Update app.py

app.py

@@ -1,5 +1,5 @@
 from base64 import b64encode
-
+import torch
 import numpy
 import torch
 from diffusers import AutoencoderKL, LMSDiscreteScheduler, UNet2DConditionModel
@@ -16,8 +16,7 @@ from transformers import CLIPTextModel, CLIPTokenizer, logging
 import os
 import numpy as np  
 
-torch.manual_seed(1)
-# if not (Path.home()/'.cache/huggingface'/'token').exists(): notebook_login()
+torch.manual_seed(24041975)
 
 # Supress some unnecessary warnings when loading the CLIPTextModel
 logging.set_verbosity_error()
@@ -172,15 +171,47 @@ def generate_with_prompt_style(prompt, style, seed = 42):
     # And generate an image with this:
     return generate_with_embs(modified_output_embeddings, text_input, seed)
 
+def contrast_loss(images):
+    variance = torch.var(images)
+    return -variance
 
-import torch
-
-# def contrast_loss(images):
-#     variance = torch.var(images)
-#     return -variance
+def blue_loss_variant(images, use_mean=False, alpha=1.0):
+    """
+    Computes the blue loss for a batch of images with an optional mean component.
+    
+    The blue loss is defined as the negative variance of the blue channel's pixel values.
+    Optionally, it can also include the mean value of the blue channel.
 
+    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.
+    use_mean (bool): If True, includes the mean of the blue channel in the loss calculation.
+    alpha (float): Weighting factor for the mean component when use_mean is True.
 
-import torch
+    Returns:
+    torch.Tensor: The blue loss, which is the negative variance of the blue channel's pixel values,
+                  optionally combined with the mean value of the blue channel.
+    """
+    # 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)
+    
+    if use_mean:
+        # Calculate the mean of the blue channel
+        mean = torch.mean(blue_channel)
+        # Combine variance and mean into the loss
+        loss = -variance + alpha * mean
+    else:
+        loss = -variance
+    
+    return loss
 
 def blue_loss(images):
     """
@@ -294,7 +325,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 = blue_loss(denoised_images) * contrast_loss_scale
+            loss = blue_loss_variant(denoised_images) * contrast_loss_scale
 
             # # Occasionally print it out
             # if i%10==0: