Spaces:
Running on Zero
Running on Zero
multimodalart HF Staff
Learn2Splat ZeroGPU demo: app, vendored optgs, prebuilt sm_120 wheels
9196862 verified | """Core `Gaussians` dataclass shared across the whole pipeline. | |
| `Gaussians` holds one scene's primitives — means, SH harmonics, opacities, scales, rotations — plus the | |
| optional per-Gaussian fields the optimizer attaches during refinement (gradients, deltas, visibility, | |
| optimization state) and a `nr_valid` count for padded scenes. Helpers (`to`, `clone`, `__getitem__`) | |
| move it across devices and batch dimensions. Initializers produce it, the optimizer refines it, and the | |
| decoder renders it. | |
| """ | |
| from dataclasses import dataclass, fields | |
| from jaxtyping import Float, Bool, Int64, BFloat16 | |
| from torch import Tensor | |
| class Gaussians: | |
| means: Float[Tensor, "batch gaussian dim"] | |
| harmonics: Float[Tensor, "batch gaussian 3 d_sh"] | |
| opacities: Float[Tensor, "batch gaussian"] | |
| scales: Float[Tensor, "batch gaussian 3"] | |
| rotations_unnorm: Float[Tensor, "batch gaussian 4"] | |
| rotations: Float[Tensor, "batch gaussian 4"] | None = None | |
| covariances: Float[Tensor, "batch gaussian dim dim"] | None = None | |
| probabilities: Float[Tensor, "batch gaussian distr"] | None = None | |
| # mask: Bool[Tensor, "batch gaussian"] | None = None | |
| sel: Int64[Tensor, "valid_gaussian_1"] | None = None | |
| filter_3D: Float[Tensor, "batch gaussian"] | None = None | |
| gradients: Float[Tensor, "batch valid_gaussian_1 total_dim"] | BFloat16[Tensor, "batch valid_gaussian_1 total_dim"] | None = None | |
| norm_gradients: Float[Tensor, "batch valid_gaussian_1 total_dim"] | BFloat16[Tensor, "batch valid_gaussian_1 total_dim"] | None = None | |
| deltas: Float[Tensor, "batch valid_gaussian_2 d_delta"] | BFloat16[Tensor, "batch valid_gaussian_2 d_delta"] | None = None # In case of predicting scale and mag, the raw deltas are 2*total_dim, cannot use SGD loss directly | |
| visibility: Float[Tensor, "batch gaussian"] | None = None # visibility information at the end of the current batch for pruning | |
| visibility_aggregator: Float[Tensor, "batch gaussian"] | None = None # aggregates visibility over epoch | |
| stores_activated: bool = True # whether scales and opacities are stored in activated form | |
| nr_valid: int = -1 # the number of valid gaussians (without padding) | |
| EXCLUDED_FROM_MASKING = {"sel", "stores_activated", "deltas", "gradients", "norm_gradients", "valid_gaussians"} # deltas are predicted to non masked values | |
| def to(self, device=None, dtype=None) -> "Gaussians": | |
| """ Move all tensors to the specified device or dtype. """ | |
| def to_with_none(tensor): | |
| if isinstance(tensor, bool): | |
| return tensor | |
| elif isinstance(tensor, int): | |
| return tensor | |
| return tensor.to(device=device, dtype=dtype) if tensor is not None else None | |
| new_tensors = {field.name: to_with_none(getattr(self, field.name)) for field in fields(self)} | |
| return Gaussians(**new_tensors) | |
| def clone(self) -> "Gaussians": | |
| """ Clone all tensors. """ | |
| # handle None and bool fields | |
| new_tensors = {} | |
| for field in fields(self): | |
| tensor = getattr(self, field.name) | |
| if isinstance(tensor, bool): | |
| new_tensors[field.name] = tensor | |
| elif isinstance(tensor, int): | |
| new_tensors[field.name] = tensor | |
| elif tensor is not None: | |
| new_tensors[field.name] = tensor.clone() | |
| else: | |
| new_tensors[field.name] = None | |
| return Gaussians(**new_tensors) | |
| # Override __getitem__ to support indexing | |
| def __getitem__(self, idx) -> "Gaussians": | |
| new_tensors = {} | |
| for field in fields(self): | |
| tensor = getattr(self, field.name) | |
| if isinstance(tensor, bool): | |
| new_tensors[field.name] = tensor | |
| elif isinstance(tensor, int): | |
| new_tensors[field.name] = tensor | |
| elif tensor is not None and field.name not in self.EXCLUDED_FROM_MASKING: | |
| new_tensors[field.name] = tensor[idx] | |
| else: | |
| new_tensors[field.name] = None | |
| return Gaussians(**new_tensors) | |
| def select_valid(self) -> "Gaussians": | |
| """Return the subset of valid (non-padding) gaussians selected by `sel`. | |
| Returns self unchanged when `sel` is None. | |
| """ | |
| if self.sel is None: | |
| return self | |
| return self[:, self.sel] | |
| def sample_subset(self, sampled_indices) -> "Gaussians": | |
| """ Randomly sample a subset of gaussians. """ | |
| total_gaussians = self.means.shape[1] | |
| sample_num = len(sampled_indices) | |
| new_tensors = {} | |
| for field in fields(self): | |
| tensor = getattr(self, field.name) | |
| if tensor is not None: | |
| if isinstance(tensor, bool): | |
| new_tensors[field.name] = tensor | |
| elif isinstance(tensor, int): | |
| new_tensors[field.name] = tensor | |
| else: | |
| new_tensors[field.name] = tensor[:, sampled_indices] | |
| else: | |
| new_tensors[field.name] = None | |
| print(f"Sampled {sample_num} / {total_gaussians} gaussians.") | |
| return Gaussians(**new_tensors) | |
| def __len__(self): | |
| return self.means.shape[1] | |
| def update_object_by_curr_mask(self, **new_values) -> "Gaussians": | |
| """ Update certain element using the current mask. """ | |
| sel = self.sel | |
| new_tensors = {} | |
| for field in fields(self): | |
| tensor = getattr(self, field.name) # [B, G, ...] | |
| if tensor is not None: | |
| if field.name in new_values: | |
| new_value = new_values[field.name] # [B, G_valid, ...] | |
| if sel is None or new_value is None or field.name in self.EXCLUDED_FROM_MASKING: | |
| tensor = new_value | |
| else: | |
| tensor = tensor.clone() | |
| tensor[:, sel, ...] = new_value | |
| new_tensors[field.name] = tensor | |
| else: | |
| if field.name in new_values: | |
| if field.name in ["deltas", "gradients", "norm_gradients"]: | |
| # Special case: allow updating deltas even if it is None | |
| new_tensors[field.name] = new_values[field.name] | |
| continue | |
| assert new_values[field.name] is None, f"Cannot update a None field! {field.name}, got {new_values[field.name]}" | |
| new_tensors[field.name] = None | |
| return Gaussians(**new_tensors) | |