Update app.py

app.py

@@ -27,13 +27,16 @@ def get_images(username):
     for prompt in prompt_list_prepend:
         image_dirs.append(stable_diffusion(prompt, fn_index=2))
 
-    # do not take captions, and only take first output 
     sd_output = []
+    
     for image_dir in image_dirs:
+        found_first_image = False
         for filename in os.listdir(image_dir):
-            if filename.endswith('.jpg'):
+            if filename.endswith('.json'):
+                continue
+            elif not found_first_image:
                 sd_output.append(os.path.join(image_dir, filename))
-                break
+                found_first_image = True
         
     return sd_output, prompt_string