Update app.py

app.py

@@ -0,0 +1,320 @@
+import gradio as gr
+import os
+import sys
+from base64 import b64encode
+
+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_animal_toys.bin','learned_embeds_fftstyle.bin',
+          'learned_embeds_midjourney_style.bin','learned_embeds_oil_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 brilliance_loss(image, target_brilliance=10):
+    # Calculate the standard deviation of color channels
+    std_dev = torch.std(image, dim=(2, 3))
+    # Calculate the mean standard deviation across the batch
+    mean_std_dev = torch.mean(std_dev)
+    # Calculate the loss as the absolute difference from the target brilliance.
+    loss = torch.abs(mean_std_dev - target_brilliance)
+    return loss
+
+
+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 = ["animal toy","fft style","mid journey","oil 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 = brilliance_loss
+  else:
+    loss_function = None
+
+  return image_generator(prompt, loss_function)
+
+description = '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="astronaut riding a cycle"),
+                            gr.Radio(["Yes", "No"], value="No"  , label="Apply Contrast Loss")],
+                    outputs=gr.Plot(label="Generated Images"), title = "Stable Diffusion", description=description)
+demo.launch()