"""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: ' · '.""" 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)