#!/usr/bin/env python3
"""Generate a demo image with PixelFlow-256."""

from pathlib import Path

import torch
from diffusers import DiffusionPipeline

REPO_ROOT = Path(__file__).resolve().parent
MODEL_DIR = REPO_ROOT / "PixelFlow-256"
OUTPUT_PATH = REPO_ROOT / "PixelFlow-256" / "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")

    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=256,
        width=256,
        num_inference_steps=[10, 10, 10, 10],
        guidance_scale=4.0,
        generator=generator,
    ).images[0]
    image.save(OUTPUT_PATH)
    print(f"Saved demo image to {OUTPUT_PATH}")


if __name__ == "__main__":
    main()