update

app.py

@@ -27,8 +27,8 @@ import yaml
 import matplotlib.pyplot as plt
 from sklearn.neighbors import NearestNeighbors
 import spaces
-# from spaces import zero
-# zero.startup()
+from spaces import zero
+zero.startup()
 
 def install_cuda_toolkit():
     # CUDA_TOOLKIT_URL = "https://developer.download.nvidia.com/compute/cuda/11.8.0/local_installers/cuda_11.8.0_520.61.05_linux.run"
@@ -51,7 +51,7 @@ def install_cuda_toolkit():
 install_cuda_toolkit()
 
 gs_path = Path(__file__).parent / "src/third-party/diff-gaussian-rasterization-w-depth"
-subprocess.check_call(["pip", "install", "-e", str(gs_path)])
+subprocess.check_call([sys.executable, "-m", "pip", "install", "-e", str(gs_path)])
 site.main()  # re-processes every *.pth in site-packages
 importlib.invalidate_caches()
 diff_gaussian_rasterization = importlib.import_module("diff_gaussian_rasterization")
@@ -113,7 +113,6 @@ def fps_new(x, enabled, n, device, random_start=False):
 
 class DynamicsVisualizer:
 
-    @spaces.GPU
     def __init__(self):
         wp.init()
 
@@ -143,7 +142,7 @@ class DynamicsVisualizer:
         cfg.sim.uniform = True
         cfg.sim.use_pv = False
 
-        device = torch.device('cuda:0')
+        device = torch.device('cuda')
 
         self.cfg = cfg
         self.device = device
@@ -156,14 +155,6 @@ class DynamicsVisualizer:
         self.dt_base = cfg.sim.dt
         self.high_freq_pred = True
 
-        gpus = [int(gpu) for gpu in cfg.gpus]
-        wp_devices = [wp.get_device(f'cuda:{gpu}') for gpu in gpus]
-        torch_devices = [torch.device(f'cuda:{gpu}') for gpu in gpus]
-        device_count = len(torch_devices)
-        assert device_count == 1
-        self.wp_device = wp_devices[0]
-        self.torch_device = torch_devices[0]
-
         seed = cfg.seed
         random.seed(seed)
         np.random.seed(seed)
@@ -172,8 +163,7 @@ class DynamicsVisualizer:
         torch.backends.cudnn.benchmark = True
 
         self.clear()
-    
-    @spaces.GPU
+
     def clear(self, clear_params=True):
         self.metadata = {}
         self.config = {}
@@ -206,7 +196,6 @@ class DynamicsVisualizer:
         self.material = None
         self.friction = None
 
-    @spaces.GPU
     def load_scaniverse(self, data_path):
 
         ### load splat params
@@ -303,7 +292,6 @@ 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)
         
@@ -369,7 +357,6 @@ 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
@@ -382,7 +369,6 @@ 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
 
@@ -405,7 +391,6 @@ 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]
@@ -442,7 +427,6 @@ 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()}
@@ -452,7 +436,6 @@ 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())
@@ -460,7 +443,6 @@ 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)
@@ -473,7 +455,6 @@ 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
@@ -484,7 +465,6 @@ 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
@@ -520,13 +500,11 @@ 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
@@ -616,7 +594,6 @@ 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']
@@ -634,7 +611,6 @@ 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']
@@ -647,7 +623,6 @@ 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']
@@ -660,7 +635,6 @@ 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']
@@ -673,7 +647,6 @@ class DynamicsVisualizer:
 
         return grippers
 
-    @spaces.GPU
     def rotate(self, params, rot_mat):
         scale = np.linalg.norm(rot_mat, axis=1, keepdims=True)
         
@@ -686,7 +659,6 @@ class DynamicsVisualizer:
         }
         return params
 
-    @spaces.GPU
     def preprocess_gs(self, params):
         if isinstance(params, dict):
             xyz = params['means3D']
@@ -736,7 +708,6 @@ 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
@@ -784,7 +755,6 @@ 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)
@@ -861,7 +831,6 @@ 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']
@@ -931,6 +900,15 @@ class DynamicsVisualizer:
 
     @spaces.GPU
     def reset(self):
+        wp.init()
+        gpus = [int(gpu) for gpu in self.cfg.gpus]
+        wp_devices = [wp.get_device(f'cuda:{gpu}') for gpu in gpus]
+        torch_devices = [torch.device(f'cuda:{gpu}') for gpu in gpus]
+        device_count = len(torch_devices)
+        assert device_count == 1
+        self.wp_device = wp_devices[0]
+        self.torch_device = torch_devices[0]
+
         params = self.preprocess_gs(self.params)
         if self.with_bg:
             self.preprocess_bg_gs()
@@ -971,7 +949,6 @@ class DynamicsVisualizer:
 
         return form_video, form_3dgs_pred
 
-    @spaces.GPU
     def run_command(self, unit_command):
 
         os.system('rm -rf ' + str(root / 'log/temp/*'))
@@ -1058,25 +1035,78 @@ class DynamicsVisualizer:
         )
         return form_video, form_3dgs_pred
     
+    @spaces.GPU
     def on_click_run_xplus(self):
+        wp.init()
+        gpus = [int(gpu) for gpu in self.cfg.gpus]
+        wp_devices = [wp.get_device(f'cuda:{gpu}') for gpu in gpus]
+        torch_devices = [torch.device(f'cuda:{gpu}') for gpu in gpus]
+        device_count = len(torch_devices)
+        assert device_count == 1
+        self.wp_device = wp_devices[0]
+        self.torch_device = torch_devices[0]
         return self.run_command([5.0, 0, 0])
 
+    @spaces.GPU
     def on_click_run_xminus(self):
+        wp.init()
+        gpus = [int(gpu) for gpu in self.cfg.gpus]
+        wp_devices = [wp.get_device(f'cuda:{gpu}') for gpu in gpus]
+        torch_devices = [torch.device(f'cuda:{gpu}') for gpu in gpus]
+        device_count = len(torch_devices)
+        assert device_count == 1
+        self.wp_device = wp_devices[0]
+        self.torch_device = torch_devices[0]
         return self.run_command([-5.0, 0, 0])
     
+    @spaces.GPU
     def on_click_run_yplus(self):
+        wp.init()
+        gpus = [int(gpu) for gpu in self.cfg.gpus]
+        wp_devices = [wp.get_device(f'cuda:{gpu}') for gpu in gpus]
+        torch_devices = [torch.device(f'cuda:{gpu}') for gpu in gpus]
+        device_count = len(torch_devices)
+        assert device_count == 1
+        self.wp_device = wp_devices[0]
+        self.torch_device = torch_devices[0]
         return self.run_command([0, 5.0, 0])
     
+    @spaces.GPU
     def on_click_run_yminus(self):
+        wp.init()
+        gpus = [int(gpu) for gpu in self.cfg.gpus]
+        wp_devices = [wp.get_device(f'cuda:{gpu}') for gpu in gpus]
+        torch_devices = [torch.device(f'cuda:{gpu}') for gpu in gpus]
+        device_count = len(torch_devices)
+        assert device_count == 1
+        self.wp_device = wp_devices[0]
+        self.torch_device = torch_devices[0]
         return self.run_command([0, -5.0, 0])
     
+    @spaces.GPU
     def on_click_run_zplus(self):
+        wp.init()
+        gpus = [int(gpu) for gpu in self.cfg.gpus]
+        wp_devices = [wp.get_device(f'cuda:{gpu}') for gpu in gpus]
+        torch_devices = [torch.device(f'cuda:{gpu}') for gpu in gpus]
+        device_count = len(torch_devices)
+        assert device_count == 1
+        self.wp_device = wp_devices[0]
+        self.torch_device = torch_devices[0]
         return self.run_command([0, 0, 5.0])
     
+    @spaces.GPU
     def on_click_run_zminus(self):
+        wp.init()
+        gpus = [int(gpu) for gpu in self.cfg.gpus]
+        wp_devices = [wp.get_device(f'cuda:{gpu}') for gpu in gpus]
+        torch_devices = [torch.device(f'cuda:{gpu}') for gpu in gpus]
+        device_count = len(torch_devices)
+        assert device_count == 1
+        self.wp_device = wp_devices[0]
+        self.torch_device = torch_devices[0]
         return self.run_command([0, 0, -5.0])
 
-    @spaces.GPU
     def launch(self, share=False):
         in_dir = root / 'log/gs/ckpts/rope_scene_1'
         batch_size = 1