code update

app.py

@@ -1,5 +1,5 @@
 from base64 import b64encode
-
+import gradio as gr
 import numpy as np
 import torch
 from diffusers import AutoencoderKL, LMSDiscreteScheduler, UNet2DConditionModel
@@ -247,7 +247,7 @@ def image_generator(prompt = "dog", loss_function=None):
   seed_values = [8,16,50,80,128]
   height = 512                        # default height of Stable Diffusion
   width = 512                         # default width of Stable Diffusion
-  num_inference_steps = 10            # Number of denoising steps
+  num_inference_steps = 1           # Number of denoising steps
   guidance_scale = 7.5                # Scale for classifier-free guidance
   num_styles = len(style_files)
 
@@ -269,3 +269,21 @@ def image_generator(prompt = "dog", loss_function=None):
 
   return display_images_in_rows(generated_sd_images, titles)
 
+
+# Create a wrapper function for show_misclassified_images()
+def image_generator_wrapper(prompt = "dog", loss_function=None):
+    prompt = string(prompt)
+    if loss_function == "Yes":
+      loss_function = vibrance_loss
+    else:
+      loss_function = None
+
+    return image_generator(prompt, loss_function)
+
+description = "Generate an image with a prompt and apply loss if you wish to"
+
+demo = gr.Interface(image_generator_wrapper,
+                    inputs=[gr.Textbox(label="Enter prompt for generating", type="str", default="dog sitting on a bench"),
+                            gr.Radio(["Yes", "No"], value="No"  , label="Apply vibrance loss")],
+                    outputs=gr.Plot(), title = "Stable Diffusion", description=description)
+demo.launch()