Update app.py

app.py

@@ -73,6 +73,9 @@ sdf_model.eval()
 set_state('Loading SDF model...')
 sdf_model.load_state_dict(load_checkpoint(sdf_name, device))
 
+stable_diffusion = gr.Blocks.load(name="spaces/runwayml/stable-diffusion-v1-5")
+
+
 set_state('')
 
 def get_sampler(model_name, txt2obj, guidance_scale):
@@ -95,14 +98,12 @@ def get_sampler(model_name, txt2obj, guidance_scale):
         )
 
 def generate_txt2img(prompt):
-    
-    stable_diffusion = gr.Blocks.load(name="spaces/runwayml/stable-diffusion-v1-5")
+
     prompt = f"“a 3d rendering of {prompt}, chair imitating an avocado, full view, white background"
     gallery_dir = stable_diffusion(prompt, fn_index=2)
-    imgs = [os.path.join(gallery_dir, img) for img in os.listdir(gallery_dir)]
-
-    return imgs[1], gr.update(visible=True)
+    imgs = [os.path.join(gallery_dir, img) for img in os.listdir(gallery_dir) if os.path.splitext(img)[1] == '.jpg']
 
+    return imgs[0], gr.update(visible=True)
 
 def generate_3D(input, model_name='base40M', guidance_scale=3.0, grid_size=32):
 
@@ -200,7 +201,7 @@ with gr.Blocks() as app:
                 btn_generate_img2obj = gr.Button(value="Generate")
 
             with gr.Tab("Text to Image to 3D"):
-                gr.Markdown("Generate an image with Stable Diffusion, then convert to 3D. Just enter the object you want to generate.")
+                gr.Markdown("Generate an image with Stable Diffusion, then convert it to 3D. Just enter the object you want to generate.")
                 prompt_sd = gr.Textbox(label="Prompt", placeholder="a 3d rendering of [your prompt], full view, white background")
                 btn_generate_txt2sd = gr.Button(value="Generate")
                 img_sd = gr.Image(label="Image")