update

asset3d_gen/models/gs_model.py

@@ -137,6 +137,7 @@ class GaussianBase:
         cls,
         path: str,
         gamma: float = 1.0,
+        device: str = "cuda",
     ) -> "GaussianBase":
         plydata = PlyData.read(path)
         xyz = torch.stack(
@@ -241,6 +242,7 @@ class GaussianBase:
             _quats=rots,
             _features_dc=features_dc,
             _features_rest=features_extra,
+            device=device,
         )
 
     def save_to_ply(
@@ -454,8 +456,9 @@ class GaussianOperator(GaussianBase):
         real_height: float = None,
         instance_pose: np.ndarray = None,
         sh_degree: int = 0,
+        device: str = "cuda",
     ) -> None:
-        gs_model = GaussianOperator.load_from_ply(in_ply, sh_degree)
+        gs_model = GaussianOperator.load_from_ply(in_ply, sh_degree, device=device)
 
         if instance_pose is not None:
             gs_model = gs_model.get_gaussians(instance_pose=instance_pose)

common.py

@@ -83,16 +83,16 @@ DELIGHT = DelightingModel()
 IMAGESR_MODEL = ImageRealESRGAN(outscale=4)
 
 
-def inverse_softplus(x):
-    return x + torch.log(-torch.expm1(-x))
+def patched_setup_functions(self):
+    def inverse_softplus(x):
+        return x + torch.log(-torch.expm1(-x))
 
-def build_covariance_from_scaling_rotation(scaling, scaling_modifier, rotation):
-    L = build_scaling_rotation(scaling_modifier * scaling, rotation)
-    actual_covariance = L @ L.transpose(1, 2)
-    symm = strip_symmetric(actual_covariance)
-    return symm
+    def build_covariance_from_scaling_rotation(scaling, scaling_modifier, rotation):
+        L = build_scaling_rotation(scaling_modifier * scaling, rotation)
+        actual_covariance = L @ L.transpose(1, 2)
+        symm = strip_symmetric(actual_covariance)
+        return symm
 
-def patched_setup_functions(self):
     if self.scaling_activation_type == "exp":
         self.scaling_activation = torch.exp
         self.inverse_scaling_activation = torch.log
@@ -514,6 +514,7 @@ def extract_3d_representations_v2(
         in_ply=gs_path,
         out_ply=aligned_gs_path,
         instance_pose=pose,
+        device="cpu",
     )
     print("debug5")
     color_path = os.path.join(user_dir, "color.png")
@@ -601,6 +602,7 @@ def extract_urdf(
         in_ply=gs_path,
         out_ply=out_gs,
         real_height=real_height,
+        device="cpu",
     )
 
     # Quality check and update .urdf file.