Update app.py

app.py

@@ -89,16 +89,16 @@ def get_style_embeddings(style_file):
 
 import torch
 
-def vibrance_loss(image):
-    # Calculate the standard deviation of color channels
-    std_dev = torch.std(image, dim=(2, 3))  # Compute standard deviation over height and width
+def sharpness_loss(image):
+    threshold = torch.tensor(0.5) 
+    std_dev = torch.std(image, dim=[2, 3])
+
     # Calculate the mean standard deviation across the batch
-    mean_std_dev = torch.mean(std_dev)
-    # You can adjust a scale factor to control the strength of vibrance regularization
-    scale_factor = 100.0
-    # Calculate the vibrance loss
-    loss = -scale_factor * mean_std_dev
-    return loss
+    mean_std_dev = torch.mean(std_dev, dim=0)
+
+    # Check if the mean standard deviation is below the threshold
+    loss = torch.mean(mean_std_dev - threshold)*50
+    return -loss
 
 
 from torchvision.transforms import ToTensor
@@ -183,6 +183,7 @@ def generate_with_embs(text_embeddings, max_length, random_seed, loss_fn = None)
         if loss_fn is not None:
             if i%2 == 0:
                 latents, custom_loss = additional_guidance(latents, scheduler, noise_pred, t, sigma, loss_fn)
+                print('loss: ',custom_loss.item())
 
         # compute the previous noisy sample x_t -> x_t-1
         latents = scheduler.step(noise_pred, t, latents).prev_sample
@@ -255,7 +256,7 @@ def image_generator(prompt = "dog", loss_function=None):
   images_without_loss = []
   images_with_loss = []
   if loss_function == "Yes":
-    loss_function = vibrance_loss
+    loss_function = sharpness_loss
   else:
     loss_function = None
 
@@ -277,7 +278,7 @@ def image_generator(prompt = "dog", loss_function=None):
 
   return display_images_in_rows(generated_sd_images, titles)
 
-description = "Generate an image with a prompt"
+description = "Generate an image with a prompt (takes time)"
 
 demo = gr.Interface(image_generator,
                     inputs=[gr.Textbox(label="prompt", type="text", value="cat fight"),