File size: 1,922 Bytes
a032108 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 |
#!/usr/bin/env python3
from diffusers import DiffusionPipeline
import torch
from diffusers.models.attention_processor import AttnAddedKVProcessor2_0, XFormersAttnAddedKVProcessor, AttnAddedKVProcessor
import time
import os
from huggingface_hub import HfApi
from pathlib import Path
api = HfApi()
prev_time = time.time()
prompt = "a picture of elon musk next to a rocket"
negative_prompt = "low quality, ugly"
pipe_prior = DiffusionPipeline.from_pretrained(
"kandinsky-community/kandinsky-2-1-prior", torch_dtype=torch.float16
)
pipe_prior.to("cuda")
t2i_pipe = DiffusionPipeline.from_pretrained("kandinsky-community/kandinsky-2-1", torch_dtype=torch.float16)
t2i_pipe.to("cuda")
t2i_pipe.unet.set_attn_processor(AttnAddedKVProcessor2_0())
t2i_pipe.unet.to(memory_format=torch.channels_last)
t2i_pipe.unet = torch.compile(t2i_pipe.unet, mode="reduce-overhead", fullgraph=True)
next_time = time.time()
print("Loading", next_time - prev_time)
prev_time = next_time
generator = torch.Generator(device="cuda").manual_seed(12)
image_embeds, negative_image_embeds = pipe_prior(prompt, negative_prompt=negative_prompt, generator=generator).to_tuple()
next_time = time.time()
print("Prior", next_time - prev_time)
prev_time = next_time
for _ in range(3):
images = t2i_pipe(prompt, image_embeds=image_embeds, negative_image_embeds=negative_image_embeds, negative_prompt=negative_prompt, num_inference_steps=50, generator=generator).images
next_time = time.time()
print("Text-to-image", next_time - prev_time)
prev_time = next_time
for i, image in enumerate(images):
path = os.path.join(Path.home(), "images", f"aa_{i}.png")
image.save(path)
api.upload_file(
path_or_fileobj=path,
path_in_repo=path.split("/")[-1],
repo_id="patrickvonplaten/images",
repo_type="dataset",
)
print(f"https://huggingface.co/datasets/patrickvonplaten/images/blob/main/aa_{i}.png")
|