test

app.py

@@ -997,9 +997,11 @@ class DynamicsVisualizer:
             clear_color=[0, 0, 0, 0],
         )
 
-        return form_video, form_3dgs_pred
+        return form_video, form_3dgs_pred, \
+            self.preprocess_metadata, self.state, self.params, \
+            self.table_params, self.gripper_params
 
-    def run_command(self, unit_command):
+    def run_command(self, unit_command, preprocess_metadata, state, params, table_params, gripper_params):
         import warp as wp
         wp.init()
         gpus = [int(gpu) for gpu in self.cfg.gpus]
@@ -1011,6 +1013,11 @@ class DynamicsVisualizer:
         self.torch_device = torch_devices[0]
         os.system('rm -rf ' + str(root / 'log/temp/*'))
         
+        self.preprocess_metadata = preprocess_metadata
+        self.state = state
+        self.params = params
+        self.table_params = table_params
+        self.gripper_params = gripper_params
         for k, v in self.preprocess_metadata.items():
             self.preprocess_metadata[k] = v.to(self.torch_device) if isinstance(v, torch.Tensor) else v
         for k, v in self.state.items():
@@ -1125,33 +1132,41 @@ class DynamicsVisualizer:
             value=root / 'log/gs/temp/gs_pred.splat',
             clear_color=[0, 0, 0, 0],
         )
-        return form_video, form_3dgs_pred
+        return form_video, form_3dgs_pred, \
+            self.preprocess_metadata, self.state, self.params, \
+            self.table_params, self.gripper_params
     
     @spaces.GPU
-    def on_click_run_xplus(self):
-        return self.run_command([5.0, 0, 0])
+    def on_click_run_xplus(self, preprocess_metadata, state, params, table_params, gripper_params):
+        return self.run_command([5.0, 0, 0], preprocess_metadata, state, params, table_params, gripper_params)
 
     @spaces.GPU
-    def on_click_run_xminus(self):
-        return self.run_command([-5.0, 0, 0])
+    def on_click_run_xminus(self, preprocess_metadata, state, params, table_params, gripper_params):
+        return self.run_command([-5.0, 0, 0], preprocess_metadata, state, params, table_params, gripper_params)
     
     @spaces.GPU
-    def on_click_run_yplus(self):
-        return self.run_command([0, 5.0, 0])
+    def on_click_run_yplus(self, preprocess_metadata, state, params, table_params, gripper_params):
+        return self.run_command([0, 5.0, 0], preprocess_metadata, state, params, table_params, gripper_params)
     
     @spaces.GPU
-    def on_click_run_yminus(self):
-        return self.run_command([0, -5.0, 0])
+    def on_click_run_yminus(self, preprocess_metadata, state, params, table_params, gripper_params):
+        return self.run_command([0, -5.0, 0], preprocess_metadata, state, params, table_params, gripper_params)
     
     @spaces.GPU
-    def on_click_run_zplus(self):
-        return self.run_command([0, 0, 5.0])
+    def on_click_run_zplus(self, preprocess_metadata, state, params, table_params, gripper_params):
+        return self.run_command([0, 0, 5.0], preprocess_metadata, state, params, table_params, gripper_params)
     
     @spaces.GPU
-    def on_click_run_zminus(self):
-        return self.run_command([0, 0, -5.0])
+    def on_click_run_zminus(self, preprocess_metadata, state, params, table_params, gripper_params):
+        return self.run_command([0, 0, -5.0], preprocess_metadata, state, params, table_params, gripper_params)
 
     def launch(self, share=False):
+        
+        preprocess_metadata = gr.State(self.preprocess_metadata)
+        state = gr.State(self.state)
+        params = gr.State(self.params)
+        table_params = gr.State(self.table_params)
+        gripper_params = gr.State(self.gripper_params)
 
         with gr.Blocks() as app:
 
@@ -1268,31 +1283,51 @@ class DynamicsVisualizer:
             # Set up callbacks
             run_reset.click(self.reset,
                     inputs=[],
-                    outputs=[form_video, form_3dgs_pred])
+                    outputs=[form_video, form_3dgs_pred, 
+                             preprocess_metadata, state, params,
+                             table_params, gripper_params])
             
             run_xplus.click(self.on_click_run_xplus,
-                    inputs=[],
-                    outputs=[form_video, form_3dgs_pred])
+                    inputs=[preprocess_metadata, state, params,
+                             table_params, gripper_params],
+                    outputs=[form_video, form_3dgs_pred, 
+                             preprocess_metadata, state, params,
+                             table_params, gripper_params])
             
             run_xminus.click(self.on_click_run_xminus,
-                    inputs=[],
-                    outputs=[form_video, form_3dgs_pred])
+                    inputs=[preprocess_metadata, state, params,
+                             table_params, gripper_params],
+                    outputs=[form_video, form_3dgs_pred, 
+                             preprocess_metadata, state, params,
+                             table_params, gripper_params])
             
             run_yplus.click(self.on_click_run_yplus,
-                    inputs=[],
-                    outputs=[form_video, form_3dgs_pred])
+                    inputs=[preprocess_metadata, state, params,
+                             table_params, gripper_params],
+                    outputs=[form_video, form_3dgs_pred, 
+                             preprocess_metadata, state, params,
+                             table_params, gripper_params])
             
             run_yminus.click(self.on_click_run_yminus,
-                    inputs=[],
-                    outputs=[form_video, form_3dgs_pred])
+                    inputs=[preprocess_metadata, state, params,
+                             table_params, gripper_params],
+                    outputs=[form_video, form_3dgs_pred, 
+                             preprocess_metadata, state, params,
+                             table_params, gripper_params])
             
             run_zplus.click(self.on_click_run_zplus,
-                    inputs=[],
-                    outputs=[form_video, form_3dgs_pred])
+                    inputs=[preprocess_metadata, state, params,
+                             table_params, gripper_params],
+                    outputs=[form_video, form_3dgs_pred, 
+                             preprocess_metadata, state, params,
+                             table_params, gripper_params])
 
             run_zminus.click(self.on_click_run_zminus,
-                    inputs=[],
-                    outputs=[form_video, form_3dgs_pred])
+                    inputs=[preprocess_metadata, state, params,
+                             table_params, gripper_params],
+                    outputs=[form_video, form_3dgs_pred, 
+                             preprocess_metadata, state, params,
+                             table_params, gripper_params])
 
             # Set up callbacks
             # run_reset_2.click(self.reset_2,