Update app.py

app.py

@@ -175,10 +175,40 @@ def generate_with_prompt_style(prompt, style, seed = 42):
 
 import torch
 
-def contrast_loss(images):
-    variance = torch.var(images)
+# def contrast_loss(images):
+#     variance = torch.var(images)
+#     return -variance
+
+
+import torch
+
+def blue_loss(images):
+    """
+    Computes the blue loss for a batch of images.
+    
+    The blue loss is defined as the negative variance of the blue channel's pixel values.
+
+    Parameters:
+    images (torch.Tensor): A batch of images. Expected shape is (N, C, H, W) where
+                           N is the batch size, C is the number of channels (3 for RGB), 
+                           H is the height, and W is the width.
+
+    Returns:
+    torch.Tensor: The blue loss, which is the negative variance of the blue channel's pixel values.
+    """
+    # Ensure the input tensor has the correct shape
+    if images.shape[1] != 3:
+        raise ValueError("Expected images with 3 channels (RGB), but got shape {}".format(images.shape))
+    
+    # Extract the blue channel (assuming the channels are in RGB order)
+    blue_channel = images[:, 2, :, :]
+    
+    # Calculate the variance of the blue channel
+    variance = torch.var(blue_channel)
+    
     return -variance
 
+
 def generate_with_prompt_style_guidance(prompt, style, seed=42):
 
     prompt = prompt + ' in style of s'
@@ -264,7 +294,7 @@ def generate_with_prompt_style_guidance(prompt, style, seed=42):
             denoised_images = vae.decode((1 / 0.18215) * latents_x0).sample / 2 + 0.5 # range (0, 1)
 
             # Calculate loss
-            loss = contrast_loss(denoised_images) * contrast_loss_scale
+            loss = blue_loss(denoised_images) * contrast_loss_scale
 
             # # Occasionally print it out
             # if i%10==0: