Surface content-safety blocks: gr.Warning + stamped output image for generate and edit routes

app.py

@@ -9,12 +9,15 @@ 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
+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"
@@ -23,6 +26,48 @@ 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 safety 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,
@@ -61,6 +106,16 @@ def generate(
         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(reason_text)
+            w, h = refs[0].size
+            return _blocked_image(reason_text, height=h, width=w)
+
         out = pipe_edit.edit(
             [prompt],
             [refs],
@@ -73,6 +128,14 @@ 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).
+    verdict = pipe_t2i.model.txt_enc.screen_text(prompt)
+    if verdict.violates:
+        reason_text = _block_reason_text(verdict)
+        gr.Warning(reason_text)
+        return _blocked_image(reason_text, height=int(height), width=int(width))
+
     img = pipe_t2i.generate(
         [prompt],
         neg_prompts=[negative_prompt or " "],