Add prompt safety filter and make guards mandatory

Run Nemotron 3.5 Content Safety on the prompt before generation (fail-closed
input filter) in addition to the post-generation image check. Remove the UI
toggles so both safety filters always run; env-var opt-outs remain for local dev.

app.py

@@ -159,6 +159,8 @@ IMAGE_GUARD_THRESHOLD = float(
 IMAGE_GUARD_OFFLOAD_T2I = os.getenv("IMAGE_GUARD_OFFLOAD_T2I", "0") == "1"
 # Opt-out: guard runs by default; set ENABLE_IMAGE_GUARD=0 or uncheck the UI box to disable.
 DEFAULT_ENABLE_IMAGE_GUARD = os.getenv("ENABLE_IMAGE_GUARD", "1") == "1"
+# Pre-generation prompt safety check (input filter), same content-safety model.
+DEFAULT_ENABLE_PROMPT_GUARD = os.getenv("ENABLE_PROMPT_GUARD", "1") == "1"
 
 GUARD_ACCESS_HELP = (
     "https://huggingface.co/nvidia/Nemotron-3.5-Content-Safety"
@@ -384,6 +386,25 @@ class T2IEngine:
             if IMAGE_GUARD_OFFLOAD_T2I:
                 self._reload_t2i_to_device()
 
+    def _moderate_prompt(self, prompt: str) -> tuple[bool, str]:
+        """Run pre-generation prompt moderation. Returns (ok, meta_or_error)."""
+        try:
+            with _suppress_tqdm_tracking():
+                guard = self._get_image_guard()
+            check = guard.check_text(prompt)
+            if not check.passed:
+                message = (
+                    "Prompt blocked by content-safety filter "
+                    f"(label={check.label})."
+                )
+                _report_guard_failure(message)
+                return False, message
+            return True, "prompt_" + _format_guard_meta(check)
+        except Exception as exc:
+            message = _guard_unavailable_message(exc)
+            _report_guard_failure(message)
+            return False, message
+
     def generate(
         self,
         prompt: str,
@@ -401,9 +422,18 @@ class T2IEngine:
         seed: int,
         micro_cond: str,
         return_animation: bool,
-        enable_image_guard: bool,
+        enable_image_guard: bool = DEFAULT_ENABLE_IMAGE_GUARD,
+        enable_prompt_guard: bool = DEFAULT_ENABLE_PROMPT_GUARD,
         progress: gr.Progress | None = None,
     ):
+        prompt_guard_meta = ""
+        if enable_prompt_guard:
+            if progress is not None:
+                progress(0.0, desc="Checking prompt…")
+            prompt_ok, prompt_guard_meta = self._moderate_prompt(prompt)
+            if not prompt_ok:
+                return None, f"ERROR: {prompt_guard_meta}"
+
         if progress is not None and self._model is None:
             progress(0.0, desc="Loading model (first run, this can take 1-2 min)…")
         with _suppress_tqdm_tracking():
@@ -461,6 +491,8 @@ class T2IEngine:
             f"sch_temp={schedule_temp} | conf={confidence_policy} | "
             f"edit_threshold={edit_threshold:.3f} | gen_time={latency:.2f}s"
         )
+        if prompt_guard_meta:
+            meta += " | " + prompt_guard_meta
 
         if progress is not None and enable_image_guard:
             progress(1.0, desc="Running safety filter…")
@@ -498,7 +530,6 @@ def generate(
     seed: int,
     micro_cond: str,
     return_animation: bool,
-    enable_image_guard: bool,
     progress: gr.Progress = gr.Progress(track_tqdm=True),
 ):
     return engine.generate(
@@ -517,7 +548,8 @@ def generate(
         seed,
         micro_cond,
         return_animation,
-        enable_image_guard,
+        enable_image_guard=DEFAULT_ENABLE_IMAGE_GUARD,
+        enable_prompt_guard=DEFAULT_ENABLE_PROMPT_GUARD,
         progress=progress,
     )
 
@@ -611,9 +643,9 @@ def build_demo() -> gr.Blocks:
                         label="Seed (-1 for random)", value=42, precision=0, scale=3
                     )
                     randomize_seed_btn = gr.Button("🎲 Randomize", scale=1)
-                enable_image_guard = gr.Checkbox(
-                    label="NSFW output filter (Nemotron 3.5 Content Safety) — uncheck to opt out",
-                    value=DEFAULT_ENABLE_IMAGE_GUARD,
+                gr.Markdown(
+                    "Safety filters (Nemotron 3.5 Content Safety) run on both the "
+                    "prompt and the generated image and cannot be disabled."
                 )
                 generate_btn = gr.Button("Generate", variant="primary")
 
@@ -681,7 +713,6 @@ def build_demo() -> gr.Blocks:
                 seed,
                 micro_cond,
                 return_animation,
-                enable_image_guard,
             ],
             outputs=[output_image, output_meta],
         )

image_guard.py

@@ -3,7 +3,7 @@
 #
 # Adapted from NVIDIA asset-harvester (asset_harvester/utils/image_guard.py).
 
-"""Image moderation utility backed by Nemotron 3.5 Content Safety."""
+"""Image and prompt moderation utility backed by Nemotron 3.5 Content Safety."""
 
 from __future__ import annotations
 
@@ -117,6 +117,50 @@ class ImageGuard:
             inference_seconds=inference_seconds,
         )
 
+    def check_text(self, prompt: str) -> ImageGuardResult:
+        """Moderate a text prompt before generation (fail-closed input filter)."""
+        self._load()
+        start_time = time.perf_counter()
+        text = self._generate_text_response(prompt)
+        inference_seconds = time.perf_counter() - start_time
+        label, score = self._parse_response(text)
+        return ImageGuardResult(
+            passed=label == "safe",
+            score=score,
+            label=label,
+            raw_response=text,
+            model_id=self.model_id,
+            inference_seconds=inference_seconds,
+        )
+
+    def _generate_text_response(self, prompt: str) -> str:
+        messages = [
+            {
+                "role": "user",
+                "content": [{"type": "text", "text": prompt}],
+            }
+        ]
+        inputs = self._processor.apply_chat_template(
+            messages,
+            add_generation_prompt=True,
+            tokenize=True,
+            return_dict=True,
+            return_tensors="pt",
+            request_categories="/categories",
+            enable_thinking=False,
+        ).to(self.device)
+
+        input_len = inputs["input_ids"].shape[-1]
+        with torch.inference_mode():
+            generation = self._model.generate(
+                **inputs,
+                max_new_tokens=_MAX_NEW_TOKENS,
+                do_sample=False,
+            )
+            generation = generation[0][input_len:]
+
+        return self._processor.decode(generation, skip_special_tokens=True).strip()
+
     def _generate_response(self, image: Image.Image) -> str:
         messages = [
             {