diffusers-parti-prompts/muse512

Dataset Card for "muse_512"

from PIL import Image import torch from muse import PipelineMuse, MaskGiTUViT from datasets import Dataset, Features from datasets import Image as ImageFeature from datasets import Value, load_dataset

device = "cuda" if torch.cuda.is_available() else "cpu"

pipe = PipelineMuse.frompretrained( transformerpath="valhalla/research-run", textencoderpath="openMUSE/clip-vit-large-patch14-text-enc", vae_path="openMUSE/vqgan-f16-8192-laion", ).to(device)

pipe.transformer = MaskGiTUViT.frompretrained("valhalla/research-run-finetuned-journeydb", revision="06bcd6ab6580a2ed3275ddfc17f463b8574457da", subfolder="emamodel").to(device) pipe.tokenizer.padtokenid = 49407

if device == "cuda": pipe.transformer.enablexformersmemoryefficientattention() pipe.text_encoder.to(torch.float16) pipe.transformer.to(torch.float16)

import PIL

def main(): print("Loading dataset...") partiprompts = loaddataset("nateraw/parti-prompts", split="train")

print("Loading pipeline...") seed = 0

device = "cuda" torch.manual_seed(0)

ckpt_id = "openMUSE/muse-512"

scale = 10

print("Running inference...") maindict = {} for i in range(len(partiprompts)): sample = parti_prompts[i] prompt = sample["Prompt"]

image = pipe( prompt, timesteps=16, negativetext=None, guidancescale=scale, temperature=(2, 0), origsize=(512, 512), cropcoords=(0, 0), aestheticscore=6, usefp16=device == "cuda", transformerseqlen=1024, use_tqdm=False, )[0]

image = image.resize((256, 256), resample=PIL.Image.Resampling.LANCZOS) imgpath = f"/home/patrick/museimages/muse512{i}.png" image.save(imgpath) maindict.update( { prompt: { "imgpath": imgpath, "Category": sample["Category"], "Challenge": sample["Challenge"], "Note": sample["Note"], "modelname": ckptid, "seed": seed, } } )

def generationfn(): for prompt in maindict: promptentry = maindict[prompt] yield { "Prompt": prompt, "Category": promptentry["Category"], "Challenge": promptentry["Challenge"], "Note": promptentry["Note"], "images": {"path": promptentry["imgpath"]}, "modelname": promptentry["modelname"], "seed": prompt_entry["seed"], }

print("Preparing HF dataset...") ds = Dataset.fromgenerator( generationfn, features=Features( Prompt=Value("string"), Category=Value("string"), Challenge=Value("string"), Note=Value("string"), images=ImageFeature(), modelname=Value("string"), seed=Value("int64"), ), ) dsid = "diffusers-parti-prompts/muse512" ds.pushtohub(ds_id)

if _name == "main_": main()