Add quality model option to unified Space demo

README.md

@@ -17,19 +17,19 @@ Efficient Native-Resolution Foundation Model for Image Generation and Editing.
 
 This Space demonstrates [Mage-Flow](https://huggingface.co/papers/2607.19064) from Microsoft, a 4B-scale image generation and editing foundation model through a single unified interface:
 
-- **No image** → the prompt generates a new image (routes to [Mage-Flow-Turbo](https://huggingface.co/microsoft/Mage-Flow-Turbo)).
-- **With an uploaded image** → the prompt is treated as an edit instruction (routes to [Mage-Flow-Edit-Turbo](https://huggingface.co/microsoft/Mage-Flow-Edit-Turbo)).
+- **No image** → text-to-image generation.
+- **With an uploaded image** → instruction-based image editing.
 
-The demo uses the Turbo variants for fast, few-step inference (4 steps, CFG 1.0).
+Choose **Mage-Flow-Turbo · Fast** (4 steps, CFG 1.0) or **Mage-Flow · Quality**. The quality variant uses 20 steps / CFG 5.0 for generation and 30 steps / CFG 5.0 for editing; uploading or removing an image updates these defaults automatically.
 
 ## Usage
 
 1. Enter a prompt.
 2. *(Optional)* Upload an image to edit — the prompt becomes an edit instruction.
-3. Click **Run**.
+3. Choose Fast or Quality, then click **Run**.
 
 ## References
 
 - Paper: [Mage-Flow: An Efficient Native-Resolution Foundation Model for Image Generation and Editing](https://huggingface.co/papers/2607.19064)
 - GitHub: [microsoft/Mage](https://github.com/microsoft/Mage)
-- Models: [microsoft/Mage-Flow-Turbo](https://huggingface.co/microsoft/Mage-Flow-Turbo), [microsoft/Mage-Flow-Edit-Turbo](https://huggingface.co/microsoft/Mage-Flow-Edit-Turbo)
+- 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)

app.py

@@ -1,10 +1,11 @@
 """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.
+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")
@@ -19,11 +20,43 @@ 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")
+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 _block_reason_text(verdict) -> str:
@@ -67,7 +100,7 @@ def _blocked_image(reason_text: str, height: int = 1024, width: int = 1024) -> I
     return img
 
 
-@spaces.GPU(duration=60)
+@spaces.GPU(duration=120)
 def generate(
     prompt: str,
     image=None,
@@ -78,13 +111,13 @@ def generate(
     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, the reference image is edited with
-    Mage-Flow-Edit-Turbo. Otherwise a new image is generated from the text
-    prompt with Mage-Flow-Turbo.
+    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).
@@ -102,6 +135,7 @@ def generate(
 
     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")]
@@ -129,6 +163,7 @@ def generate(
     # 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).
+    pipe_t2i = _get_pipe("t2i", model_variant)
     verdict = pipe_t2i.model.txt_enc.screen_text(prompt)
     if verdict.violates:
         reason_text = _block_reason_text(verdict)
@@ -160,7 +195,9 @@ with gr.Blocks(css=CSS) as demo:
             "# 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) & "
+            "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)"
@@ -179,6 +216,11 @@ with gr.Blocks(css=CSS) as demo:
                     )
                     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",
@@ -232,7 +274,10 @@ with gr.Blocks(css=CSS) as demo:
             cache_mode="lazy",
         )
 
-    inputs = [prompt, image, negative_prompt, steps, cfg, height, width, max_size, seed]
+    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",
     )