| |
| 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 |
|
|
