microsoft/mage-flow

1"""Mage-Flow: Efficient Native-Resolution Foundation Model for Image Generation and Editing.2 3Gradio Space demo with a single unified interface: image presence selects4editing vs. generation, while the model control selects fast vs. quality.5"""6import os7 8# Use flash_attention_2 for the HF text encoder (flash_attn is installed via wheel)9os.environ.setdefault("VF_HF_ATTN_IMPL", "flash_attention_2")10 11import spaces  # MUST be first (after env setup)12import gradio as gr13from PIL import Image14 15from mage_flow.pipeline import MageFlowPipeline16 17MODEL_VARIANTS = {18    "turbo": {19        "t2i": "microsoft/Mage-Flow-Turbo", "edit": "microsoft/Mage-Flow-Edit-Turbo",20        "t2i_steps": 4, "edit_steps": 4, "cfg": 1.0,21    },22    "quality": {23        "t2i": "microsoft/Mage-Flow", "edit": "microsoft/Mage-Flow-Edit",24        "t2i_steps": 20, "edit_steps": 30, "cfg": 5.0,25    },26}27 28# ZeroGPU requires every model to be placed on CUDA at module scope: the backend29# registers the weights, offloads them to disk at startup, and streams them into30# VRAM for each @spaces.GPU call. Loading inside the GPU function instead would31# charge every switch to the caller's GPU quota and is not carried across the32# forked GPU workers.33PIPES = {34    (task, variant): MageFlowPipeline.from_pretrained(spec[task], device="cuda")35    for variant, spec in MODEL_VARIANTS.items()36    for task in ("t2i", "edit")37}38 39 40def _recommended(variant: str, image):41    spec = MODEL_VARIANTS[variant]42    return (spec["edit_steps"] if image is not None else spec["t2i_steps"], spec["cfg"])43 44 45@spaces.GPU(duration=120)46def generate(47    prompt: str,48    image=None,49    negative_prompt: str = " ",50    steps: int = 4,51    cfg: float = 1.0,52    height: int = 1024,53    width: int = 1024,54    max_size: int = 1024,55    seed: int = 42,56    model_variant: str = "turbo",57    progress=gr.Progress(track_tqdm=True),58):59    """Generate or edit an image with Mage-Flow.60 61    If ``image`` is provided, route to the selected edit model; otherwise route62    to the selected text-to-image model.63 64    Args:65        prompt: Text description (generation) or edit instruction (editing).66        image: Optional reference image. When given, routes to the edit model.67        negative_prompt: What to avoid in the result.68        steps: Number of denoising steps (Turbo uses 4).69        cfg: Classifier-free guidance scale (Turbo uses 1.0).70        height: Output image height for text-to-image (multiple of 16).71        width: Output image width for text-to-image (multiple of 16).72        max_size: Longest side of edited output (0 = keep source resolution).73        seed: Random seed for reproducibility.74    """75    if not (prompt or "").strip():76        raise gr.Error("Prompt is empty.")77 78    if image is not None:79        # Route to the edit model when an image is provided.80        pipe_edit = PIPES[("edit", model_variant)]81        if isinstance(image, str):82            image = Image.open(image)83        refs = [image.convert("RGB")]84 85        # Content-safety gate: blocked requests return a blank image.86        verdict = pipe_edit.model.txt_enc.screen_edit(prompt, refs)87        if verdict.violates:88            w, h = refs[0].size89            return Image.new("RGB", (w, h), (255, 255, 255))90 91        out = pipe_edit.edit(92            [prompt],93            [refs],94            neg_prompts=[negative_prompt or " "],95            seeds=[int(seed)],96            steps=int(steps),97            cfg=float(cfg),98            max_size=int(max_size) if max_size else None,99        )[0]100        return out101 102    # No image: route to the text-to-image model.103    # Content-safety gate: blocked requests return a blank image.104    pipe_t2i = PIPES[("t2i", model_variant)]105    verdict = pipe_t2i.model.txt_enc.screen_text(prompt)106    if verdict.violates:107        return Image.new("RGB", (int(width), int(height)), (255, 255, 255))108 109    img = pipe_t2i.generate(110        [prompt],111        neg_prompts=[negative_prompt or " "],112        seeds=[int(seed)],113        steps=int(steps),114        cfg=float(cfg),115        heights=[int(height)],116        widths=[int(width)],117    )[0]118    return img119 120 121ASSETS_DIR = os.path.join(os.path.dirname(__file__), "mage_flow", "assets")122 123CSS = """124#col-container { margin: 0 auto; max-width: 1100px; }125.dark .gradio-container { color: var(--body-text-color); }126"""127 128with gr.Blocks(css=CSS) as demo:129    with gr.Column(elem_id="col-container"):130        gr.Markdown(131            "# Mage-Flow\n"132            "Efficient Native-Resolution Foundation Model for Image Generation and Editing. "133            "Enter a prompt to generate an image, or upload an image to edit it.\n\n"134            "Models: [Mage-Flow](https://huggingface.co/microsoft/Mage-Flow), "135            "[Mage-Flow-Turbo](https://huggingface.co/microsoft/Mage-Flow-Turbo), "136            "[Mage-Flow-Edit](https://huggingface.co/microsoft/Mage-Flow-Edit), "137            "[Mage-Flow-Edit-Turbo](https://huggingface.co/microsoft/Mage-Flow-Edit-Turbo) | "138            "[Paper](https://huggingface.co/papers/2607.19064) | "139            "[GitHub](https://github.com/microsoft/Mage)"140        )141 142        with gr.Row():143            with gr.Column(scale=1):144                with gr.Row():145                    prompt = gr.Textbox(146                        label="Prompt",147                        show_label=False,148                        max_lines=3,149                        placeholder="Describe an image to generate, or an edit instruction for an uploaded image",150                        container=False,151                        scale=4,152                    )153                    run_btn = gr.Button("Run", variant="primary", scale=1)154 155                model_variant = gr.Radio(156                    [("Mage-Flow-Turbo · Fast", "turbo"), ("Mage-Flow · Quality", "quality")],157                    value="turbo", label="Model",158                )159 160                with gr.Accordion("Input image (optional — enables editing)", open=True):161                    image = gr.Image(162                        type="pil",163                        label="Input image",164                        show_label=False,165                        height=300,166                    )167 168                with gr.Accordion("Advanced Settings", open=False):169                    negative_prompt = gr.Textbox(label="Negative prompt", value=" ", lines=1)170                    with gr.Row():171                        steps = gr.Slider(1, 50, value=4, step=1, label="Steps")172                        cfg = gr.Slider(1.0, 10.0, value=1.0, step=0.5, label="CFG")173                    with gr.Row():174                        height = gr.Slider(256, 1536, value=1024, step=16, label="Height (text→image)")175                        width = gr.Slider(256, 1536, value=1024, step=16, label="Width (text→image)")176                    max_size = gr.Slider(177                        0, 1536, value=1024, step=16,178                        label="Max output side for editing (0 = keep source size)",179                    )180                    seed = gr.Number(value=42, precision=0, label="Seed")181 182            with gr.Column(scale=1):183                result = gr.Image(type="pil", label="Output", height=560)184 185        gr.Markdown("### Text → Image examples")186        gr.Examples(187            examples=[188                ["A close-up portrait of an elderly Hausa man with deep wrinkles, wearing a traditional hat, soft natural lighting, ultra realistic."],189                ["A serene mountain landscape at sunset, with snow-capped peaks reflecting golden light, photorealistic."],190                ["A cute robot playing a guitar in a neon-lit cyberpunk city, digital art style."],191            ],192            inputs=[prompt],193            outputs=result,194            fn=generate,195            cache_examples=True,196            cache_mode="lazy",197        )198 199        gr.Markdown("### Image editing examples")200        gr.Examples(201            examples=[202                ["change the background to a city street", os.path.join(ASSETS_DIR, "dog.jpg")],203                ["make it look like a painting", os.path.join(ASSETS_DIR, "cuisine.jpg")],204                ["add a hat to the person", os.path.join(ASSETS_DIR, "portrait.jpg")],205            ],206            inputs=[prompt, image],207            outputs=result,208            fn=generate,209            cache_examples=True,210            cache_mode="lazy",211        )212 213    model_variant.change(_recommended, [model_variant, image], [steps, cfg], api_name=False)214    image.change(_recommended, [model_variant, image], [steps, cfg], api_name=False)215 216    inputs = [prompt, image, negative_prompt, steps, cfg, height, width, max_size, seed, model_variant]217    run_btn.click(lambda: None, None, result).then(218        generate, inputs, result, api_name="generate",219    )220    prompt.submit(lambda: None, None, result).then(221        generate, inputs, result, api_name=False,222    )223 224if __name__ == "__main__":225    demo.launch(theme=gr.themes.Citrus(), mcp_server=True, show_error=True)226