diffusers/tools

1#!/usr/bin/env python32from PIL import Image  3import torch4from muse import PipelineMuse, MaskGiTUViT5from datasets import Dataset, Features6from datasets import Image as ImageFeature7from datasets import Value, load_dataset8 9device = "cuda" if torch.cuda.is_available() else "cpu"10 11pipe = PipelineMuse.from_pretrained(12    transformer_path="valhalla/research-run",13    text_encoder_path="openMUSE/clip-vit-large-patch14-text-enc",14    vae_path="openMUSE/vqgan-f16-8192-laion",15).to(device)16 17# pipe.transformer = MaskGiTUViT.from_pretrained("valhalla/research-run-finetuned-journeydb", revision="06bcd6ab6580a2ed3275ddfc17f463b8574457da", subfolder="ema_model").to(device)18pipe.transformer = MaskGiTUViT.from_pretrained("valhalla/muse-research-run", subfolder="ema_model").to(device)19pipe.tokenizer.pad_token_id = 4940720 21if device == "cuda":22    pipe.transformer.enable_xformers_memory_efficient_attention()23    pipe.text_encoder.to(torch.float16)24    pipe.transformer.to(torch.float16)25 26 27import PIL28 29 30def main():31    print("Loading dataset...")32    parti_prompts = load_dataset("nateraw/parti-prompts", split="train")33 34    print("Loading pipeline...")35    seed = 036 37    device = "cuda"38    torch.manual_seed(0)39 40    ckpt_id = "openMUSE/muse-256"41 42    scale = 1043 44    print("Running inference...")45    main_dict = {}46    for i in range(len(parti_prompts)):47        sample = parti_prompts[i]48        prompt = sample["Prompt"]49 50        image = pipe(51            prompt,52            timesteps=16,53            negative_text=None,54            guidance_scale=scale,55            temperature=(2, 0),56            orig_size=(256, 256),57            crop_coords=(0, 0),58            aesthetic_score=6,59            use_fp16=device == "cuda",60            transformer_seq_len=256,61            use_tqdm=False,62        )[0]63 64        image = image.resize((256, 256), resample=PIL.Image.Resampling.LANCZOS)65        img_path = f"/home/patrick/muse_images/muse_256_{i}.png"66        image.save(img_path)67        main_dict.update(68            {69                prompt: {70                    "img_path": img_path,71                    "Category": sample["Category"],72                    "Challenge": sample["Challenge"],73                    "Note": sample["Note"],74                    "model_name": ckpt_id,75                    "seed": seed,76                }77            }78        )79 80    def generation_fn():81        for prompt in main_dict:82            prompt_entry = main_dict[prompt]83            yield {84                "Prompt": prompt,85                "Category": prompt_entry["Category"],86                "Challenge": prompt_entry["Challenge"],87                "Note": prompt_entry["Note"],88                "images": {"path": prompt_entry["img_path"]},89                "model_name": prompt_entry["model_name"],90                "seed": prompt_entry["seed"],91            }92 93    print("Preparing HF dataset...")94    ds = Dataset.from_generator(95        generation_fn,96        features=Features(97            Prompt=Value("string"),98            Category=Value("string"),99            Challenge=Value("string"),100            Note=Value("string"),101            images=ImageFeature(),102            model_name=Value("string"),103            seed=Value("int64"),104        ),105    )106    ds_id = "diffusers-parti-prompts/muse_256"107    ds.push_to_hub(ds_id)108 109 110if __name__ == "__main__":111    main()112