Upload app.py

app.py

@@ -175,6 +175,126 @@ def contrast_loss(images):
     variance = torch.var(images)
     return -variance
 
+
+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 ymca_loss(images, weights=(1.0, 1.0, 1.0, 1.0)):
+    """
+    Computes the YMCA loss for a batch of images.
+    
+    The YMCA loss is a custom loss function combining the mean value of the Y (luminance) channel,
+    the mean value of the M (magenta) channel, the variance of the C (cyan) channel, and the 
+    absolute sum of the A (alpha) 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 (assumed 4 for RGBA), 
+                           H is the height, and W is the width.
+    weights (tuple): A tuple of four floats representing the weights for each component of the loss
+                     (default is (1.0, 1.0, 1.0, 1.0)).
+
+    Returns:
+    torch.Tensor: The YMCA loss, combining the specified components.
+    """
+    # Ensure the input tensor has the correct shape
+    if images.shape[1] != 4:
+        raise ValueError("Expected images with 4 channels (RGBA), but got shape {}".format(images.shape))
+    
+    # Extract the RGBA channels
+    R = images[:, 0, :, :]
+    G = images[:, 1, :, :]
+    B = images[:, 2, :, :]
+    A = images[:, 3, :, :]
+    
+    # Convert RGB to Y (luminance) channel
+    Y = 0.299 * R + 0.587 * G + 0.114 * B
+    
+    # Convert RGB to M (magenta) channel
+    M = 1 - G
+    
+    # Convert RGB to C (cyan) channel
+    C = 1 - R
+    
+    # Compute the mean of the Y channel
+    mean_Y = torch.mean(Y)
+    
+    # Compute the mean of the M channel
+    mean_M = torch.mean(M)
+    
+    # Compute the variance of the C channel
+    variance_C = torch.var(C)
+    
+    # Compute the absolute sum of the A channel
+    abs_sum_A = torch.sum(torch.abs(A))
+    
+    # Combine the components with the given weights
+    loss = (weights[0] * mean_Y) + (weights[1] * mean_M) - (weights[2] * variance_C) + (weights[3] * abs_sum_A)
+    
+    return loss
+
+
+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.
+
+    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 generate_with_prompt_style_guidance(prompt, style, seed,num_inference_steps,guidance_scale):
 
     prompt = prompt + ' in style of s'
@@ -260,7 +380,7 @@ def generate_with_prompt_style_guidance(prompt, style, seed,num_inference_steps,
             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 = ymca_loss(denoised_images) * contrast_loss_scale
 
             # # Occasionally print it out
             # if i%10==0:
@@ -291,9 +411,10 @@ dict_styles = {
     'Oil Painting':'styles/learned_embeds_oil.bin',  
 }
 
-def inference(prompt, seed, style,num_inference_steps,guidance_scale): 
+def inference(prompt, seed, style,num_inference_steps,guidance_scale,loss_function): 
     
     if prompt is not None and style is not None and seed is not None: 
+        print(loss_function)
         style = dict_styles[style]
         torch.manual_seed(seed)
         result = generate_with_prompt_style_guidance(prompt, style,seed,num_inference_steps,guidance_scale)
@@ -323,7 +444,7 @@ demo = gr.Interface(inference,
                                 step=8,
                                 label="Select Guidance Scale",
                                 interactive=True,
-                                )
+                                ),gr.Radio(["contrast", "blue-original", "blue-modified","ymca_loss"], label="loss-function", info="loss-function"),
                               ],
                     outputs = [
                               gr.Image(label="Stable Diffusion Output"),