Spaces:
Running on Zero
Running on Zero
File size: 2,570 Bytes
904bf44 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 | #!/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())
|