Initial commit

app.py

@@ -0,0 +1,199 @@
+import gradio as gr
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import torchvision
+from torchvision import transforms as tfms
+import torchvision.models as models
+
+from PIL import Image
+import numpy as np
+from diffusers import LMSDiscreteScheduler, DiffusionPipeline
+
+import random
+import os
+import subprocess
+
+from matplotlib import pyplot as plt
+from pathlib import Path
+from torch import autocast
+from tqdm.auto import tqdm
+
+# Set device
+torch_device = "cuda" if torch.cuda.is_available() else "cpu"
+
+# Load a pre-trained VGG model (you can use other models as well)
+vgg_model = models.vgg16(pretrained=True).features
+vgg_model = vgg_model.to(torch_device)
+
+# Create a new model that extracts features from the chosen layers
+feature_extractor = nn.Sequential()
+for name, layer in vgg_model._modules.items():
+    if name == '0':  # Stop at the 0th layer
+        break
+    feature_extractor.add_module(name, layer)
+feature_extractor = feature_extractor.to(torch_device)
+
+pretrained_model_name_or_path = "segmind/tiny-sd"
+pipe = DiffusionPipeline.from_pretrained(
+    pretrained_model_name_or_path,
+    torch_dtype=torch.float32
+).to(torch_device)
+
+# The noise scheduler
+scheduler = LMSDiscreteScheduler(beta_start=0.00085, beta_end=0.012, beta_schedule="scaled_linear", num_train_timesteps=1000)
+
+concept_dict={'anime_bg_v2':('sd-concepts-library/anime-background-style-v2','<anime-background-style-v2>',31),
+              'birb':('sd-concepts-library/birb-style','<birb-style>',32),
+              'depthmap':('sd-concepts-library/depthmap','<depthmap>',33),
+              'gta5_artwork':('sd-concepts-library/gta5-artwork','<gta5_artwork>',34),
+              'midjourney':('sd-concepts-library/midjourney-style','<midjourney-style>',35),
+              'beetlejuice':('sd-concepts-library/beetlejuice-cartoon-style','<beetlejuice-cartoon>',36)}
+
+cache_style_list = []
+
+def transform_pattern_image(pattern_image):
+    preprocess = tfms.Compose([
+        tfms.Resize((320, 320)),
+        tfms.ToTensor(),
+    ])
+    tfms_pattern_image = preprocess(pattern_image).unsqueeze(0)
+    return tfms_pattern_image
+
+def load_required_style(style):
+    for concept, value in concept_dict.items():
+        if style in concept:
+            concept_key = value[1]
+            concept_seed = value[2]
+            if style not in cache_style_list:
+                pipe.load_textual_inversion(value[0])
+                cache_style_list.append(style)
+            break
+    return concept_key, concept_seed
+
+def pil_to_latent(input_im):
+    # Single image -> single latent in a batch (so size 1, 4, 64, 64)
+    with torch.no_grad():
+        latent = pipe.vae.encode(tfms.ToTensor()(input_im).unsqueeze(0).to(torch_device)*2-1) # Note scaling
+    return 0.18215 * latent.latent_dist.sample()  # [1, 4, 64, 64]
+
+def latents_to_pil(latents):
+    # bath of latents -> list of images
+    latents = (1 / 0.18215) * latents
+    with torch.no_grad():
+        image = pipe.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 perceptual_loss(images, pattern):
+    """
+    This function calculates the perceptual loss between the output image and the target image.
+
+    Parameters:
+    """
+    criterion = nn.MSELoss()
+    mse_loss = criterion(images, pattern)
+    return mse_loss
+
+#Generating image with the modified embeddings with pattern loss guidance and saving the images to steps/{concept} folder
+def generate_with_embs_pattern_loss(prompt, concept_seed, tfm_pattern_image, num_inf_steps):
+    height = 320                        # default height of Stable Diffusion
+    width = 320                         # default width of Stable Diffusion
+    num_inference_steps = num_inf_steps # Number of denoising steps
+    guidance_scale = 8                  # Scale for classifier-free guidance
+    generator = torch.manual_seed(concept_seed) # Seed generator to create the inital latent noise
+    batch_size = 1
+    pattern_loss_scale = 20
+
+    text_input = pipe.tokenizer(prompt, padding="max_length", max_length=pipe.tokenizer.model_max_length, truncation=True, return_tensors="pt")
+    input_ids = text_input.input_ids.to(torch_device)
+    with torch.no_grad():
+        text_embeddings = pipe.text_encoder(text_input.input_ids.to(torch_device))[0]
+
+    max_length = text_input.input_ids.shape[-1]
+    uncond_input = pipe.tokenizer([""] * batch_size, padding="max_length", max_length=max_length, return_tensors="pt")
+    with torch.no_grad():
+        uncond_embeddings = pipe.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, pipe.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)):
+        # 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 = pipe.unet(latent_model_input, t, encoder_hidden_states=text_embeddings)["sample"]
+
+        # perform CFG (Classifier Free Guidance)
+        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 ###
+        if (i%3 == 0):
+            # Requires grad on the latents
+            latents = latents.detach().requires_grad_()
+
+            # Get the predicted x0:
+            latents_x0 = latents - sigma * noise_pred
+            # latents_x0 = scheduler.step(noise_pred, t, latents).pred_original_sample
+
+            # Decode to image space
+            denoised_images = pipe.vae.decode((1 / 0.18215) * latents_x0).sample / 2 + 0.5 # range (0, 1)
+            # Calculate loss
+            denoised_images_extr = feature_extractor(denoised_images)
+            reference_img_extr = feature_extractor(tfm_pattern_image)
+            loss = perceptual_loss(denoised_images_extr, reference_img_extr) * pattern_loss_scale
+            # Get gradient
+            cond_grad = torch.autograd.grad(loss, latents)[0]
+
+            # Modify the latents based on this gradient
+            latents = latents.detach() - cond_grad * sigma**2
+
+        # Now step with scheduler. compute the previous noisy sample x_t -> x_t-1
+        latents = scheduler.step(noise_pred, t, latents).prev_sample
+
+    return latents
+
+def generate_image(prompt, pattern_image, style, num_inf_steps):
+    tfm_pattern_image = transform_pattern_image(pattern_image)  # Transform the pattern image to be fed to feature extractor
+    tfm_pattern_image = tfm_pattern_image.to(torch_device)
+    if style == "no-style":
+        concept_seed = 40
+        main_prompt = str(prompt)
+    else:
+        concept_key, concept_seed = load_required_style(style)
+        main_prompt = f"{str(prompt)} in the style of {concept_key}"
+    latents = generate_with_embs_pattern_loss(main_prompt, concept_seed, tfm_pattern_image, num_inf_steps)
+    generated_image = latents_to_pil(latents)[0]
+    return generated_image
+
+def gradio_fn(prompt, pattern_image, style, num_inf_steps):
+    output_pil_image = generate_image(prompt, pattern_image, style, num_inf_steps)
+    return output_pil_image
+
+demo = gr.Interface(fn=gradio_fn,
+                    inputs=[gr.Textbox(info="Example prompt: 'A toddler gazing at sky'"),
+                            gr.Image(type="pil", height=224, width=224, info='Sample image to emulate the pattern'),
+                            gr.Radio(["anime","birb","depthmap","gta5","midjourney","beetlejuice","no-style"], label="Style",
+                                     info="Choose the style in which image to be made"),
+                            gr.Slider(50, 100, value=50, label="Num_inference_steps", info="Choose between 50 & 100")],
+                    outputs=gr.Image(height=320, width=320),
+                    title="ImageAlchemy using Stable Diffusion",
+                    description="- Stable Diffusion model that generates single image to fit \
+                                  (a) given text prompt (b) given reference image and (c) selected style.")
+
+demo.launch(share=True)

requirements.txt

@@ -0,0 +1,7 @@
+transformers==4.34.1
+diffusers==0.21.4
+ftfy==6.1.1
+accelerate==0.23.0
+scipy
+torch==2.1.0
+torchvision==0.16.0