demo_inference_t2i.py

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

from pathlib import Path

import torch
from diffusers import DiffusionPipeline

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

    generator = torch.Generator(device="cuda").manual_seed(42)
    image = pipe(
        prompt="A golden retriever playing in a sunny garden",
        height=1024,
        width=1024,
        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()