Upload app.py

app.py

@@ -202,7 +202,6 @@ def blue_loss(images):
     
     return -variance
 
-import torch
 
 def ymca_loss(images, weights=(1.0, 1.0, 1.0, 1.0)):
     """
@@ -263,6 +262,82 @@ def ymca_loss(images, weights=(1.0, 1.0, 1.0, 1.0)):
     return loss
 
 
+
+def rgb_to_cmyk(images):
+    """
+    Converts an RGB image tensor to CMYK.
+
+    Parameters:
+    images (torch.Tensor): A batch of images in RGB format. Expected shape is (N, 3, H, W).
+
+    Returns:
+    torch.Tensor: A tensor containing the CMYK channels.
+    """
+    R = images[:, 0, :, :]
+    G = images[:, 1, :, :]
+    B = images[:, 2, :, :]
+
+    # Convert RGB to CMY
+    C = 1 - R
+    M = 1 - G
+    Y = 1 - B
+
+    # Convert CMY to CMYK
+    K = torch.min(torch.min(C, M), Y)
+    C = (C - K) / (1 - K + 1e-8)
+    M = (M - K) / (1 - K + 1e-8)
+    Y = (Y - K) / (1 - K + 1e-8)
+
+    CMYK = torch.stack([C, M, Y, K], dim=1)
+    return CMYK
+
+def cymk_loss(images, weights=(1.0, 1.0, 1.0, 1.0)):
+    """
+    Computes the CYMK loss for a batch of images.
+    
+    The CYMK loss is a custom loss function combining the variance of the Cyan channel,
+    the mean value of the Yellow channel, the variance of the Magenta channel, and the 
+    absolute sum of the Black channel.
+
+    Parameters:
+    images (torch.Tensor): A batch of images. Expected shape is (N, 3, H, W) for RGB input.
+    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 CYMK loss, combining the specified components.
+    """
+    # 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))
+    
+    # Convert RGB to CMYK
+    cmyk_images = rgb_to_cmyk(images)
+    
+    # Extract CMYK channels
+    C = cmyk_images[:, 0, :, :]
+    M = cmyk_images[:, 1, :, :]
+    Y = cmyk_images[:, 2, :, :]
+    K = cmyk_images[:, 3, :, :]
+    
+    # Compute the variance of the C channel
+    variance_C = torch.var(C)
+    
+    # Compute the mean of the Y channel
+    mean_Y = torch.mean(Y)
+    
+    # Compute the variance of the M channel
+    variance_M = torch.var(M)
+    
+    # Compute the absolute sum of the K channel
+    abs_sum_K = torch.sum(torch.abs(K))
+    
+    # Combine the components with the given weights
+    loss = (weights[0] * variance_C) + (weights[1] * mean_Y) + (weights[2] * variance_M) + (weights[3] * abs_sum_K)
+    
+    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.
@@ -301,7 +376,7 @@ def blue_loss_variant(images, use_mean=False, alpha=1.0):
     
     return loss
 
-def generate_with_prompt_style_guidance(prompt, style, seed,num_inference_steps,guidance_scale):
+def generate_with_prompt_style_guidance(prompt, style, seed,num_inference_steps,guidance_scale,loss_function):
 
     prompt = prompt + ' in style of s'
     
@@ -386,7 +461,19 @@ 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 = ymca_loss(denoised_images) * contrast_loss_scale
+            # "contrast", "blue_original", "blue_modified","ymca_loss","cymk_loss"
+            if loss_function == "contrast":
+                loss = contrast_loss(denoised_images) * contrast_loss_scale
+            elif loss_function == "blue_original":
+                loss = blue_loss(denoised_images) * contrast_loss_scale
+            elif loss_function == "blue_modified":
+                loss = blue_loss_variant(denoised_images) * contrast_loss_scale
+            elif loss_function == "ymca_loss":
+                loss = ymca_loss(denoised_images) * contrast_loss_scale
+            elif loss_function == "cymk_loss":
+                loss = cymk_loss(denoised_images) * contrast_loss_scale
+            else :
+                loss = ymca_loss(denoised_images) * contrast_loss_scale
 
             # # Occasionally print it out
             # if i%10==0:
@@ -423,7 +510,7 @@ def inference(prompt, seed, style,num_inference_steps,guidance_scale,loss_functi
         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)
+        result = generate_with_prompt_style_guidance(prompt, style,seed,num_inference_steps,guidance_scale,loss_function)
         return np.array(result)
     else:
         return None
@@ -450,7 +537,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"),
+                                ),gr.Radio(["contrast", "blue_original", "blue_modified","ymca_loss","cymk_loss"], label="loss-function", info="loss-function" , value="ymca_loss"),
                               ],
                     outputs = [
                               gr.Image(label="Stable Diffusion Output"),