update code

app.py

@@ -1,5 +1,4 @@
 import gradio as gr
-import importlib
 import sys
 import site
 from PIL import Image
@@ -10,12 +9,15 @@ import random
 import math
 import hydra
 import numpy as np
+
+# import importlib
 # import torch
 # import torch.nn as nn
 # print(torch.__version__)
 # print(torch.version.cuda)
 # import torch.backends.cudnn
 # from torch.utils.data import DataLoader
+
 import glob
 import os
 import subprocess
@@ -25,6 +27,7 @@ import copy
 import yaml
 import matplotlib.pyplot as plt
 from sklearn.neighbors import NearestNeighbors
+
 import spaces
 from spaces import zero
 zero.startup()
@@ -65,43 +68,46 @@ zero.startup()
 # importlib.invalidate_caches()
 # torch = importlib.import_module("torch")
 
-import torch
-import torch.nn as nn
-from torch.utils.data import DataLoader
-import kornia
-import warp as wp
+# import torch
+# import torch.nn as nn
+# from torch.utils.data import DataLoader
+# import kornia
+# import warp as wp
 
-print(torch.__version__)
-print(torch.version.cuda)
+# print(torch.__version__)
+# print(torch.version.cuda)
 
 import sys
 sys.path.insert(0, str(Path(__file__).parent / "src"))
 sys.path.append(str(Path(__file__).parent / "src" / "experiments"))
 
-from pgnd.sim import Friction, CacheDiffSimWithFrictionBatch, StaticsBatch, CollidersBatch
-from pgnd.material import PGNDModel
-from pgnd.utils import Logger, get_root, mkdir
-from pgnd.ffmpeg import make_video
+# from pgnd.sim import Friction, CacheDiffSimWithFrictionBatch, StaticsBatch, CollidersBatch
+# from pgnd.material import PGNDModel
+# from pgnd.utils import Logger, get_root, mkdir
+# from pgnd.ffmpeg import make_video
 
-from real_world.utils.render_utils import interpolate_motions
-from real_world.gs.helpers import setup_camera
-from real_world.gs.convert import save_to_splat, read_splat
+# from real_world.utils.render_utils import interpolate_motions
+# from real_world.gs.helpers import setup_camera
+# from real_world.gs.convert import save_to_splat, read_splat
 
-from diff_gaussian_rasterization import GaussianRasterizer
-from diff_gaussian_rasterization import GaussianRasterizationSettings as Camera
+# from diff_gaussian_rasterization import GaussianRasterizer
+# from diff_gaussian_rasterization import GaussianRasterizationSettings as Camera
 
 root = Path(__file__).parent / "src" / "experiments"
 
 
 def quat2mat(quat):
+    import kornia
     return kornia.geometry.conversions.quaternion_to_rotation_matrix(quat)
 
 
 def mat2quat(mat):
+    import kornia
     return kornia.geometry.conversions.rotation_matrix_to_quaternion(mat)
 
 
 def fps(x, enabled, n, device, random_start=False):
+    import torch
     from dgl.geometry import farthest_point_sampler
     assert torch.diff(enabled * 1.0).sum() in [0.0, -1.0]
     start_idx = random.randint(0, enabled.sum() - 1) if random_start else 0
@@ -139,10 +145,10 @@ class DynamicsVisualizer:
         cfg.sim.uniform = True
         cfg.sim.use_pv = False
 
-        device = torch.device('cuda')
+        # device = torch.device('cuda')
 
         self.cfg = cfg
-        self.device = device
+        # self.device = device
         self.k_rel = 8  # knn for relations
         self.k_wgt = 16  # knn for weights
         self.with_bg = True
@@ -155,7 +161,7 @@ class DynamicsVisualizer:
         seed = cfg.seed
         random.seed(seed)
         np.random.seed(seed)
-        torch.manual_seed(seed)
+        # torch.manual_seed(seed)
         # torch.autograd.set_detect_anomaly(True)
         # torch.backends.cudnn.benchmark = True
 
@@ -424,7 +430,7 @@ class DynamicsVisualizer:
             'global_translation': global_translation,
         }
 
-    @torch.no_grad
+    # @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()}
         w, h = self.metadata['w'], self.metadata['h']
@@ -502,7 +508,7 @@ class DynamicsVisualizer:
         sim = CacheDiffSimWithFrictionBatch(self.cfg, num_steps, 1, self.wp_device, requires_grad=True)
         self.sim = sim
 
-    @torch.no_grad
+    # @torch.no_grad
     def step(self):
         cfg = self.cfg
         batch_size = 1
@@ -897,6 +903,28 @@ 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
+        
+        from pgnd.sim import Friction, CacheDiffSimWithFrictionBatch, StaticsBatch, CollidersBatch
+        from pgnd.material import PGNDModel
+        from pgnd.utils import Logger, get_root, mkdir
+        from pgnd.ffmpeg import make_video
+
+        from real_world.utils.render_utils import interpolate_motions
+        from real_world.gs.helpers import setup_camera
+        from real_world.gs.convert import save_to_splat, read_splat
+
+        from diff_gaussian_rasterization import GaussianRasterizer
+        from diff_gaussian_rasterization import GaussianRasterizationSettings as Camera
+        
+        device = torch.device('cuda:0')
+        self.device = device
+        torch.manual_seed(self.cfg.seed)
+        
         wp.init()
         gpus = [int(gpu) for gpu in self.cfg.gpus]
         wp_devices = [wp.get_device(f'cuda:{gpu}') for gpu in gpus]
@@ -1041,7 +1069,6 @@ class DynamicsVisualizer:
     
     @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]
@@ -1053,7 +1080,6 @@ class DynamicsVisualizer:
 
     @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]
@@ -1065,7 +1091,6 @@ class DynamicsVisualizer:
     
     @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]
@@ -1077,7 +1102,6 @@ class DynamicsVisualizer:
     
     @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]
@@ -1089,7 +1113,6 @@ class DynamicsVisualizer:
     
     @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]
@@ -1101,7 +1124,6 @@ class DynamicsVisualizer:
     
     @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]