refactor interface to improve session management and state handling for 3D model extraction

app.py

@@ -219,7 +219,6 @@ def infer(prompt, seed=DEFAULT_SEED,
         return None, f"Error generating image: {str(e)}"
  
  
-# Format: [prompt, system_prompt]
 examples = [
     "a backpack for kids, flower style",
     "medieval flip flops",
@@ -420,7 +419,14 @@ def create_interface():
         model_status = "ℹ️ Models will be loaded on demand"
 
     with gr.Blocks(css=css) as demo:
+        # Set up session management
+        demo.load(start_session)
+        demo.unload(end_session)
+
         gr.Info(model_status)
+        
+        # State for storing 3D model data - moved to the top level inside Blocks context
+        output_state = gr.State(None)
 
         with gr.Column(elem_id="col-container"):
             gr.Markdown("# Text to Product\nUsing Mistral-7B-Instruct-v0.3 + FLUX.1-dev + Trellis")
@@ -511,6 +517,13 @@ def create_interface():
                     mesh_simplify = gr.Slider(0.9, 0.98, label="Simplify", value=0.95, step=0.01)
                     texture_size = gr.Slider(512, 2048, label="Texture Size", value=1024, step=512)
             
+            with gr.Row():
+                extract_glb_btn = gr.Button("Extract GLB", interactive=False)
+                extract_gs_btn = gr.Button("Extract Gaussian", interactive=False)
+            gr.Markdown("""
+                        *NOTE: Gaussian file can be very large (~50MB), it will take a while to display and download.*
+                        """)
+
             output_buf = gr.State()
 
             # Examples section - simplified version that only updates the prompt fields
@@ -537,17 +550,41 @@ def create_interface():
             outputs=[generated_image, message_box]
         )
 
+        # Updated to use output_state
         gr.on(
             triggers=[gen3d_button.click],
             fn=image_to_3d,
             inputs=[generated_image, seed, ss_guidance_strength, ss_sampling_steps, slat_guidance_strength, slat_sampling_steps],
-            outputs=[output_buf, video_output],
+            outputs=[output_state, video_output],
+        ).success(
+            # Update button states after successful 3D generation
+            lambda: (gr.Button.update(interactive=True), gr.Button.update(interactive=True)),
+            outputs=[extract_glb_btn, extract_gs_btn]
+        )
+        
+        # Add handlers for GLB and Gaussian extraction
+        gr.on(
+            triggers=[extract_glb_btn.click],
+            fn=extract_glb,
+            inputs=[output_state, mesh_simplify, texture_size],
+            outputs=[model_output, download_glb]
+        ).success(
+            lambda path: gr.DownloadButton.update(interactive=True, value=path),
+            inputs=[model_output],
+            outputs=[download_glb]
+        )
+
+        gr.on(
+            triggers=[extract_gs_btn.click],
+            fn=extract_gaussian,
+            inputs=[output_state],
+            outputs=[model_output, download_gs]
+        ).success(
+            lambda path: gr.DownloadButton.update(interactive=True, value=path),
+            inputs=[model_output],
+            outputs=[download_gs]
         )
 
-    # Handlers
-    demo.load(start_session)
-    demo.unload(end_session)
-    
     return demo