chore: catch error

app.py

@@ -178,53 +178,58 @@ def _default_render_cameras(batch_size: int = 1):
 
 @spaces.GPU
 def generate_mesh(image, source_size=512, render_size=384, mesh_size=512, export_mesh=False, export_video=False, fps=30):
-    image = preprocess_image(image, source_size).to(model_wrapper.device)
-    source_camera = _default_source_camera(batch_size=1).to(model_wrapper.device)
-
-    with torch.no_grad():
-        planes = model_wrapper.forward(image, source_camera)
-
-        if export_mesh:
-            grid_out = model_wrapper.model.synthesizer.forward_grid(planes=planes, grid_size=mesh_size)
-            vtx, faces = mcubes.marching_cubes(grid_out['sigma'].float().squeeze(0).squeeze(-1).cpu().numpy(), 1.0)
-            vtx = vtx / (mesh_size - 1) * 2 - 1
-            vtx_tensor = torch.tensor(vtx, dtype=torch.float32, device=model_wrapper.device).unsqueeze(0)
-            vtx_colors = model_wrapper.model.synthesizer.forward_points(planes, vtx_tensor)['rgb'].float().squeeze(0).cpu().numpy()
-            vtx_colors = (vtx_colors * 255).astype(np.uint8)
-            mesh = trimesh.Trimesh(vertices=vtx, faces=faces, vertex_colors=vtx_colors)
-
-            mesh_path = "xiaoxis_mesh.obj"
-            mesh.export(mesh_path, 'obj')
-
-            return None, mesh_path
-
-        if export_video:
-            render_cameras = _default_render_cameras(batch_size=1).to(model_wrapper.device)
-            frames = []
-            chunk_size = 1
-            for i in range(0, render_cameras.shape[1], chunk_size):
-                frame_chunk = model_wrapper.model.synthesizer(
-                    planes,
-                    render_cameras[:, i:i + chunk_size],
-                    render_size,
-                    render_size,
-                    0,
-                    0
-                )
-                frames.append(frame_chunk['images_rgb'])
-
-            frames = torch.cat(frames, dim=1)
-            frames = frames.squeeze(0)
-            frames = (frames.permute(0, 2, 3, 1).cpu().numpy() * 255).astype(np.uint8)
-
-            video_path = "xiaoxis_video.mp4"
-            imageio.mimwrite(video_path, frames, fps=fps)
-
-            return None, video_path
-
-        return planes, None
-
-    return None, None
+    try:
+        image = preprocess_image(image, source_size).to(model_wrapper.device)
+        source_camera = _default_source_camera(batch_size=1).to(model_wrapper.device)
+    
+        with torch.no_grad():
+            planes = model_wrapper.forward(image, source_camera)
+    
+            if export_mesh:
+                grid_out = model_wrapper.model.synthesizer.forward_grid(planes=planes, grid_size=mesh_size)
+                vtx, faces = mcubes.marching_cubes(grid_out['sigma'].float().squeeze(0).squeeze(-1).cpu().numpy(), 1.0)
+                vtx = vtx / (mesh_size - 1) * 2 - 1
+                vtx_tensor = torch.tensor(vtx, dtype=torch.float32, device=model_wrapper.device).unsqueeze(0)
+                vtx_colors = model_wrapper.model.synthesizer.forward_points(planes, vtx_tensor)['rgb'].float().squeeze(0).cpu().numpy()
+                vtx_colors = (vtx_colors * 255).astype(np.uint8)
+                mesh = trimesh.Trimesh(vertices=vtx, faces=faces, vertex_colors=vtx_colors)
+    
+                mesh_path = "xiaoxis_mesh.obj"
+                mesh.export(mesh_path, 'obj')
+    
+                return None, mesh_path
+    
+            if export_video:
+                render_cameras = _default_render_cameras(batch_size=1).to(model_wrapper.device)
+                frames = []
+                chunk_size = 1
+                for i in range(0, render_cameras.shape[1], chunk_size):
+                    frame_chunk = model_wrapper.model.synthesizer(
+                        planes,
+                        render_cameras[:, i:i + chunk_size],
+                        render_size,
+                        render_size,
+                        0,
+                        0
+                    )
+                    frames.append(frame_chunk['images_rgb'])
+    
+                frames = torch.cat(frames, dim=1)
+                frames = frames.squeeze(0)
+                frames = (frames.permute(0, 2, 3, 1).cpu().numpy() * 255).astype(np.uint8)
+    
+                video_path = "xiaoxis_video.mp4"
+                imageio.mimwrite(video_path, frames, fps=fps)
+    
+                return None, video_path
+    
+            return planes, None
+    
+        return None, None
+
+    except Exception as e:
+        print(f"Error: {str(e)}")
+        return None, None
 
 def step_1_generate_planes(image):
     planes, _ = generate_mesh(image)