Update app.py

app.py

@@ -565,16 +565,25 @@ def unpack_state(state: dict) -> Tuple[SparseTensor, SparseTensor, int]:
     return shape_slat, tex_slat, state['res']
 
 
-@spaces.GPU(duration=120)
-def image_to_3d(
+@spaces.GPU(duration=180)
+def generate_and_extract(
     image: Image.Image,
     req: gr.Request,
     progress=gr.Progress(track_tqdm=True),
-) -> str:
+) -> Tuple[str, str, str]:
+    """
+    Combined function: Generate 3D from image AND extract GLB in one GPU session.
+    This avoids issues with chaining multiple @spaces.GPU functions.
+    """
+    user_dir = os.path.join(TMP_DIR, str(req.session_hash))
+    os.makedirs(user_dir, exist_ok=True)
+    
     # Hardcoded values
     seed = np.random.randint(0, MAX_SEED)
-    resolution = "1024"
+    decimation_target = 300000
+    texture_size = 4096
     
+    # === STAGE 1: Generate 3D ===
     outputs, latents = pipeline.run(
         image,
         seed=seed,
@@ -602,11 +611,11 @@ def image_to_3d(
     )
     mesh = outputs[0]
     mesh.simplify(16777216)
+    
+    # Render preview images
     images = render_utils.render_snapshot(mesh, resolution=1024, r=2, fov=36, nviews=STEPS, envmap=envmap)
-    state = pack_state(latents)
-    torch.cuda.empty_cache()
     
-    # Build HTML
+    # Build preview HTML
     images_html = ""
     for m_idx, mode in enumerate(MODES):
         for s_idx in range(STEPS):
@@ -621,7 +630,7 @@ def image_to_3d(
         active_class = "active" if idx == DEFAULT_MODE else ""
         btns_html += f'<img src="{mode["icon_base64"]}" class="mode-btn {active_class}" onclick="selectMode({idx})" title="{mode["name"]}">'
     
-    full_html = f"""
+    preview_html = f"""
     <div class="previewer-container">
         <div class="display-row">{images_html}</div>
         <div class="mode-row">{btns_html}</div>
@@ -631,23 +640,11 @@ def image_to_3d(
     </div>
     """
     
-    return state, full_html
-
-
-@spaces.GPU(duration=120)
-def extract_glb(
-    state: dict,
-    req: gr.Request,
-    progress=gr.Progress(track_tqdm=True),
-) -> Tuple[str, str]:
-    # Hardcoded values
-    decimation_target = 300000
-    texture_size = 4096
-    
-    user_dir = os.path.join(TMP_DIR, str(req.session_hash))
-    shape_slat, tex_slat, res = unpack_state(state)
+    # === STAGE 2: Extract GLB ===
+    shape_slat, tex_slat, res = latents
     mesh = pipeline.decode_latent(shape_slat, tex_slat, res)[0]
     mesh.simplify(16777216)
+    
     glb = o_voxel.postprocess.to_glb(
         vertices=mesh.vertices,
         faces=mesh.faces,
@@ -663,13 +660,16 @@ def extract_glb(
         remesh_project=0,
         use_tqdm=True,
     )
+    
     now = datetime.now()
     timestamp = now.strftime("%Y-%m-%dT%H%M%S") + f".{now.microsecond // 1000:03d}"
-    os.makedirs(user_dir, exist_ok=True)
     glb_path = os.path.join(user_dir, f'sample_{timestamp}.glb')
     glb.export(glb_path, extension_webp=True)
+    
     torch.cuda.empty_cache()
-    return glb_path, glb_path
+    
+    # Return: preview_html, glb_path (for viewer), glb_path (for download)
+    return preview_html, glb_path, glb_path
 
 
 # ═══════════════════════════════════════════════════════════════
@@ -720,9 +720,7 @@ with gr.Blocks(theme=gr.themes.Base(primary_hue="indigo"), delete_cache=(600, 60
                     download_btn = gr.DownloadButton("Download GLB", elem_classes=["primary-btn"], size="lg")
     
     # Footer
-    gr.HTML('<div class="footer-note">Generation includes automatic GLB extraction. This may take 60+ seconds total.</div>')
-                    
-    output_buf = gr.State()
+    gr.HTML('<div class="footer-note">Generation includes automatic GLB extraction. This may take 90+ seconds total.</div>')
 
     # Event Handlers
     demo.load(start_session)
@@ -734,19 +732,13 @@ with gr.Blocks(theme=gr.themes.Base(primary_hue="indigo"), delete_cache=(600, 60
         outputs=[image_prompt],
     )
 
-    # Generate 3D now automatically chains to Extract GLB
+    # Single GPU call: Generate 3D + Extract GLB
     generate_btn.click(
-        lambda: gr.Walkthrough(selected=0), outputs=walkthrough
-    ).then(
-        image_to_3d,
+        generate_and_extract,
         inputs=[image_prompt],
-        outputs=[output_buf, preview_output],
+        outputs=[preview_output, glb_output, download_btn],
     ).then(
         lambda: gr.Walkthrough(selected=1), outputs=walkthrough
-    ).then(
-        extract_glb,
-        inputs=[output_buf],
-        outputs=[glb_output, download_btn],
     )