diffusers/tools

1#!/usr/bin/env python32from diffusers import IFBasePipeline, IFSuperResolutionPipeline, UNet2DConditionModel3from transformers import T5EncoderModel, T5Tokenizer4import torch5import gc6import os7from pathlib import Path8 9prompt = '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"'10 11model_id = "diffusers/if"12model_id = "/home/patrick/if"13 14# T515t5_tok = T5Tokenizer.from_pretrained(model_id, subfolder="tokenizer", torch_dtype=torch.float16, variant="fp16", use_safetensors=True)16 17t5 = T5EncoderModel.from_pretrained(model_id, subfolder="text_encoder", torch_dtype=torch.float16, variant="fp16", low_cpu_mem_usage=True)18t5.cuda()19 20prompt = prompt.lower().strip()  # make sure everything is lower-cased21with torch.no_grad():22    inputs = t5_tok(prompt, max_length=77, return_tensors="pt", truncation=True, padding="max_length").to("cuda")23    prompt_embeds = t5(**inputs).last_hidden_state24 25with torch.no_grad():26    inputs = t5_tok("", max_length=77, return_tensors="pt", truncation=True, padding="max_length").to("cuda")27    neg_prompt_embeds = t5(**inputs).last_hidden_state28 29del t530torch.cuda.empty_cache()31gc.collect()32 33generator = torch.Generator("cuda").manual_seed(0)34 35# Stage 136pipe = IFBasePipeline.from_pretrained(model_id, text_encoder=None, torch_dtype=torch.float16, variant="fp16")37pipe.to("cuda")38 39image = pipe(prompt_embeds=prompt_embeds, negative_prompt_embeds=neg_prompt_embeds, output_type="pt", num_inference_steps=100, generator=generator).images40 41# save_image42pil_image = pipe.numpy_to_pil(pipe.decode_latents(image))[0]43pil_image.save(os.path.join(Path.home(), "images", "if_I_0.png"))44 45# offload46del pipe47torch.cuda.empty_cache()48gc.collect()49 50# Stage 251unet = UNet2DConditionModel.from_pretrained(model_id, subfolder="super_res_1_unet", torch_dtype=torch.float16)52pipe = IFSuperResolutionPipeline.from_pretrained(model_id, unet=unet, text_encoder=None, torch_dtype=torch.float16, variant="fp16")53pipe.to("cuda")54 55image = pipe(image=image, prompt_embeds=prompt_embeds, negative_prompt_embeds=neg_prompt_embeds, num_inference_steps=50, noise_level=250, output_type="pt", generator=generator).images56 57# save_image58pil_image = pipe.numpy_to_pil(pipe.decode_latents(image))[0]59pil_image.save(os.path.join(Path.home(), "images", "if_II_0.png"))60 61# offload62del pipe63torch.cuda.empty_cache()64gc.collect()65 66# Stage 367unet = UNet2DConditionModel.from_pretrained(model_id, subfolder="super_res_2_unet", torch_dtype=torch.float16)68pipe = IFSuperResolutionPipeline.from_pretrained(model_id, unet=unet, text_encoder=None, torch_dtype=torch.float16, variant="fp16")69pipe.to("cuda")70 71image = pipe(image=image, prompt_embeds=prompt_embeds, negative_prompt_embeds=neg_prompt_embeds, num_inference_steps=40, noise_level=0, output_type="pt", generator=generator).images72 73# save image74pil_image = pipe.numpy_to_pil(pipe.decode_latents(image))[0]75pil_image.save(os.path.join(Path.home(), "images", "if_III_0.png"))76 77# offload78del pipe79torch.cuda.empty_cache()80gc.collect()81