code update

app.py

@@ -246,7 +246,7 @@ def display_images_in_rows(images_with_titles, titles):
                 axes[r, c].imshow(image)
                 axes[r, c].axis('off')
 
-    plt.show()
+    return plt
 
 def image_generator(prompt = "dog", loss_function=None):
     
@@ -276,11 +276,10 @@ def image_generator(prompt = "dog", loss_function=None):
     titles = ["Arcane Style", "Birb Style", "Dr Strange Style", "Midjourney Style", "Oil Style"]
     
     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.append((images_without_loss[i], titles[i])) 
+        if images_with_loss != []:
+          generated_sd_images.append((images_with_loss[i], titles[i])) 
+            
     return display_images_in_rows(generated_sd_images, titles)