code update

app.py

@@ -249,7 +249,12 @@ def display_images_in_rows(images_with_titles, titles):
     plt.show()
 
 def image_generator(prompt = "dog", loss_function=None):
-
+    
+    if loss_function == "Yes":
+      loss_function = vibrance_loss
+    else:
+      loss_function = None
+        
   images_without_loss = []
   images_with_loss = []
 
@@ -290,7 +295,7 @@ def image_generator_wrapper(prompt = "dog", loss_function=None):
 
 description = "Generate an image with a prompt and apply loss if you wish to"
 
-demo = gr.Interface(image_generator_wrapper,
+demo = gr.Interface(image_generator,
                     inputs=[gr.Textbox(label="Enter prompt for generating", type="text", value="dog sitting on a bench"),
                             gr.Radio(["Yes", "No"], value="No"  , label="Apply vibrance loss")],
                     outputs=gr.Plot(label="Outputs"), title = "Stable Diffusion", description=description)