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
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 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
import numpy as np
from PIL import Image
import torch
from scipy.stats import wasserstein_distance
def exposure_loss(image, target_exposure = 3):
# Calculate the brightness (exposure) of the image.
image_brightness = torch.mean(image)
# Calculate the loss as the absolute difference from the target exposure.
loss = torch.abs(image_brightness - target_exposure)
return loss
def color_diversity_loss(images):
# Calculate color diversity by measuring the variance of color channels (R, G, B).
color_variance = torch.var(images, dim=(2, 3), keepdim=True)
# Sum the color variances for each channel to get the total color diversity.
total_color_diversity = torch.sum(color_variance, dim=1)
return total_color_diversity
def sharpness_loss(images):
# Apply the Laplacian filter to the images to measure sharpness.
laplacian_filter = torch.Tensor([[-1, -1, -1],
[-1, 8, -1],
[-1, -1, -1]]).view(1, 1, 3, 3).to(images.device)
# Expand the filter to match the number of channels in the input image.
laplacian_filter = laplacian_filter.expand(-1, images.shape[1], -1, -1)
# Apply the convolution operation.
laplacian = torch.abs(F.conv2d(images, laplacian_filter))
# Calculate sharpness as the negative of the Laplacian variance.
sharpness = torch.var(laplacian)
return sharpness
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 = color_diversity_loss
else:
loss_function = None
return image_generator(prompt, loss_function)
icon_html = ''
title = f"""
{icon_html} Image Generation using Stable Diffusion
"""
description = f"""
{icon_html}
Embedding New Styles Into Stable Diffusion