Spaces:
Running on Zero
Running on Zero
multimodalart HF Staff
Learn2Splat ZeroGPU demo: app, vendored optgs, prebuilt sm_120 wheels
9196862 verified | """Shared orchestration for the CUDA-rasterizer decoders (inria, fastgs). | |
| Both backends share identical forward/depth plumbing — flattening B×V cameras, | |
| handling Gaussians vs GaussiansModule, optional view chunking, and reshaping the | |
| rasterizer outputs back to (B, V). Only the rasterizer call itself differs, so | |
| that is the single overridable hook (``_raster`` / ``_raster_depth``). | |
| This module imports no rasterizer backend, so importing it never requires | |
| ``diff_gaussian_rasterization`` or ``diff_gaussian_rasterization_fastgs``. | |
| """ | |
| from typing import TypeVar | |
| import torch | |
| from einops import rearrange, repeat | |
| from jaxtyping import Float | |
| from torch import Tensor | |
| from tqdm import tqdm | |
| from ...dataset import DatasetCfg | |
| from ...scene_trainer.gaussian_module import GaussiansModule | |
| from ..types import Gaussians | |
| from .decoder import Decoder, DecoderOutput, DepthRenderingMode | |
| T = TypeVar("T") | |
| class SplattingCUDADecoder(Decoder[T]): | |
| """Base for the inria/fastgs splatting decoders. Subclasses implement only the two | |
| backend hooks below; everything else (camera flattening, output reshaping) is shared.""" | |
| background_color: Float[Tensor, "3"] | |
| def __init__(self, cfg: T, dataset_cfg: DatasetCfg) -> None: | |
| super().__init__(cfg, dataset_cfg) | |
| self.register_buffer( | |
| "background_color", | |
| torch.tensor(dataset_cfg.background_color, dtype=torch.float32), | |
| persistent=False, | |
| ) | |
| # --- backend hooks ----------------------------------------------------- | |
| # Both receive flat (B*V) camera tensors and return flat outputs. | |
| def _raster(self, ext, intr, near, far, image_shape, bg, means, covars, shs, opacities, | |
| scales, rotations_wxyz, means2d_out, means2d_abs_out=None): | |
| """Renders, writing screen-space means into ``means2d_out`` ([bv,N,2]) so its gradient | |
| is reachable via autograd.grad. Returns (images [bv,3,H,W], radii [bv,N]). | |
| ``means2d_abs_out`` is the optional FastGS Abs-GS leaf (cols [2:] of its screen tensor); | |
| backends that do not produce it ignore the argument.""" | |
| raise NotImplementedError | |
| def _raster_depth(self, ext, intr, near, far, image_shape, means, covars, opacities, mode): | |
| """Returns depth [bv, H, W].""" | |
| raise NotImplementedError | |
| def _produces_abs_grad(self) -> bool: | |
| """Whether the backend exposes the FastGS abs-gradient (cols [2:]) via ``means2d_abs_out``. | |
| Only the FastGS decoder overrides this to True; gsplat/inria stay False.""" | |
| return False | |
| # --- shared Gaussian tensor prep -------------------------------------- | |
| def _prepare_flat_gaussians(self, gaussians: Gaussians | GaussiansModule, b: int, v: int): | |
| """Flatten Gaussian params to (B*V) rasterizer layout. Returns | |
| (means, shs, opacities, scales, rotations_wxyz, covars); scales/rotations are None when | |
| ``use_covariances`` (covars supplied instead) and vice-versa. Shared by forward and the | |
| FastGS metric-counts render so they build identical inputs.""" | |
| bv = b * v | |
| scales = rotations_wxyz = covars = None | |
| if isinstance(gaussians, GaussiansModule): | |
| means = repeat(gaussians.means, "g xyz -> bv g xyz", bv=bv) | |
| shs = repeat(gaussians.harmonics, "g c d -> bv g c d", bv=bv) | |
| opacities = repeat(gaussians.opacities, "g -> bv g", bv=bv) | |
| if self.cfg.use_covariances: | |
| covars = repeat(gaussians.covariances, "g i j -> bv g i j", bv=bv) | |
| else: | |
| scales = repeat(gaussians.scales, "g d -> bv g d", bv=bv) | |
| # gaussians.rotations is xyzw post-normalization; the rasterizer wants wxyz. | |
| rotations_wxyz = repeat(gaussians.rotations[:, [3, 0, 1, 2]], "g d -> bv g d", bv=bv) | |
| elif isinstance(gaussians, Gaussians): | |
| means = repeat(gaussians.means, "b g xyz -> (b v) g xyz", v=v) | |
| shs = repeat(gaussians.harmonics, "b g c d -> (b v) g c d", v=v) | |
| opacities = repeat(gaussians.opacities, "b g -> (b v) g", v=v) | |
| if self.cfg.use_covariances: | |
| if gaussians.covariances is None: | |
| raise ValueError("use_covariances=true but gaussians.covariances is None.") | |
| covars = repeat(gaussians.covariances, "b g i j -> (b v) g i j", v=v) | |
| else: | |
| _scales = gaussians.scales if gaussians.stores_activated else torch.exp(gaussians.scales) | |
| scales = repeat(_scales, "b g d -> (b v) g d", v=v) | |
| # Normalize rotations_unnorm here (the always-present grad leaf) rather than | |
| # using gaussians.rotations: the learned optimizer precomputes the latter under | |
| # torch.no_grad(), so the screen-space loss must reach rotations_unnorm directly | |
| # (else autograd.grad sees it as unused). Numerically == gaussians.rotations; | |
| # xyzw -> wxyz for the rasterizer. Mirrors the gsplat decoder. | |
| rot = torch.nn.functional.normalize(gaussians.rotations_unnorm, dim=-1) | |
| rotations_wxyz = repeat(rot[..., [3, 0, 1, 2]], "b g d -> (b v) g d", v=v) | |
| if not gaussians.stores_activated: | |
| opacities = torch.sigmoid(opacities) | |
| else: | |
| raise ValueError(f"Unknown gaussians type: {type(gaussians)}") | |
| return means, shs, opacities, scales, rotations_wxyz, covars | |
| # --- shared forward ---------------------------------------------------- | |
| def forward( | |
| self, | |
| gaussians: Gaussians | GaussiansModule, | |
| extrinsics: Float[Tensor, "batch view 4 4"], | |
| intrinsics: Float[Tensor, "batch view 3 3"], | |
| near: Float[Tensor, "batch view"], | |
| far: Float[Tensor, "batch view"], | |
| image_shape: tuple[int, int], | |
| depth_mode: DepthRenderingMode | None = None, | |
| return_radii: bool = False, | |
| iter_batch_size: int = -1, | |
| to_cpu: bool = False, | |
| ) -> DecoderOutput: | |
| b, v, _, _ = extrinsics.shape | |
| bv = b * v | |
| # Flatten camera params to (B*V) | |
| flat_ext = rearrange(extrinsics, "b v i j -> (b v) i j") | |
| flat_int = rearrange(intrinsics, "b v i j -> (b v) i j") | |
| flat_near = rearrange(near, "b v -> (b v)") | |
| flat_far = rearrange(far, "b v -> (b v)") | |
| flat_bg = repeat(self.background_color, "c -> (b v) c", b=b, v=v) | |
| # Prepare Gaussian tensors in flat (B*V) format | |
| means, shs, opacities, scales, rotations_wxyz, covars = self._prepare_flat_gaussians(gaussians, b, v) | |
| # Single [B, V, N, 2] screen-space-means leaf. Each view's rasterizer call consumes a | |
| # slice of it (via means2d_flat), so the 2D gradient is reachable as | |
| # torch.autograd.grad(loss, out.means2d) — uniformly with the gsplat decoder, and | |
| # without a .backward()/.grad pass. Returned as-is (no reshape, which would detach it). | |
| n_gauss = means.shape[1] | |
| means2d = torch.zeros((b, v, n_gauss, 2), dtype=means.dtype, device=means.device, | |
| requires_grad=True) | |
| means2d_flat = means2d.reshape(bv, n_gauss, 2) | |
| # FastGS only: a second leaf for the abs-gradient (cols [2:] of its screen tensor). | |
| means2d_abs = means2d_abs_flat = None | |
| if self._produces_abs_grad(): | |
| means2d_abs = torch.zeros((b, v, n_gauss, 2), dtype=means.dtype, device=means.device, | |
| requires_grad=True) | |
| means2d_abs_flat = means2d_abs.reshape(bv, n_gauss, 2) | |
| def _render_flat(s: slice): | |
| return self._raster( | |
| flat_ext[s], flat_int[s], flat_near[s], flat_far[s], image_shape, flat_bg[s], | |
| means[s], covars[s] if covars is not None else None, shs[s], opacities[s], | |
| scales[s] if scales is not None else None, | |
| rotations_wxyz[s] if rotations_wxyz is not None else None, | |
| means2d_flat[s], | |
| means2d_abs_flat[s] if means2d_abs_flat is not None else None, | |
| ) | |
| if iter_batch_size < 0: | |
| imgs, radii_flat = _render_flat(slice(None)) | |
| if to_cpu: | |
| imgs = imgs.detach().cpu() | |
| radii_flat = radii_flat.detach().cpu() | |
| else: | |
| all_imgs, all_radii = [], [] | |
| for i in tqdm(range(0, bv, iter_batch_size), desc="Rendering in batches"): | |
| s = slice(i, min(i + iter_batch_size, bv)) | |
| imgs_c, rad_c = _render_flat(s) | |
| if to_cpu: | |
| imgs_c = imgs_c.detach().cpu() | |
| rad_c = rad_c.detach().cpu() | |
| all_imgs.append(imgs_c) | |
| all_radii.append(rad_c) | |
| imgs = torch.cat(all_imgs, dim=0) | |
| radii_flat = torch.cat(all_radii, dim=0) | |
| # Reshape (B*V) → (B, V) | |
| color = rearrange(imgs, "(b v) c h w -> b v c h w", b=b, v=v) | |
| radii_bv = rearrange(radii_flat, "(b v) n -> b v n", b=b, v=v) | |
| means2d_bv = means2d.detach().cpu() if to_cpu else means2d # [B, V, N, 2] | |
| means2d_abs_bv = None | |
| if means2d_abs is not None: | |
| means2d_abs_bv = means2d_abs.detach().cpu() if to_cpu else means2d_abs # [B, V, N, 2] | |
| # Expand scalar radii [B, V, N] → [B, V, N, 2] to match gsplat interface | |
| radii_out = radii_bv.unsqueeze(-1).expand(-1, -1, -1, 2).contiguous() | |
| visibility_filter = radii_bv > 0 # [B, V, N] | |
| depth = ( | |
| self._render_depth(gaussians, extrinsics, intrinsics, near, far, image_shape, depth_mode) | |
| if depth_mode is not None | |
| else None | |
| ) | |
| return DecoderOutput( | |
| color=color, | |
| depth=depth, | |
| accumulated_alpha=None, | |
| means2d=means2d_bv, | |
| means2d_abs=means2d_abs_bv, | |
| radii=radii_out, | |
| visibility_filter=visibility_filter, | |
| ) | |
| # --- shared depth ------------------------------------------------------ | |
| def _render_depth( | |
| self, | |
| gaussians: Gaussians | GaussiansModule, | |
| extrinsics: Float[Tensor, "batch view 4 4"], | |
| intrinsics: Float[Tensor, "batch view 3 3"], | |
| near: Float[Tensor, "batch view"], | |
| far: Float[Tensor, "batch view"], | |
| image_shape: tuple[int, int], | |
| mode: DepthRenderingMode = "depth", | |
| ) -> Float[Tensor, "batch view height width"]: | |
| b, v, _, _ = extrinsics.shape | |
| if isinstance(gaussians, GaussiansModule): | |
| means = repeat(gaussians.means, "g xyz -> (b v) g xyz", b=b, v=v) | |
| covars = repeat(gaussians.covariances, "g i j -> (b v) g i j", b=b, v=v) | |
| opacities = repeat(gaussians.opacities, "g -> (b v) g", b=b, v=v) | |
| else: | |
| means = repeat(gaussians.means, "b g xyz -> (b v) g xyz", v=v) | |
| covars = repeat(gaussians.covariances, "b g i j -> (b v) g i j", v=v) | |
| opacities = repeat(gaussians.opacities, "b g -> (b v) g", v=v) | |
| if not gaussians.stores_activated: | |
| opacities = torch.sigmoid(opacities) | |
| result = self._raster_depth( | |
| rearrange(extrinsics, "b v i j -> (b v) i j"), | |
| rearrange(intrinsics, "b v i j -> (b v) i j"), | |
| rearrange(near, "b v -> (b v)"), | |
| rearrange(far, "b v -> (b v)"), | |
| image_shape, | |
| means, | |
| covars, | |
| opacities, | |
| mode, | |
| ) | |
| return rearrange(result, "(b v) h w -> b v h w", b=b, v=v) | |