diffusers/tools

1#!/usr/bin/env python32from diffusers import DiffusionPipeline3import torch4 5run_compile = True  # Set True / False6 7pipe = DiffusionPipeline.from_pretrained("DeepFloyd/IF-I-M-v1.0", variant="fp16", text_encoder=None, torch_dtype=torch.float16)8pipe.to("cuda")9pipe_2 = DiffusionPipeline.from_pretrained("DeepFloyd/IF-II-M-v1.0", variant="fp16", text_encoder=None, torch_dtype=torch.float16)10pipe_2.to("cuda")11pipe_3 = DiffusionPipeline.from_pretrained("stabilityai/stable-diffusion-x4-upscaler", torch_dtype=torch.float16)12pipe_3.to("cuda")13 14 15pipe.unet.to(memory_format=torch.channels_last)16pipe_2.unet.to(memory_format=torch.channels_last)17pipe_3.unet.to(memory_format=torch.channels_last)18 19if run_compile:20    pipe.unet = torch.compile(pipe.unet, mode="reduce-overhead", fullgraph=True)21    pipe_2.unet = torch.compile(pipe_2.unet, mode="reduce-overhead", fullgraph=True)22    pipe_3.unet = torch.compile(pipe_3.unet, mode="reduce-overhead", fullgraph=True)23 24prompt = "the blue hulk"25 26prompt_embeds = torch.randn((1, 2, 4096), dtype=torch.float16)27neg_prompt_embeds = torch.randn((1, 2, 4096), dtype=torch.float16)28 29for _ in range(3):30    image = pipe(prompt_embeds=prompt_embeds, negative_prompt_embeds=neg_prompt_embeds, output_type="pt").images31    image_2 = pipe_2(image=image, prompt_embeds=prompt_embeds, negative_prompt_embeds=neg_prompt_embeds, output_type="pt").images32    image_3 = pipe_3(prompt=prompt, image=image, noise_level=100).images33