Add Gradio T2I demo with Nemotron 3.5 guard, model write-up, and figures

- Gradio Space app for nvidia/NL-Diffusion-Image (text-to-image)
- NSFW guard via nvidia/Nemotron-3.5-Content-Safety (EU-accessible)
- About-the-model section with sanitized announcement content + assets
- NVIDIA-green theme; PNG/GIF assets via Git LFS; ignore internal/local artifacts

.gitattributes

@@ -4,6 +4,11 @@
 *.bz2 filter=lfs diff=lfs merge=lfs -text
 *.ckpt filter=lfs diff=lfs merge=lfs -text
 *.ftz filter=lfs diff=lfs merge=lfs -text
+*.gif filter=lfs diff=lfs merge=lfs -text
+*.png filter=lfs diff=lfs merge=lfs -text
+*.jpg filter=lfs diff=lfs merge=lfs -text
+*.jpeg filter=lfs diff=lfs merge=lfs -text
+*.webp filter=lfs diff=lfs merge=lfs -text
 *.gz filter=lfs diff=lfs merge=lfs -text
 *.h5 filter=lfs diff=lfs merge=lfs -text
 *.joblib filter=lfs diff=lfs merge=lfs -text

.gitignore

@@ -0,0 +1,9 @@
+# Internal announcement source (do not publish)
+Nemotron Labs Diffusion Image - IPR.zip
+announcement_extract/
+
+# Local artifacts
+__pycache__/
+*.pyc
+log.txt
+test_hub_output.webp

app.py

@@ -0,0 +1,663 @@
+"""NL-Diffusion-Image Gradio demo for Hugging Face Spaces (and local testing).
+
+Local test with private Hub model:
+
+    conda activate lavida
+    export HF_TOKEN=hf_...
+    export MODEL_ID=nvidia/NL-Diffusion-Image
+    python app.py
+
+Post-generation NSFW guard (on by default — opt out with ENABLE_IMAGE_GUARD=0):
+
+    export ENABLE_IMAGE_GUARD=0
+    python app.py
+"""
+
+from __future__ import annotations
+
+import gc
+import os
+import tempfile
+import time
+from pathlib import Path
+from typing import Any
+
+ASSETS_DIR = Path(__file__).resolve().parent / "assets"
+
+
+def _asset(name: str) -> str:
+    return str(ASSETS_DIR / name)
+
+import gradio as gr
+import imageio.v3 as iio
+import torch
+from PIL import ImageDraw
+from transformers import AutoModelForCausalLM, PreTrainedTokenizerFast
+
+from image_guard import (
+    DEFAULT_IMAGE_GUARD_MODEL_ID,
+    DEFAULT_IMAGE_GUARD_THRESHOLD,
+    ImageGuard,
+)
+
+try:
+    import spaces
+except ImportError:
+
+    class _SpacesStub:
+        @staticmethod
+        def GPU(*args, **kwargs):
+            def decorator(fn):
+                return fn
+
+            if args and callable(args[0]):
+                return args[0]
+            return decorator
+
+    spaces = _SpacesStub()
+
+os.environ.setdefault("DEBUG_FIX_PADDING", "1")
+os.environ.setdefault("NOT_ALWASY_DO_2DPOOL", "1")
+
+if "CUDA_HOME" not in os.environ:
+    _local_cuda = "/lustre/fsw/portfolios/llmservice/users/gheinrich/cuda/cuda_12.4"
+    if os.path.isdir(_local_cuda):
+        os.environ["CUDA_HOME"] = _local_cuda
+
+MODEL_ID = os.getenv("MODEL_ID", "nvidia/NL-Diffusion-Image")
+DEVICE = os.getenv("DEVICE", "cuda")
+IMAGE_GUARD_MODEL_ID = os.getenv("IMAGE_GUARD_MODEL_ID", DEFAULT_IMAGE_GUARD_MODEL_ID)
+IMAGE_GUARD_THRESHOLD = float(
+    os.getenv("IMAGE_GUARD_THRESHOLD", str(DEFAULT_IMAGE_GUARD_THRESHOLD))
+)
+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"
+
+GUARD_ACCESS_HELP = (
+    "https://huggingface.co/nvidia/Nemotron-3.5-Content-Safety"
+)
+
+
+def _guard_unavailable_message(exc: Exception) -> str:
+    text = str(exc).lower()
+    if "gated repo" in text or "403" in text or "authorized list" in text:
+        return (
+            "NSFW filter is enabled but Nemotron 3.5 Content Safety is not accessible. "
+            f"See {GUARD_ACCESS_HELP}, ensure HF_TOKEN has read access, "
+            "or uncheck 'NSFW output filter' to opt out."
+        )
+    return f"NSFW filter is enabled but Nemotron 3.5 Content Safety failed: {exc}"
+
+
+def _report_guard_failure(message: str) -> None:
+    """Surface guard failures in the Gradio UI without breaking output components."""
+    gr.Warning(message)
+    print(f"GUARD ERROR: {message}", flush=True)
+
+# Defaults aligned with nemotron-diffusion-omni/gradio_t2i_demo.py
+DEFAULT_PROMPT = (
+    "A full-body shot of hyper-realistic female cyborg, human facial skin seamlessly integrated "
+    "with a glossy white mechanical head shell. Features a realistic human ear, blue eyes. bright, "
+    "outdoor, background with blue sky, illuminated by striking bright white studio lighting, "
+    "casting soft shadows. Cyberpunk aesthetic, high-tech minimalism, shot on 85mm lens, "
+    "photorealistic, Unreal Engine 5 render, intricately detailed, 8k resolution, high dynamic "
+    "range, chest with whit armor plate, cute, beautiful, sexy, glossy surface, reflective, "
+    "Artstation, pixiv, no hair, 3D render, stylized eyesz"
+)
+DEFAULT_MICRO_COND = (
+    "ORIGINAL WIDTH : 1024; ORIGINAL HEIGHT : 1024; TOP : 0; LEFT : 0; "
+    "SCORE : 6.520; HPS: 3.220"
+)
+
+DEFAULT_GENERATION_CONFIG: dict[str, Any] = {
+    "guidance_scale": 5.0,
+    "n_steps": 64,
+    "shift": 5,
+    "schedule": "shift",
+    "alg_temp": 1.0,
+    "dynamic_temperature": False,
+    "min_temperature": 0.01,
+    "schedule_temp": "linear",
+    "temperature": 0.86,
+    "confidence_policy": "mmada",
+    "micro_cond": DEFAULT_MICRO_COND,
+    "edit_threshold": 0.6,
+    "block_policy": 2,
+    "is_legacy": False,
+    "use_cache": False,
+}
+
+
+def n_tokens_from_resolution(image_resolution: int) -> int:
+    return (image_resolution // 16) * (image_resolution // 16)
+
+
+def process_gif(image_list):
+    if not image_list:
+        return None
+
+    with tempfile.NamedTemporaryFile(suffix=".gif", delete=False) as tmp_file:
+        gif_path = tmp_file.name
+
+    frames = []
+    total_frames = len(image_list)
+    for i, img in enumerate(image_list):
+        frame = img.resize((400, 400))
+        draw = ImageDraw.Draw(frame)
+        text = f"Frame: {i + 1} / {total_frames}"
+        x, y = 15, 15
+        for dx, dy in [(-1, -1), (1, -1), (-1, 1), (1, 1)]:
+            draw.text((x + dx, y + dy), text, fill="black")
+        draw.text((x, y), text, fill="white")
+        frames.append(frame)
+
+    duration = [1000 / 20] * len(frames)
+    duration[-1] = 2000
+    iio.imwrite(gif_path, frames, extension=".gif", duration=duration, loop=0)
+    return gif_path
+
+
+def process_webp(pil_image):
+    with tempfile.NamedTemporaryFile(suffix=".webp", delete=False) as tmp_file:
+        webp_path = tmp_file.name
+    pil_image.save(webp_path, "webp", quality=95)
+    return webp_path
+
+
+def load_release_model_and_tokenizer(model_id: str, device: str):
+    hf_token = os.getenv("HF_TOKEN")
+    tokenizer = PreTrainedTokenizerFast.from_pretrained(model_id, token=hf_token)
+    if tokenizer.pad_token_id is None:
+        tokenizer.pad_token_id = tokenizer.eos_token_id
+        tokenizer.pad_token = tokenizer.eos_token
+
+    model = AutoModelForCausalLM.from_pretrained(
+        model_id,
+        trust_remote_code=True,
+        torch_dtype=torch.bfloat16,
+        low_cpu_mem_usage=False,
+        token=hf_token,
+    )
+    model.to(device)
+    model.eval()
+    model.requires_grad_(False)
+    model.config.dlm_paradigm = "bidirectional"
+    return tokenizer, model
+
+
+def _format_guard_meta(result) -> str:
+    return (
+        f"guard={result.model_id} | label={result.label} | "
+        f"unsafe_score={result.score:.3f} | guard_time={result.inference_seconds:.2f}s"
+    )
+
+
+class T2IEngine:
+    def __init__(self, model_id: str, device: str = "cuda") -> None:
+        self.model_id = model_id
+        self.device = device
+        self._tokenizer = None
+        self._model = None
+        self._image_guard: ImageGuard | None = None
+
+    def _lazy_load(self) -> None:
+        if self._model is not None and self._tokenizer is not None:
+            return
+        print(f"Loading model from {self.model_id} ...", flush=True)
+        self._tokenizer, self._model = load_release_model_and_tokenizer(
+            self.model_id, self.device
+        )
+        print("Model loaded.", flush=True)
+
+    def _get_image_guard(self) -> ImageGuard:
+        if self._image_guard is None:
+            print(f"Loading image guard from {IMAGE_GUARD_MODEL_ID} ...", flush=True)
+            self._image_guard = ImageGuard(
+                model_id=IMAGE_GUARD_MODEL_ID,
+                threshold=IMAGE_GUARD_THRESHOLD,
+                device=self.device,
+            )
+        return self._image_guard
+
+    def _offload_t2i_to_cpu(self) -> None:
+        if self._model is not None:
+            self._model.to("cpu")
+            gc.collect()
+            if torch.cuda.is_available():
+                torch.cuda.empty_cache()
+
+    def _reload_t2i_to_device(self) -> None:
+        if self._model is not None:
+            self._model.to(self.device)
+
+    def _moderate_output(
+        self,
+        result,
+        return_animation: bool,
+        enable_image_guard: bool,
+    ) -> tuple[bool, str]:
+        """Run post-generation moderation. Returns (ok, meta_suffix_or_error_message)."""
+        if not enable_image_guard:
+            return True, ""
+
+        if IMAGE_GUARD_OFFLOAD_T2I:
+            self._offload_t2i_to_cpu()
+
+        try:
+            guard = self._get_image_guard()
+            frames = result if return_animation else [result]
+            guard_parts = []
+            for frame_idx, frame in enumerate(frames):
+                check = guard.check_image(frame)
+                guard_parts.append(_format_guard_meta(check))
+                if not check.passed:
+                    message = (
+                        "Generated image blocked by NSFW filter "
+                        f"(frame {frame_idx + 1}/{len(frames)}, "
+                        f"unsafe_score={check.score:.3f}, threshold={IMAGE_GUARD_THRESHOLD})."
+                    )
+                    _report_guard_failure(message)
+                    return False, message
+            return True, " | " + guard_parts[0] if guard_parts else ""
+        except Exception as exc:
+            message = _guard_unavailable_message(exc)
+            _report_guard_failure(message)
+            return False, message
+        finally:
+            if IMAGE_GUARD_OFFLOAD_T2I:
+                self._reload_t2i_to_device()
+
+    def generate(
+        self,
+        prompt: str,
+        image_resolution: int,
+        guidance_scale: float,
+        temperature: float,
+        n_steps: int,
+        schedule: str,
+        shift: int,
+        confidence_policy: str,
+        schedule_temp: str,
+        alg_temp: float,
+        dynamic_temperature: bool,
+        min_temperature: float,
+        edit_threshold: float,
+        seed: int,
+        micro_cond: str,
+        use_cache: bool,
+        is_legacy: bool,
+        return_animation: bool,
+        enable_image_guard: bool,
+    ):
+        self._lazy_load()
+
+        gen_cfg = dict(DEFAULT_GENERATION_CONFIG)
+        gen_cfg.update(
+            micro_cond=micro_cond,
+            guidance_scale=guidance_scale,
+            temperature=temperature,
+            edit_threshold=edit_threshold,
+            n_steps=int(n_steps),
+            schedule=schedule,
+            shift=int(shift),
+            confidence_policy=confidence_policy,
+            schedule_temp=schedule_temp,
+            alg_temp=alg_temp,
+            dynamic_temperature=dynamic_temperature,
+            min_temperature=min_temperature,
+            block_policy=2,
+            is_legacy=is_legacy,
+            use_cache=use_cache,
+        )
+
+        if seed < 0:
+            seed = int(torch.seed() % (2**31 - 1))
+        torch.manual_seed(int(seed))
+
+        n_tokens = n_tokens_from_resolution(int(image_resolution))
+
+        t0 = time.time()
+        with torch.no_grad():
+            with torch.inference_mode():
+                result = self._model.text_to_image(
+                    prompt,
+                    tokenizer=self._tokenizer,
+                    **gen_cfg,
+                    image_resolution=int(image_resolution),
+                    n_tokens=n_tokens,
+                    disable_tqdm=True,
+                    return_intermediate_steps=return_animation,
+                )
+        latency = time.time() - t0
+
+        meta = (
+            f"model={self.model_id} | seed={seed} | res={image_resolution} | "
+            f"n_tokens={n_tokens} | steps={n_steps} | "
+            f"cfg={guidance_scale:.2f} | temp={temperature:.3f} | "
+            f"schedule={schedule} | shift={shift} | alg_temp={alg_temp:.2f} | "
+            f"dyn_temp={dynamic_temperature} | min_temp={min_temperature:.3f} | "
+            f"sch_temp={schedule_temp} | conf={confidence_policy} | "
+            f"edit_threshold={edit_threshold:.3f} | gen_time={latency:.2f}s"
+        )
+
+        guard_ok, guard_meta = self._moderate_output(
+            result, return_animation, enable_image_guard
+        )
+        if not guard_ok:
+            return None, f"ERROR: {guard_meta}\n\n{meta}"
+
+        meta += guard_meta
+
+        if return_animation:
+            return process_gif(result), meta
+        return process_webp(result), meta
+
+
+engine = T2IEngine(model_id=MODEL_ID, device=DEVICE)
+
+
+@spaces.GPU(duration=240)
+def generate(
+    prompt: str,
+    image_resolution: int,
+    guidance_scale: float,
+    temperature: float,
+    n_steps: int,
+    schedule: str,
+    shift: int,
+    confidence_policy: str,
+    schedule_temp: str,
+    alg_temp: float,
+    dynamic_temperature: bool,
+    min_temperature: float,
+    edit_threshold: float,
+    seed: int,
+    micro_cond: str,
+    use_cache: bool,
+    is_legacy: bool,
+    return_animation: bool,
+    enable_image_guard: bool,
+):
+    return engine.generate(
+        prompt,
+        image_resolution,
+        guidance_scale,
+        temperature,
+        n_steps,
+        schedule,
+        shift,
+        confidence_policy,
+        schedule_temp,
+        alg_temp,
+        dynamic_temperature,
+        min_temperature,
+        edit_threshold,
+        seed,
+        micro_cond,
+        use_cache,
+        is_legacy,
+        return_animation,
+        enable_image_guard,
+    )
+
+
+def build_demo() -> gr.Blocks:
+    nvidia_green = gr.themes.Color(
+        c50="#f3f9e6",
+        c100="#e3f1c2",
+        c200="#cfe88f",
+        c300="#b6dc56",
+        c400="#97c61f",
+        c500="#76b900",
+        c600="#69a600",
+        c700="#548400",
+        c800="#3f6300",
+        c900="#2a4200",
+        c950="#1a2900",
+    )
+    theme = gr.themes.Soft(primary_hue=nvidia_green, secondary_hue=nvidia_green)
+    with gr.Blocks(title="Nemotron Labs Diffusion Image", theme=theme) as demo:
+        gr.Markdown(
+            "# Nemotron Labs Diffusion Image\n\n"
+            "NL-Diffusion-Image generates high-resolution images via **masked discrete diffusion** "
+            "over tokenized image patches. Each image is encoded into discrete tokens "
+            "(131K codebook), and generation proceeds through iterative parallel unmasking—similar "
+            "to diffusion LLMs. The model is fine-tuned from "
+            "[Nemotron-Labs-Diffusion](https://huggingface.co/nvidia/Nemotron-Labs-Diffusion-8B) "
+            "with two key additions:\n\n"
+            "- **Token editing** — revise already-unmasked tokens during inference.\n"
+            "- **Grouped Cross-Entropy (GCE)** — efficient training with large vocabularies.\n\n"
+            "This aligns image generation with LLM training and inference infrastructure, "
+            "making it highly scalable.\n\n"
+            f"**Model:** [`{MODEL_ID}`](https://huggingface.co/nvidia/NL-Diffusion-Image)"
+        )
+
+        gr.Markdown(
+            "| GenEval | DPG | HPSv3 | Speed vs EMU3.5 |\n"
+            "|:---:|:---:|:---:|:---:|\n"
+            "| **0.90** | **86.9** | **10.76** | **42.4× faster** |"
+        )
+
+        gr.Markdown("## Generate an image")
+
+        with gr.Row():
+            with gr.Column(scale=2):
+                prompt = gr.Textbox(label="Prompt", lines=4, value=DEFAULT_PROMPT)
+                micro_cond = gr.Textbox(label="Micro Cond", lines=2, value=DEFAULT_MICRO_COND)
+                with gr.Row():
+                    image_resolution = gr.Dropdown(
+                        choices=[256, 512, 1024],
+                        value=1024,
+                        label="Image Resolution",
+                    )
+                    n_steps = gr.Slider(
+                        minimum=1, maximum=128, value=64, step=1, label="Diffusion Steps"
+                    )
+                with gr.Row():
+                    guidance_scale = gr.Slider(
+                        minimum=1.0, maximum=10.0, value=5.0, step=0.1, label="Guidance Scale"
+                    )
+                    temperature = gr.Slider(
+                        minimum=0.05, maximum=1.5, value=0.86, step=0.01, label="Temperature"
+                    )
+                with gr.Row():
+                    schedule = gr.Dropdown(choices=["shift"], value="shift", label="Schedule")
+                    shift = gr.Slider(minimum=0, maximum=16, value=5, step=1, label="Shift")
+                    confidence_policy = gr.Dropdown(
+                        choices=["mask_git", "mmada", "stratified"],
+                        value="mmada",
+                        label="Confidence Policy",
+                    )
+                with gr.Row():
+                    schedule_temp = gr.Dropdown(
+                        choices=["linear", "cosine2", "shift", "exp"],
+                        value="linear",
+                        label="Schedule Temp",
+                    )
+                    alg_temp = gr.Slider(
+                        minimum=0.1, maximum=3.0, value=1.0, step=0.1, label="Alg Temp"
+                    )
+                    dynamic_temperature = gr.Checkbox(label="Dynamic Temp", value=False)
+                with gr.Row():
+                    min_temperature = gr.Slider(
+                        minimum=0.0, maximum=1.0, value=0.01, step=0.01, label="Min Temp"
+                    )
+                    edit_threshold = gr.Slider(
+                        minimum=-1.0, maximum=1.0, value=0.6, step=0.01, label="Edit Threshold"
+                    )
+                    seed = gr.Number(label="Seed (-1 for random)", value=42, precision=0)
+                with gr.Row():
+                    use_cache = gr.Checkbox(label="Use Cache", value=False)
+                    is_legacy = gr.Checkbox(label="Legacy Mode", value=False)
+                    return_animation = gr.Checkbox(
+                        label="Return Animation (resized to 400x400 for preview)",
+                        value=False,
+                    )
+                    enable_image_guard = gr.Checkbox(
+                        label="NSFW output filter (Nemotron 3.5 Content Safety) — uncheck to opt out",
+                        value=DEFAULT_ENABLE_IMAGE_GUARD,
+                    )
+
+                generate_btn = gr.Button("Generate", variant="primary")
+
+            with gr.Column(scale=3):
+                output_image = gr.Image(label="Generated Image", type="filepath")
+                output_meta = gr.Textbox(label="Generation Info", lines=6)
+
+        generate_btn.click(
+            fn=generate,
+            inputs=[
+                prompt,
+                image_resolution,
+                guidance_scale,
+                temperature,
+                n_steps,
+                schedule,
+                shift,
+                confidence_policy,
+                schedule_temp,
+                alg_temp,
+                dynamic_temperature,
+                min_temperature,
+                edit_threshold,
+                seed,
+                micro_cond,
+                use_cache,
+                is_legacy,
+                return_animation,
+                enable_image_guard,
+            ],
+            outputs=[output_image, output_meta],
+        )
+
+        with gr.Accordion("About the model", open=True):
+            gr.Markdown(
+                "_Masked Discrete Diffusion · Text-to-Image Synthesis · Token Editing · "
+                "Grouped Cross-Entropy (GCE) · High-Resolution Image Generation_"
+            )
+
+            gr.Markdown("### Sample outputs")
+            gr.Gallery(
+                value=[
+                    _asset("demo_1.gif"),
+                    _asset("demo_2.gif"),
+                    _asset("demo_3.gif"),
+                ],
+                columns=3,
+                height="auto",
+                object_fit="contain",
+                show_label=False,
+                show_download_button=True,
+            )
+
+            gr.Markdown(
+                "### Generation speed\n\n"
+                "Side-by-side at 1024×1024. **Left:** NL-Diffusion-Image. "
+                "**Right:** EMU3.5 (autoregressive). NL-Diffusion-Image is **42.4× faster** "
+                "while scoring higher on GenEval."
+            )
+            gr.Image(
+                value=_asset("speed_comparison.gif"),
+                show_label=False,
+                interactive=False,
+                show_download_button=True,
+            )
+
+            with gr.Row(equal_height=True):
+                with gr.Column(scale=1):
+                    gr.Markdown(
+                        "### Architecture\n\n"
+                        "16×16 image patches are encoded with a pretrained discrete tokenizer "
+                        "from EMU3.5 (128K codebook). The Nemotron-Labs-Diffusion vocabulary is "
+                        "expanded with randomly initialized embeddings and fine-tuned on "
+                        "image/caption pairs."
+                    )
+                with gr.Column(scale=1):
+                    gr.Image(
+                        value=_asset("architecture.png"),
+                        show_label=False,
+                        interactive=False,
+                        show_download_button=True,
+                    )
+
+            with gr.Row(equal_height=True):
+                with gr.Column(scale=1):
+                    gr.Markdown(
+                        "### Benchmarks\n\n"
+                        "State-of-the-art among discrete image generators at 1024px text-to-image, "
+                        "surpassing prior masked image generators on quality while remaining "
+                        "dramatically faster than autoregressive baselines."
+                    )
+                with gr.Column(scale=1):
+                    gr.Image(
+                        value=_asset("benchmarks.png"),
+                        show_label=False,
+                        interactive=False,
+                        show_download_button=True,
+                    )
+
+            gr.Markdown("### Key findings")
+
+            with gr.Row(equal_height=True):
+                with gr.Column(scale=1):
+                    gr.Markdown(
+                        "**Token editing for self-correction**\n\n"
+                        "Token editing lets the model iteratively refine outputs during inference, "
+                        "correcting artifacts and improving texture detail."
+                    )
+                with gr.Column(scale=1):
+                    gr.Image(
+                        value=_asset("token_editing.png"),
+                        show_label=False,
+                        interactive=False,
+                        show_download_button=True,
+                    )
+
+            with gr.Row(equal_height=True):
+                with gr.Column(scale=1):
+                    gr.Markdown(
+                        "**Grouped Cross-Entropy (GCE)**\n\n"
+                        "GCE alleviates codebook sparsity by supervising semantically close "
+                        "non-top-1 tokens in embedding space.\n\n"
+                        "A fused GCE operator cuts peak VRAM from 25.2 GB to 16.1 GB and latency "
+                        "from 44.14 ms to 20.04 ms versus an eager implementation."
+                    )
+                with gr.Column(scale=1):
+                    gr.Image(
+                        value=_asset("gce_objective.png"),
+                        show_label=False,
+                        interactive=False,
+                        show_download_button=True,
+                    )
+
+            with gr.Row(equal_height=True):
+                with gr.Column(scale=1):
+                    gr.Markdown(
+                        "**Few-step generation**\n\n"
+                        "Unlike continuous flow-matching models that predict blurry mean fields at "
+                        "low step counts, NL-Diffusion-Image produces reasonable quality in as few "
+                        "as 4 steps without distillation."
+                    )
+                with gr.Column(scale=1):
+                    gr.Image(
+                        value=_asset("few_step_generation.png"),
+                        show_label=False,
+                        interactive=False,
+                        show_download_button=True,
+                    )
+
+            gr.Markdown(
+                "**Future work:** extend the model to unified vision generation and understanding."
+            )
+
+    return demo
+
+
+demo = build_demo()
+
+if __name__ == "__main__":
+    demo.queue(default_concurrency_limit=1).launch(
+        server_name=os.getenv("HOST", "0.0.0.0"),
+        server_port=int(os.getenv("PORT", "7860")),
+    )

image_guard.py

@@ -0,0 +1,186 @@
+# 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 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 _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)}")

requirements.txt

@@ -0,0 +1,12 @@
+torch
+transformers>=4.57.1
+accelerate
+safetensors
+huggingface_hub
+gradio
+spaces
+einops
+diffusers
+tqdm
+imageio
+Pillow

test_hub_load.py

@@ -0,0 +1,93 @@
+#!/usr/bin/env python3
+"""Quick smoke test: load private Hub model and generate one image.
+
+Defaults match nemotron-diffusion-omni/gradio_t2i_demo.py:
+  image_resolution=1024, n_tokens=4096, is_legacy=False, 64 steps
+  NSFW filter on by default (set ENABLE_IMAGE_GUARD=0 to opt out)
+
+Usage on a GPU node:
+
+    conda activate lavida
+    export HF_TOKEN=hf_...
+    export MODEL_ID=nvidia/NL-Diffusion-Image
+    python test_hub_load.py
+
+Opt out of guard for local runs:
+
+    ENABLE_IMAGE_GUARD=0 python test_hub_load.py
+"""
+
+from __future__ import annotations
+
+import os
+import sys
+import time
+
+import torch
+
+from app import DEFAULT_MICRO_COND, DEFAULT_PROMPT, T2IEngine
+
+MODEL_ID = os.getenv("MODEL_ID", "nvidia/NL-Diffusion-Image")
+OUTPUT_PATH = os.getenv("OUTPUT_PATH", "test_hub_output.webp")
+RESOLUTION = int(os.getenv("TEST_RESOLUTION", "1024"))
+
+
+def main() -> int:
+    if not torch.cuda.is_available():
+        print("ERROR: CUDA is not available. Run this on a GPU node.", file=sys.stderr)
+        return 1
+
+    if not os.getenv("HF_TOKEN") and not os.path.isdir(MODEL_ID):
+        print(
+            "ERROR: Set HF_TOKEN to load the private Hub model, "
+            f"or set MODEL_ID to a local checkpoint directory.",
+            file=sys.stderr,
+        )
+        return 1
+
+    print(f"CUDA device: {torch.cuda.get_device_name()}")
+    print(f"Model: {MODEL_ID}")
+    print(f"Resolution: {RESOLUTION} (gradio_t2i_demo.py default: 1024, is_legacy=False)")
+
+    engine = T2IEngine(model_id=MODEL_ID, device="cuda")
+
+    t0 = time.time()
+    image_path, meta = engine.generate(
+        prompt=os.getenv("TEST_PROMPT", DEFAULT_PROMPT),
+        image_resolution=RESOLUTION,
+        guidance_scale=5.0,
+        temperature=0.86,
+        n_steps=int(os.getenv("TEST_STEPS", "64")),
+        schedule="shift",
+        shift=5,
+        confidence_policy="mmada",
+        schedule_temp="linear",
+        alg_temp=1.0,
+        dynamic_temperature=False,
+        min_temperature=0.01,
+        edit_threshold=0.6,
+        seed=42,
+        micro_cond=os.getenv("TEST_MICRO_COND", DEFAULT_MICRO_COND),
+        use_cache=False,
+        is_legacy=False,
+        return_animation=False,
+        enable_image_guard=os.getenv("ENABLE_IMAGE_GUARD", "1") == "1",
+    )
+    elapsed = time.time() - t0
+
+    if image_path is None:
+        print(meta, file=sys.stderr)
+        print(f"Total wall time: {elapsed:.2f}s")
+        return 1
+
+    import shutil
+
+    shutil.copy(image_path, OUTPUT_PATH)
+    print(f"Saved {OUTPUT_PATH}")
+    print(meta)
+    print(f"Total wall time: {elapsed:.2f}s")
+    return 0
+
+
+if __name__ == "__main__":
+    raise SystemExit(main())