Spaces:
Running on Zero
Running on Zero
| """Mage-Flow: Efficient Native-Resolution Foundation Model for Image Generation and Editing. | |
| Gradio Space demo with a single unified interface: image presence selects | |
| editing vs. generation, while the model control selects fast vs. quality. | |
| """ | |
| import gc | |
| import os | |
| import threading | |
| # 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 spaces # MUST be first (after env setup) | |
| import torch | |
| import gradio as gr | |
| from PIL import Image | |
| from mage_flow.pipeline import MageFlowPipeline | |
| MODEL_VARIANTS = { | |
| "turbo": { | |
| "t2i": "microsoft/Mage-Flow-Turbo", "edit": "microsoft/Mage-Flow-Edit-Turbo", | |
| "t2i_steps": 4, "edit_steps": 4, "cfg": 1.0, | |
| }, | |
| "quality": { | |
| "t2i": "microsoft/Mage-Flow", "edit": "microsoft/Mage-Flow-Edit", | |
| "t2i_steps": 20, "edit_steps": 30, "cfg": 5.0, | |
| }, | |
| } | |
| _pipe_slots = { | |
| "t2i": {"variant": "turbo", "pipe": MageFlowPipeline.from_pretrained(MODEL_VARIANTS["turbo"]["t2i"], device="cuda")}, | |
| "edit": {"variant": "turbo", "pipe": MageFlowPipeline.from_pretrained(MODEL_VARIANTS["turbo"]["edit"], device="cuda")}, | |
| } | |
| _pipe_lock = threading.Lock() | |
| def _get_pipe(task: str, variant: str): | |
| """Keep one loaded variant per task, matching the original two-pipeline footprint.""" | |
| with _pipe_lock: | |
| slot = _pipe_slots.get(task) | |
| if slot and slot["variant"] == variant: | |
| return slot["pipe"] | |
| if slot: | |
| del _pipe_slots[task] | |
| del slot | |
| gc.collect() | |
| torch.cuda.empty_cache() | |
| pipe = MageFlowPipeline.from_pretrained(MODEL_VARIANTS[variant][task], device="cuda") | |
| _pipe_slots[task] = {"variant": variant, "pipe": pipe} | |
| return pipe | |
| def _recommended(variant: str, image): | |
| spec = MODEL_VARIANTS[variant] | |
| return (spec["edit_steps"] if image is not None else spec["t2i_steps"], spec["cfg"]) | |
| 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, | |
| model_variant: str = "turbo", | |
| progress=gr.Progress(track_tqdm=True), | |
| ): | |
| """Generate or edit an image with Mage-Flow. | |
| If ``image`` is provided, route to the selected edit model; otherwise route | |
| to the selected text-to-image model. | |
| 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. | |
| pipe_edit = _get_pipe("edit", model_variant) | |
| if isinstance(image, str): | |
| image = Image.open(image) | |
| refs = [image.convert("RGB")] | |
| # Content-safety gate: blocked requests return a blank image. | |
| verdict = pipe_edit.model.txt_enc.screen_edit(prompt, refs) | |
| if verdict.violates: | |
| w, h = refs[0].size | |
| return Image.new("RGB", (w, h), (255, 255, 255)) | |
| 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: blocked requests return a blank image. | |
| pipe_t2i = _get_pipe("t2i", model_variant) | |
| verdict = pipe_t2i.model.txt_enc.screen_text(prompt) | |
| if verdict.violates: | |
| return Image.new("RGB", (int(width), int(height)), (255, 255, 255)) | |
| 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](https://huggingface.co/microsoft/Mage-Flow), " | |
| "[Mage-Flow-Turbo](https://huggingface.co/microsoft/Mage-Flow-Turbo), " | |
| "[Mage-Flow-Edit](https://huggingface.co/microsoft/Mage-Flow-Edit), " | |
| "[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) | |
| model_variant = gr.Radio( | |
| [("Mage-Flow-Turbo · Fast", "turbo"), ("Mage-Flow · Quality", "quality")], | |
| value="turbo", label="Model", | |
| ) | |
| 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", | |
| ) | |
| model_variant.change(_recommended, [model_variant, image], [steps, cfg], api_name=False) | |
| image.change(_recommended, [model_variant, image], [steps, cfg], api_name=False) | |
| inputs = [prompt, image, negative_prompt, steps, cfg, height, width, max_size, seed, model_variant] | |
| 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) | |