refactor session handling to use session-specific directories and remove unused extraction functions

app.py

@@ -55,13 +55,13 @@ _trellis_pipeline = None
 
 
 def start_session(req: gr.Request):
-    user_dir = os.path.join(TMP_DIR, "temp_output")
-    # user_dir = os.path.join(TMP_DIR, str(req.session_hash))
+    # user_dir = os.path.join(TMP_DIR, "temp_output")
+    user_dir = os.path.join(TMP_DIR, str(req.session_hash))
     os.makedirs(user_dir, exist_ok=True)
     
 def end_session(req: gr.Request):
-    user_dir = os.path.join(TMP_DIR, "temp_output")
-    # user_dir = os.path.join(TMP_DIR, str(req.session_hash))
+    # user_dir = os.path.join(TMP_DIR, "temp_output")
+    user_dir = os.path.join(TMP_DIR, str(req.session_hash))
     shutil.rmtree(user_dir)
 
 def preprocess_image(image: Image.Image) -> Image.Image:
@@ -321,6 +321,7 @@ def image_to_3d(
     ss_sampling_steps: int,
     slat_guidance_strength: float,
     slat_sampling_steps: int,
+    req: gr.Request,
 ) -> Tuple[dict, str]:
     try:
         # Load the Trellis pipeline
@@ -347,8 +348,9 @@ def image_to_3d(
                 "cfg_strength": slat_guidance_strength,
             },
         )
-        temp_dir = os.path.join(TMP_DIR, "temp_output")
-        
+        # temp_dir = os.path.join(TMP_DIR, "temp_output")
+        temp_dir = os.path.join(TMP_DIR, str(req.session_hash))
+
         video = render_utils.render_video(outputs['gaussian'][0], num_frames=120)['color']
         video_geo = render_utils.render_video(outputs['mesh'][0], num_frames=120)['normal']
         video = [np.concatenate([video[i], video_geo[i]], axis=1) for i in range(len(video))]
@@ -364,51 +366,6 @@ def image_to_3d(
         return None, f"Error generating 3D model: {str(e)}"
 
 
-@spaces.GPU(duration=90)
-def extract_glb(
-    state: dict,
-    mesh_simplify: float,
-    texture_size: int,
-) -> Tuple[str, str]:
-    """
-    Extract a GLB file from the 3D model.
-
-    Args:
-        state (dict): The state of the generated 3D model.
-        mesh_simplify (float): The mesh simplification factor.
-        texture_size (int): The texture resolution.
-
-    Returns:
-        str: The path to the extracted GLB file.
-    """
-    temp_dir = os.path.join(TMP_DIR, "temp_output")
-    gs, mesh = unpack_state(state)
-    glb = postprocessing_utils.to_glb(gs, mesh, simplify=mesh_simplify, texture_size=texture_size, verbose=False)
-    glb_path = os.path.join(temp_dir, 'sample.glb')
-    glb.export(glb_path)
-    torch.cuda.empty_cache()
-    return glb_path, glb_path
-
-
-@spaces.GPU
-def extract_gaussian(state: dict) -> Tuple[str, str]:
-    """
-    Extract a Gaussian file from the 3D model.
-
-    Args:
-        state (dict): The state of the generated 3D model.
-
-    Returns:
-        str: The path to the extracted Gaussian file.
-    """
-    temp_dir = os.path.join(TMP_DIR, "temp_output")
-    gs, _ = unpack_state(state)
-    gaussian_path = os.path.join(temp_dir, 'sample.ply')
-    gs.save_ply(gaussian_path)
-    torch.cuda.empty_cache()
-    return gaussian_path, gaussian_path
-
-
 # Create a combined function that handles the whole pipeline from example to image
 # This version gets the parameters from the UI components
 @spaces.GPU()
@@ -465,17 +422,10 @@ def create_interface():
             visual_button = gr.Button("Create visual with Flux")
             generated_image = gr.Image(show_label=False, format="png", image_mode="RGBA", type="pil", height=300)
 
-            preprocessed_button = gr.Button("Preprocess image")
-            preprocessed_image = gr.Image(show_label=False)
-
             gen3d_button = gr.Button("Create 3D visual with Trellis")
             video_output = gr.Video(label="Generated 3D Asset", autoplay=True, loop=True, height=300)
             # model_output = LitModel3D(label="Extracted GLB/Gaussian", exposure=10.0, height=300)
             
-            with gr.Row():
-                download_glb = gr.DownloadButton(label="Download GLB", interactive=False)
-                download_gs = gr.DownloadButton(label="Download Gaussian", interactive=False)  
-    
             message_box = gr.Textbox(
                 label="Status Messages", 
                 interactive=False,
@@ -531,18 +481,8 @@ def create_interface():
                         slat_guidance_strength = gr.Slider(0.0, 10.0, label="Guidance Strength", value=3.0, step=0.1)
                         slat_sampling_steps = gr.Slider(1, 50, label="Sampling Steps", value=12, step=1)
 
-                with gr.Tab("GLB Extraction Settings"):
-                    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()
+            # output_buf = gr.State()
 
             gr.Examples(
                 examples=examples,
@@ -566,13 +506,6 @@ def create_interface():
             outputs=[generated_image, message_box]
         )
 
-        gr.on(
-            triggers=[preprocessed_button.click],
-            fn=preprocess_image,
-            inputs=[generated_image],
-            outputs=[preprocessed_image]
-        )
-
         gr.on(
             triggers=[gen3d_button.click],
             fn=image_to_3d,