app.py

"""Mage-Flow: Efficient Native-Resolution Foundation Model for Image Generation and Editing.

Gradio Space demo with a single unified interface: enter a prompt and optionally
upload an image. If an image is provided, the request is routed to the
Mage-Flow-Edit-Turbo model; otherwise it is routed to the Mage-Flow-Turbo model.
"""
import os

# Use flash_attention_2 for the HF text encoder (flash_attn is installed via wheel)
os.environ.setdefault("VF_HF_ATTN_IMPL", "flash_attention_2")

import textwrap

import spaces  # MUST be first (after env setup)
import torch
import gradio as gr
from PIL import Image, ImageDraw, ImageFont

from mage_flow.pipeline import MageFlowPipeline
from mage_flow.models.modules.mage_text import CATEGORY_DISPLAY

T2I_MODEL = "microsoft/Mage-Flow-Turbo"
EDIT_MODEL = "microsoft/Mage-Flow-Edit-Turbo"

pipe_t2i = MageFlowPipeline.from_pretrained(T2I_MODEL, device="cuda")
pipe_edit = MageFlowPipeline.from_pretrained(EDIT_MODEL, device="cuda")


def _block_reason_text(verdict) -> str:
    """Human-readable block reason: '<category> · <explanation>'."""
    cat = ", ".join(
        CATEGORY_DISPLAY.get(c, c) for c in verdict.categories
    ) or "policy violation"
    reason = (verdict.reason or "").strip()
    return f"{cat} · {reason}" if reason else cat


def _blocked_image(reason_text: str, height: int = 1024, width: int = 1024) -> Image.Image:
    """Neutral background image with the safety-filter message drawn on it."""
    img = Image.new("RGB", (int(width), int(height)), color=(245, 245, 245))
    draw = ImageDraw.Draw(img)
    try:
        font = ImageFont.load_default(size=28)
    except TypeError:  # older Pillow: load_default() takes no size
        font = ImageFont.load_default()

    message = (
        "The classifiers included by default in the Microsoft Mage Flow model flagged the following:\n\n" + reason_text
    )

    # Wrap each paragraph to a column width that fits the image.
    max_chars = max(20, int(width / 16))
    lines = []
    for para in message.split("\n"):
        if not para:
            lines.append("")
            continue
        lines.extend(textwrap.wrap(para, width=max_chars))

    # Vertically center the text block.
    line_h = 38
    total_h = line_h * len(lines)
    y = max(20, (int(height) - total_h) // 2)
    for line in lines:
        draw.text((40, y), line, fill=(20, 20, 20), font=font)
        y += line_h
    return img


@spaces.GPU(duration=60)
def generate(
    prompt: str,
    image=None,
    negative_prompt: str = " ",
    steps: int = 4,
    cfg: float = 1.0,
    height: int = 1024,
    width: int = 1024,
    max_size: int = 1024,
    seed: int = 42,
    progress=gr.Progress(track_tqdm=True),
):
    """Generate or edit an image with Mage-Flow.

    If ``image`` is provided, the reference image is edited with
    Mage-Flow-Edit-Turbo. Otherwise a new image is generated from the text
    prompt with Mage-Flow-Turbo.

    Args:
        prompt: Text description (generation) or edit instruction (editing).
        image: Optional reference image. When given, routes to the edit model.
        negative_prompt: What to avoid in the result.
        steps: Number of denoising steps (Turbo uses 4).
        cfg: Classifier-free guidance scale (Turbo uses 1.0).
        height: Output image height for text-to-image (multiple of 16).
        width: Output image width for text-to-image (multiple of 16).
        max_size: Longest side of edited output (0 = keep source resolution).
        seed: Random seed for reproducibility.
    """
    if not (prompt or "").strip():
        raise gr.Error("Prompt is empty.")

    if image is not None:
        # Route to the edit model when an image is provided.
        if isinstance(image, str):
            image = Image.open(image)
        refs = [image.convert("RGB")]

        # Content-safety gate: surface a block to the user instead of failing
        # silently (gr.Warning toast + a stamped output image).
        verdict = pipe_edit.model.txt_enc.screen_edit(prompt, refs)
        if verdict.violates:
            reason_text = _block_reason_text(verdict)
            gr.Warning("The classifiers included by default in the Microsoft Mage Flow model flagged the following: " + reason_text)
            w, h = refs[0].size
            return _blocked_image(reason_text, height=h, width=w)

        out = pipe_edit.edit(
            [prompt],
            [refs],
            neg_prompts=[negative_prompt or " "],
            seeds=[int(seed)],
            steps=int(steps),
            cfg=float(cfg),
            max_size=int(max_size) if max_size else None,
        )[0]
        return out

    # No image: route to the text-to-image model.
    # Content-safety gate: surface a block to the user instead of failing
    # silently (gr.Warning toast + a stamped output image).
    verdict = pipe_t2i.model.txt_enc.screen_text(prompt)
    if verdict.violates:
        reason_text = _block_reason_text(verdict)
        gr.Warning("The classifiers included by default in the Microsoft Mage Flow model flagged the following: " + reason_text)
        return _blocked_image(reason_text, height=int(height), width=int(width))

    img = pipe_t2i.generate(
        [prompt],
        neg_prompts=[negative_prompt or " "],
        seeds=[int(seed)],
        steps=int(steps),
        cfg=float(cfg),
        heights=[int(height)],
        widths=[int(width)],
    )[0]
    return img


ASSETS_DIR = os.path.join(os.path.dirname(__file__), "mage_flow", "assets")

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

with gr.Blocks(css=CSS) as demo:
    with gr.Column(elem_id="col-container"):
        gr.Markdown(
            "# Mage-Flow\n"
            "Efficient Native-Resolution Foundation Model for Image Generation and Editing. "
            "Enter a prompt to generate an image, or upload an image to edit it.\n\n"
            "Models: [Mage-Flow-Turbo](https://huggingface.co/microsoft/Mage-Flow-Turbo) & "
            "[Mage-Flow-Edit-Turbo](https://huggingface.co/microsoft/Mage-Flow-Edit-Turbo) | "
            "[Paper](https://huggingface.co/papers/2607.19064) | "
            "[GitHub](https://github.com/microsoft/Mage)"
        )

        with gr.Row():
            with gr.Column(scale=1):
                with gr.Row():
                    prompt = gr.Textbox(
                        label="Prompt",
                        show_label=False,
                        max_lines=3,
                        placeholder="Describe an image to generate, or an edit instruction for an uploaded image",
                        container=False,
                        scale=4,
                    )
                    run_btn = gr.Button("Run", variant="primary", scale=1)

                with gr.Accordion("Input image (optional — enables editing)", open=True):
                    image = gr.Image(
                        type="pil",
                        label="Input image",
                        show_label=False,
                        height=300,
                    )

                with gr.Accordion("Advanced Settings", open=False):
                    negative_prompt = gr.Textbox(label="Negative prompt", value=" ", lines=1)
                    with gr.Row():
                        steps = gr.Slider(1, 50, value=4, step=1, label="Steps")
                        cfg = gr.Slider(1.0, 10.0, value=1.0, step=0.5, label="CFG")
                    with gr.Row():
                        height = gr.Slider(256, 1536, value=1024, step=16, label="Height (text→image)")
                        width = gr.Slider(256, 1536, value=1024, step=16, label="Width (text→image)")
                    max_size = gr.Slider(
                        0, 1536, value=1024, step=16,
                        label="Max output side for editing (0 = keep source size)",
                    )
                    seed = gr.Number(value=42, precision=0, label="Seed")

            with gr.Column(scale=1):
                result = gr.Image(type="pil", label="Output", height=560)

        gr.Markdown("### Text → Image examples")
        gr.Examples(
            examples=[
                ["A close-up portrait of an elderly Hausa man with deep wrinkles, wearing a traditional hat, soft natural lighting, ultra realistic."],
                ["A serene mountain landscape at sunset, with snow-capped peaks reflecting golden light, photorealistic."],
                ["A cute robot playing a guitar in a neon-lit cyberpunk city, digital art style."],
            ],
            inputs=[prompt],
            outputs=result,
            fn=generate,
            cache_examples=True,
            cache_mode="lazy",
        )

        gr.Markdown("### Image editing examples")
        gr.Examples(
            examples=[
                ["change the background to a city street", os.path.join(ASSETS_DIR, "dog.jpg")],
                ["make it look like a painting", os.path.join(ASSETS_DIR, "cuisine.jpg")],
                ["add a hat to the person", os.path.join(ASSETS_DIR, "portrait.jpg")],
            ],
            inputs=[prompt, image],
            outputs=result,
            fn=generate,
            cache_examples=True,
            cache_mode="lazy",
        )

    inputs = [prompt, image, negative_prompt, steps, cfg, height, width, max_size, seed]
    run_btn.click(lambda: None, None, result).then(
        generate, inputs, result, api_name="generate",
    )
    prompt.submit(lambda: None, None, result).then(
        generate, inputs, result, api_name=False,
    )

if __name__ == "__main__":
    demo.launch(theme=gr.themes.Citrus(), mcp_server=True, show_error=True)