Update app.py

app.py

@@ -300,39 +300,35 @@ def display_images_in_rows(images_with_titles, titles):
     return fig
     # plt.show()
 
+def image_generator(prompt = "dog", loss_function=None):
+  images_without_loss = []
+  images_with_loss = []
 
-def image_generator(prompt="cat", loss_function=None):
-    images_without_loss = []
-    images_with_loss = []
+  for i in range(num_styles):
+      generated_img = generate_image_custom_style(prompt,style_num = i,random_seed = seed_values[i],custom_loss_fn = None)
+      images_without_loss.append(generated_img)
+      if loss_function:
+        generated_img = generate_image_custom_style(prompt,style_num = i,random_seed = seed_values[i],custom_loss_fn = loss_function)
+        images_with_loss.append(generated_img)
 
-    for i in range(num_styles):
-        generated_img = generate_image_custom_style(prompt, style_num=i, random_seed=seed_values[i], custom_loss_fn=None)
-        images_without_loss.append(generated_img)
+  generated_sd_images = []
+  titles = ["animal toy", "fft style", "mid journey", "oil style", "Space style"]
 
-        if loss_function:
-            if loss_function == "Exposure":
-                generated_img = generate_image_custom_style(prompt, style_num=i, random_seed=seed_values[i], custom_loss_fn=exposure_loss)
-            elif loss_function == "Color Diversity":
-                generated_img = generate_image_custom_style(prompt, style_num=i, random_seed=seed_values[i], custom_loss_fn=color_diversity_loss)
-            elif loss_function == "Sharpness":
-                generated_img = generate_image_custom_style(prompt, style_num=i, random_seed=seed_values[i], custom_loss_fn=sharpness_loss)
-            elif loss_function == "Brilliance":
-                generated_img = generate_image_custom_style(prompt, style_num=i, random_seed=seed_values[i], custom_loss_fn=brilliance_loss)
-            images_with_loss.append(generated_img)
+  for i in range(len(titles)):
+    generated_sd_images.append((images_without_loss[i], titles[i])) 
+    if images_with_loss != []:
+      generated_sd_images.append((images_with_loss[i], titles[i])) 
 
-    generated_sd_images = []
-    titles = ["animal toy", "fft style", "mid journey", "oil style", "Space style"]
+  return display_images_in_rows(generated_sd_images, titles)
 
-    for i in range(len(titles)):
-        generated_sd_images.append((images_without_loss[i], titles[i])) 
-        if images_with_loss != []:
-          generated_sd_images.append((images_with_loss[i], titles[i])) 
+# Create a wrapper function for show_misclassified_images()
+def image_generator_wrapper(prompt = "dog", loss_function=None):
+  if loss_function == "Yes":
+    loss_function = color_diversity_loss
+  else:
+    loss_function = None
 
-    return display_images_in_rows(generated_sd_images, titles)
-
-# Create a wrapper function for image_generator()
-def image_generator_wrapper(prompt="cat", loss_function=None):
-    return image_generator(prompt, loss_function)
+  return image_generator(prompt, loss_function)
 
 icon_html = '<i class="fas fa-chart-bar"></i>'
 title = f"""
@@ -345,21 +341,12 @@ description = f"""
     {icon_html}
     <span style="margin-left: 10px;">
         <p><strong>Embedding New Styles Into Stable Diffusion</strong></p>
-        <p>Following are some Losses tried</p>
-        <ul>
-            <li>exposure : It helps control the overall exposure of generated images. It ensures that the contrast of the generated images align with the desired aesthetic, preventing overexposure or underexposure</li>
-            <li>Brilliance:  Brilliance loss is a loss function that emphasizes the brilliance or luminance of specific image components, such as highlights. It can be used to highlight or enhance certain aspects of the generated artwork, adding a touch of brilliance or radiance to the final image.</li>
-            <li>color diversity: Color diversity loss encourages the model to produce images with a wider range of colors and hues. It helps create visually diverse and vibrant artworks by minimizing color repetition and promoting a rich color palette in the generated images</li>
-            <li>sharpness: Sharpness loss is used to enhance the level of detail and clarity in generated images. It encourages the model to produce crisp and well-defined visual elements, leading to sharper and more realistic results.</li>
-        </ul>
     </span>
 </div>
 """
 
 demo = gr.Interface(image_generator_wrapper,
-                    inputs=[gr.Textbox(label="Enter prompt for generating Image", type="text", value="A ballerina cat dancing in space"),
-                            gr.Radio(["None", "Exposure", "Color Diversity", "Sharpness", "Brilliance"], value="None", label="Select Loss")],
-                    outputs=gr.Plot(label="Generated Images"),
-                    title=title,
-                    description=description)
+                    inputs=[gr.Textbox(label="Enter prompt for generation", type="text", value="A ballerina cat dancing in space"),
+                            gr.Radio(["Yes", "No"], value="No"  , label="Apply color diversity Loss")],
+                    outputs=gr.Plot(label="Generated Images"), title = title, description=description)
 demo.launch()