Add quality model option and streamline blocked-content handling

Adds a Fast/Quality model selector (Mage-Flow-Turbo vs Mage-Flow) with per-task lazy pipeline swapping and recommended step/CFG defaults. Blocked requests now return a plain blank image.

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,73 +1,62 @@
 """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")
 
-import textwrap
-
 import spaces  # MUST be first (after env setup)
 import torch
 import gradio as gr
-from PIL import Image, ImageDraw, ImageFont
+from PIL import Image
 
 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)
+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"])
+
+
+@spaces.GPU(duration=120)
 def generate(
     prompt: str,
     image=None,
@@ -78,13 +67,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,18 +91,16 @@ 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")]
 
-        # Content-safety gate: surface a block to the user instead of failing
-        # silently (gr.Warning toast + a stamped output image).
+        # Content-safety gate: blocked requests return a blank 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)
+            return Image.new("RGB", (w, h), (255, 255, 255))
 
         out = pipe_edit.edit(
             [prompt],
@@ -127,13 +114,11 @@ def generate(
         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).
+    # 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:
-        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))
+        return Image.new("RGB", (int(width), int(height)), (255, 255, 255))
 
     img = pipe_t2i.generate(
         [prompt],
@@ -160,7 +145,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 +166,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 +224,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",
     )

mage_flow/pipeline.py

@@ -286,7 +286,6 @@ def generate_images(model, prompts, neg_prompts=None, seeds=None, steps=30, cfg=
         verdict = model.txt_enc.screen_text(prompts[i])
         if verdict.violates:
             h_, w_ = _make_divisible_by_16(heights[i]), _make_divisible_by_16(widths[i])
-            print(verdict.banner())
             results[i] = make_refusal_image(verdict, height=h_, width=w_)
             continue
         active.append(i)
@@ -477,7 +476,6 @@ def generate_edits(model, prompts, ref_images, neg_prompts=None, seeds=None, ste
         verdict = model.txt_enc.screen_edit(prompts[i], pils_per_sample[i])
         if verdict.violates:
             h_, w_ = res_hw[i]
-            print(verdict.banner())
             results[i] = make_refusal_image(verdict, height=h_, width=w_)
             continue
         active.append(i)