BiliSakura/JiT-diffusers

1#!/usr/bin/env python32"""Generate a demo image with JiT-H-32."""3 4from pathlib import Path5 6import torch7from diffusers import DiffusionPipeline, FlowMatchHeunDiscreteScheduler8 9REPO_ROOT = Path(__file__).resolve().parent10MODEL_DIR = REPO_ROOT / "JiT-H-32"11OUTPUT_PATH = REPO_ROOT / "demo.png"12 13 14def main() -> None:15    pipe = DiffusionPipeline.from_pretrained(16        str(MODEL_DIR),17        custom_pipeline=str(MODEL_DIR / "pipeline.py"),18        trust_remote_code=True,19        torch_dtype=torch.bfloat16,20    )21    pipe.scheduler = FlowMatchHeunDiscreteScheduler.from_config(pipe.scheduler.config, shift=4.0)22    pipe.to("cuda")23    pipe.set_progress_bar_config(disable=False)24 25    print(pipe.id2label[207])26    print(pipe.get_label_ids("golden retriever"))27 28    generator = torch.Generator(device="cuda").manual_seed(42)29    image = pipe(30        class_labels="golden retriever",31        num_inference_steps=50,32        guidance_scale=2.3,33        generator=generator,34    ).images[0]35    image.save(OUTPUT_PATH)36    print(f"Saved demo image to {OUTPUT_PATH}")37 38 39if __name__ == "__main__":40    main()41