"""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 random
import tempfile
import time
from contextlib import contextmanager
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()


# The released model drives its denoising loop with a tqdm bar over a
# length-less iterable (`enumerate(...)` plus a separate `total=`). Gradio's
# track_tqdm ignores that explicit total and the bar has no description, so the
# UI shows a bogus "Downloading (incomplete total...)" indicator. This shim
# backfills the missing length/description onto the tracked bar, but only while
# a generation is in flight (so download bars during model load are untouched).
_DENOISE_PROGRESS = {"active": False, "total": None, "desc": "Generating image"}


def _install_progress_shim() -> None:
    try:
        from gradio import helpers as _gr_helpers
    except Exception:
        return
    if getattr(_gr_helpers.Progress.tqdm, "_denoise_shim", False):
        return
    _orig_tqdm = _gr_helpers.Progress.tqdm

    def _patched_tqdm(self, *args, **kwargs):
        out = _orig_tqdm(self, *args, **kwargs)
        try:
            if _DENOISE_PROGRESS["active"] and getattr(self, "iterables", None):
                ti = self.iterables[-1]
                if getattr(ti, "length", None) in (None, 0) and _DENOISE_PROGRESS["total"]:
                    ti.length = _DENOISE_PROGRESS["total"]
                if not getattr(ti, "desc", None):
                    ti.desc = _DENOISE_PROGRESS["desc"]
        except Exception:
            pass
        return out

    _patched_tqdm._denoise_shim = True
    _gr_helpers.Progress.tqdm = _patched_tqdm


_install_progress_shim()


@contextmanager
def _suppress_tqdm_tracking():
    """Stop track_tqdm from capturing tqdm bars (e.g. Hugging Face download
    bars during model loading), so only our explicit status message shows."""
    try:
        from gradio.context import LocalContext
    except Exception:
        yield
        return
    token = LocalContext.progress.set(None)
    try:
        yield
    finally:
        try:
            LocalContext.progress.reset(token)
        except Exception:
            pass

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

def resolve_hf_token() -> str | None:
    """Return an HF token from any of the common env var names (or None)."""
    for var in ("HF_TOKEN", "HUGGING_FACE_HUB_TOKEN", "HUGGINGFACEHUB_API_TOKEN"):
        value = os.getenv(var)
        if value:
            return value.strip()
    return None


HF_TOKEN = resolve_hf_token()


def _log_token_status() -> None:
    if HF_TOKEN:
        masked = f"{HF_TOKEN[:4]}…{HF_TOKEN[-2:]}" if len(HF_TOKEN) > 6 else "set"
        print(f"HF token detected ({masked}).", flush=True)
    else:
        print(
            "WARNING: no HF token found (HF_TOKEN / HUGGING_FACE_HUB_TOKEN). "
            "Private models will fail to load with a 404. "
            "Add an HF_TOKEN secret in the Space Settings with read access.",
            flush=True,
        )


_log_token_status()

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


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

EXAMPLE_PROMPTS = [
    "A photorealistic portrait of an astronaut riding a horse on the moon, "
    "golden hour lighting, 85mm lens, ultra detailed, sharp focus",
    "A cozy bookstore cafe interior, warm lighting, hanging plants, wooden shelves, "
    "cinematic, highly detailed",
    "A majestic snow leopard standing on a rocky cliff, national geographic "
    "photography, crisp fur detail, soft bokeh background",
    "A futuristic city skyline at dusk, neon signs, wet streets with reflections, "
    "cyberpunk aesthetic, 8k, cinematic lighting",
    "A steaming bowl of ramen with a soft-boiled egg, studio food photography, "
    "shallow depth of field, rich colors",
    "An oil painting of a lighthouse on a stormy coast, dramatic clouds, crashing "
    "waves, impressionist style",
]

CITATION_BIBTEX = """@article{li2026nemotron,
  title  = {Nemotron-Labs-Diffusion-Image: Advancing Masked Discrete Diffusion
            for High-Resolution Image Synthesis},
  author = {Li, Shufan and Heinrich, Greg and Ye, Hanrong and Fu, Yonggan and
            Grover, Aditya and Kautz, Jan and Molchanov, Pavlo},
  journal = {arXiv preprint arXiv:2606.29814},
  year   = {2026}
}"""

DEFAULT_GENERATION_CONFIG: dict[str, Any] = {
    "guidance_scale": 5.0,
    "n_steps": 64,
    "shift": 5,
    "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,
    "is_legacy": 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 = HF_TOKEN
    if hf_token is None and not os.path.isdir(model_id):
        raise RuntimeError(
            f"Cannot load '{model_id}': no HF token found. "
            "This is a private repo — add an HF_TOKEN secret in the Space Settings "
            "(Settings → Variables and secrets) using a token with read access to "
            f"{model_id}."
        )
    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,
                hf_token=HF_TOKEN,
            )
        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:
            with _suppress_tqdm_tracking():
                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 _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,
        image_resolution: int,
        guidance_scale: float,
        temperature: float,
        n_steps: int,
        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,
        return_animation: 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():
            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),
            shift=int(shift),
            confidence_policy=confidence_policy,
            schedule_temp=schedule_temp,
            alg_temp=alg_temp,
            dynamic_temperature=dynamic_temperature,
            min_temperature=min_temperature,
        )

        if seed < 0:
            seed = int(torch.seed() % (2**31 - 1))
        torch.manual_seed(int(seed))

        n_tokens = n_tokens_from_resolution(int(image_resolution))

        if progress is not None:
            progress(0.0, desc="Generating…")

        _DENOISE_PROGRESS["total"] = int(n_steps)
        _DENOISE_PROGRESS["active"] = progress is not None
        t0 = time.time()
        try:
            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=progress is None,
                        return_intermediate_steps=return_animation,
                    )
        finally:
            _DENOISE_PROGRESS["active"] = False
        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"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"
        )
        if prompt_guard_meta:
            meta += " | " + prompt_guard_meta

        if progress is not None and enable_image_guard:
            progress(1.0, desc="Running safety filter…")

        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,
    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,
    return_animation: bool,
    progress: gr.Progress = gr.Progress(track_tqdm=True),
):
    return engine.generate(
        prompt,
        image_resolution,
        guidance_scale,
        temperature,
        n_steps,
        shift,
        confidence_policy,
        schedule_temp,
        alg_temp,
        dynamic_temperature,
        min_temperature,
        edit_threshold,
        seed,
        micro_cond,
        return_animation,
        enable_image_guard=DEFAULT_ENABLE_IMAGE_GUARD,
        enable_prompt_guard=DEFAULT_ENABLE_PROMPT_GUARD,
        progress=progress,
    )


def make_theme() -> gr.themes.Base:
    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",
    )
    return gr.themes.Soft(
        primary_hue=nvidia_green,
        secondary_hue=nvidia_green,
        font=[
            gr.themes.GoogleFont("Inter"),
            "ui-sans-serif",
            "system-ui",
            "sans-serif",
        ],
        font_mono=[
            gr.themes.GoogleFont("JetBrains Mono"),
            "ui-monospace",
            "monospace",
        ],
    )


# Gradio 6.0 moved `theme` from Blocks(...) to launch(...).
_GRADIO_MAJOR = int(gr.__version__.split(".")[0])


def build_demo() -> gr.Blocks:
    theme = make_theme()
    blocks_kwargs: dict[str, Any] = {"title": "Nemotron Labs Diffusion Image"}
    if _GRADIO_MAJOR < 6:
        blocks_kwargs["theme"] = theme
    with gr.Blocks(**blocks_kwargs) 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"
            "📄 [Paper (arXiv:2606.29814)](https://arxiv.org/abs/2606.29814) · "
            "🤗 [Model](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)
                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():
                    seed = gr.Number(
                        label="Seed (-1 for random)", value=42, precision=0, scale=3
                    )
                    randomize_seed_btn = gr.Button("🎲 Randomize", scale=1)
                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")

                with gr.Accordion("Advanced settings", open=False):
                    micro_cond = gr.Textbox(
                        label="Micro Cond", lines=2, value=DEFAULT_MICRO_COND
                    )
                    with gr.Row():
                        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",
                        )
                    return_animation = gr.Checkbox(
                        label="Return Animation (resized to 400x400 for preview)",
                        value=False,
                    )

            with gr.Column(scale=3):
                output_image = gr.Image(label="Generated Image", type="filepath")
                output_meta = gr.Textbox(label="Generation Info", lines=6)

        gr.Examples(examples=EXAMPLE_PROMPTS, inputs=[prompt], label="Example prompts")

        randomize_seed_btn.click(
            fn=lambda: random.randint(0, 2**31 - 1), inputs=None, outputs=seed
        )

        generate_btn.click(
            fn=generate,
            inputs=[
                prompt,
                image_resolution,
                guidance_scale,
                temperature,
                n_steps,
                shift,
                confidence_policy,
                schedule_temp,
                alg_temp,
                dynamic_temperature,
                min_temperature,
                edit_threshold,
                seed,
                micro_cond,
                return_animation,
            ],
            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,
            )

            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,            )

            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,                    )

            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,                    )

            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,                    )

            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,                    )

            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,                    )

            gr.Markdown(
                "**Future work:** extend the model to unified vision generation and understanding."
            )

            gr.Markdown("### Citation")
            gr.Code(value=CITATION_BIBTEX, language=None, label="BibTeX")

    return demo


demo = build_demo()

if __name__ == "__main__":
    launch_kwargs: dict[str, Any] = {
        "server_name": os.getenv("HOST", "0.0.0.0"),
        "server_port": int(os.getenv("PORT", "7860")),
    }
    if _GRADIO_MAJOR >= 6:
        launch_kwargs["theme"] = make_theme()
    demo.queue(default_concurrency_limit=1).launch(**launch_kwargs)