fixed version bug

app.py

@@ -26,9 +26,9 @@ def main():
 
     cfg = OmegaConf.load(model_cfg_path)
     model = GaussianPredictor(cfg)
-    device = torch.device("cuda:0")
-    model.to(device)
+    device = torch.device(device)
     model.load_model(model_path)
+    model.to(device)
 
     pad_border_fn = TT.Pad((cfg.dataset.pad_border_aug, cfg.dataset.pad_border_aug))
     to_tensor = TT.ToTensor()
@@ -150,7 +150,7 @@ def main():
         )
 
     demo.queue(max_size=1)
-    demo.launch(share=True)
+    demo.launch()
 
 
 if __name__ == "__main__":

flash3d/networks/gaussian_predictor.py

@@ -235,7 +235,7 @@ class GaussianPredictor(nn.Module):
             for ckpt in ckpts[num_ckpts:]:
                 ckpt.unlink()
 
-    def load_model(self, weights_path, optimizer=None):
+    def load_model(self, weights_path, optimizer=None, device='cpu'):
         """Load model(s) from disk
         """
         weights_path = Path(weights_path)
@@ -246,7 +246,7 @@ class GaussianPredictor(nn.Module):
             return
 
         logging.info(f"Loading weights from {weights_path}...")
-        state_dict = torch.load(weights_path)
+        state_dict = torch.load(weights_path, map_location=torch.device(device))
         if "version" in state_dict and state_dict["version"] == "1.0":
             new_dict = {}
             for k, v in state_dict["model"].items():

flash3d/unidepth/utils/geometric.py

@@ -48,7 +48,6 @@ def generate_rays(
     return ray_directions, angles
 
 
-@torch.jit.script
 def spherical_zbuffer_to_euclidean(spherical_tensor: torch.Tensor) -> torch.Tensor:
     theta = spherical_tensor[..., 0]  # Extract polar angle
     phi = spherical_tensor[..., 1]  # Extract azimuthal angle
@@ -68,7 +67,6 @@ def spherical_zbuffer_to_euclidean(spherical_tensor: torch.Tensor) -> torch.Tens
     return euclidean_tensor
 
 
-@torch.jit.script
 def spherical_to_euclidean(spherical_tensor: torch.Tensor) -> torch.Tensor:
     theta = spherical_tensor[..., 0]  # Extract polar angle
     phi = spherical_tensor[..., 1]  # Extract azimuthal angle
@@ -84,7 +82,6 @@ def spherical_to_euclidean(spherical_tensor: torch.Tensor) -> torch.Tensor:
     return euclidean_tensor
 
 
-@torch.jit.script
 def euclidean_to_spherical(spherical_tensor: torch.Tensor) -> torch.Tensor:
     x = spherical_tensor[..., 0]  # Extract polar angle
     y = spherical_tensor[..., 1]  # Extract azimuthal angle
@@ -100,7 +97,6 @@ def euclidean_to_spherical(spherical_tensor: torch.Tensor) -> torch.Tensor:
     return euclidean_tensor
 
 
-@torch.jit.script
 def euclidean_to_spherical_zbuffer(euclidean_tensor: torch.Tensor) -> torch.Tensor:
     pitch = torch.asin(euclidean_tensor[..., 1])
     yaw = torch.atan2(euclidean_tensor[..., 0], euclidean_tensor[..., -1])
@@ -109,7 +105,6 @@ def euclidean_to_spherical_zbuffer(euclidean_tensor: torch.Tensor) -> torch.Tens
     return euclidean_tensor
 
 
-@torch.jit.script
 def unproject_points(
     depth: torch.Tensor, camera_intrinsics: torch.Tensor
 ) -> torch.Tensor:
@@ -152,7 +147,6 @@ def unproject_points(
     return unprojected_points
 
 
-@torch.jit.script
 def project_points(
     points_3d: torch.Tensor,
     intrinsic_matrix: torch.Tensor,
@@ -200,7 +194,6 @@ def project_points(
     return mean_depth_maps.reshape(-1, 1, *image_shape)  # (B, 1, H, W)
 
 
-@torch.jit.script
 def downsample(data: torch.Tensor, downsample_factor: int = 2):
     N, _, H, W = data.shape
     data = data.view(
@@ -220,7 +213,6 @@ def downsample(data: torch.Tensor, downsample_factor: int = 2):
     return data
 
 
-@torch.jit.script
 def flat_interpolate(
     flat_tensor: torch.Tensor,
     old: Tuple[int, int],

pre-requirements.txt

@@ -1,5 +1,5 @@
 --extra-index-url https://download.pytorch.org/whl/cu118
-torch==2.2.2
+torch
 torchvision
 torchaudio
 xformers==0.0.25.post1

requirements.txt

@@ -13,4 +13,5 @@ plyfile
 omegaconf
 jaxtyping
 gradio
-spaces
+spaces
+opencv-python