Add TestCode

app.py

@@ -15,6 +15,9 @@ 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')
@@ -103,28 +106,32 @@ def text_to_3d(
         dict: The information of the generated 3D model.
         str: The path to the video of the 3D model.
     """
-    user_dir = os.path.join(TMP_DIR, str(req.session_hash))
-    outputs = pipeline.run(
-        prompt,
-        seed=seed,
-        formats=["gaussian", "mesh"],
-        sparse_structure_sampler_params={
-            "steps": ss_sampling_steps,
-            "cfg_strength": ss_guidance_strength,
-        },
-        slat_sampler_params={
-            "steps": slat_sampling_steps,
-            "cfg_strength": slat_guidance_strength,
-        },
-    )
-    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))]
-    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])
-    torch.cuda.empty_cache()
-    return state, video_path
+    try:
+        user_dir = os.path.join(TMP_DIR, str(req.session_hash))
+        outputs = pipeline.run(
+            prompt,
+            seed=seed,
+            formats=["gaussian", "mesh"],
+            sparse_structure_sampler_params={
+                "steps": ss_sampling_steps,
+                "cfg_strength": ss_guidance_strength,
+            },
+            slat_sampler_params={
+                "steps": slat_sampling_steps,
+                "cfg_strength": slat_guidance_strength,
+            },
+        )
+        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))]
+        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])
+        torch.cuda.empty_cache()
+        return state, video_path
+    except Exception as e:
+        et, ev, tb = sys.exc_info()
+        traceback.print_tb(tb)
 
 
 @spaces.GPU(duration=90)