Update app.py

app.py

@@ -1,3 +1,4 @@
+import os
 from PIL import Image
 import torch
 
@@ -94,11 +95,13 @@ def get_sampler(model_name, txt2obj, guidance_scale):
         )
 
 def generate_txt2img(prompt):
-    prompt = f"“a 3d rendering of {prompt}, chair imitating an avocado, full view, white background"
+    
     stable_diffusion = gr.Blocks.load(name="spaces/runwayml/stable-diffusion-v1-5")
-    image = stable_diffusion(prompt, fn_index=2)[0]
+    prompt = f"“a 3d rendering of {prompt}, chair imitating an avocado, full view, white background"
+    gallery_dir = stable_diffusion(prompt, fn_index=2)[0]
+    imgs = [os.path.join(gallery_dir, img) for img in os.listdir(gallery_dir)]
 
-    return image, gr.update(visible=True)
+    return imgs[1], gr.update(visible=True)
 
 
 def generate_3D(input, model_name='base40M', guidance_scale=3.0, grid_size=32):