revert bck to memory error

app.py

@@ -13,7 +13,6 @@ from easydict import EasyDict as edict
 from trellis.pipelines import TrellisTextTo3DPipeline
 from trellis.representations import Gaussian, MeshExtractResult
 from trellis.utils import render_utils, postprocessing_utils
-import joblib # Added for saving/loading state
 
 import traceback
 import sys
@@ -90,7 +89,7 @@ def text_to_3d(
     slat_guidance_strength: float,
     slat_sampling_steps: int,
     req: gr.Request,
-) -> Tuple[str, str, str]:
+) -> Tuple[dict, str, dict]:
     """
     Convert an text prompt to a 3D model.
     Args:
@@ -101,9 +100,9 @@ def text_to_3d(
         slat_guidance_strength (float): The guidance strength for structured latent generation.
         slat_sampling_steps (int): The number of sampling steps for structured latent generation.
     Returns:
-        str: Path to the saved state file.
-        str: Path to the generated video.
-        str: Path to the saved state file (for internal buffer).
+        dict: The information of the generated 3D model.
+        str: The path to the video of the 3D model.
+        dict: The state of the generated 3D model.
     """
     user_dir = os.path.join(TMP_DIR, str(req.session_hash))
     os.makedirs(user_dir, exist_ok=True)
@@ -126,70 +125,34 @@ def text_to_3d(
     video_path = os.path.join(user_dir, 'sample.mp4')
     imageio.mimsave(video_path, video, fps=15)
     state = pack_state(outputs['gaussian'][0], outputs['mesh'][0])
-    
-    # Save state to file
-    state_file_path = os.path.join(user_dir, f'state_{seed}.joblib')
-    try:
-        joblib.dump(state, state_file_path)
-        print(f"[Trellis] State saved to {state_file_path}")
-    except Exception as e:
-        print(f"Error saving state to {state_file_path}: {e}")
-        # Decide how to handle error - maybe return None or raise?
-        # For now, let's allow it to proceed but log the error
-        state_file_path = None # Indicate failure
-
     torch.cuda.empty_cache()
-    # Return state file path for API, video path for Video, and state path again for internal buffer
-    # Return None for path if saving failed
-    return state_file_path, video_path, state_file_path
+    # Return state for JSON, video path for Video, and state again for internal buffer
+    return state, video_path, state
 
 
 @spaces.GPU(duration=90)
 def extract_glb(
-    state_file_path: str, # Changed input from state: dict
+    state: dict,
     mesh_simplify: float,
     texture_size: int,
     req: gr.Request,
 ) -> Tuple[str, str]:
     """
-    Extract a GLB file from the 3D model state file.
+    Extract a GLB file from the 3D model.
     Args:
-        state_file_path (str): Path to the file containing the state.
+        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.
-        str: The path to the extracted GLB file (for download button).
     """
-    if not state_file_path or not os.path.exists(state_file_path):
-        print(f"Error: State file path invalid or file not found: {state_file_path}")
-        # Return dummy paths or raise an error
-        return None, None 
-
     user_dir = os.path.join(TMP_DIR, str(req.session_hash))
     os.makedirs(user_dir, exist_ok=True)
-    
-    # Load state from file
-    try:
-        state = joblib.load(state_file_path)
-        print(f"[Trellis] State loaded from {state_file_path}")
-    except Exception as e:
-        print(f"Error loading state from {state_file_path}: {e}")
-        return None, None
-
     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(user_dir, 'sample.glb')
     glb.export(glb_path)
     torch.cuda.empty_cache()
-
-    # Optional: Clean up the state file after use
-    try:
-        os.remove(state_file_path)
-        print(f"[Trellis] Cleaned up state file: {state_file_path}")
-    except OSError as e:
-        print(f"Error removing state file {state_file_path}: {e.strerror}")
-
     return glb_path, glb_path
 
 
@@ -215,8 +178,8 @@ output_buf = gr.State()
 video_output = gr.Video(label="Generated 3D Asset", autoplay=True, loop=True, height=300)
 model_output = gr.Model3D(label="Extracted GLB/Gaussian", height=300)
 
-# Change hidden JSON to hidden Textbox for the state file path
-state_output_path_textbox = gr.Textbox(visible=False, label="State File Path Output")
+# Add a hidden JSON output for the state object for API calls
+state_output_json = gr.JSON(visible=False, label="State JSON Output")
 
 with gr.Blocks(delete_cache=(600, 600)) as demo:
     gr.Markdown("""
@@ -275,8 +238,8 @@ with gr.Blocks(delete_cache=(600, 600)) as demo:
     ).then(
         text_to_3d,
         inputs=[text_prompt, seed, ss_guidance_strength, ss_sampling_steps, slat_guidance_strength, slat_sampling_steps],
-        # Output state path to hidden Textbox, video to Video, state path to internal buffer
-        outputs=[state_output_path_textbox, video_output, output_buf],
+        # Output state to hidden JSON first, then video to visible component, then state to internal buffer
+        outputs=[state_output_json, video_output, output_buf],
     ).then(
         lambda: tuple([gr.Button(interactive=True), gr.Button(interactive=True)]),
         outputs=[extract_glb_btn, extract_gs_btn],
@@ -289,7 +252,6 @@ with gr.Blocks(delete_cache=(600, 600)) as demo:
 
     extract_glb_btn.click(
         extract_glb,
-        # Input state path from internal buffer (assuming it holds the path now)
         inputs=[output_buf, mesh_simplify, texture_size],
         outputs=[model_output, download_glb],
     ).then(
@@ -299,8 +261,7 @@ with gr.Blocks(delete_cache=(600, 600)) as demo:
     
     extract_gs_btn.click(
         extract_gaussian,
-        # This likely needs adjustment too if it relies on output_buf holding the state dict
-        inputs=[output_buf], 
+        inputs=[output_buf],
         outputs=[model_output, download_gs],
     ).then(
         lambda: gr.Button(interactive=True),
@@ -344,11 +305,11 @@ api_text_to_3d = gr.Interface(
 # --- API-only endpoint for GLB extraction --- 
 # Explicitly defines state input as JSON for server calls.
 api_extract_glb = gr.Interface(
-    fn=lambda state_file_path, mesh_simplify, texture_size: extract_glb(
-        state_file_path, mesh_simplify, texture_size, gr.Request()
+    fn=lambda state, mesh_simplify, texture_size: extract_glb(
+        state, mesh_simplify, texture_size, gr.Request()
     ),
     inputs=[
-        gr.Textbox(label="State File Path"), # Expect state file path as string
+        gr.JSON(label="State Object"), # Expect state as JSON
         gr.Slider(0.9, 0.98, label="Simplify", value=0.95, step=0.01),
         gr.Slider(512, 2048, label="Texture Size", value=1024, step=512)
     ],