Update app.py

app.py

@@ -1,4 +1,3 @@
-import gradio as gr
 import PIL
 import torch
 import numpy as np
@@ -13,7 +12,7 @@ torch_device        = "cuda" if torch.cuda.is_available() else "mps" if torch.ba
 height, width       = 512, 512
 guidance_scale      = 8
 loss_scale          = 200
-num_inference_steps = 10
+num_inference_steps = 50
 
 
 model_path = "CompVis/stable-diffusion-v1-4"
@@ -42,109 +41,168 @@ styles_mapping = {
 
 # Define seeds for all the styles
 seed_list = [11, 56, 110, 65, 5, 29, 47]
-# Optimized loss computation functions
+
+# Loss Function based on Edge Detection
 def edge_detection(image):
     channels = image.shape[1]
-    kernels = torch.tensor([[-1, -2, -1], [0, 0, 0], [1, 2, 1],
-                            [-1, 0, 1], [-2, 0, 2], [-1, 0, 1]], device=image.device).float()
-    kernels = kernels.view(2, 1, 3, 3).repeat(channels, 1, 1, 1)
-    padded_image = F.pad(image, (1, 1, 1, 1), mode='replicate')
-    edge = F.conv2d(padded_image, kernels, groups=channels)
-    return torch.sqrt(edge[:, :channels]**2 + edge[:, channels:]**2)
-
-def compute_loss(original_image, loss_type: str):
+
+    # Define the kernels for Edge Detection
+    ed_x = torch.tensor([[-1, 0, 1], [-2, 0, 2], [-1, 0, 1]], dtype=torch.float32).unsqueeze(0).unsqueeze(0)
+    ed_y = torch.tensor([[1, 2, 1], [0, 0, 0], [-1, -2, -1]], dtype=torch.float32).unsqueeze(0).unsqueeze(0)
+
+    # Replicate the Edge detection kernels for each channel
+    ed_x = ed_x.repeat(channels, 1, 1, 1).to(image.device)
+    ed_y = ed_y.repeat(channels, 1, 1, 1).to(image.device)
+
+    # ed_x = ed_x.to(torch.float16)
+    # ed_y = ed_y.to(torch.float16)
+
+    # Convolve the image with the Edge detection kernels
+    conv_ed_x = F.conv2d(image, ed_x, padding=1, groups=channels)
+    conv_ed_y = F.conv2d(image, ed_y, padding=1, groups=channels)
+
+    # Combine the x and y gradients after convolution
+    ed_value = torch.sqrt(conv_ed_x**2 + conv_ed_y**2)
+
+    return ed_value
+
+def edge_loss(image):
+    ed_value = edge_detection(image)
+    ed_capped = (ed_value > 0.5).to(torch.float32)
+    return F.mse_loss(ed_value, ed_capped)
+
+def compute_loss(original_image, loss_type):
+
     if loss_type == 'blue':
-        return torch.abs(original_image[:,2] - 0.9).mean()
+        # blue loss
+        # [:,2] -> all images in batch, only the blue channel
+        error = torch.abs(original_image[:,2] - 0.9).mean()
     elif loss_type == 'edge':
-        ed_value = edge_detection(original_image)
-        return F.mse_loss(ed_value, (ed_value > 0.5).float())
+        # edge loss
+        error = edge_loss(original_image)        
     elif loss_type == 'contrast':
-        transformed_image = TF.adjust_contrast(original_image, contrast_factor=2.0)
-        return torch.abs(transformed_image - original_image).mean()
+        # RGB to Gray loss
+        transformed_image = T.functional.adjust_contrast(original_image, contrast_factor = 2)
+        error = torch.abs(transformed_image - original_image).mean()
     elif loss_type == 'brightness':
-        transformed_image = TF.adjust_brightness(original_image, brightness_factor=2.0)
-        return torch.abs(transformed_image - original_image).mean()
+        # brightnesss loss
+        transformed_image = T.functional.adjust_brightness(original_image, brightness_factor = 2)
+        error = torch.abs(transformed_image - original_image).mean()
     elif loss_type == 'sharpness':
-        transformed_image = TF.adjust_sharpness(original_image, sharpness_factor=2.0)
-        return torch.abs(transformed_image - original_image).mean()
+        # sharpness loss
+        transformed_image = T.functional.adjust_sharpness(original_image, sharpness_factor = 2)
+        error = torch.abs(transformed_image - original_image).mean()
     elif loss_type == 'saturation':
-        transformed_image = TF.adjust_saturation(original_image, saturation_factor=10.0)
-        return torch.abs(transformed_image - original_image).mean()
+        # saturation loss
+        transformed_image = T.functional.adjust_saturation(original_image, saturation_factor = 10)
+        error = torch.abs(transformed_image - original_image).mean()
     else:
-        return torch.tensor(0.0, device=original_image.device)
+        print("error. Loss not defined")
+
+    return error
+
+
+def get_examples():
+   examples = [
+      ['A bird sitting on a tree', 'Midjourney', 'edge', 5],
+      ['Cats fighting on the road', 'Marc Allante', 'brightness', 65],
+      ['A mouse with the head of a puppy', 'Hitokomoru Style', 'contrast', 110],
+      ['A woman with a smiling face in front of an Italian Pizza', 'Hanfu Anime', 'brightness', 29],
+      ['A campfire (oil on canvas)', 'Birb Style', 'blue', 47],
+   ]
+   return(examples)
+
+
+def latents_to_pil(latents):
+    # bath of latents -> list of images
+    latents = (1 / 0.18215) * latents
+    with torch.no_grad():
+        image = sd_pipeline.vae.decode(latents).sample
+    image = (image / 2 + 0.5).clamp(0, 1) # 0 to 1
+    image = image.detach().cpu().permute(0, 2, 3, 1).numpy()
+    image = (image * 255).round().astype("uint8")
+    return Image.fromarray(image[0])
+
+
+def show_image(prompt, concept, guidance_type):
+
+  for idx, sd in enumerate(styles_mapping.keys()):
+    if(sd == concept):
+      break
+  seed =  seed_list[idx]
+  prompt = f"{prompt} in the style of {styles_mapping[sd]}"
+  styled_image_without_loss = latents_to_pil(generate_image(seed, prompt, guidance_type, loss_flag=False))
+  styled_image_with_loss = latents_to_pil(generate_image(seed, prompt, guidance_type, loss_flag=True))
+  return([styled_image_without_loss, styled_image_with_loss])
+
 
-# Optimized generate_image function
 def generate_image(seed, prompt, loss_type, loss_flag=False):
-    generator = torch.manual_seed(seed)
-    batch_size = 1
 
-    text_embeddings = sd_pipeline._encode_prompt(prompt, sd_pipeline.device, 1, True)
+    generator           = torch.manual_seed(seed)
+    batch_size          = 1
+
+    # scheduler
+    scheduler    = LMSDiscreteScheduler(beta_start = 0.00085, beta_end = 0.012, beta_schedule = "scaled_linear", num_train_timesteps = 1000)
+    scheduler.set_timesteps(num_inference_steps)
+    scheduler.timesteps = scheduler.timesteps.to(torch.float32)
+
+    # text embeddings of the prompt
+    text_input = sd_pipeline.tokenizer(prompt, padding='max_length', max_length = sd_pipeline.tokenizer.model_max_length, truncation= True, return_tensors="pt")
+    input_ids = text_input.input_ids.to(torch_device)
+
+    with torch.no_grad():
+        text_embeddings = sd_pipeline.text_encoder(text_input.input_ids.to(torch_device))[0]
+
+    max_length = text_input.input_ids.shape[-1]
+    uncond_input = sd_pipeline.tokenizer(
+          [""] * batch_size, padding="max_length", max_length= max_length, return_tensors="pt"
+    )
+
+    with torch.no_grad():
+        uncond_embeddings = sd_pipeline.text_encoder(uncond_input.input_ids.to(torch_device))[0]
 
+    text_embeddings = torch.cat([uncond_embeddings,text_embeddings]) # shape: 2,77,768
+
+    # random latent
     latents = torch.randn(
-        (batch_size, sd_pipeline.unet.config.in_channels, height // 8, width // 8),
-        generator=generator,
-    ).to(sd_pipeline.device)
+        (batch_size, sd_pipeline.unet.config.in_channels, height// 8, width //8),
+        generator = generator,
+    ) .to(torch.float32)
+
 
-    latents = latents * sd_pipeline.scheduler.init_noise_sigma
+    latents = latents.to(torch_device)
+    latents = latents * scheduler.init_noise_sigma
 
-    sd_pipeline.scheduler.set_timesteps(num_inference_steps)
+    for i, t in tqdm(enumerate(scheduler.timesteps), total = len(scheduler.timesteps)):
 
-    for i, t in enumerate(tqdm(sd_pipeline.scheduler.timesteps)):
         latent_model_input = torch.cat([latents] * 2)
-        latent_model_input = sd_pipeline.scheduler.scale_model_input(latent_model_input, t)
+        sigma = scheduler.sigmas[i]
+        latent_model_input = scheduler.scale_model_input(latent_model_input, t)
 
         with torch.no_grad():
-            noise_pred = sd_pipeline.unet(latent_model_input, t, encoder_hidden_states=text_embeddings).sample
+            noise_pred = sd_pipeline.unet(latent_model_input.to(torch.float32), t, encoder_hidden_states=text_embeddings)["sample"]
 
         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_flag and i % 5 == 0:
+        if loss_flag and i%5 == 0:
+
             latents = latents.detach().requires_grad_()
-            latents_x0 = sd_pipeline.scheduler.step(noise_pred, t, latents).prev_sample
-            with torch.no_grad():
-                denoised_images = sd_pipeline.vae.decode((1 / 0.18215) * latents_x0).sample / 2 + 0.5
+            # the following line alone does not work, it requires change to reduce step only once
+            # hence commenting it out
+            #latents_x0 = scheduler.step(noise_pred,t, latents).pred_original_sample 
+            latents_x0 = latents - sigma * noise_pred
+
+            # use vae to decode the image
+            denoised_images = sd_pipeline.vae.decode((1/ 0.18215) * latents_x0).sample / 2 + 0.5 # range(0,1)
 
             loss = compute_loss(denoised_images, loss_type) * loss_scale
-            print(f"Step {i}, Loss: {loss.item():.4f}")
+            #loss = loss.to(torch.float16)
+            print(f"{i} loss {loss}")
 
             cond_grad = torch.autograd.grad(loss, latents)[0]
-            latents = latents.detach() - cond_grad * sd_pipeline.scheduler.sigmas[i] ** 2
+            latents = latents.detach() - cond_grad * sigma**2
 
-        latents = sd_pipeline.scheduler.step(noise_pred, t, latents).prev_sample
+        latents = scheduler.step(noise_pred,t, latents).prev_sample
 
-    return latents
-
-def generate_image(prompt, style, guidance_type):
-    styled_prompt = f"{prompt} in the style of {styles_mapping[style]}"
-    seed = torch.randint(0, 1000000, (1,)).item()
-    latents = generate_image(seed, styled_prompt, guidance_type, loss_flag=True)
-    with torch.no_grad():
-        image = sd_pipeline.decode_latents(latents)
-    image = sd_pipeline.numpy_to_pil(image)[0]
-    return image
-
-def get_examples():
-    examples = [
-        ["A bird sitting on a tree", "Midjourney", "edge"],
-        ["Cats fighting on the road", "Marc Allante", "brightness"],
-        ["A mouse with the head of a puppy", "Hitokomoru Style", "contrast"],
-        ["A woman with a smiling face in front of an Italian Pizza", "Hanfu Anime", "brightness"],
-        ["A campfire (oil on canvas)", "Birb Style", "blue"],
-    ]
-    return examples
-
-iface = gr.Interface(
-    fn=generate_image,
-    inputs=[
-        gr.Textbox(label="Prompt"),
-        gr.Dropdown(list(styles_mapping.keys()), label="Style"),
-        gr.Dropdown(["blue", "edge", "contrast", "brightness", "sharpness", "saturation"], label="Guidance Type"),
-    ],
-    outputs=gr.Image(label="Generated Image"),
-    title="Stable Diffusion with Custom Styles",
-    description="Generate images using a custom Stable Diffusion model with various styles and guidance types.",
-    examples=get_examples(),
-)
-
-iface.launch()
+    return latents