refactor session management and update event handlers for improved clarity and functionality

app.py

@@ -430,17 +430,17 @@ 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)
+        # Set up session management - COMMENT THESE OUT FOR TESTING
+        # 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
+        # State for storing 3D model data
         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")
+            gr.Markdown("# Text to Product\nUsing Mistral-7B + FLUX.1-dev + Trellis")
             
             prompt = gr.Text(
                 show_label=False,
@@ -497,8 +497,8 @@ def create_interface():
                     
                 with gr.Tab("Flux"):
                     # Flux settings
-                    seed = gr.Slider(label="Seed", minimum=0, maximum=MAX_SEED, step=1, value=DEFAULT_SEED)
-                    randomize_seed = gr.Checkbox(label="Randomize seed", value=DEFAULT_RANDOMIZE_SEED)
+                    flux_seed = gr.Slider(label="Seed", minimum=0, maximum=MAX_SEED, step=1, value=DEFAULT_SEED)
+                    flux_randomize_seed = gr.Checkbox(label="Randomize seed", value=DEFAULT_RANDOMIZE_SEED)
                     
                     with gr.Row():
                         width = gr.Slider(label="Width", minimum=256, maximum=MAX_IMAGE_SIZE, step=32, value=DEFAULT_WIDTH)
@@ -513,8 +513,8 @@ def create_interface():
                     )
                 
                 with gr.Tab("3D Generation Settings"):
-                    seed = gr.Slider(0, MAX_SEED, label="Seed", value=0, step=1)
-                    randomize_seed = gr.Checkbox(label="Randomize Seed", value=True)
+                    trellis_seed = gr.Slider(0, MAX_SEED, label="Seed", value=0, step=1)
+                    trellis_randomize_seed = gr.Checkbox(label="Randomize Seed", value=True)
                     gr.Markdown("Stage 1: Sparse Structure Generation")
                     with gr.Row():
                         ss_guidance_strength = gr.Slider(0.0, 10.0, label="Guidance Strength", value=7.5, step=0.1)
@@ -539,38 +539,37 @@ def create_interface():
 
             # Examples section - simplified version that only updates the prompt fields
             gr.Examples(
-                examples=examples,  # Now just a list of prompts
+                examples=examples,
                 fn=process_example_pipeline,
-                inputs=[prompt],  # Add system_prompt as input
-                outputs=[refined_prompt, message_box],  # Don't output image
+                inputs=[prompt],
+                outputs=[refined_prompt, message_box],
                 cache_examples=True,
             )
 
-        # Event handlers
+        # Event handlers - Fixed to use the renamed components
         gr.on(
             triggers=[prompt_button.click, prompt.submit],
             fn=refine_prompt,
-            inputs=[prompt, system_prompt],  # Add system_prompt as input
+            inputs=[prompt, system_prompt],
             outputs=[refined_prompt, message_box]
         )
 
         gr.on(
             triggers=[visual_button.click],
             fn=infer,
-            inputs=[refined_prompt, seed, randomize_seed, width, height, num_inference_steps],
+            inputs=[refined_prompt, flux_seed, flux_randomize_seed, width, height, num_inference_steps],
             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],
+            inputs=[generated_image, trellis_seed, ss_guidance_strength, ss_sampling_steps, slat_guidance_strength, slat_sampling_steps],
             outputs=[output_state, video_output],
-        ).success(
+        ).then(
             # 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]
+            lambda: (gr.Button.update(interactive=True), gr.Button.update(interactive=True), "3D model generated successfully"),
+            outputs=[extract_glb_btn, extract_gs_btn, message_box]
         )
         
         # Add handlers for GLB and Gaussian extraction
@@ -579,10 +578,10 @@ def create_interface():
             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),
+        ).then(
+            lambda path: (gr.DownloadButton.update(interactive=True, value=path), "GLB extraction completed"),
             inputs=[model_output],
-            outputs=[download_glb]
+            outputs=[download_glb, message_box]
         )
 
         gr.on(
@@ -590,10 +589,10 @@ def create_interface():
             fn=extract_gaussian,
             inputs=[output_state],
             outputs=[model_output, download_gs]
-        ).success(
-            lambda path: gr.DownloadButton.update(interactive=True, value=path),
+        ).then(
+            lambda path: (gr.DownloadButton.update(interactive=True, value=path), "Gaussian extraction completed"),
             inputs=[model_output],
-            outputs=[download_gs]
+            outputs=[download_gs, message_box]
         )
 
     return demo