test

app.py

@@ -925,10 +925,6 @@ class DynamicsVisualizer:
 
     @spaces.GPU
     def reset(self):
-        # import torch
-        # import torch.nn as nn
-        # from torch.utils.data import DataLoader
-        # import kornia
         import warp as wp
         
         device = torch.device('cuda:0')
@@ -977,7 +973,7 @@ class DynamicsVisualizer:
 
         form_video = gr.Video(
             label='Predicted video',
-            value=root / f'log/gs/temp/form_video.mp4',
+            value=root / f'log/gs/temp/form_video_init.mp4',
             format='mp4',
             width=self.width,
             height=self.height,
@@ -992,7 +988,15 @@ class DynamicsVisualizer:
         return form_video, form_3dgs_pred
 
     def run_command(self, unit_command):
-
+        import warp as wp
+        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]
         os.system('rm -rf ' + str(root / 'log/temp/*'))
 
         # im_list = []
@@ -1079,68 +1083,26 @@ class DynamicsVisualizer:
     
     @spaces.GPU
     def on_click_run_xplus(self):
-        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):
-        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):
-        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):
-        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):
-        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):
-        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])
 
     def launch(self, share=False):