Update app.py

app.py

@@ -42,174 +42,82 @@ styles_mapping = {
 
 # Define seeds for all the styles
 seed_list = [11, 56, 110, 65, 5, 29, 47]
-
-# Loss Function based on Edge Detection
+# Optimized loss computation functions
 def edge_detection(image):
     channels = image.shape[1]
-
-    # 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):
-
+    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)
+
+@torch.jit.script
+def compute_loss(original_image, loss_type: str):
     if loss_type == 'blue':
-        # blue loss
-        # [:,2] -> all images in batch, only the blue channel
-        error = torch.abs(original_image[:,2] - 0.9).mean()
+        return torch.abs(original_image[:,2] - 0.9).mean()
     elif loss_type == 'edge':
-        # edge loss
-        error = edge_loss(original_image)        
+        ed_value = edge_detection(original_image)
+        return F.mse_loss(ed_value, (ed_value > 0.5).float())
     elif loss_type == 'contrast':
-        # RGB to Gray loss
-        transformed_image = T.functional.adjust_contrast(original_image, contrast_factor = 2)
-        error = torch.abs(transformed_image - original_image).mean()
+        transformed_image = T.functional.adjust_contrast(original_image, contrast_factor=2)
+        return torch.abs(transformed_image - original_image).mean()
     elif loss_type == 'brightness':
-        # brightnesss loss
-        transformed_image = T.functional.adjust_brightness(original_image, brightness_factor = 2)
-        error = torch.abs(transformed_image - original_image).mean()
+        transformed_image = T.functional.adjust_brightness(original_image, brightness_factor=2)
+        return torch.abs(transformed_image - original_image).mean()
     elif loss_type == 'sharpness':
-        # sharpness loss
-        transformed_image = T.functional.adjust_sharpness(original_image, sharpness_factor = 2)
-        error = torch.abs(transformed_image - original_image).mean()
+        transformed_image = T.functional.adjust_sharpness(original_image, sharpness_factor=2)
+        return torch.abs(transformed_image - original_image).mean()
     elif loss_type == 'saturation':
-        # saturation loss
-        transformed_image = T.functional.adjust_saturation(original_image, saturation_factor = 10)
-        error = torch.abs(transformed_image - original_image).mean()
+        transformed_image = T.functional.adjust_saturation(original_image, saturation_factor=10)
+        return torch.abs(transformed_image - original_image).mean()
     else:
-        print("error. Loss not defined")
-
-    return error
-
-
-
-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
-
-# Existing functions (latents_to_pil, show_image, generate_image)
-# ... (Copy all the existing functions here)
-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])
-
+        return torch.tensor(0.0, device=original_image.device)
 
+# Optimized generate_image function
+@torch.cuda.amp.autocast()
 def generate_image(seed, prompt, loss_type, loss_flag=False):
+    generator = torch.manual_seed(seed)
+    batch_size = 1
 
-    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]
+    text_embeddings = sd_pipeline._encode_prompt(prompt, sd_pipeline.device, 1, True)
 
-    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(torch.float32)
-
+        (batch_size, sd_pipeline.unet.config.in_channels, height // 8, width // 8),
+        generator=generator,
+    ).to(sd_pipeline.device)
 
-    latents = latents.to(torch_device)
-    latents = latents * scheduler.init_noise_sigma
+    latents = latents * sd_pipeline.scheduler.init_noise_sigma
 
-    for i, t in tqdm(enumerate(scheduler.timesteps), total = len(scheduler.timesteps)):
+    sd_pipeline.scheduler.set_timesteps(num_inference_steps)
 
+    for i, t in enumerate(tqdm(sd_pipeline.scheduler.timesteps)):
         latent_model_input = torch.cat([latents] * 2)
-        sigma = scheduler.sigmas[i]
-        latent_model_input = scheduler.scale_model_input(latent_model_input, t)
+        latent_model_input = sd_pipeline.scheduler.scale_model_input(latent_model_input, t)
 
         with torch.no_grad():
-            noise_pred = sd_pipeline.unet(latent_model_input.to(torch.float32), t, encoder_hidden_states=text_embeddings)["sample"]
+            noise_pred = sd_pipeline.unet(latent_model_input, 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_()
-            # 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)
+            latents_x0 = sd_pipeline.scheduler.step(noise_pred, t, latents).pred_original_sample
+            with torch.no_grad():
+                denoised_images = sd_pipeline.vae.decode((1 / 0.18215) * latents_x0).sample / 2 + 0.5
 
             loss = compute_loss(denoised_images, loss_type) * loss_scale
-            #loss = loss.to(torch.float16)
-            print(f"{i} loss {loss}")
+            print(f"Step {i}, Loss: {loss.item():.4f}")
 
             cond_grad = torch.autograd.grad(loss, latents)[0]
-            latents = latents.detach() - cond_grad * sigma**2
+            latents = latents.detach() - cond_grad * sd_pipeline.scheduler.sigmas[i] ** 2
 
-        latents = scheduler.step(noise_pred,t, latents).prev_sample
+        latents = sd_pipeline.scheduler.step(noise_pred, t, latents).prev_sample
 
     return latents
 
+
 # Gradio interface function
 def generate_images(prompt, style, guidance_type):
     images = show_image(prompt, style, guidance_type)