import torch from transformers import CLIPTextModel, CLIPTokenizer from diffusers import AutoencoderKL, UNet2DConditionModel, LMSDiscreteScheduler from tqdm.auto import tqdm from torch import autocast from PIL import Image from matplotlib import pyplot as plt import numpy from torchvision import transforms as tfms import shutil # For video display: import cv2 import os from utils import color_loss,latents_to_pil,pil_to_latent,sketch_loss # Set device torch_device = "cpu" vae = AutoencoderKL.from_pretrained("CompVis/stable-diffusion-v1-4", subfolder="vae") # Load the tokenizer and text encoder to tokenize and encode the text. tokenizer = CLIPTokenizer.from_pretrained("openai/clip-vit-large-patch14") text_encoder = CLIPTextModel.from_pretrained("openai/clip-vit-large-patch14") # The UNet model for generating the latents. unet = UNet2DConditionModel.from_pretrained("CompVis/stable-diffusion-v1-4", subfolder="unet") # The noise scheduler scheduler = LMSDiscreteScheduler(beta_start=0.00085, beta_end=0.012, beta_schedule="scaled_linear", num_train_timesteps=1000) vae = vae.to(torch_device) text_encoder = text_encoder.to(torch_device) unet = unet.to(torch_device) scheduler.set_timesteps(15) def generate_mixed_image(prompt1, prompt2,num_inference_steps=50,noised_image=False): mix_factor = 0.4 #@param height = 512 # default height of Stable Diffusion width = 512 # default width of Stable Diffusion num_inference_steps = num_inference_steps #@param # Number of denoising steps guidance_scale = 8 # Scale for classifier-free guidance generator = torch.manual_seed(32) # Seed generator to create the inital latent noise batch_size = 1 # Prep text # Embed both prompts text_input1 = tokenizer([prompt1], padding="max_length", max_length=tokenizer.model_max_length, truncation=True, return_tensors="pt") with torch.no_grad(): text_embeddings1 = text_encoder(text_input1.input_ids.to(torch_device))[0] text_input2 = tokenizer([prompt2], padding="max_length", max_length=tokenizer.model_max_length, truncation=True, return_tensors="pt") with torch.no_grad(): text_embeddings2 = text_encoder(text_input2.input_ids.to(torch_device))[0] # Take the average text_embeddings = (text_embeddings1*mix_factor + \ text_embeddings2*(1-mix_factor)) # And the uncond. input as before: max_length = max(text_input1.input_ids.shape[-1],text_input2.input_ids.shape[-1]) uncond_input = tokenizer( [""] * batch_size, padding="max_length", max_length=max_length, return_tensors="pt" ) with torch.no_grad(): uncond_embeddings = text_encoder(uncond_input.input_ids.to(torch_device))[0] text_embeddings = torch.cat([uncond_embeddings, text_embeddings]) # Prep Scheduler scheduler.set_timesteps(num_inference_steps) # Prep latents latents = torch.randn( (batch_size, unet.in_channels, height // 8, width // 8), generator=generator, ) latents = latents.to(torch_device) latents = latents * scheduler.sigmas[0] # Need to scale to match k # Loop with autocast("cuda"): for i, t in tqdm(enumerate(scheduler.timesteps)): # expand the latents if we are doing classifier-free guidance to avoid doing two forward passes. latent_model_input = torch.cat([latents] * 2) sigma = scheduler.sigmas[i] latent_model_input = latent_model_input / ((sigma**2 + 1) ** 0.5) # predict the noise residual with torch.no_grad(): noise_pred = unet(latent_model_input, t, encoder_hidden_states=text_embeddings)["sample"] # perform guidance noise_pred_uncond, noise_pred_text = noise_pred.chunk(2) noise_pred = noise_pred_uncond + guidance_scale * (noise_pred_text - noise_pred_uncond) # compute the previous noisy sample x_t -> x_t-1 latents = scheduler.step(noise_pred, i, latents)["prev_sample"] if noised_image: output = generate_noised_version_of_image(latents_to_pil(latents,vae)[0]) else: output = latents_to_pil(latents,vae)[0] return output def generate_image(prompt,num_inference_steps=50,color_postprocessing=False,postporcessing_color=None,color_loss_scale=40,noised_image=False): #@title Store the predicted outputs and next frame for later viewing #prompt = 'A campfire (oil on canvas)' # height = 512 # default height of Stable Diffusion width = 512 # default width of Stable Diffusion num_inference_steps = num_inference_steps # # Number of denoising steps guidance_scale = 8 # # Scale for classifier-free guidance generator = torch.manual_seed(32) # Seed generator to create the inital latent noise batch_size = 1 # Define the directory name directory_name = "steps" # Check if the directory exists, and if so, delete it if os.path.exists(directory_name): shutil.rmtree(directory_name) #Create the directory os.makedirs(directory_name) # Prep text text_input = tokenizer([prompt], padding="max_length", max_length=tokenizer.model_max_length, truncation=True, return_tensors="pt") with torch.no_grad(): text_embeddings = text_encoder(text_input.input_ids.to(torch_device))[0] # And the uncond. input as before: max_length = text_input.input_ids.shape[-1] uncond_input = tokenizer( [""] * batch_size, padding="max_length", max_length=max_length, return_tensors="pt" ) with torch.no_grad(): uncond_embeddings = text_encoder(uncond_input.input_ids.to(torch_device))[0] text_embeddings = torch.cat([uncond_embeddings, text_embeddings]) # Prep Scheduler scheduler.set_timesteps(num_inference_steps) # Prep latents latents = torch.randn( (batch_size, unet.in_channels, height // 8, width // 8), generator=generator, ) latents = latents.to(torch_device) latents = latents * scheduler.sigmas[0] # Need to scale to match k # Loop with autocast("cuda"): for i, t in tqdm(enumerate(scheduler.timesteps)): # expand the latents if we are doing classifier-free guidance to avoid doing two forward passes. latent_model_input = torch.cat([latents] * 2) sigma = scheduler.sigmas[i] latent_model_input = latent_model_input / ((sigma**2 + 1) ** 0.5) # predict the noise residual with torch.no_grad(): noise_pred = unet(latent_model_input, t, encoder_hidden_states=text_embeddings)["sample"] # perform CFG noise_pred_uncond, noise_pred_text = noise_pred.chunk(2) noise_pred = noise_pred_uncond + guidance_scale * (noise_pred_text - noise_pred_uncond) #### ADDITIONAL GUIDANCE ### # Requires grad on the latents if color_postprocessing: latents = latents.detach().requires_grad_() # Get the predicted x0: latents_x0 = latents - sigma * noise_pred # Decode to image space denoised_images = vae.decode((1 / 0.18215) * latents_x0) / 2 + 0.5 # (0, 1) # Calculate loss #loss = sketch_loss(denoised_images) * color_loss_scale loss = color_loss(denoised_images,postporcessing_color) * color_loss_scale if i%10==0: print(i, 'loss:', loss.item()) # Get gradient cond_grad = -torch.autograd.grad(loss, latents)[0] # Modify the latents based on this gradient latents = latents.detach() + cond_grad * sigma**2 ### And saving as before ### # Get the predicted x0: latents_x0 = latents - sigma * noise_pred im_t0 = latents_to_pil(latents_x0,vae)[0] # And the previous noisy sample x_t -> x_t-1 latents = scheduler.step(noise_pred, i, latents)["prev_sample"] im_next = latents_to_pil(latents,vae)[0] # Combine the two images and save for later viewing im = Image.new('RGB', (1024, 512)) im.paste(im_next, (0, 0)) im.paste(im_t0, (512, 0)) im.save(f'steps/{i:04}.jpeg') else: latents = scheduler.step(noise_pred, i, latents)["prev_sample"] if noised_image: output = generate_noised_version_of_image(latents_to_pil(latents,vae)[0]) else: output = latents_to_pil(latents,vae)[0] return output def generate_noised_version_of_image(pil_image): # View a noised version encoded = pil_to_latent(pil_image,vae) noise = torch.randn_like(encoded) # Random noise timestep = 150 # i.e. equivalent to that at 150/1000 training steps encoded_and_noised = scheduler.add_noise(encoded, noise, timestep) return latents_to_pil(encoded_and_noised,vae)[0] # Display # if __name__ == "__main__": # prompt = 'A campfire (oil on canvas)' # color_loss_scale = 40 # color_postprocessing = False # pil_image = generate_mixed_image("a dog", "a cat") # #pil_image = generate_image(prompt,color_postprocessing,color_loss_scale) # #pil_image = generate_noised_version_of_image(Image.open('output.png').resize((512, 512))) # pil_image.save("output1.png")