use gsplat

app.py

@@ -192,6 +192,12 @@ def run_depth_estimation(
 ) -> Image.Image:
     rgb_image = image_prompts["image"].convert("RGB")
 
+    import torch
+    print(torch.__version__)
+    print(torch.version.cuda)
+    print(torch.__file__)
+    print(torch.cuda.is_available())
+
     from threeDFixer.datasets.utils import (
         normalize_vertices,
         project2ply
@@ -325,12 +331,6 @@ def run_generation(
     t_rescale: float = 3.0,
     work_space: dict = None,
 ):
-    
-    try:
-        import diff_gaussian_rasterization
-    except ModuleNotFoundError:
-        install_mipsplatting()
-    
     if work_space is None:
         raise gr.Error("Please run step 1 and step 2 first.")
     required_keys = ["dir", "depth_mask", "depth", "K", "c2w", "trans", "scale"]

requirements.txt

@@ -40,3 +40,4 @@ pydantic==2.10.6
 kaolin==0.18.0
 flash-attn==2.8.3+pt2.8.0cu129
 nvdiffrast==0.4.0+253ac4fpt2.8.0cu129
+git+https://github.com/nerfstudio-project/gsplat.git

threeDFixer/renderers/gaussian_render.py

@@ -47,97 +47,71 @@ def intrinsics_to_projection(
     return ret
 
 
-def render(viewpoint_camera, pc : Gaussian, pipe, bg_color : torch.Tensor, scaling_modifier = 1.0, override_color = None):
-    """
-    Render the scene. 
-    
-    Background tensor (bg_color) must be on GPU!
-    """
+def render(viewpoint_camera, pc, pipe, bg_color: torch.Tensor, scaling_modifier=1.0, override_color=None):
     # lazy import
-    if 'GaussianRasterizer' not in globals():
-        from diff_gaussian_rasterization import GaussianRasterizer, GaussianRasterizationSettings
-    
-    # Create zero tensor. We will use it to make pytorch return gradients of the 2D (screen-space) means
-    screenspace_points = torch.zeros_like(pc.get_xyz, dtype=pc.get_xyz.dtype, requires_grad=True, device="cuda") + 0
-    try:
-        screenspace_points.retain_grad()
-    except:
-        pass
-    # Set up rasterization configuration
+    if "rasterization" not in globals():
+        from gsplat import rasterization
+
     tanfovx = math.tan(viewpoint_camera.FoVx * 0.5)
     tanfovy = math.tan(viewpoint_camera.FoVy * 0.5)
-    
-    kernel_size = pipe.kernel_size
-    subpixel_offset = torch.zeros((int(viewpoint_camera.image_height), int(viewpoint_camera.image_width), 2), dtype=torch.float32, device="cuda")
-
-    raster_settings = GaussianRasterizationSettings(
-        image_height=int(viewpoint_camera.image_height),
-        image_width=int(viewpoint_camera.image_width),
-        tanfovx=tanfovx,
-        tanfovy=tanfovy,
-        kernel_size=kernel_size,
-        subpixel_offset=subpixel_offset,
-        bg=bg_color,
-        scale_modifier=scaling_modifier,
-        viewmatrix=viewpoint_camera.world_view_transform,
-        projmatrix=viewpoint_camera.full_proj_transform,
-        sh_degree=pc.active_sh_degree,
-        campos=viewpoint_camera.camera_center,
-        prefiltered=False,
-        debug=pipe.debug
+
+    focal_length_x = viewpoint_camera.image_width / (2 * tanfovx)
+    focal_length_y = viewpoint_camera.image_height / (2 * tanfovy)
+
+    K = torch.tensor(
+        [
+            [focal_length_x, 0, viewpoint_camera.image_width / 2.0],
+            [0, focal_length_y, viewpoint_camera.image_height / 2.0],
+            [0, 0, 1],
+        ],
+        device=pc.get_xyz.device,
+        dtype=torch.float32,
     )
-    
-    rasterizer = GaussianRasterizer(raster_settings=raster_settings)
 
     means3D = pc.get_xyz
-    means2D = screenspace_points
     opacity = pc.get_opacity
+    scales = pc.get_scaling * scaling_modifier
+    rotations = pc.get_rotation
 
-    # If precomputed 3d covariance is provided, use it. If not, then it will be computed from
-    # scaling / rotation by the rasterizer.
-    scales = None
-    rotations = None
-    cov3D_precomp = None
-    if pipe.compute_cov3D_python:
-        cov3D_precomp = pc.get_covariance(scaling_modifier)
-    else:
-        scales = pc.get_scaling
-        rotations = pc.get_rotation
-
-    # If precomputed colors are provided, use them. Otherwise, if it is desired to precompute colors
-    # from SHs in Python, do it. If not, then SH -> RGB conversion will be done by rasterizer.
-    shs = None
-    colors_precomp = None
-    if override_color is None:
-        if pipe.convert_SHs_python:
-            shs_view = pc.get_features.transpose(1, 2).view(-1, 3, (pc.max_sh_degree+1)**2)
-            dir_pp = (pc.get_xyz - viewpoint_camera.camera_center.repeat(pc.get_features.shape[0], 1))
-            dir_pp_normalized = dir_pp/dir_pp.norm(dim=1, keepdim=True)
-            sh2rgb = eval_sh(pc.active_sh_degree, shs_view, dir_pp_normalized)
-            colors_precomp = torch.clamp_min(sh2rgb + 0.5, 0.0)
-        else:
-            shs = pc.get_features
+    if override_color is not None:
+        colors = override_color          # [N, 3]
+        sh_degree = None
     else:
-        colors_precomp = override_color
-
-    # Rasterize visible Gaussians to image, obtain their radii (on screen). 
-    rendered_image, radii = rasterizer(
-        means3D = means3D,
-        means2D = means2D,
-        shs = shs,
-        colors_precomp = colors_precomp,
-        opacities = opacity,
-        scales = scales,
-        rotations = rotations,
-        cov3D_precomp = cov3D_precomp
+        colors = pc.get_features         # [N, K, 3]
+        sh_degree = pc.active_sh_degree
+
+    viewmat = viewpoint_camera.world_view_transform.transpose(0, 1)
+
+    render_colors, render_alphas, info = rasterization(
+        means=means3D,                        # [N, 3]
+        quats=rotations,                      # [N, 4]
+        scales=scales,                        # [N, 3]
+        opacities=opacity.squeeze(-1),        # [N]
+        colors=colors,
+        viewmats=viewmat[None],               # [1, 4, 4]
+        Ks=K[None],                           # [1, 3, 3]
+        backgrounds=bg_color[None],
+        width=int(viewpoint_camera.image_width),
+        height=int(viewpoint_camera.image_height),
+        packed=False,
+        sh_degree=sh_degree,
+        rasterize_mode='antialiased'
     )
 
-    # Those Gaussians that were frustum culled or had a radius of 0 were not visible.
-    # They will be excluded from value updates used in the splitting criteria.
-    return edict({"render": rendered_image,
-            "viewspace_points": screenspace_points,
-            "visibility_filter" : radii > 0,
-            "radii": radii})
+    rendered_image = render_colors[0].permute(2, 0, 1)
+    radii = info["radii"].squeeze(0)
+
+    try:
+        info["means2d"].retain_grad()
+    except Exception:
+        pass
+
+    return edict({
+        "render": rendered_image,
+        "viewspace_points": info["means2d"],
+        "visibility_filter": radii > 0,
+        "radii": radii,
+    })
 
 
 class GaussianRenderer: