diffusers/tools

1#!/usr/bin/env python32import torch3from datasets import Dataset, Features4from datasets import Image as ImageFeature5from datasets import Value, load_dataset6from diffusers import AutoPipelineForText2Image7 8import PIL9 10 11def main():12    print("Loading dataset...")13    parti_prompts = load_dataset("nateraw/parti-prompts", split="train")14 15    print("Loading pipeline...")16    seed = 017 18    device = "cuda"19    generator = torch.Generator(device).manual_seed(seed)20    dtype = torch.float1621 22    ckpt_id = "warp-diffusion/wuerstchen"23 24    pipeline =  AutoPipelineForText2Image.from_pretrained(25        ckpt_id, torch_dtype=dtype26    ).to(device)27 28    pipeline.prior_prior = torch.compile(pipeline.prior_prior, mode="reduce-overhead", fullgraph=True)29    pipeline.decoder = torch.compile(pipeline.decoder, mode="reduce-overhead", fullgraph=True)30 31    print("Running inference...")32    main_dict = {}33    for i in range(len(parti_prompts)):34        sample = parti_prompts[i]35        prompt = sample["Prompt"]36 37        image = pipeline(38            prompt=prompt,39            height=1024,40            width=1024,41            prior_guidance_scale=4.0,42            decoder_guidance_scale=0.0,43            generator=generator,44        ).images[0]45 46        image = image.resize((256, 256), resample=PIL.Image.Resampling.LANCZOS)47        img_path = f"wuerstchen_{i}.png"48        image.save(img_path)49        main_dict.update(50            {51                prompt: {52                    "img_path": img_path,53                    "Category": sample["Category"],54                    "Challenge": sample["Challenge"],55                    "Note": sample["Note"],56                    "model_name": ckpt_id,57                    "seed": seed,58                }59            }60        )61 62    def generation_fn():63        for prompt in main_dict:64            prompt_entry = main_dict[prompt]65            yield {66                "Prompt": prompt,67                "Category": prompt_entry["Category"],68                "Challenge": prompt_entry["Challenge"],69                "Note": prompt_entry["Note"],70                "images": {"path": prompt_entry["img_path"]},71                "model_name": prompt_entry["model_name"],72                "seed": prompt_entry["seed"],73            }74 75    print("Preparing HF dataset...")76    ds = Dataset.from_generator(77        generation_fn,78        features=Features(79            Prompt=Value("string"),80            Category=Value("string"),81            Challenge=Value("string"),82            Note=Value("string"),83            images=ImageFeature(),84            model_name=Value("string"),85            seed=Value("int64"),86        ),87    )88    ds_id = "diffusers-parti-prompts/wuerstchen"89    ds.push_to_hub(ds_id)90 91 92if __name__ == "__main__":93    main()94