Upload app.py

app.py

@@ -388,7 +388,7 @@ def generate_with_prompt_style_guidance(prompt, style, seed,num_inference_steps,
     guidance_scale = guidance_scale #               # Scale for classifier-free guidance
     generator = torch.manual_seed(seed)   # Seed generator to create the inital latent noise
     batch_size = 1
-    contrast_loss_scale = 200 #
+    
 
     # Prep text
     text_input = tokenizer([prompt], padding="max_length", max_length=tokenizer.model_max_length, truncation=True, return_tensors="pt")
@@ -463,17 +463,22 @@ def generate_with_prompt_style_guidance(prompt, style, seed,num_inference_steps,
             # Calculate loss
             # "contrast", "blue_original", "blue_modified","ymca_loss","cymk_loss"
             if loss_function == "contrast":
-                loss = contrast_loss(denoised_images) * contrast_loss_scale
+                loss_scale = 200 #
+                loss = contrast_loss(denoised_images) * loss_scale
             elif loss_function == "blue_original":
-                loss = blue_loss(denoised_images) * contrast_loss_scale
+                loss_scale = 200 #
+                loss = blue_loss(denoised_images) * 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
+                loss_scale = 200 #
+                loss = blue_loss_variant(denoised_images) * loss_scale
+            elif loss_function == "ymca":
+                loss_scale = 200 #
+                loss = ymca_loss(denoised_images) * loss_scale
+            elif loss_function == "cmyk":
+                loss_scale = 10 #
+                loss = cymk_loss(denoised_images) * loss_scale
             else :
-                loss = ymca_loss(denoised_images) * contrast_loss_scale
+                loss = ymca_loss(denoised_images) * loss_scale
 
             # # Occasionally print it out
             # if i%10==0:
@@ -537,7 +542,7 @@ demo = gr.Interface(inference,
                                 step=8,
                                 label="Select Guidance Scale",
                                 interactive=True,
-                                ),gr.Radio(["contrast", "blue_original", "blue_modified","ymca_loss","cymk_loss"], label="loss-function", info="loss-function" , value="ymca_loss"),
+                                ),gr.Radio(["contrast", "blue_original", "blue_modified","ymca","cmyk"], label="loss-function", info="loss-function" , value="ymca_loss"),
                               ],
                     outputs = [
                               gr.Image(label="Stable Diffusion Output"),