update bug fix -- @cpu

app.py

@@ -11,7 +11,7 @@ import base64
 import json
 from io import BytesIO
 from PIL import Image
-from transformers import pipeline as hf_pipeline
+import ncii_guard
 
 MAX_SEED = np.iinfo(np.int32).max
 LANCZOS = getattr(Image, "Resampling", Image).LANCZOS
@@ -52,37 +52,35 @@ try:
 except Exception as e:
     print(f"Warning: Could not set FA3 processor: {e}")
 
-# ── NCII safety guard (CPU-only, loaded at startup) ──────────────────────────
-NCII_MODEL_ID = "hfmlsoc/ncii-guard-v02"
+# ── NCII safety guard (pure CPU subprocess, never touches spaces/CUDA) ──────
 NCII_UNSAFE_LABEL = "ncii"
 NCII_THRESHOLD = 0.5
 NCII_BLOCK_MESSAGE = "You entered prompt is NCII (non-consensual intimate imagery) and your request will not be processed. Try with Safe Prompts."
 
-print("Loading NCII safety guard model (CPU-only)...")
+print("Starting NCII safety guard subprocess (CPU-only)…")
 try:
-    # device=-1 forces CPU inference — never touches the GPU
-    ncii_guard = hf_pipeline("text-classification", model=NCII_MODEL_ID, device=-1)
-    print("NCII guard loaded successfully on CPU.")
+    ncii_guard.start()
+    print("NCII guard subprocess ready.")
 except Exception as e:
-    ncii_guard = None
-    print(f"Warning: Could not load NCII guard model ({NCII_MODEL_ID}): {e}")
+    print(f"Warning: Could not start NCII guard subprocess: {e}")
 
 
 def check_ncii_safety(prompt_text):
-    """Run NCII classification on CPU. Returns (is_unsafe, label, score)."""
-    if ncii_guard is None or not prompt_text or not prompt_text.strip():
+    """Ask the CPU-side subprocess. No GPU is requested by this call.
+
+    Returns: (is_unsafe: bool, label: str, score: float)
+    """
+    if not prompt_text or not prompt_text.strip():
         return False, "unknown", 0.0
     try:
-        out = ncii_guard(prompt_text.strip())
-        if isinstance(out, list) and len(out) > 0:
-            top = out[0]
-            label = str(top.get("label", "")).lower()
-            score = float(top.get("score", 0.0))
-            is_unsafe = (label == NCII_UNSAFE_LABEL) and (score >= NCII_THRESHOLD)
-            return is_unsafe, label, score
+        out = ncii_guard.classify(prompt_text.strip())
+        label = str(out.get("label", "")).lower()
+        score = float(out.get("score", 0.0))
+        is_unsafe = (label == NCII_UNSAFE_LABEL) and (score >= NCII_THRESHOLD)
+        return is_unsafe, label, score
     except Exception as e:
         print(f"NCII guard inference error: {e}")
-    return False, "unknown", 0.0
+        return False, "unknown", 0.0
 
 
 # ── Examples Config ───────────────────────────────────────────────────────────
@@ -95,7 +93,6 @@ EXAMPLES_CONFIG = [
     {"images": ["examples/10.jpg", "examples/11.png"], "prompt": "Replace the current clothing with the clothing from the reference image 2. Keep the person's face, hairstyle, body pose, background, lighting, and camera angle unchanged. Ensure the new outfit fits naturally with realistic fabric texture, proper shadows, folds, and accurate proportions. Match the lighting, color tone, and overall style for a seamless and high-quality result."},
 ]
 
-
 def make_thumb_b64(path, max_dim=220):
     if not os.path.exists(path):
         return ""
@@ -109,7 +106,6 @@ def make_thumb_b64(path, max_dim=220):
         print(f"Thumbnail error for {path}: {e}")
         return ""
 
-
 def encode_full_image(path):
     if not os.path.exists(path):
         return ""
@@ -123,7 +119,6 @@ def encode_full_image(path):
         print(f"Encode error for {path}: {e}")
         return ""
 
-
 def build_client_config():
     """Static config consumed by the frontend: example cards."""
     examples = []
@@ -138,12 +133,10 @@ def build_client_config():
         "examples": examples,
     }
 
-
 print("Building client config (example thumbnails)…")
 CLIENT_CONFIG = build_client_config()
 print(f"Built config with {len(EXAMPLES_CONFIG)} examples.")
 
-
 def b64_to_pil_list(b64_json_str):
     if not b64_json_str or b64_json_str.strip() in ("", "[]"):
         return []
@@ -166,13 +159,11 @@ def b64_to_pil_list(b64_json_str):
             print(f"Error decoding image: {e}")
     return pil_images
 
-
 def pil_to_b64_png(image: Image.Image) -> str:
     buf = BytesIO()
     image.save(buf, format="PNG")
     return f"data:image/png;base64,{base64.b64encode(buf.getvalue()).decode()}"
 
-
 def update_dimensions_on_upload(image):
     if image is None:
         return 1024, 1024
@@ -190,31 +181,22 @@ def update_dimensions_on_upload(image):
 app = Server(title="FireRed-Image-Edit-1.0-Fast")
 
 
-# ── CPU-only safety check endpoint ───────────────────────────────────────────
-# This endpoint is NOT decorated with @spaces.GPU, so it runs entirely on CPU
-# and never triggers GPU allocation. The frontend calls this BEFORE submitting
-# the GPU job, so blocked prompts are caught early without wasting GPU resources.
 @app.api(name="check_safety", queue=False)
 def check_safety(prompt: str) -> dict:
-    """Check if a prompt is safe (NCII guard). Runs on CPU only — no GPU allocation.
+    """Pure-CPU safety check. Runs BEFORE any @spaces.GPU resource is requested.
 
-    Returns {"safe": bool, "label": str, "score": float, "message": str}.
+    Returns:
+        {"status": "ok"}                       — prompt is safe, proceed
+        {"status": "blocked", "message": ...}  — NCII detected
     """
     if not prompt or not prompt.strip():
-        return {"safe": True, "label": "empty", "score": 0.0, "message": ""}
-
+        return {"status": "ok"}
     is_unsafe, label, score = check_ncii_safety(prompt)
     if is_unsafe:
-        return {
-            "safe": False,
-            "label": label,
-            "score": score,
-            "message": NCII_BLOCK_MESSAGE,
-        }
-    return {"safe": True, "label": label, "score": score, "message": ""}
+        return {"status": "blocked", "message": NCII_BLOCK_MESSAGE}
+    return {"status": "ok", "label": label, "score": score}
 
 
-# ── GPU image-edit endpoint ──────────────────────────────────────────────────
 @app.mcp.tool(name="edit_image")
 @app.api(name="edit_image")
 @spaces.GPU(size="xlarge")
@@ -228,11 +210,8 @@ def infer(
 ) -> dict:
     """Edit one or more images with FireRed-Image-Edit-1.1.
 
-    Pre-condition: The frontend must call /check_safety BEFORE this endpoint
-    to ensure the prompt is not blocked by the NCII guard. This function assumes
-    the prompt has already been validated and focuses solely on GPU inference.
-
-    Returns {"image": <base64 PNG data URL>, "seed": <seed used>, "status": "success"}.
+    The frontend is expected to have already called /check_safety and
+    only submit here when safe. This is a defense-in-depth re-check.
     """
     gc.collect()
     torch.cuda.empty_cache()
@@ -243,6 +222,13 @@ def infer(
     if not prompt or prompt.strip() == "":
         raise gr.Error("Please enter an edit prompt.")
 
+    # ── Defense-in-depth NCII re-check inside the GPU worker ──
+    is_unsafe, _, _ = check_ncii_safety(prompt)
+    if is_unsafe:
+        gc.collect()
+        torch.cuda.empty_cache()
+        return {"image": "", "seed": seed, "status": "blocked", "message": NCII_BLOCK_MESSAGE}
+
     if randomize_seed:
         seed = random.randint(0, MAX_SEED)