update

app.py

@@ -113,6 +113,7 @@ def fps_new(x, enabled, n, device, random_start=False):
 
 class DynamicsVisualizer:
 
+    @spaces.GPU
     def __init__(self):
         wp.init()
 
@@ -172,6 +173,7 @@ class DynamicsVisualizer:
 
         self.clear()
     
+    @spaces.GPU
     def clear(self, clear_params=True):
         self.metadata = {}
         self.config = {}
@@ -204,6 +206,7 @@ class DynamicsVisualizer:
         self.material = None
         self.friction = None
 
+    @spaces.GPU
     def load_scaniverse(self, data_path):
 
         ### load splat params
@@ -300,7 +303,7 @@ class DynamicsVisualizer:
             # self.state['prev_key_pos_timestamp'] = torch.zeros(1).to(self.device).to(torch.float32)
             self.state['gripper_radius'] = cfg.model.gripper_radius
 
-
+    @spaces.GPU
     def load_params(self, params_path, remove_low_opa=True, remove_black=False):
         pts, colors, scales, quats, opacities = read_splat(params_path)
         
@@ -366,7 +369,7 @@ class DynamicsVisualizer:
         self.state['clip_bound'] = torch.tensor([self.cfg.model.clip_bound], dtype=torch.float32)
         self.state['enabled'] = torch.ones(n_particles, dtype=torch.bool)
 
-
+    @spaces.GPU
     def set_camera(self, w, h, intr, w2c=None, R=None, t=None, near=0.01, far=100.0):
         if w2c is None:
             assert R is not None and t is not None
@@ -379,6 +382,7 @@ class DynamicsVisualizer:
         }
         self.config = {'near': near, 'far': far}
     
+    @spaces.GPU
     def load_eef(self, grippers=None, eef_t=None):
         assert self.state['prev_key_pos'] is None
 
@@ -401,6 +405,7 @@ class DynamicsVisualizer:
             # self.state['prev_key_pos_timestamp'] = torch.zeros(1).to(self.device).to(torch.float32) + 0.001
             self.state['gripper_radius'] = self.cfg.model.gripper_radius
 
+    @spaces.GPU
     def load_preprocess_metadata(self, p_x_orig):
         cfg = self.cfg
         dx = cfg.sim.num_grids[-1]
@@ -437,6 +442,7 @@ class DynamicsVisualizer:
             'global_translation': global_translation,
         }
 
+    @spaces.GPU
     @torch.no_grad
     def render(self, render_data, cam_id, bg=[0.7, 0.7, 0.7]):
         render_data = {k: v.to(self.device) for k, v in render_data.items()}
@@ -446,6 +452,7 @@ class DynamicsVisualizer:
         im, _, depth, = GaussianRasterizer(raster_settings=cam)(**render_data)
         return im, depth
     
+    @spaces.GPU
     def knn_relations(self, bones):
         k = self.k_rel
         knn = NearestNeighbors(n_neighbors=k+1, algorithm='kd_tree').fit(bones.detach().cpu().numpy())
@@ -453,6 +460,7 @@ class DynamicsVisualizer:
         indices = indices[:, 1:]  # exclude self
         return indices
     
+    @spaces.GPU
     def knn_weights_brute(self, bones, pts):
         k = self.k_wgt
         dist = torch.norm(pts[:, None] - bones, dim=-1)  # (n_pts, n_bones)
@@ -465,6 +473,7 @@ class DynamicsVisualizer:
         weights_all[torch.arange(pts.shape[0])[:, None], indices] = weights
         return weights_all
     
+    @spaces.GPU
     def update_camera(self, k, w2c, w=None, h=None, near=0.01, far=100.0):
         self.metadata['k'] = k
         self.metadata['w2c'] = w2c
@@ -475,6 +484,7 @@ class DynamicsVisualizer:
         self.config['near'] = near
         self.config['far'] = far
     
+    @spaces.GPU
     def init_model(self, batch_size, num_steps, num_particles, ckpt_path=None):
         self.cfg.sim.num_steps = num_steps
         cfg = self.cfg
@@ -510,11 +520,13 @@ class DynamicsVisualizer:
         self.material = material
         self.friction = friction
 
+    @spaces.GPU
     def reload_model(self, num_steps):  # only change num_steps
         self.cfg.sim.num_steps = num_steps
         sim = CacheDiffSimWithFrictionBatch(self.cfg, num_steps, 1, self.wp_device, requires_grad=True)
         self.sim = sim
 
+    @spaces.GPU
     @torch.no_grad
     def step(self):
         cfg = self.cfg
@@ -604,6 +616,7 @@ class DynamicsVisualizer:
         self.state['sub_pos'] = None
         # self.state['sub_pos_timestamps'] = None
 
+    @spaces.GPU
     def preprocess_x(self, p_x):  # viewer frame to model frame (not data frame)
         R = self.preprocess_metadata['R']
         R_viewer = self.preprocess_metadata['R_viewer']
@@ -621,6 +634,7 @@ class DynamicsVisualizer:
 
         return p_x
     
+    @spaces.GPU
     def preprocess_gripper(self, grippers):  # viewer frame to model frame (not data frame)
         R = self.preprocess_metadata['R']
         R_viewer = self.preprocess_metadata['R_viewer']
@@ -633,6 +647,7 @@ class DynamicsVisualizer:
 
         return grippers
     
+    @spaces.GPU
     def inverse_preprocess_x(self, p_x):  # model frame (not data frame) to viewer frame
         R = self.preprocess_metadata['R']
         R_viewer = self.preprocess_metadata['R_viewer']
@@ -645,6 +660,7 @@ class DynamicsVisualizer:
 
         return p_x
     
+    @spaces.GPU
     def inverse_preprocess_gripper(self, grippers):  # model frame (not data frame) to viewer frame
         R = self.preprocess_metadata['R']
         R_viewer = self.preprocess_metadata['R_viewer']
@@ -657,55 +673,20 @@ class DynamicsVisualizer:
 
         return grippers
 
-    def preprocess_gs(self, params):
-        if isinstance(params, dict):
-            xyz = params['means3D']
-            rgb = params['rgb_colors']
-            quat = torch.nn.functional.normalize(params['unnorm_rotations'])
-            opa = torch.sigmoid(params['logit_opacities'])
-            scales = torch.exp(params['log_scales'])
-        else:
-            assert isinstance(params, tuple)
-            xyz, rgb, quat, opa, scales = params
-        
-        quat = torch.nn.functional.normalize(quat, dim=-1)
-
-        # transform
-        R = self.preprocess_metadata['R']
-        R_viewer = self.preprocess_metadata['R_viewer']
-        scale = self.preprocess_metadata['scale']
-        global_translation = self.preprocess_metadata['global_translation']
-
-        mat = quat2mat(quat)
-        mat = R @ mat
-        xyz = xyz @ R.T
-        xyz = xyz * scale
-        xyz += global_translation
-        quat = mat2quat(mat)
-        scales = scales * scale
-
-        # viewer-specific transform (flip y and z)
-        # model frame to viewer frame
-        xyz = xyz @ R_viewer.T
-        quat = mat2quat(R_viewer @ quat2mat(quat))
-
-        t_viewer = -xyz.mean(dim=0)
-        t_viewer[2] = 0
-        xyz += t_viewer
-        print('Overwriting t_viewer to be the planar mean of the object')
-        self.preprocess_metadata['t_viewer'] = t_viewer
-
-        if isinstance(params, dict):
-            params['means3D'] = xyz
-            params['rgb_colors'] = rgb
-            params['unnorm_rotations'] = quat
-            params['logit_opacities'] = opa
-            params['log_scales'] = torch.log(scales)
-        else:
-            params = xyz, rgb, quat, opa, scales
+    @spaces.GPU
+    def rotate(self, params, rot_mat):
+        scale = np.linalg.norm(rot_mat, axis=1, keepdims=True)
         
+        params = {
+            'means3D': pts,
+            'rgb_colors': params['rgb_colors'],
+            'log_scales': params['log_scales'],
+            'unnorm_rotations': quats,
+            'logit_opacities': params['logit_opacities'],
+        }
         return params
 
+    @spaces.GPU
     def preprocess_gs(self, params):
         if isinstance(params, dict):
             xyz = params['means3D']
@@ -755,6 +736,7 @@ class DynamicsVisualizer:
         
         return params
     
+    @spaces.GPU
     def preprocess_bg_gs(self):
         t_pts, t_colors, t_scales, t_quats, t_opacities = self.table_params
         g_pts, g_colors, g_scales, g_quats, g_opacities = self.gripper_params
@@ -802,6 +784,7 @@ class DynamicsVisualizer:
         self.table_params = t_pts, t_colors, t_scales, t_quats, t_opacities
         self.gripper_params = g_pts, g_colors, g_scales, g_quats, g_opacities
 
+    @spaces.GPU
     def update_rendervar(self, rendervar):
         p_x = self.state['x']
         p_x_viewer = self.inverse_preprocess_x(p_x)
@@ -878,6 +861,7 @@ class DynamicsVisualizer:
 
         return rendervar, rendervar_full
 
+    @spaces.GPU
     def reset_state(self, params, visualize_image=False, init=False):
         xyz_0 = params['means3D']
         rgb_0 = params['rgb_colors']
@@ -945,12 +929,7 @@ class DynamicsVisualizer:
 
         return rendervar_init
 
-
-
-
-
-
-
+    @spaces.GPU
     def reset(self):
         params = self.preprocess_gs(self.params)
         if self.with_bg:
@@ -992,6 +971,7 @@ class DynamicsVisualizer:
 
         return form_video, form_3dgs_pred
 
+    @spaces.GPU
     def run_command(self, unit_command):
 
         os.system('rm -rf ' + str(root / 'log/temp/*'))
@@ -1095,21 +1075,9 @@ class DynamicsVisualizer:
     
     def on_click_run_zminus(self):
         return self.run_command([0, 0, -5.0])
-    
-    def rotate(self, params, rot_mat):
-        scale = np.linalg.norm(rot_mat, axis=1, keepdims=True)
-        
-        params = {
-            'means3D': pts,
-            'rgb_colors': params['rgb_colors'],
-            'log_scales': params['log_scales'],
-            'unnorm_rotations': quats,
-            'logit_opacities': params['logit_opacities'],
-        }
-        return params
 
+    @spaces.GPU
     def launch(self, share=False):
-
         in_dir = root / 'log/gs/ckpts/rope_scene_1'
         batch_size = 1
         num_steps = 1