Update app.py

app.py

@@ -14,18 +14,24 @@ from trellis.pipelines import TrellisTextTo3DPipeline
 from trellis.representations import Gaussian, MeshExtractResult
 from trellis.utils import render_utils, postprocessing_utils
 
+import traceback
+import sys
+
 MAX_SEED = np.iinfo(np.int32).max
 TMP_DIR = os.path.join(os.path.dirname(os.path.abspath(__file__)), 'tmp')
 os.makedirs(TMP_DIR, exist_ok=True)
 
+
 def start_session(req: gr.Request):
     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, str(req.session_hash))
     shutil.rmtree(user_dir)
 
+
 def pack_state(gs: Gaussian, mesh: MeshExtractResult) -> dict:
     return {
         'gaussian': {
@@ -42,6 +48,7 @@ def pack_state(gs: Gaussian, mesh: MeshExtractResult) -> dict:
         },
     }
     
+    
 def unpack_state(state: dict) -> Tuple[Gaussian, edict, str]:
     gs = Gaussian(
         aabb=state['gaussian']['aabb'],
@@ -64,6 +71,7 @@ def unpack_state(state: dict) -> Tuple[Gaussian, edict, str]:
     
     return gs, mesh
 
+
 def get_seed(randomize_seed: bool, seed: int) -> int:
     return np.random.randint(0, MAX_SEED) if randomize_seed else seed
 
@@ -100,7 +108,7 @@ def text_to_3d(
     torch.cuda.empty_cache()
     return state, video_path
 
-@spaces.GPU
+@spaces.GPU(duration=90)
 def extract_glb(
     state: dict,
     mesh_simplify: float,
@@ -115,6 +123,16 @@ def extract_glb(
     torch.cuda.empty_cache()
     return glb_path, glb_path
 
+
+@spaces.GPU
+def extract_gaussian(state: dict, req: gr.Request) -> Tuple[str, str]:
+    user_dir = os.path.join(TMP_DIR, str(req.session_hash))
+    gs, _ = unpack_state(state)
+    gaussian_path = os.path.join(user_dir, 'sample.ply')
+    gs.save_ply(gaussian_path)
+    torch.cuda.empty_cache()
+    return gaussian_path, gaussian_path
+
 with gr.Blocks(delete_cache=(600, 600)) as demo:
     gr.Markdown("""
     # UTPL - Conversión de Texto a objetos 3D usando IA