# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 # # Adapted from NVIDIA asset-harvester (asset_harvester/utils/image_guard.py). """Image and prompt moderation utility backed by Nemotron 3.5 Content Safety.""" from __future__ import annotations import base64 import gc import io import os import re import time from dataclasses import asdict, dataclass from pathlib import Path import numpy as np import torch from PIL import Image from transformers import AutoProcessor, Gemma3ForConditionalGeneration DEFAULT_IMAGE_GUARD_MODEL_ID = "nvidia/Nemotron-3.5-Content-Safety" DEFAULT_IMAGE_GUARD_THRESHOLD = 0.5 _MAX_NEW_TOKENS = 100 _MODERATION_PROMPT = ( "Evaluate whether this image is safe to display in a public text-to-image demo." ) _USER_SAFETY_RE = re.compile(r"^\s*user safety:\s*(safe|unsafe)\s*$", re.IGNORECASE) class ImageGuardError(RuntimeError): """Raised when the content-safety model returns an unexpected moderation result.""" @dataclass class ImageGuardResult: passed: bool score: float label: str raw_response: str model_id: str inference_seconds: float def to_dict(self) -> dict: return asdict(self) class ImageGuard: """Lazy-loading wrapper around Nemotron 3.5 Content Safety.""" def __init__( self, model_id: str = DEFAULT_IMAGE_GUARD_MODEL_ID, threshold: float = DEFAULT_IMAGE_GUARD_THRESHOLD, device: str | torch.device | None = None, dtype: torch.dtype | None = None, hf_token: str | None = None, ) -> None: if device is None: device = "cuda" if torch.cuda.is_available() else "cpu" self.device = torch.device(device) if dtype is None: if self.device.type == "cuda" and torch.cuda.is_bf16_supported(): dtype = torch.bfloat16 elif self.device.type == "cuda": dtype = torch.float16 else: dtype = torch.float32 self.dtype = dtype self.model_id = model_id self.threshold = threshold self.hf_token = hf_token or os.getenv("HF_TOKEN") self._processor = None self._model = None def _load(self) -> None: if self._processor is not None and self._model is not None: return self._processor = AutoProcessor.from_pretrained( self.model_id, token=self.hf_token ) self._model = Gemma3ForConditionalGeneration.from_pretrained( self.model_id, torch_dtype=self.dtype, token=self.hf_token, ).to(self.device) self._model.eval() def load(self) -> None: self._load() def unload(self) -> None: if self._model is not None: self._model.to("cpu") self._processor = None self._model = None gc.collect() if torch.cuda.is_available(): torch.cuda.empty_cache() def check_image(self, image: str | Path | Image.Image | np.ndarray) -> ImageGuardResult: self._load() image_pil = self._coerce_image(image) start_time = time.perf_counter() text = self._generate_response(image_pil) inference_seconds = time.perf_counter() - start_time label, score = self._parse_response(text) return ImageGuardResult( passed=label == "safe" and score < self.threshold, score=score, label=label, raw_response=text, model_id=self.model_id, 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 = [ { "role": "user", "content": [ self._image_to_message_content(image), {"type": "text", "text": _MODERATION_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() @staticmethod def _image_to_message_content(image: Image.Image) -> dict[str, str]: img_bytes = io.BytesIO() image.save(img_bytes, format="JPEG") return { "type": "image", "image": base64.b64encode(img_bytes.getvalue()).decode("utf-8"), } @staticmethod def _parse_response(text: str) -> tuple[str, float]: for line in text.splitlines(): match = _USER_SAFETY_RE.match(line) if match: label = match.group(1).lower() score = 1.0 if label == "unsafe" else 0.0 return label, score normalized = text.strip().lower() if "user safety: unsafe" in normalized or normalized.startswith("unsafe"): return "unsafe", 1.0 if "user safety: safe" in normalized or normalized.startswith("safe"): return "safe", 0.0 raise ImageGuardError(f"Unexpected image guard response: {text!r}") @staticmethod def _coerce_image(image: str | Path | Image.Image | np.ndarray) -> Image.Image: if isinstance(image, Image.Image): return image.convert("RGB") if isinstance(image, (str, Path)): return Image.open(image).convert("RGB") if isinstance(image, np.ndarray): if image.ndim == 2: image = np.stack([image, image, image], axis=-1) return Image.fromarray(image.astype(np.uint8)).convert("RGB") raise TypeError(f"Unsupported image type: {type(image)}")