Add lazy loading for cuda libs

demo/gs_train.py

@@ -6,28 +6,10 @@ import uuid
 from tqdm.auto import tqdm
 import gradio as gr
 import importlib.util
+from dataclasses import dataclass, field
 
-# Add the path to the gaussian-splatting repository
-gaussian_splatting_path = 'wild-gaussian-splatting/gaussian-splatting/'
-sys.path.append(gaussian_splatting_path)
-
-# Import necessary modules from the gaussian-splatting directory
-from utils.loss_utils import l1_loss, ssim
-from diff_gaussian_rasterization import GaussianRasterizationSettings, GaussianRasterizer
-from scene import Scene, GaussianModel
-from utils.general_utils import safe_state
-from utils.image_utils import psnr
-
-# Dynamically import the train module from the gaussian-splatting directory
-train_spec = importlib.util.spec_from_file_location("gaussian_splatting_train", os.path.join(gaussian_splatting_path, "train.py"))
-gaussian_splatting_train = importlib.util.module_from_spec(train_spec)
-train_spec.loader.exec_module(gaussian_splatting_train)
-
-# Import the necessary functions from the dynamically loaded module
-prepare_output_and_logger = gaussian_splatting_train.prepare_output_and_logger
-training_report = gaussian_splatting_train.training_report
+import spaces
 
-from dataclasses import dataclass, field
 
 @dataclass
 class PipelineParams:
@@ -78,6 +60,7 @@ class TrainingArgs:
     checkpoint_iterations: list[int] = field(default_factory=lambda: [7_000, 15_000, 30_000])
     start_checkpoint: str = None
 
+@spaces.GPU(duration=600)
 def train(
     data_source_path, sh_degree, model_path, images, resolution, white_background, data_device, eval,
     convert_SHs_python, compute_cov3D_python, debug,
@@ -86,6 +69,29 @@ def train(
     percent_dense, lambda_dssim, densification_interval, opacity_reset_interval,
     densify_from_iter, densify_until_iter, densify_grad_threshold, random_background
 ):
+
+    # Add the path to the gaussian-splatting repository
+    if 'GaussianRasterizer' not in globals():
+        gaussian_splatting_path = 'wild-gaussian-splatting/gaussian-splatting/'
+        sys.path.append(gaussian_splatting_path)
+
+        # Import necessary modules from the gaussian-splatting directory
+        from utils.loss_utils import l1_loss, ssim
+        from diff_gaussian_rasterization import GaussianRasterizationSettings, GaussianRasterizer
+        from scene import Scene, GaussianModel
+        from utils.general_utils import safe_state
+        from utils.image_utils import psnr
+        from utils.graphics_utils import focal2fov, fov2focal, getProjectionMatrix
+
+        # Dynamically import the train module from the gaussian-splatting directory
+        train_spec = importlib.util.spec_from_file_location("gaussian_splatting_train", os.path.join(gaussian_splatting_path, "train.py"))
+        gaussian_splatting_train = importlib.util.module_from_spec(train_spec)
+        train_spec.loader.exec_module(gaussian_splatting_train)
+
+        # Import the necessary functions from the dynamically loaded module
+        prepare_output_and_logger = gaussian_splatting_train.prepare_output_and_logger
+        training_report = gaussian_splatting_train.training_report
+
     print(data_source_path)
     # Create instances of the parameter dataclasses
     dataset = ModelParams(
@@ -221,7 +227,6 @@ def train(
 
 
     from os import makedirs
-    from utils.graphics_utils import focal2fov, fov2focal, getProjectionMatrix
     import torchvision
     import subprocess
 

demo/mast3r_demo.py

@@ -29,18 +29,6 @@ from dust3r.utils.device import to_numpy
 from dust3r.viz import add_scene_cam, CAM_COLORS, OPENGL, pts3d_to_trimesh, cat_meshes
 from dust3r.demo import get_args_parser as dust3r_get_args_parser
 
-sys.path.append(os.path.join(os.path.dirname(__file__), '../wild-gaussian-splatting/gaussian-splatting'))
-sys.path.append(os.path.join(os.path.dirname(__file__), '../wild-gaussian-splatting/src'))
-# from colmap_dataset_utils import (
-#     inv,
-#     init_filestructure,
-#     save_images_masks,
-#     save_cameras,
-#     save_imagestxt,
-#     save_pointcloud,
-#     save_pointcloud_with_normals
-# )
-
 import matplotlib.pyplot as pl
 
 import torch
@@ -151,27 +139,40 @@ def get_3D_model_from_scene(silent, scene_state, min_conf_thr=2, as_pointcloud=F
     return _convert_scene_output_to_glb(outfile, rgbimg, pts3d, msk, focals, cams2world, as_pointcloud=as_pointcloud,
                                         transparent_cams=transparent_cams, cam_size=cam_size, silent=silent)
 
-# def save_colmap_scene(scene, save_dir, min_conf_thr=2, clean_depth=False):
-#     cam2world = scene.get_im_poses().detach().cpu().numpy()
-#     world2cam = inv(cam2world) #
-#     principal_points = scene.get_principal_points().detach().cpu().numpy()
-#     focals = scene.get_focals().detach().cpu().numpy()[..., None]
-#     imgs = np.array(scene.imgs)
-
-#     pts3d, _, confs = scene.get_dense_pts3d(clean_depth=clean_depth)
-#     pts3d = [i.detach().reshape(imgs[0].shape) for i in pts3d] #
-
-#     masks = to_numpy([c > min_conf_thr for c in to_numpy(confs)])
-
-#     # move
-#     mask_images = True
-
-#     save_path, images_path, masks_path, sparse_path = init_filestructure(save_dir)
-#     save_images_masks(imgs, masks, images_path, masks_path, mask_images)
-#     save_cameras(focals, principal_points, sparse_path, imgs_shape=imgs.shape)
-#     save_imagestxt(world2cam, sparse_path)
-#     save_pointcloud_with_normals(imgs, pts3d, masks, sparse_path)
-#     return save_path
+def save_colmap_scene(scene, save_dir, min_conf_thr=2, clean_depth=False):
+    if 'save_pointcloud_with_normals' not in globals():
+        sys.path.append(os.path.join(os.path.dirname(__file__), '../wild-gaussian-splatting/gaussian-splatting'))
+        sys.path.append(os.path.join(os.path.dirname(__file__), '../wild-gaussian-splatting/src'))
+        from colmap_dataset_utils import (
+            inv,
+            init_filestructure,
+            save_images_masks,
+            save_cameras,
+            save_imagestxt,
+            save_pointcloud,
+            save_pointcloud_with_normals
+        )
+
+    cam2world = scene.get_im_poses().detach().cpu().numpy()
+    world2cam = inv(cam2world) #
+    principal_points = scene.get_principal_points().detach().cpu().numpy()
+    focals = scene.get_focals().detach().cpu().numpy()[..., None]
+    imgs = np.array(scene.imgs)
+
+    pts3d, _, confs = scene.get_dense_pts3d(clean_depth=clean_depth)
+    pts3d = [i.detach().reshape(imgs[0].shape) for i in pts3d] #
+
+    masks = to_numpy([c > min_conf_thr for c in to_numpy(confs)])
+
+    # move
+    mask_images = True
+
+    save_path, images_path, masks_path, sparse_path = init_filestructure(save_dir)
+    save_images_masks(imgs, masks, images_path, masks_path, mask_images)
+    save_cameras(focals, principal_points, sparse_path, imgs_shape=imgs.shape)
+    save_imagestxt(world2cam, sparse_path)
+    save_pointcloud_with_normals(imgs, pts3d, masks, sparse_path)
+    return save_path
 
 @spaces.GPU(duration=300)
 def get_reconstructed_scene(outdir, model, device, silent, image_size, current_scene_state,
@@ -222,7 +223,7 @@ def get_reconstructed_scene(outdir, model, device, silent, image_size, current_s
     os.makedirs(base_colmapdata_dir, exist_ok=True)
     colmap_data_dir = get_next_dir(base_colmapdata_dir)
     # 
-    # save_colmap_scene(scene, colmap_data_dir, min_conf_thr, clean_depth)
+    save_colmap_scene(scene, colmap_data_dir, min_conf_thr, clean_depth)
     
     if current_scene_state is not None and \
        current_scene_state.outfile_name is not None: