import gradio as gr import os import sys from base64 import b64encode import numpy import numpy as np import torch from diffusers import AutoencoderKL, LMSDiscreteScheduler, UNet2DConditionModel from matplotlib import pyplot as plt from pathlib import Path from PIL import Image from torch import autocast from torchvision import transforms as tfms from tqdm.auto import tqdm from transformers import CLIPTextModel, CLIPTokenizer, logging import os import cv2 import torchvision.transforms as T torch.manual_seed(1) logging.set_verbosity_error() torch_device = "cuda" if torch.cuda.is_available() else "cpu" # Load the autoencoder vae = AutoencoderKL.from_pretrained("CompVis/stable-diffusion-v1-4", subfolder='vae') # Load 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") # Unet model for generating latents unet = UNet2DConditionModel.from_pretrained("CompVis/stable-diffusion-v1-4", subfolder='unet') # Noise scheduler scheduler = LMSDiscreteScheduler(beta_start=0.00085, beta_end=0.012, beta_schedule="scaled_linear", num_train_timesteps=1000) # Move everything to GPU vae = vae.to(torch_device) text_encoder = text_encoder.to(torch_device) unet = unet.to(torch_device) def get_output_embeds(input_embeddings): # CLIP's text model uses causal mask, so we prepare it here: bsz, seq_len = input_embeddings.shape[:2] causal_attention_mask = text_encoder.text_model._build_causal_attention_mask(bsz, seq_len, dtype=input_embeddings.dtype) # Getting the output embeddings involves calling the model with passing output_hidden_states=True # so that it doesn't just return the pooled final predictions: encoder_outputs = text_encoder.text_model.encoder( inputs_embeds=input_embeddings, attention_mask=None, # We aren't using an attention mask so that can be None causal_attention_mask=causal_attention_mask.to(torch_device), output_attentions=None, output_hidden_states=True, # We want the output embs not the final output return_dict=None, ) # We're interested in the output hidden state only output = encoder_outputs[0] # There is a final layer norm we need to pass these through output = text_encoder.text_model.final_layer_norm(output) # And now they're ready! return output # Prep Scheduler def set_timesteps(scheduler, num_inference_steps): scheduler.set_timesteps(num_inference_steps) scheduler.timesteps = scheduler.timesteps.to(torch.float32) # minor fix to ensure MPS compatibility, fixed in diffusers PR 3925 style_files = ['learned_embeds_birb_style.bin','learned_embeds_cute_game_style.bin', 'learned_embeds_manga_style.bin','learned_embeds_midjourney_style.bin','learned_embeds_space_style.bin'] seed_values = [8,16,50,80,128] height = 512 # default height of Stable Diffusion width = 512 # default width of Stable Diffusion num_inference_steps = 5 # Number of denoising steps guidance_scale = 7.5 # Scale for classifier-free guidance num_styles = len(style_files) def get_style_embeddings(style_file): style_embed = torch.load(style_file) style_name = list(style_embed.keys())[0] return style_embed[style_name] def get_EOS_pos_in_prompt(prompt): return len(prompt.split())+1 import torch.nn.functional as F """ def gradient_loss(images): # Compute gradient magnitude using Sobel filters. gradient_x = F.conv2d(images, torch.Tensor([[-1, 0, 1], [-2, 0, 2], [-1, 0, 1]]).view(1, 1, 3, 3).to(images.device)) gradient_y = F.conv2d(images, torch.Tensor([[-1, -2, -1], [0, 0, 0], [1, 2, 1]]).view(1, 1, 3, 3).to(images.device)) gradient_magnitude = torch.sqrt(gradient_x**2 + gradient_y**2) return gradient_magnitude.mean() """ from torchvision.transforms import ToTensor def pil_to_latent(input_im): # Single image -> single latent in a batch (so size 1, 4, 64, 64) with torch.no_grad(): latent = vae.encode(tfms.ToTensor()(input_im).unsqueeze(0).to(torch_device)*2-1) # Note scaling return 0.18215 * latent.latent_dist.sample() def latents_to_pil(latents): # bath of latents -> list of images latents = (1 / 0.18215) * latents with torch.no_grad(): image = vae.decode(latents).sample image = (image / 2 + 0.5).clamp(0, 1) image = image.detach().cpu().permute(0, 2, 3, 1).numpy() images = (image * 255).round().astype("uint8") pil_images = [Image.fromarray(image) for image in images] return pil_images def additional_guidance(latents, scheduler, noise_pred, t, sigma, custom_loss_fn, custom_loss_scale): #### ADDITIONAL GUIDANCE ### # Requires grad on the latents 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).sample / 2 + 0.5 # range (0, 1) # Calculate loss loss = custom_loss_fn(denoised_images) * custom_loss_scale # Get gradient cond_grad = torch.autograd.grad(loss, latents, allow_unused=False)[0] # Modify the latents based on this gradient latents = latents.detach() - cond_grad * sigma**2 return latents, loss def generate_with_embs(text_embeddings, max_length, random_seed, loss_fn = None, custom_loss_scale=1.0): height = 512 # default height of Stable Diffusion width = 512 # default width of Stable Diffusion num_inference_steps = 5 # Number of denoising steps guidance_scale = 7.5 # Scale for classifier-free guidance generator = torch.manual_seed(random_seed) # Seed generator to create the inital latent noise batch_size = 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 set_timesteps(scheduler, 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.init_noise_sigma # Loop for i, t in tqdm(enumerate(scheduler.timesteps), total=len(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 = scheduler.scale_model_input(latent_model_input, t) # 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) if loss_fn is not None: if i%2 == 0: latents, custom_loss = additional_guidance(latents, scheduler, noise_pred, t, sigma, loss_fn, custom_loss_scale) print(i, 'loss:', custom_loss.item()) # compute the previous noisy sample x_t -> x_t-1 latents = scheduler.step(noise_pred, t, latents).prev_sample return latents_to_pil(latents)[0] def generate_image_custom_style(prompt, style_num=None, random_seed=41, custom_loss_fn = None, custom_loss_scale=1.0): eos_pos = get_EOS_pos_in_prompt(prompt) style_token_embedding = None if style_num: style_token_embedding = get_style_embeddings(style_files[style_num]) # tokenize text_input = tokenizer(prompt, padding="max_length", max_length=tokenizer.model_max_length, truncation=True, return_tensors="pt") max_length = text_input.input_ids.shape[-1] input_ids = text_input.input_ids.to(torch_device) # get token embeddings token_emb_layer = text_encoder.text_model.embeddings.token_embedding token_embeddings = token_emb_layer(input_ids) # Append style token towards the end of the sentence embeddings if style_token_embedding is not None: token_embeddings[-1, eos_pos, :] = style_token_embedding # combine with pos embs pos_emb_layer = text_encoder.text_model.embeddings.position_embedding position_ids = text_encoder.text_model.embeddings.position_ids[:, :77] position_embeddings = pos_emb_layer(position_ids) input_embeddings = token_embeddings + position_embeddings # Feed through to get final output embs modified_output_embeddings = get_output_embeds(input_embeddings) # And generate an image with this: generated_image = generate_with_embs(modified_output_embeddings, max_length, random_seed, custom_loss_fn, custom_loss_scale) return generated_image def show_images(images_list): # Let's visualize the four channels of this latent representation: fig, axs = plt.subplots(1, len(images_list), figsize=(16, 4)) for c in range(len(images_list)): axs[c].imshow(images_list[c]) plt.show() def invert_loss(gen_image): inverter = T.RandomInvert(p=1.0) inverted_img = inverter(gen_image) #loss = torch.abs(gen_image - inverted_img).sum() loss = torch.nn.functional.mse_loss(gen_image[:,0], gen_image[:,2]) + torch.nn.functional.mse_loss(gen_image[:,2], gen_image[:,1]) + torch.nn.functional.mse_loss(gen_image[:,0], gen_image[:,1]) return loss def contrast_loss(images): # Calculate the variance of pixel values as a measure of contrast. variance = torch.var(images) return -variance def blue_loss(images): # How far are the blue channel values to 0.9: error = torch.abs(images[:,2] - 0.9).mean() # [:,2] -> all images in batch, only the blue channel return error def display_images_in_rows(images_with_titles, titles): num_images = len(images_with_titles) rows = 5 # Display 5 rows always columns = 1 if num_images == 5 else 2 # Use 1 column if there are 5 images, otherwise 2 columns fig, axes = plt.subplots(rows, columns + 1, figsize=(15, 5 * rows)) # Add an extra column for titles for r in range(rows): # Add the title on the extreme left in the middle of each picture axes[r, 0].text(0.5, 0.5, titles[r], ha='center', va='center') axes[r, 0].axis('off') # Add "Without Loss" label above the first column and "With Loss" label above the second column (if applicable) if columns == 2: axes[r, 1].set_title("Without Loss", pad=10) axes[r, 2].set_title("With Loss", pad=10) for c in range(1, columns + 1): index = r * columns + c - 1 if index < num_images: image, _ = images_with_titles[index] axes[r, c].imshow(image) axes[r, c].axis('off') return fig # plt.show() def image_generator(prompt = "dog", loss_function=None): images_without_loss = [] images_with_loss = [] for i in range(num_styles): generated_img = generate_image_custom_style(prompt,style_num = i,random_seed = seed_values[i],custom_loss_fn = None) images_without_loss.append(generated_img) if loss_function: generated_img = generate_image_custom_style(prompt,style_num = i,random_seed = seed_values[i],custom_loss_fn = loss_function) images_with_loss.append(generated_img) generated_sd_images = [] titles = ["Birb Style","Cute Game Style","Manga Style","Mid Journey Style","Space Style"] for i in range(len(titles)): generated_sd_images.append((images_without_loss[i], titles[i])) if images_with_loss != []: generated_sd_images.append((images_with_loss[i], titles[i])) return display_images_in_rows(generated_sd_images, titles) # Create a wrapper function for show_misclassified_images() def image_generator_wrapper(prompt = "dog", loss_function=None): if loss_function == "Yes": loss_function = contrast_loss else: loss_function = None return image_generator(prompt, loss_function) description = '(Team Project EE267) Stable Diffusion is a generative artificial intelligence (generative AI) model that produces unique photorealistic images from text and image prompts.' title = 'Image Generation using Stable Diffusion' demo = gr.Interface(image_generator_wrapper, inputs=[gr.Textbox(label="Enter prompt for generation", type="text", value="a dog walking on moon"), gr.Radio(["Yes", "No"], value="No" , label="Custom Loss Function")], outputs=gr.Plot(label="Generated Images"), title = "EE 267 Team Project - Stable Diffusion", description=description) demo.launch()