Update utils.py

utils.py

@@ -1,4 +1,4 @@
-'''from base64 import b64encode
+from base64 import b64encode
 
 import numpy
 import torch
@@ -118,9 +118,9 @@ def get_output_embeds(input_embeddings):
     # And now they're ready!
     return output
 
-def blue_loss(images):
+def blue_loss(images, contrast_perc=80):
     # How far the pixels are from +80% contrast:
-    contrast = 230 # it ranges from -255 to +255
+    contrast = 255*contrast_perc // 100 # it ranges from -255 to +255
     contrast_scale_factor = (259 * (contrast + 255)) / (255 * (259 - contrast))
     cimgs = (contrast_scale_factor * (images - 0.5) + 0.5 )
     cimgs = torch.where(cimgs > 1.0, 1.0, cimgs)
@@ -131,7 +131,7 @@ def blue_loss(images):
     return error
 
 # Generating an image with these modified embeddings
-def generate_with_embs(text_input, text_embeddings, output=None, generator=None, additional_guidance=False):
+def generate_with_embs(text_input, text_embeddings, output=None, generator=None, contrast_loss=False, contrast_perc=0):
     height = 512  # default height of Stable Diffusion
     width = 512  # default width of Stable Diffusion
     num_inference_steps = 30  # Number of denoising steps
@@ -177,8 +177,8 @@ def generate_with_embs(text_input, text_embeddings, output=None, generator=None,
         noise_pred = noise_pred_uncond + guidance_scale * (noise_pred_text - noise_pred_uncond)
 
         #### ADDITIONAL GUIDANCE ###
-        if additional_guidance:
-            blue_loss_scale = 80
+        if contrast_loss:
+            blue_loss_scale = 70
             if i % 5 == 0:
                 # Requires grad on the latents
                 latents = latents.detach().requires_grad_()
@@ -191,7 +191,7 @@ def generate_with_embs(text_input, text_embeddings, output=None, generator=None,
                 denoised_images = vae.decode((1 / 0.18215) * latents_x0).sample / 2 + 0.5  # range (0, 1)
 
                 # Calculate loss
-                loss = blue_loss(denoised_images) * blue_loss_scale
+                loss = blue_loss(denoised_images, contrast_perc=contrast_perc) * blue_loss_scale
 
                 # Occasionally print it out
                 if i % 10 == 0:
@@ -218,4 +218,3 @@ token_emb_layer = text_encoder.text_model.embeddings.token_embedding
 
 pos_emb_layer = text_encoder.text_model.embeddings.position_embedding
 #pos_emb_layer
-'''