diffusers/tools

1#!/usr/bin/env python32from diffusers import DiffusionPipeline3from diffusers.utils import pt_to_pil4import torch5 6# stage 17stage_1 = DiffusionPipeline.from_pretrained("DeepFloyd/IF-I-M-v1.0", variant="fp16", torch_dtype=torch.float16)8stage_1.enable_xformers_memory_efficient_attention()  # remove line if torch.__version__ >= 2.0.09stage_1.enable_model_cpu_offload()10 11# stage 212stage_2 = DiffusionPipeline.from_pretrained(13    "DeepFloyd/IF-II-M-v1.0", text_encoder=None, variant="fp16", torch_dtype=torch.float1614)15stage_2.enable_xformers_memory_efficient_attention()  # remove line if torch.__version__ >= 2.0.016stage_2.enable_model_cpu_offload()17 18# stage 319safety_modules = {"feature_extractor": stage_1.feature_extractor, "safety_checker": stage_1.safety_checker, "watermarker": stage_1.watermarker}20stage_3 = DiffusionPipeline.from_pretrained("stabilityai/stable-diffusion-x4-upscaler", **safety_modules, torch_dtype=torch.float16)21stage_3.enable_xformers_memory_efficient_attention()  # remove line if torch.__version__ >= 2.0.022stage_3.enable_model_cpu_offload()23 24prompt = 'a photo of a kangaroo wearing an orange hoodie and blue sunglasses standing in front of the eiffel tower holding a sign that says "very deep learning"'25 26# text embeds27prompt_embeds, negative_embeds = stage_1.encode_prompt(prompt)28 29generator = torch.manual_seed(0)30 31image = stage_1(prompt_embeds=prompt_embeds, negative_prompt_embeds=negative_embeds, generator=generator, output_type="pt").images32pt_to_pil(image)[0].save("./if_stage_I.png")33 34image = stage_2(35    image=image, prompt_embeds=prompt_embeds, negative_prompt_embeds=negative_embeds, generator=generator, output_type="pt"36).images37pt_to_pil(image)[0].save("./if_stage_II.png")38 39image = stage_3(prompt=prompt, image=image, generator=generator, noise_level=100).images40image[0].save("./if_stage_III.png")41 42