app.py

import os

os.environ.setdefault("PYTORCH_CUDA_ALLOC_CONF", "expandable_segments:True")

import spaces  # MUST come before torch / any CUDA-touching import

import random
import time

import gradio as gr
import torch
from huggingface_hub import hf_hub_download
from diffusers import Flux2KleinPipeline, Flux2Transformer2DModel

DTYPE = torch.bfloat16
BASE_REPO = "black-forest-labs/FLUX.2-klein-4B"
CKPT_REPO = "shy0423/MGFlow"
CKPT_FILE = "Checkpoints/T2I/FLUX2-klein-4B_COCO-joint.pth"
TRANSFORMER_CONFIG = {
    "attention_head_dim": 128,
    "axes_dims_rope": [32, 32, 32, 32],
    "eps": 1e-06,
    "guidance_embeds": False,
    "in_channels": 128,
    "joint_attention_dim": 7680,
    "mlp_ratio": 3.0,
    "num_attention_heads": 24,
    "num_layers": 5,
    "num_single_layers": 20,
    "out_channels": None,
    "patch_size": 1,
    "rope_theta": 2000,
    "timestep_guidance_channels": 256,
}

MAX_SEED = 2**32 - 1


def _load_mgflow_transformer() -> Flux2Transformer2DModel:
    """Build the FLUX.2 [klein] 4B transformer and swap in the MGFlow post-trained weights.

    The released MGFlow checkpoint is a torch.save state dict whose keys are the
    diffusers `Flux2Transformer2DModel` keys prefixed with `transformer.` and stored
    in float32. We build the model on the meta device, stream the checkpoint with
    mmap, cast each tensor to bfloat16, and assign — so we never hold two full
    copies of the 4B checkpoint in RAM and never download the base transformer
    weights at all.
    """
    ckpt_path = hf_hub_download(CKPT_REPO, CKPT_FILE)
    print(f"[mgflow] checkpoint at {ckpt_path}", flush=True)
    ckpt = torch.load(ckpt_path, map_location="cpu", mmap=True, weights_only=True)
    raw = ckpt["model"] if "model" in ckpt else ckpt
    print(f"[mgflow] checkpoint step: {ckpt.get('step', 'n/a')}", flush=True)

    with torch.device("meta"):
        model = Flux2Transformer2DModel.from_config(TRANSFORMER_CONFIG)
    state = {k[len("transformer."):]: v.to(DTYPE) for k, v in raw.items() if k.startswith("transformer.")}
    model.load_state_dict(state, strict=True, assign=True)
    del ckpt, raw, state
    print("[mgflow] transformer weights loaded", flush=True)
    return model


print("[mgflow] loading pipeline (text encoder, vae, scheduler)...", flush=True)
pipe = Flux2KleinPipeline.from_pretrained(
    BASE_REPO,
    transformer=_load_mgflow_transformer(),
    torch_dtype=DTYPE,
)
pipe.to("cuda")
print("[mgflow] pipeline ready on cuda", flush=True)


@spaces.GPU(duration=20)
def generate(
    prompt: str,
    seed: int = 0,
    randomize_seed: bool = True,
    width: int = 512,
    height: int = 512,
    progress=gr.Progress(track_tqdm=True),
):
    """Generate an image from a text prompt in a single step.

    MGFlow post-trains FLUX.2 [klein] 4B into a one-step generator: one forward
    pass of the transformer (1 NFE, no classifier-free guidance) produces the
    image, matching the paper's evaluation protocol at 512x512.

    Args:
        prompt: text description of the image to generate.
        seed: RNG seed for reproducibility.
        randomize_seed: draw a fresh seed instead of using the one above.
        width: output image width in pixels.
        height: output image height in pixels.
    """
    if randomize_seed:
        seed = random.randint(0, MAX_SEED)
    generator = torch.Generator(device="cuda").manual_seed(int(seed))
    t0 = time.perf_counter()
    image = pipe(
        prompt=prompt,
        width=int(width),
        height=int(height),
        num_inference_steps=1,
        guidance_scale=1.0,
        max_sequence_length=512,
        generator=generator,
    ).images[0]
    elapsed = time.perf_counter() - t0
    print(f"[mgflow] generated in {elapsed:.2f}s (seed {seed})", flush=True)
    return image, gr.update(value=int(seed))


CSS = """
#col-container { max-width: 1100px; margin: 0 auto; }
.dark .gradio-container { color: var(--body-text-color); }
main, .gradio-container, .gradio-container .fillable { max-width: 1100px !important; margin: 0 auto !important; }
"""

with gr.Blocks() as demo:
    with gr.Column(elem_id="col-container"):
        gr.Markdown(
            """
            # MGFlow · 512×512 in One Step

            MGFlow post-trains FLUX.2 [klein] 4B into a **512×512 one-step generator** (**1 NFE, no CFG**).
            The released **[joint image–text checkpoint](https://huggingface.co/shy0423/MGFlow/blob/main/Checkpoints/T2I/FLUX2-klein-4B_COCO-joint.pth)** achieves **state-of-the-art GenEval (0.900) and
            PickScore (21.98) among one-step models post-trained from this backbone**.

            [Paper](https://huggingface.co/papers/2609.35763) · [Code](https://github.com/shihaoyang0423/MGFlow) · [Checkpoints](https://huggingface.co/shy0423/MGFlow) · [Project page](https://shihaoyang0423.github.io/MGFlow-website/)
            """
        )
        with gr.Row():
            prompt = gr.Textbox(
                label="Prompt",
                show_label=False,
                placeholder="Describe the image you want…",
                lines=2,
                scale=4,
                container=False,
            )
            run = gr.Button("Generate", variant="primary", scale=1)
        output = gr.Image(label="Generated image (1 step)", height=512)

        with gr.Accordion("Advanced settings", open=False):
            with gr.Row():
                seed = gr.Number(label="Seed", value=0, precision=0, minimum=0, maximum=MAX_SEED)
                randomize_seed = gr.Checkbox(label="Randomize seed", value=True)
            with gr.Row():
                width = gr.Slider(label="Width", minimum=256, maximum=1024, step=32, value=512)
                height = gr.Slider(label="Height", minimum=256, maximum=1024, step=32, value=512)

        gr.Examples(
            examples=[
                ["A fluffy orange cat sitting on a blue cushion, soft window light, simple background, realistic photography."],
                ["A red ceramic mug on a light wooden table, soft morning light, clean background, product photography."],
                ["A ripe peach on a white ceramic plate, pale pink background, soft studio lighting, still-life photography."],
                ["A close-up portrait of a red panda, blurred green foliage, soft daylight, wildlife photography."],
                ["A white lighthouse on a grassy coastal hill, calm blue sea, clear sky, watercolor illustration."],
                ["A single oak tree on a gentle grassy hill, clear blue sky, warm afternoon light, simple watercolor illustration."],
            ],
            inputs=[prompt],
            outputs=[output, seed],
            fn=generate,
            cache_examples=True,
            cache_mode="lazy",
        )

    run.click(
        generate,
        inputs=[prompt, seed, randomize_seed, width, height],
        outputs=[output, seed],
        api_name="generate",
    )
    prompt.submit(
        generate,
        inputs=[prompt, seed, randomize_seed, width, height],
        outputs=[output, seed],
        api_name="generate",
    )

demo.launch(theme=gr.themes.Citrus(), css=CSS, mcp_server=True)