Update app.py

app.py

@@ -324,11 +324,11 @@ def image_generator(prompt="cat", loss_function=None):
     titles = ["animal toy", "fft style", "mid journey", "oil style", "Space 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 generated_sd_images
+    return display_images_in_rows(generated_sd_images, titles)
 
 # Create a wrapper function for image_generator()
 def image_generator_wrapper(prompt="dog", selected_loss="None"):