demo_inference.py

#!/usr/bin/env python3
"""Generate a demo image with NiT-XL."""

from pathlib import Path

import torch
from diffusers import DiffusionPipeline

REPO_ROOT = Path(__file__).resolve().parent
MODEL_DIR = REPO_ROOT / "NiT-XL"
OUTPUT_PATH = REPO_ROOT / "demo.png"


def main() -> None:
    pipe = DiffusionPipeline.from_pretrained(
        str(MODEL_DIR),
        local_files_only=True,
        custom_pipeline=str(MODEL_DIR / "pipeline.py"),
        trust_remote_code=True,
        torch_dtype=torch.bfloat16,
    )
    pipe.to("cuda")
    pipe.set_progress_bar_config(disable=False)

    print(pipe.id2label[207])
    print(pipe.get_label_ids("golden retriever"))

    generator = torch.Generator(device="cuda").manual_seed(42)
    image = pipe(
        class_labels="golden retriever",
        height=512,
        width=512,
        num_inference_steps=250,
        guidance_scale=2.05,
        guidance_interval=(0.0, 0.7),
        generator=generator,
    ).images[0]
    image.save(OUTPUT_PATH)
    print(f"Saved demo image to {OUTPUT_PATH}")


if __name__ == "__main__":
    main()