|
|
|
from diffusers import DiffusionPipeline |
|
import torch |
|
|
|
run_compile = False |
|
|
|
pipe = DiffusionPipeline.from_pretrained("DeepFloyd/IF-I-M-v1.0", variant="fp16", text_encoder=None, torch_dtype=torch.float16) |
|
pipe.to("cuda") |
|
pipe_2 = DiffusionPipeline.from_pretrained("DeepFloyd/IF-II-M-v1.0", variant="fp16", text_encoder=None, torch_dtype=torch.float16) |
|
pipe_2.to("cuda") |
|
pipe_3 = DiffusionPipeline.from_pretrained("stabilityai/stable-diffusion-x4-upscaler", torch_dtype=torch.float16) |
|
pipe_3.to("cuda") |
|
|
|
|
|
pipe.unet.to(memory_format=torch.channels_last) |
|
pipe_2.unet.to(memory_format=torch.channels_last) |
|
pipe_3.unet.to(memory_format=torch.channels_last) |
|
|
|
if run_compile: |
|
pipe.unet = torch.compile(pipe.unet, mode="reduce-overhead", fullgraph=True) |
|
pipe_2.unet = torch.compile(pipe_2.unet, mode="reduce-overhead", fullgraph=True) |
|
pipe_3.unet = torch.compile(pipe_3.unet, mode="reduce-overhead", fullgraph=True) |
|
|
|
prompt = "the blue hulk" |
|
|
|
prompt_embeds = torch.randn((1, 2, 4096), dtype=torch.float16) |
|
neg_prompt_embeds = torch.randn((1, 2, 4096), dtype=torch.float16) |
|
|
|
for _ in range(3): |
|
image = pipe(prompt_embeds=prompt_embeds, negative_prompt_embeds=neg_prompt_embeds, output_type="pt").images |
|
image_2 = pipe_2(image=image, prompt_embeds=prompt_embeds, negative_prompt_embeds=neg_prompt_embeds, output_type="pt").images |
|
image_3 = pipe_3(prompt=prompt, image=image, noise_level=100).images |
|
|
