Update app.py

Key changes based on the paper's implementation:

Split the processing into clear stages with memory cleanup between each:

Stage 1: Generate sparse structure
Stage 2: Generate video preview in batches
GLB generation when requested


Added batched processing for video frame generation:

Process 30 frames at a time instead of all 120 at once
Clear CUDA cache after each batch


Added explicit memory management:

torch.cuda.empty_cache() calls at key points
Explicit deletion of large temporary data
Clear video data after saving


Separated high-quality and reduced GLB generation into distinct functions
Added progress visibility with verbose=True for GLB generation to track progress

This should help prevent GPU timeouts by:

Breaking up large operations into smaller chunks
Managing memory more efficiently
Cleaning up resources between steps

app.py

@@ -35,13 +35,6 @@ def end_session(req: gr.Request):
 def preprocess_image(image: Image.Image) -> Tuple[str, Image.Image]:
     """
     Preprocess the input image.
-
-    Args:
-        image (Image.Image): The input image.
-
-    Returns:
-        str: uuid of the trial.
-        Image.Image: The preprocessed image.
     """
     processed_image = pipeline.preprocess_image(image)
     return processed_image
@@ -102,11 +95,11 @@ def image_to_3d(
     req: gr.Request,
 ) -> Tuple[dict, str]:
     """
-    Convert an image to a 3D model with memory management.
+    Convert an image to a 3D model.
     """
     user_dir = os.path.join(TMP_DIR, str(req.session_hash))
     
-    # Generate base outputs
+    # First stage: Generate sparse structure
     outputs = pipeline.run(
         image,
         seed=seed,
@@ -121,68 +114,67 @@ def image_to_3d(
             "cfg_strength": slat_guidance_strength,
         },
     )
-
-    # Clear CUDA cache after model generation
-    torch.cuda.empty_cache()
     
-    # Generate video preview in smaller batches
-    video = []
-    video_geo = []
+    # Clear CUDA cache after structure generation
+    torch.cuda.empty_cache()
+
+    # Second stage: Generate video preview in batches
+    video_frames = []
+    video_geo_frames = []
     batch_size = 30  # Process 30 frames at a time
     num_frames = 120
     
     for i in range(0, num_frames, batch_size):
         end_idx = min(i + batch_size, num_frames)
-        curr_frames = end_idx - i
-        
-        # Generate color frames
         batch_frames = render_utils.render_video(
-            outputs['gaussian'][0],
-            num_frames=curr_frames,
+            outputs['gaussian'][0], 
+            num_frames=end_idx - i,
             start_frame=i
         )['color']
-        video.extend(batch_frames)
+        video_frames.extend(batch_frames)
         
-        # Generate geometry frames
         batch_geo = render_utils.render_video(
-            outputs['mesh'][0],
-            num_frames=curr_frames,
+            outputs['mesh'][0], 
+            num_frames=end_idx - i,
             start_frame=i
         )['normal']
-        video_geo.extend(batch_geo)
+        video_geo_frames.extend(batch_geo)
         
         # Clear cache after each batch
         torch.cuda.empty_cache()
-
-    # Combine and save video
-    video = [np.concatenate([video[i], video_geo[i]], axis=1) for i in range(len(video))]
+    
+    # Combine frames and save video
+    video = [np.concatenate([video_frames[i], video_geo_frames[i]], axis=1) 
+            for i in range(len(video_frames))]
     trial_id = str(uuid.uuid4())
     video_path = os.path.join(user_dir, f"{trial_id}.mp4")
     imageio.mimsave(video_path, video, fps=15)
     
-    # Clear memory
+    # Clear video data
+    del video_frames
+    del video_geo_frames
     del video
-    del video_geo
     torch.cuda.empty_cache()
     
-    # Pack state and return
+    # Pack state
     state = pack_state(outputs['gaussian'][0], outputs['mesh'][0], trial_id)
     return state, video_path
 
 @spaces.GPU
-def export_full_quality_glb(
+def extract_high_quality_glb(
     state: dict,
     req: gr.Request,
 ) -> Tuple[str, str]:
     """
-    Export a full-quality GLB file with memory management.
+    Extract a high-quality GLB file with memory optimization.
     """
     user_dir = os.path.join(TMP_DIR, str(req.session_hash))
     gs, mesh, trial_id = unpack_state(state)
     
-    # Clear cache before starting
+    # Clear cache before GLB generation
     torch.cuda.empty_cache()
     
+    # Process mesh in original quality (no reduction)
     glb = postprocessing_utils.to_glb(
         gs,
         mesh,
@@ -192,37 +184,51 @@ def export_full_quality_glb(
         texture_size=2048,  # Maximum texture resolution
         verbose=True  # Show progress
     )
+    
     glb_path = os.path.join(user_dir, f"{trial_id}_full.glb")
     glb.export(glb_path)
     
-    # Clear cache after finishing
+    # Final cleanup
     torch.cuda.empty_cache()
     return glb_path, glb_path
 
 @spaces.GPU
-def extract_glb(
+def extract_reduced_glb(
     state: dict,
     mesh_simplify: float,
     texture_size: int,
     req: gr.Request,
 ) -> Tuple[str, str]:
     """
-    Extract a GLB file from the 3D model.
+    Extract a reduced-quality GLB file with memory optimization.
     """
     user_dir = os.path.join(TMP_DIR, str(req.session_hash))
     gs, mesh, trial_id = unpack_state(state)
-    glb = postprocessing_utils.to_glb(gs, mesh, simplify=mesh_simplify, texture_size=texture_size, verbose=False)
-    glb_path = os.path.join(user_dir, f"{trial_id}.glb")
+    
+    # Clear cache before GLB generation
+    torch.cuda.empty_cache()
+    
+    glb = postprocessing_utils.to_glb(
+        gs,
+        mesh,
+        simplify=mesh_simplify,
+        texture_size=texture_size,
+        verbose=True
+    )
+    glb_path = os.path.join(user_dir, f"{trial_id}_reduced.glb")
     glb.export(glb_path)
+    
+    # Final cleanup
+    torch.cuda.empty_cache()
     return glb_path, glb_path
 
 with gr.Blocks(delete_cache=(600, 600)) as demo:
     gr.Markdown("""
     ## Image to 3D Asset with [TRELLIS](https://trellis3d.github.io/)
-    * Upload an image and click "Generate" to create a 3D asset. If the image has alpha channel, it be used as the mask. Otherwise, we use `rembg` to remove the background.
+    * Upload an image and click "Generate" to create a 3D asset
     * After generation:
-        * Click "Download Full-Quality GLB" for maximum quality
-        * Or use GLB Extraction Settings for a reduced size version
+        * Click "Extract Full GLB" for maximum quality (no reduction)
+        * Or use GLB Extraction Settings for a reduced version
     """)
     
     with gr.Row():
@@ -242,12 +248,13 @@ with gr.Blocks(delete_cache=(600, 600)) as demo:
                     slat_sampling_steps = gr.Slider(1, 500, label="Sampling Steps", value=12, step=1)
 
             generate_btn = gr.Button("Generate")
+            extract_full_btn = gr.Button("Extract Full GLB", interactive=False)
             
             with gr.Accordion(label="GLB Extraction Settings", open=False):
                 mesh_simplify = gr.Slider(0.0, 0.98, label="Simplify", value=0.95, step=0.01)
                 texture_size = gr.Slider(512, 2048, label="Texture Size", value=1024, step=512)
             
-            extract_glb_btn = gr.Button("Extract GLB", interactive=False)
+            extract_reduced_btn = gr.Button("Extract Reduced GLB", interactive=False)
 
         with gr.Column():
             video_output = gr.Video(label="Generated 3D Asset", autoplay=True, loop=True, height=300)
@@ -258,7 +265,7 @@ with gr.Blocks(delete_cache=(600, 600)) as demo:
             
     output_buf = gr.State()
 
-    # Example images at the bottom of the page
+    # Example images
     with gr.Row():
         examples = gr.Examples(
             examples=[
@@ -291,12 +298,12 @@ with gr.Blocks(delete_cache=(600, 600)) as demo:
         inputs=[image_prompt, seed, ss_guidance_strength, ss_sampling_steps, slat_guidance_strength, slat_sampling_steps],
         outputs=[output_buf, video_output],
     ).then(
-        lambda: [gr.Button(interactive=True), gr.Button(interactive=True), gr.Button(interactive=False)],
-        outputs=[download_full, extract_glb_btn, download_reduced],
+        lambda: [gr.Button(interactive=True), gr.Button(interactive=True)],
+        outputs=[extract_full_btn, extract_reduced_btn],
     )
 
-    download_full.click(
-        export_full_quality_glb,
+    extract_full_btn.click(
+        extract_high_quality_glb,
         inputs=[output_buf],
         outputs=[model_output, download_full],
     ).then(
@@ -304,8 +311,8 @@ with gr.Blocks(delete_cache=(600, 600)) as demo:
         outputs=[download_full],
     )
 
-    extract_glb_btn.click(
-        extract_glb,
+    extract_reduced_btn.click(
+        extract_reduced_glb,
         inputs=[output_buf, mesh_simplify, texture_size],
         outputs=[model_output, download_reduced],
     ).then(