Add PE-T2I/PE-I2I prompt enhancement nodes and NCII guard for image-input requests

app.py

@@ -1,30 +1,44 @@
 """Qwen-Image 2.1 — Gradio Workflow on ZeroGPU.
 
 A node-based canvas (gr.Workflow) exposing the diffusers QwenImage21Pipeline:
-  - text_to_image: prompt -> image
-  - edit_image:    condition image + instruction -> edited image (chained after
-                   text_to_image, or fed from an uploaded image reference)
 
-The pipeline is placed on `cuda` at module level: outside @spaces.GPU functions
-PyTorch runs in CUDA emulation mode, and a real ZeroGPU is attached only while
-a decorated function executes.
+  - enhance_prompt_t2i: prompt -> rewritten prompt   (Qwen/Qwen-Image-2.1-PE-T2I)
+  - text_to_image:      prompt -> image              (Qwen/Qwen-Image-2.1)
+  - enhance_prompt_i2i: image + instruction -> rewritten instruction
+                                                     (Qwen/Qwen-Image-2.1-PE-I2I)
+  - edit_image:         condition image + instruction -> edited image
+
+Image-input requests pass through an NCII prompt classifier
+(hfmlsoc/ncii-guard-v02) before any rewriting or editing runs.
+
+The diffusion pipeline stays resident on `cuda` (loaded at module level, as
+ZeroGPU requires). The 9B prompt-rewriting models don't fit alongside it on a
+48GB card, so they live in CPU RAM and are moved onto the GPU only inside
+their @spaces.GPU calls. The 270M guard classifier runs on CPU.
 """
 
 import base64
+import io
+import json
 import os
 import urllib.parse
 import urllib.request
 
 import gradio as gr
-from gradio_client import utils as client_utils
-from gradio.utils import get_upload_folder
 import spaces
 import torch
+from gradio_client import utils as client_utils
+from gradio.utils import get_upload_folder
+from huggingface_hub import hf_hub_download
 from PIL import Image
 
 from diffusers import QwenImage21Pipeline
 
 MODEL_ID = os.environ.get("QWEN_IMAGE_MODEL", "Qwen/Qwen-Image-2.1")
+PE_T2I_ID = "Qwen/Qwen-Image-2.1-PE-T2I"
+PE_I2I_ID = "Qwen/Qwen-Image-2.1-PE-I2I"
+GUARD_ID = "hfmlsoc/ncii-guard-v02"
+GUARD_THRESHOLD = 0.5
 
 # The model repo is private/gated. Per the Hub auth docs, the HF_TOKEN
 # environment variable is used implicitly for all Hub requests and takes
@@ -39,14 +53,46 @@ if not os.environ.get("HF_TOKEN"):
 pipe = QwenImage21Pipeline.from_pretrained(MODEL_ID, dtype=torch.bfloat16)
 pipe.to("cuda")
 
+# --- Prompt rewriting models (resident in CPU RAM, moved to GPU per call) ---
+from transformers import (
+    AutoModelForCausalLM,
+    AutoModelForImageTextToText,
+    AutoProcessor,
+    AutoTokenizer,
+)
+
+pe_t2i_tokenizer = AutoTokenizer.from_pretrained(PE_T2I_ID)
+pe_t2i = AutoModelForCausalLM.from_pretrained(PE_T2I_ID, dtype=torch.bfloat16).eval()
+pe_t2i_system = open(hf_hub_download(PE_T2I_ID, "system_prompt.txt")).read().strip()
+
+pe_i2i_processor = AutoProcessor.from_pretrained(PE_I2I_ID)
+pe_i2i = AutoModelForImageTextToText.from_pretrained(PE_I2I_ID, dtype=torch.bfloat16).eval()
+pe_i2i_system = open(hf_hub_download(PE_I2I_ID, "system_prompt.txt")).read().strip()
+
+# --- NCII guard for image-input requests (runs on CPU) ---
+from transformers import AutoModelForSequenceClassification
+
+guard_tokenizer = AutoTokenizer.from_pretrained(GUARD_ID)  # carries the normalizer
+guard = AutoModelForSequenceClassification.from_pretrained(GUARD_ID).eval()
+
+
+def _check_prompt_guard(prompt: str) -> None:
+    """Reject image-editing prompts the NCII classifier flags. The error is
+    deliberately generic and does not say which classifier fired."""
+    batch = guard_tokenizer(
+        [prompt], truncation=True, max_length=256, padding=True, return_tensors="pt"
+    )
+    with torch.no_grad():
+        prob = torch.softmax(guard(**batch).logits.float(), dim=-1)[0, 1].item()
+    if prob >= GUARD_THRESHOLD:
+        raise gr.Error("prompt invalid based on our classifiers, try again")
+
 
 def _to_pil(image) -> Image.Image:
     """Accept whatever the canvas hands a bound function for an image port:
     a PIL image, a local path, a /gradio_api/file= reference, an http(s) or
     data: URL, or a file dict carrying any of those. Mirrors _file_ref in
     gradio.workflow — canvas file values carry only `url`, no `path`."""
-    import io
-
     if isinstance(image, Image.Image):
         return image
     if isinstance(image, dict):
@@ -80,24 +126,99 @@ def _steps(value, default: int = 40) -> int:
     return default if value is None else int(value)
 
 
-@spaces.GPU(duration=120)
+def _parse_rewrite(gen: str, fallback: str) -> str:
+    """Split the PE model's <think> block from its JSON answer and return the
+    rewritten prompt. Fall back to the original prompt if parsing fails."""
+    _, _, answer = gen.partition("</think>")
+    try:
+        return json.loads(answer.strip()).get("rewritten_prompt") or fallback
+    except (json.JSONDecodeError, AttributeError):
+        return fallback
+
+
+@spaces.GPU(duration=300)
+def enhance_prompt_t2i(prompt: str) -> str:
+    """Rewrite a brief text-to-image request into a detailed English prompt
+    with Qwen-Image-2.1-PE-T2I."""
+    text = pe_t2i_tokenizer.apply_chat_template(
+        [{"role": "system", "content": pe_t2i_system},
+         {"role": "user", "content": prompt}],
+        tokenize=False, add_generation_prompt=True, enable_thinking=True,
+    )
+    pe_t2i.to("cuda")
+    try:
+        inputs = pe_t2i_tokenizer(text, return_tensors="pt").to("cuda")
+        with torch.no_grad():
+            out = pe_t2i.generate(
+                **inputs, max_new_tokens=4096,
+                do_sample=True, temperature=1.0, top_p=0.95, top_k=20,
+            )
+        gen = pe_t2i_tokenizer.decode(
+            out[0, inputs["input_ids"].shape[1]:], skip_special_tokens=True
+        )
+    finally:
+        pe_t2i.to("cpu")
+        torch.cuda.empty_cache()
+    return _parse_rewrite(gen, prompt)
+
+
+@spaces.GPU(duration=300)
+def enhance_prompt_i2i(image, instruction: str) -> str:
+    """Rewrite an image-editing instruction against the condition image with
+    Qwen-Image-2.1-PE-I2I."""
+    _check_prompt_guard(instruction)
+    pil_image = _to_pil(image).convert("RGB")
+    messages = [
+        {"role": "system", "content": [{"type": "text", "text": pe_i2i_system}]},
+        {"role": "user", "content": [
+            {"type": "image", "image": pil_image},
+            {"type": "text", "text": instruction},
+        ]},
+    ]
+    pe_i2i.to("cuda")
+    try:
+        inputs = pe_i2i_processor.apply_chat_template(
+            messages, add_generation_prompt=True, tokenize=True,
+            return_dict=True, return_tensors="pt", enable_thinking=True,
+        ).to("cuda")
+        with torch.no_grad():
+            out = pe_i2i.generate(
+                **inputs, max_new_tokens=4096,
+                do_sample=True, temperature=1.0, top_p=0.95, top_k=20,
+            )
+        gen = pe_i2i_processor.tokenizer.decode(
+            out[0, inputs["input_ids"].shape[1]:], skip_special_tokens=True
+        )
+    finally:
+        pe_i2i.to("cpu")
+        torch.cuda.empty_cache()
+    return _parse_rewrite(gen, instruction)
+
+
+@spaces.GPU(duration=180)
 def text_to_image(prompt: str, steps: int = 40) -> dict:
     """Generate an image from a text prompt with Qwen-Image 2.1."""
     image = pipe(prompt, num_inference_steps=_steps(steps)).images[0]
     return _save(image)
 
 
-@spaces.GPU(duration=120)
+@spaces.GPU(duration=180)
 def edit_image(image, instruction: str, steps: int = 40) -> dict:
     """Edit a condition image following an instruction (image-conditioned
     generation with Qwen-Image 2.1)."""
+    _check_prompt_guard(instruction)
     edited = pipe(instruction, image=_to_pil(image), num_inference_steps=_steps(steps)).images[0]
     return _save(edited)
 
 
 demo = gr.Workflow(
     graph=os.path.join(os.path.dirname(__file__), "workflow.json"),
-    bind={"text_to_image": text_to_image, "edit_image": edit_image},
+    bind={
+        "enhance_prompt_t2i": enhance_prompt_t2i,
+        "enhance_prompt_i2i": enhance_prompt_i2i,
+        "text_to_image": text_to_image,
+        "edit_image": edit_image,
+    },
 )
 
 if __name__ == "__main__":

workflow.json

@@ -9,24 +9,9 @@
       "asset_type": "text",
       "x": 40,
       "y": 60,
-      "inputs": [
-        {
-          "id": "in",
-          "label": "Text",
-          "type": "text"
-        }
-      ],
-      "outputs": [
-        {
-          "id": "out",
-          "label": "Text",
-          "type": "text",
-          "default_value": "A capybara wearing a wizard hat, oil painting"
-        }
-      ],
-      "data": {
-        "out": "A capybara wearing a wizard hat, oil painting"
-      }
+      "data": {"out": "A capybara wearing a wizard hat, oil painting"},
+      "inputs": [{"id": "in", "label": "Text", "type": "text"}],
+      "outputs": [{"id": "out", "label": "Text", "type": "text"}]
     },
     {
       "id": "ref_image",
@@ -35,21 +20,9 @@
       "asset_type": "image",
       "x": 40,
       "y": 380,
-      "inputs": [
-        {
-          "id": "in",
-          "label": "Image",
-          "type": "image"
-        }
-      ],
-      "outputs": [
-        {
-          "id": "out",
-          "label": "Image",
-          "type": "image"
-        }
-      ],
-      "data": {}
+      "data": {},
+      "inputs": [{"id": "in", "label": "Image", "type": "image"}],
+      "outputs": [{"id": "out", "label": "Image", "type": "image"}]
     },
     {
       "id": "ref_instruction",
@@ -58,60 +31,53 @@
       "asset_type": "text",
       "x": 40,
       "y": 640,
-      "inputs": [
-        {
-          "id": "in",
-          "label": "Text",
-          "type": "text"
-        }
-      ],
-      "outputs": [
-        {
-          "id": "out",
-          "label": "Text",
-          "type": "text",
-          "default_value": "Move it to a snowy mountain top"
-        }
-      ],
-      "data": {
-        "out": "Move it to a snowy mountain top"
-      }
+      "data": {"out": "Move it to a snowy mountain top"},
+      "inputs": [{"id": "in", "label": "Text", "type": "text"}],
+      "outputs": [{"id": "out", "label": "Text", "type": "text"}]
     }
   ],
   "operators": [
+    {
+      "id": "op_enhance_t2i",
+      "label": "Enhance Prompt (PE-T2I)",
+      "role": "operator",
+      "kind": "fn",
+      "fn": "enhance_prompt_t2i",
+      "x": 420,
+      "y": 60,
+      "inputs": [
+        {"id": "prompt", "label": "Prompt", "type": "text", "required": true}
+      ],
+      "outputs": [{"id": "out_0", "label": "Rewritten Prompt", "type": "text", "output_index": 0}]
+    },
     {
       "id": "op_t2i",
       "label": "Qwen-Image 2.1 Text-to-Image",
       "role": "operator",
       "kind": "fn",
       "fn": "text_to_image",
-      "x": 420,
+      "x": 800,
       "y": 60,
+      "data": {"steps": 40},
       "inputs": [
-        {
-          "id": "prompt",
-          "label": "Prompt",
-          "type": "text",
-          "required": true
-        },
-        {
-          "id": "steps",
-          "label": "Steps",
-          "type": "number",
-          "default_value": 40
-        }
+        {"id": "prompt", "label": "Prompt", "type": "text", "required": true},
+        {"id": "steps", "label": "Steps", "type": "number"}
       ],
-      "outputs": [
-        {
-          "id": "out_0",
-          "label": "Image",
-          "type": "image",
-          "output_index": 0
-        }
+      "outputs": [{"id": "out_0", "label": "Image", "type": "image", "output_index": 0}]
+    },
+    {
+      "id": "op_enhance_i2i",
+      "label": "Enhance Instruction (PE-I2I)",
+      "role": "operator",
+      "kind": "fn",
+      "fn": "enhance_prompt_i2i",
+      "x": 420,
+      "y": 480,
+      "inputs": [
+        {"id": "image", "label": "Image", "type": "image", "required": true},
+        {"id": "instruction", "label": "Instruction", "type": "text", "required": true}
       ],
-      "data": {
-        "steps": 40
-      }
+      "outputs": [{"id": "out_0", "label": "Rewritten Instruction", "type": "text", "output_index": 0}]
     },
     {
       "id": "op_edit",
@@ -119,127 +85,119 @@
       "role": "operator",
       "kind": "fn",
       "fn": "edit_image",
-      "x": 820,
-      "y": 320,
+      "x": 800,
+      "y": 480,
+      "data": {"steps": 40},
       "inputs": [
-        {
-          "id": "image",
-          "label": "Image",
-          "type": "image",
-          "required": true
-        },
-        {
-          "id": "instruction",
-          "label": "Instruction",
-          "type": "text",
-          "required": true
-        },
-        {
-          "id": "steps",
-          "label": "Steps",
-          "type": "number",
-          "default_value": 40
-        }
+        {"id": "image", "label": "Image", "type": "image", "required": true},
+        {"id": "instruction", "label": "Instruction", "type": "text", "required": true},
+        {"id": "steps", "label": "Steps", "type": "number"}
       ],
-      "outputs": [
-        {
-          "id": "out_0",
-          "label": "Image",
-          "type": "image",
-          "output_index": 0
-        }
-      ],
-      "data": {
-        "steps": 40
-      }
+      "outputs": [{"id": "out_0", "label": "Image", "type": "image", "output_index": 0}]
     }
   ],
   "subjects": [
+    {
+      "id": "sub_prompt",
+      "label": "Rewritten Prompt",
+      "role": "subject",
+      "asset_type": "text",
+      "x": 800,
+      "y": 300,
+      "inputs": [{"id": "in", "label": "Text", "type": "text"}],
+      "outputs": [{"id": "out", "label": "Text", "type": "text"}]
+    },
     {
       "id": "sub_generated",
       "label": "Generated Image",
       "role": "subject",
       "asset_type": "image",
-      "x": 820,
+      "x": 1180,
       "y": 60,
-      "inputs": [
-        {
-          "id": "in",
-          "label": "Image",
-          "type": "image"
-        }
-      ],
-      "outputs": [
-        {
-          "id": "out",
-          "label": "Image",
-          "type": "image"
-        }
-      ]
+      "inputs": [{"id": "in", "label": "Image", "type": "image"}],
+      "outputs": [{"id": "out", "label": "Image", "type": "image"}]
+    },
+    {
+      "id": "sub_edit_prompt",
+      "label": "Rewritten Instruction",
+      "role": "subject",
+      "asset_type": "text",
+      "x": 800,
+      "y": 720,
+      "inputs": [{"id": "in", "label": "Text", "type": "text"}],
+      "outputs": [{"id": "out", "label": "Text", "type": "text"}]
     },
     {
       "id": "sub_edited",
       "label": "Edited Image",
       "role": "subject",
       "asset_type": "image",
-      "x": 1220,
-      "y": 320,
-      "inputs": [
-        {
-          "id": "in",
-          "label": "Image",
-          "type": "image"
-        }
-      ],
-      "outputs": [
-        {
-          "id": "out",
-          "label": "Image",
-          "type": "image"
-        }
-      ]
+      "x": 1180,
+      "y": 480,
+      "inputs": [{"id": "in", "label": "Image", "type": "image"}],
+      "outputs": [{"id": "out", "label": "Image", "type": "image"}]
     }
   ],
   "edges": [
     {
       "id": "e1",
-      "from_node_id": "ref_prompt",
-      "from_port_id": "out",
-      "to_node_id": "op_t2i",
-      "to_port_id": "prompt",
+      "from_node_id": "ref_prompt", "from_port_id": "out",
+      "to_node_id": "op_enhance_t2i", "to_port_id": "prompt",
       "type": "text"
     },
     {
       "id": "e2",
-      "from_node_id": "op_t2i",
-      "from_port_id": "out_0",
-      "to_node_id": "sub_generated",
-      "to_port_id": "in",
-      "type": "image"
+      "from_node_id": "op_enhance_t2i", "from_port_id": "out_0",
+      "to_node_id": "op_t2i", "to_port_id": "prompt",
+      "type": "text"
     },
     {
-      "id": "e4",
-      "from_node_id": "ref_instruction",
-      "from_port_id": "out",
-      "to_node_id": "op_edit",
-      "to_port_id": "instruction",
+      "id": "e3",
+      "from_node_id": "op_enhance_t2i", "from_port_id": "out_0",
+      "to_node_id": "sub_prompt", "to_port_id": "in",
       "type": "text"
     },
+    {
+      "id": "e4",
+      "from_node_id": "op_t2i", "from_port_id": "out_0",
+      "to_node_id": "sub_generated", "to_port_id": "in",
+      "type": "image"
+    },
     {
       "id": "e5",
-      "from_node_id": "op_edit",
-      "from_port_id": "out_0",
-      "to_node_id": "sub_edited",
-      "to_port_id": "in",
+      "from_node_id": "ref_image", "from_port_id": "out",
+      "to_node_id": "op_enhance_i2i", "to_port_id": "image",
       "type": "image"
     },
     {
-      "id": "e3",
-      "from_node_id": "ref_image",
-      "from_port_id": "out",
-      "to_node_id": "op_edit",
-      "to_port_id": "image",
+      "id": "e6",
+      "from_node_id": "ref_instruction", "from_port_id": "out",
+      "to_node_id": "op_enhance_i2i", "to_port_id": "instruction",
+      "type": "text"
+    },
+    {
+      "id": "e7",
+      "from_node_id": "op_enhance_i2i", "from_port_id": "out_0",
+      "to_node_id": "op_edit", "to_port_id": "instruction",
+      "type": "text"
+    },
+    {
+      "id": "e8",
+      "from_node_id": "op_enhance_i2i", "from_port_id": "out_0",
+      "to_node_id": "sub_edit_prompt", "to_port_id": "in",
+      "type": "text"
+    },
+    {
+      "id": "e9",
+      "from_node_id": "ref_image", "from_port_id": "out",
+      "to_node_id": "op_edit", "to_port_id": "image",
+      "type": "image"
+    },
+    {
+      "id": "e10",
+      "from_node_id": "op_edit", "from_port_id": "out_0",
+      "to_node_id": "sub_edited", "to_port_id": "in",
       "type": "image"
     }
   ]
-}
+}