[Admin] Add NCII prompt guard for image-conditioned generations

app.py

@@ -80,6 +80,18 @@ LOADED_IN: float | None = None
 LORA_STATUS: str | None = None
 
 
+def check_prompt(prompt: str) -> None:
+    """The NCII guard, on requests carrying a keyframe — the edit-on-a-real-photo case. It runs before the
+    conditioner call and the denoise booking, so a refused prompt costs no GPU time on either half. The
+    classifier lives in `ncii_guard`'s spawned subprocess — in the main process it kills every GPU worker."""
+    import ncii_guard
+
+    flag = ncii_guard.classify(prompt)
+    if flag["label"] == "ncii":
+        print(f"[guard] prompt refused (ncii {flag['score']:.2f}): {prompt!r}", flush=True)
+        raise Exception("This prompt was flagged by a content filter and wasn't run.")
+
+
 def status() -> str:
     if LOAD_ERROR:
         return LOAD_ERROR
@@ -318,6 +330,8 @@ def generate(prompt, image_path=None, last_image_path=None, canvas=DEFAULT_CANVA
         raise Exception("The denoiser is still loading.")
     if not prompt or not prompt.strip():
         raise Exception("MiniMax-H3 always takes a prompt, keyframes or not.")
+    if image_path or last_image_path:
+        check_prompt(prompt)
 
     from PIL import Image, ImageOps
 
@@ -438,6 +452,9 @@ def homepage():
         return f.read()
 
 
+import ncii_guard
+
+ncii_guard.start()
 load_models()
 
 if __name__ == "__main__":

ncii_guard.py

@@ -0,0 +1,90 @@
+"""The NCII prompt guard, in its own process.
+
+[`hfmlsoc/ncii-light-guard-v01`](https://huggingface.co/hfmlsoc/ncii-light-guard-v01) is a 270M CPU text
+classifier scoring the NCII risk of an edit prompt. It cannot live in the main process: with it loaded there,
+every subsequent `@spaces.GPU` worker dies at `worker_init` with `RuntimeError: No CUDA GPUs are available` —
+the fork inherits whatever CUDA driver state the classifier's torch activity left behind, and a factory reboot
+does not clear it. It cannot be a `multiprocessing.spawn` child either: spawn re-imports the parent's main
+module, and on a Space that main module is `app.py` — the child would re-run the whole startup, `start()`
+included. So the classifier runs this file as a plain subprocess — a fresh interpreter that never sees `spaces`
+— and answers over stdin/stdout, one JSON object per line.
+"""
+
+from __future__ import annotations
+
+import json
+import os
+import select
+import subprocess
+import sys
+import threading
+
+GUARD_REPO = "hfmlsoc/ncii-light-guard-v01"
+
+_lock = threading.Lock()
+_process: subprocess.Popen | None = None
+
+
+def _read(timeout: float) -> dict:
+    readable, _, _ = select.select([_process.stdout], [], [], timeout)
+    if not readable:
+        raise TimeoutError(f"the guard did not answer within {timeout}s")
+    line = _process.stdout.readline()
+    if not line:
+        raise EOFError("the guard process died")
+    return json.loads(line)
+
+
+def _spawn() -> None:
+    global _process
+    _process = subprocess.Popen(
+        [sys.executable, os.path.abspath(__file__)],
+        stdin=subprocess.PIPE,
+        stdout=subprocess.PIPE,
+        text=True,
+        bufsize=1,
+    )
+    # Generous: a cold cache downloads the checkpoint first.
+    assert _read(300.0) == {"status": "ready"}
+
+
+def start() -> None:
+    """Launch the worker and block until its model is up. Called once at startup; `classify` revives it if it dies."""
+    with _lock:
+        _spawn()
+
+
+def classify(prompt: str, timeout: float = 60.0) -> dict:
+    """`{'label': 'safe' | 'ncii', 'score': ...}` for one prompt, replacing a dead or wedged worker once."""
+    with _lock:
+        for attempt in (0, 1):
+            try:
+                if _process is None or _process.poll() is not None:
+                    _spawn()
+                _process.stdin.write(json.dumps({"prompt": prompt}) + "\n")
+                _process.stdin.flush()
+                return _read(timeout)
+            except Exception:
+                if attempt:
+                    raise
+                if _process is not None and _process.poll() is None:
+                    _process.kill()
+
+
+def _serve() -> None:
+    """The child: plain torch on CPU. The protocol keeps the real stdout to itself — everything else
+    (download progress, warnings) is pushed over to stderr so it cannot corrupt a reply."""
+    protocol = os.fdopen(os.dup(1), "w", buffering=1)
+    os.dup2(2, 1)
+
+    from transformers import pipeline
+
+    classifier = pipeline("text-classification", model=GUARD_REPO, device="cpu")
+    protocol.write(json.dumps({"status": "ready"}) + "\n")
+    for line in sys.stdin:
+        result = classifier(json.loads(line)["prompt"], truncation=True)[0]
+        protocol.write(json.dumps({"label": result["label"], "score": float(result["score"])}) + "\n")
+
+
+if __name__ == "__main__":
+    _serve()