test_hub_load.py

#!/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")),
        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),
        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())