image_guard.py

# 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)}")