ffgs 0.1.0.dev0__py3-none-any.whl

ffgs/__init__.py

@@ -0,0 +1,39 @@
+"""Standard inference for feed-forward 3D Gaussian Splatting models.
+
+Posed images in (`Views`), Gaussians out (in the model's frame, or the caller's
+world frame with `out.to_world()`), one `GSPipeline.from_pretrained` for every
+registered model. See README.md.
+
+Depends on torch, torchvision, numpy, pillow, huggingface_hub and safetensors;
+gsplat is imported only when rendering (`pip install ffgs[render]`).
+"""
+
+import importlib.metadata as _metadata
+
+from .geometry import ModelFrame
+from .pipeline import GSOutput, GSPipeline
+from .processor import ImageFitConfig, Prepared, Processor
+from .registry import ModelSpec, get_model_spec, register_model
+from .render import render
+from .types import Cameras, Gaussians, Views
+
+try:
+    __version__ = _metadata.version("ffgs")
+except _metadata.PackageNotFoundError:  # a source tree that is not installed
+    __version__ = "0+unknown"
+
+__all__ = [
+    "Cameras",
+    "GSOutput",
+    "GSPipeline",
+    "Gaussians",
+    "ImageFitConfig",
+    "ModelFrame",
+    "ModelSpec",
+    "Prepared",
+    "Processor",
+    "Views",
+    "get_model_spec",
+    "register_model",
+    "render",
+]

ffgs/dynamic.py

@@ -0,0 +1,239 @@
+"""Model code shipped inside a model repo, named by `config.json:auto_map`.
+
+The convention of transformers / diffusers: a repo whose model is not installed
+lists its own classes,
+
+    "auto_map": {"model": "modeling_foo.FooModel",
+                 "processor": "processing_foo.FooProcessor"}
+
+and the `.py` files next to `config.json` are imported only when the caller passes
+`trust_remote_code=True`. The files are imported as one package per directory
+(Hub snapshot or local directory), so relative imports between them work.
+`export_code` writes such a directory from an installed package
+(`GSPipeline.save_pretrained(..., include_code=True)`).
+"""
+
+from __future__ import annotations
+
+import ast
+import hashlib
+import importlib
+import importlib.machinery
+import importlib.util
+import shutil
+import sys
+from pathlib import Path
+from types import ModuleType
+from typing import Any
+
+from .processor import Processor
+from .registry import ModelSpec
+
+AUTO_MAP_KEYS = ("model", "processor")
+
+# Parent package of every imported code directory.
+_PACKAGE = "ffgs_remote_code"
+
+
+def parse_auto_map(auto_map: Any) -> dict[str, tuple[str, str]]:
+    """`{"model": "module.Class", ...}` -> `{"model": ("module", "Class"), ...}`."""
+    if not isinstance(auto_map, dict) or set(auto_map) != set(AUTO_MAP_KEYS):
+        raise ValueError(
+            f"config.json:auto_map must map exactly {list(AUTO_MAP_KEYS)} to "
+            f'"module.Class" references, got {auto_map!r}'
+        )
+    parsed = {}
+    for key, reference in auto_map.items():
+        if not isinstance(reference, str) or "." not in reference:
+            raise ValueError(
+                f'auto_map["{key}"] must be "module.Class", got {reference!r}'
+            )
+        if "--" in reference:
+            raise ValueError(
+                f'auto_map["{key}"] = {reference!r} points at another repo; '
+                "ffgs only runs code from the repo being loaded, copy it there"
+            )
+        module, name = reference.rsplit(".", 1)
+        parsed[key] = (module, name)
+    return parsed
+
+
+def spec_from_auto_map(
+    model_type: str, auto_map: Any, pretrained: str | Path, **hub_kwargs: Any
+) -> ModelSpec:
+    """Import the classes `auto_map` names from `pretrained` (runs its code)."""
+    parsed = parse_auto_map(auto_map)
+    package = _import_package(_code_directory(pretrained, **hub_kwargs))
+    classes = {
+        key: _load_class(package, module, name, pretrained)
+        for key, (module, name) in parsed.items()
+    }
+    if not issubclass(classes["processor"], Processor):
+        raise TypeError(
+            f"auto_map processor {classes['processor'].__name__} of {pretrained} "
+            "does not subclass ffgs.Processor"
+        )
+    return ModelSpec(model_type, classes["model"], classes["processor"])
+
+
+def save_code(
+    model_cls: type, processor_cls: type, save_directory: Path
+) -> dict[str, str] | None:
+    """Copy the repo code `model_cls` / `processor_cls` came from (every `.py`
+    file of its directory) into `save_directory` and return the matching
+    auto_map; None for installed code."""
+    remote = [_is_remote(cls) for cls in (model_cls, processor_cls)]
+    if not any(remote):
+        return None
+    if not all(remote) or model_cls.__module__.split(".")[1] != (
+        processor_cls.__module__.split(".")[1]
+    ):
+        raise ValueError(
+            "cannot save a model and processor that come from different code "
+            "sources as one repo"
+        )
+    package = ".".join(model_cls.__module__.split(".")[:2])
+    source = Path(sys.modules[package].__path__[0])
+    # All of them, as they are downloaded: a module the model imports only when
+    # it runs is not imported yet.
+    for file in sorted(p for p in source.rglob("*.py") if "__pycache__" not in p.parts):
+        target = save_directory / file.relative_to(source)
+        if target.exists() and target.samefile(file):
+            continue
+        target.parent.mkdir(parents=True, exist_ok=True)
+        shutil.copyfile(file, target)
+    return {
+        key: f"{cls.__module__[len(package) + 1:]}.{cls.__qualname__}"
+        for key, cls in zip(AUTO_MAP_KEYS, (model_cls, processor_cls))
+    }
+
+
+def export_code(
+    model_cls: type, processor_cls: type, save_directory: Path
+) -> dict[str, str]:
+    """Copy the installed package holding `model_cls` and `processor_cls` (every
+    `.py` file under it) into `save_directory` and return the matching auto_map,
+    so the directory loads with `trust_remote_code=True` where that package is not
+    installed.
+
+    The package is the innermost one holding both classes. Its code may import
+    ffgs and third-party packages, and itself only relatively: an absolute import
+    of its own distribution, or a relative one reaching above the package, would
+    not resolve in the copy and is refused.
+    """
+    package = _common_package(model_cls, processor_cls)
+    top = package.split(".")[0]
+    if top == "ffgs":
+        raise ValueError(
+            f"{model_cls.__name__} is built into ffgs; its directories load "
+            "without shipping code"
+        )
+    source = Path(sys.modules[package].__path__[0])
+    files = sorted(p for p in source.rglob("*.py") if "__pycache__" not in p.parts)
+    problems = [
+        problem for file in files for problem in _outside_imports(file, source, top)
+    ]
+    if problems:
+        raise ValueError(
+            f"cannot ship package {package!r} as model code: it imports code "
+            "that is not in it (import within the package relatively):\n  "
+            + "\n  ".join(problems)
+        )
+    for file in files:
+        target = save_directory / file.relative_to(source)
+        target.parent.mkdir(parents=True, exist_ok=True)
+        shutil.copyfile(file, target)
+    return {
+        key: f"{cls.__module__[len(package) + 1:]}.{cls.__qualname__}"
+        for key, cls in zip(AUTO_MAP_KEYS, (model_cls, processor_cls))
+    }
+
+
+def _common_package(*classes: type) -> str:
+    paths = []
+    for cls in classes:
+        module = sys.modules[cls.__module__]
+        name = module.__name__
+        paths.append(name if hasattr(module, "__path__") else name.rpartition(".")[0])
+    common = []
+    if all(paths):
+        for names in zip(*(path.split(".") for path in paths)):
+            if len(set(names)) > 1:
+                break
+            common.append(names[0])
+    if not common:
+        raise ValueError(
+            "the model and processor classes must live in one package to be "
+            f"shipped as model code, got {[cls.__module__ for cls in classes]}"
+        )
+    return ".".join(common)
+
+
+def _outside_imports(file: Path, source: Path, top: str) -> list[str]:
+    """Imports in `file` that would not resolve once `source` is copied alone."""
+    depth = len(file.parent.relative_to(source).parts)
+    problems = []
+    for node in ast.walk(ast.parse(file.read_text(), filename=str(file))):
+        if isinstance(node, ast.ImportFrom) and node.level:
+            if node.level - 1 > depth:
+                where = f"{file.relative_to(source)}:{node.lineno}"
+                problems.append(f"{where}: {'.' * node.level}{node.module or ''}")
+            continue
+        if isinstance(node, ast.ImportFrom):
+            names = [node.module]
+        elif isinstance(node, ast.Import):
+            names = [alias.name for alias in node.names]
+        else:
+            continue
+        for name in names:
+            if name.split(".")[0] == top:
+                problems.append(f"{file.relative_to(source)}:{node.lineno}: {name}")
+    return problems
+
+
+def _is_remote(cls: type) -> bool:
+    return cls.__module__.startswith(_PACKAGE + ".")
+
+
+def _code_directory(pretrained: str | Path, **hub_kwargs: Any) -> Path:
+    local = Path(pretrained)
+    if local.is_dir():
+        return local.resolve()
+    from huggingface_hub import snapshot_download
+
+    return Path(
+        snapshot_download(str(pretrained), allow_patterns=["*.py"], **hub_kwargs)
+    )
+
+
+def _import_package(directory: Path) -> str:
+    """Register `directory` as a package (its `__init__.py` is not run)."""
+    if _PACKAGE not in sys.modules:
+        root = importlib.machinery.ModuleSpec(_PACKAGE, None, is_package=True)
+        root.submodule_search_locations = []
+        sys.modules[_PACKAGE] = importlib.util.module_from_spec(root)
+    digest = hashlib.sha256(str(directory).encode()).hexdigest()[:16]
+    name = f"{_PACKAGE}.m{digest}"
+    if name not in sys.modules:
+        spec = importlib.machinery.ModuleSpec(name, None, is_package=True)
+        spec.submodule_search_locations = [str(directory)]
+        module: ModuleType = importlib.util.module_from_spec(spec)
+        sys.modules[name] = module
+        setattr(sys.modules[_PACKAGE], f"m{digest}", module)
+    return name
+
+
+def _load_class(package: str, module: str, name: str, pretrained: str | Path) -> type:
+    try:
+        imported = importlib.import_module(f"{package}.{module}")
+    except ModuleNotFoundError as err:
+        if err.name != f"{package}.{module}":
+            raise
+        raise ModuleNotFoundError(
+            f"auto_map names module {module!r} but {module.replace('.', '/')}.py "
+            f"is not in {pretrained}"
+        ) from None
+    cls = getattr(imported, name, None)
+    if not isinstance(cls, type):
+        raise ImportError(f"{module}.py of {pretrained} defines no class {name!r}")
+    return cls

ffgs/geometry.py

@@ -0,0 +1,418 @@
+"""The model's working frame, and moving cameras / Gaussians in and out of it.
+
+Feed-forward models are trained on normalised cameras: here, every pose relative to
+the first context view, then translations scaled by a constant `scene_scale`. A
+`ModelFrame` records that normalisation per sample, so inputs go in with exactly the
+training arithmetic, and predicted Gaussians can be taken to the caller's world
+frame by one `Similarity` shared by every model (SH rotated to any degree).
+"""
+
+from __future__ import annotations
+
+import math
+from dataclasses import dataclass, replace
+
+import torch
+
+from .types import Cameras, Gaussians
+
+# Tolerance of R R^T = I for a rotation taken from caller c2w (fp32, maybe parsed
+# from text): loose enough for rounded poses, tight enough to catch any scale.
+_ORTHONORMAL_ATOL = 1e-3
+
+
+def _check_rotation(r: torch.Tensor, what: str) -> None:
+    """Raise unless every r [.., 3, 3] is a proper rotation."""
+    r = r.double()
+    if not torch.isfinite(r).all():
+        raise ValueError(f"{what}: the rotation is not finite")
+    eye = torch.eye(3, dtype=r.dtype, device=r.device)
+    error = (r @ r.transpose(-1, -2) - eye).abs().amax(dim=(-2, -1))
+    if (error > _ORTHONORMAL_ATOL).any():
+        raise ValueError(
+            f"{what}: the rotation is not orthonormal (non-uniform scale or "
+            f"shear; max |R R^T - I| = {error.max().item():.2e})"
+        )
+    if (torch.linalg.det(r) < 0).any():
+        raise ValueError(
+            f"{what}: the rotation has det < 0 (a reflection, e.g. left-handed "
+            "c2w axes)"
+        )
+
+
+@dataclass
+class PoseResiduals:
+    """How far the cameras a fitted frame maps to the world land from the given
+    ones, per view [B, V]."""
+
+    rotation_deg: torch.Tensor  # angle between the mapped and the given rotation
+    center_error: torch.Tensor  # distance between the camera centres, world units
+    # center_error over the RMS distance of the given centres from their mean.
+    relative_center_error: torch.Tensor
+
+
+@dataclass
+class ModelFrame:
+    """world -> model: x_m = scale * inverse(anchor_c2w) @ x_w, per sample.
+
+    `scale` is one number for the batch, or a [B] tensor (fitted frames).
+    `residuals` is set on frames fitted to given poses (`ModelFrame.fit`).
+    """
+
+    anchor_c2w: torch.Tensor  # [B, 4, 4], world frame
+    scale: float | torch.Tensor
+    residuals: PoseResiduals | None = None
+
+    @classmethod
+    def fit(cls, predicted_c2w: torch.Tensor, given_c2w: torch.Tensor) -> ModelFrame:
+        """The frame that maps a model's predicted cameras `predicted_c2w` onto
+        the caller's `given_c2w` (both [B, V, 4, 4], the same views).
+
+        Closed form over all V >= 2 views, in fp64: the rotation is the chordal
+        mean of R_given R_predicted^T (SVD, projected to det +1); with it fixed,
+        scale and translation are the least-squares fit of the camera centres.
+        """
+        pred, given = predicted_c2w.double(), given_c2w.to(predicted_c2w).double()
+        if pred.shape != given.shape:
+            raise ValueError(
+                f"predicted cameras {tuple(pred.shape)} and given c2w "
+                f"{tuple(given.shape)} disagree"
+            )
+        if pred.shape[1] < 2:
+            raise ValueError(
+                f"fitting a frame to given poses needs at least 2 views, got "
+                f"{pred.shape[1]}"
+            )
+        _check_rotation(given[..., :3, :3], "given c2w")
+        _check_rotation(pred[..., :3, :3], "predicted c2w")
+        for what, c2w in (("given", given), ("predicted", pred)):
+            if not torch.isfinite(c2w[..., :3, 3]).all():
+                raise ValueError(f"{what} c2w: the camera centres are not finite")
+        r_pred, r_given = pred[..., :3, :3], given[..., :3, :3]
+        c_pred, c_given = pred[..., :3, 3], given[..., :3, 3]
+        u, _, vh = torch.linalg.svd((r_given @ r_pred.transpose(-1, -2)).sum(dim=1))
+        flip = torch.ones_like(u[..., 0])
+        flip[..., -1] = torch.sign(torch.linalg.det(u @ vh))
+        rotation = u @ torch.diag_embed(flip) @ vh  # [B, 3, 3]
+
+        d_pred = c_pred - c_pred.mean(dim=1, keepdim=True)
+        d_given = c_given - c_given.mean(dim=1, keepdim=True)
+        spread = (d_pred**2).sum(dim=(1, 2))
+        extent = (d_given**2).sum(dim=-1).mean(dim=1).sqrt()  # RMS, given
+        if (spread <= 1e-12 * (c_pred**2).sum(dim=(1, 2)).clamp_min(1e-300)).any():
+            raise ValueError("the predicted camera centres coincide: no scale to fit")
+        rotated = torch.einsum("bij,bvj->bvi", rotation, d_pred)
+        scale = (d_given * rotated).sum(dim=(1, 2)) / spread  # model -> world
+        if (scale <= 0).any():
+            raise ValueError(
+                "the given poses do not fit the predicted cameras (fitted scale "
+                f"{scale.tolist()} <= 0)"
+            )
+        translation = c_given.mean(dim=1) - scale[:, None] * torch.einsum(
+            "bij,bj->bi", rotation, c_pred.mean(dim=1)
+        )
+
+        delta = r_given @ (rotation[:, None] @ r_pred).transpose(-1, -2)
+        cos = (delta.diagonal(dim1=-2, dim2=-1).sum(-1) - 1) / 2
+        skew = delta - delta.transpose(-1, -2)  # atan2: accurate at small angles
+        sin = torch.stack([skew[..., 2, 1], skew[..., 0, 2], skew[..., 1, 0]], -1)
+        sin = sin.norm(dim=-1) / 2
+        mapped_c = scale[:, None, None] * torch.einsum("bij,bvj->bvi", rotation, c_pred)
+        error = (mapped_c + translation[:, None] - c_given).norm(dim=-1)
+        residuals = PoseResiduals(
+            rotation_deg=torch.rad2deg(torch.atan2(sin, cos)).float(),
+            center_error=error.float(),
+            relative_center_error=(error / extent[:, None].clamp_min(1e-300)).float(),
+        )
+        anchor = torch.eye(4, dtype=torch.float64, device=pred.device)
+        anchor = anchor.repeat(pred.shape[0], 1, 1)
+        anchor[:, :3, :3] = rotation
+        anchor[:, :3, 3] = translation
+        return cls(
+            anchor_c2w=anchor.float(),
+            scale=(1 / scale).float(),
+            residuals=residuals,
+        )
+
+    def c2w_to_model(self, c2w: torch.Tensor) -> torch.Tensor:
+        """World c2w [B, V, 4, 4] -> model frame.
+
+        The dataset's arithmetic, op for op (`camera_normalization` followed by the
+        in-place translation scale), so the result is bit-identical to training.
+        """
+        if c2w.shape[0] != self.anchor_c2w.shape[0]:
+            raise ValueError(
+                f"{c2w.shape[0]} camera samples for a frame of "
+                f"{self.anchor_c2w.shape[0]}"
+            )
+        out = []
+        for anchor, poses in zip(self.anchor_c2w, c2w):
+            canonical = torch.eye(4, dtype=torch.float32, device=anchor.device)[None]
+            norm = torch.bmm(canonical, torch.inverse(anchor[None]))
+            out.append(torch.bmm(norm.repeat(poses.shape[0], 1, 1), poses))
+        model = torch.stack(out)
+        if isinstance(self.scale, torch.Tensor):
+            model[..., :3, 3] *= self.scale.to(model)[:, None, None]
+        elif self.scale != 1.0:
+            model[..., :3, 3] *= self.scale
+        return model
+
+    def cameras_to_model(self, cameras: Cameras) -> Cameras:
+        """World-frame cameras -> model frame (fp32, on the frame's device)."""
+        c2w = cameras.c2w.float().to(self.anchor_c2w.device)
+        return replace(cameras, c2w=self.c2w_to_model(c2w))
+
+    def model_to_world(self) -> Similarity:
+        """model -> world: x_w = anchor_c2w @ (x_m / scale)."""
+        anchor = self.anchor_c2w
+        if isinstance(self.scale, torch.Tensor):
+            scale = 1.0 / self.scale.to(anchor)
+        else:
+            scale = torch.full_like(anchor[:, 0, 0], 1.0 / self.scale)
+        return Similarity(
+            rotation=anchor[:, :3, :3], scale=scale, translation=anchor[:, :3, 3]
+        )
+
+
+@dataclass
+class Similarity:
+    """x -> scale * rotation @ x + translation, per sample.
+
+    The one transform every model's Gaussians go through between frames (the
+    public entry is `GSOutput.to_world`). Only rotations, uniform scales and
+    translations keep a Gaussian a Gaussian with the same opacity and its colour
+    an SH field, so anything else is rejected.
+    """
+
+    rotation: torch.Tensor  # [B, 3, 3], proper (det +1)
+    scale: torch.Tensor  # [B], > 0
+    translation: torch.Tensor  # [B, 3]
+
+    def __post_init__(self) -> None:
+        _check_rotation(self.rotation, "not a similarity")
+        if not (torch.isfinite(self.scale) & (self.scale > 0)).all():
+            raise ValueError(f"scale must be finite and > 0, got {self.scale}")
+        if not torch.isfinite(self.translation).all():
+            raise ValueError(f"translation must be finite, got {self.translation}")
+
+    def apply_cameras(self, cameras: Cameras) -> Cameras:
+        """Camera poses moved with the scene: rotations by `rotation`, centres
+        as points. Intrinsics are unchanged."""
+        c2w = cameras.c2w
+        rotation = self.rotation.to(c2w)[:, None]
+        out = c2w.clone()
+        out[..., :3, :3] = rotation @ c2w[..., :3, :3]
+        centers = (c2w[..., :3, 3] * self.scale.to(c2w)[:, None, None])[..., None]
+        out[..., :3, 3] = (rotation @ centers)[..., 0]
+        out[..., :3, 3] += self.translation.to(c2w)[:, None]
+        return replace(cameras, c2w=out)
+
+    def apply(self, gaussians: Gaussians) -> Gaussians:
+        """Each attribute as the transform acts on it: means fully, scales by
+        `scale`, quats by `rotation`, SH of degree >= 1 rotated; opacities and RGB /
+        degree-0 colours unchanged. Computed in the dtype of `gaussians`."""
+        g = gaussians
+        rotation = self.rotation.to(g.means)
+        scale = self.scale.to(g.means)
+        means = (g.means * scale[:, None, None]) @ rotation.transpose(-1, -2)
+        means = means + self.translation.to(g.means)[:, None]
+        quats = quat_multiply(matrix_to_quat(rotation)[:, None], g.quats)
+        colors = g.colors
+        if g.sh_degree is not None and g.sh_degree >= 1:
+            colors = rotate_sh(colors, rotation, g.sh_degree)
+        return replace(
+            g,
+            means=means,
+            scales=g.scales * scale[:, None, None],
+            quats=quats,
+            colors=colors,
+        )
+
+
+def anchored_frame(views_c2w: torch.Tensor, scale: float) -> ModelFrame:
+    """The frame anchored to the first context view of each sample."""
+    return ModelFrame(anchor_c2w=views_c2w[:, 0].float(), scale=float(scale))
+
+
+def matrix_to_quat(rotation: torch.Tensor) -> torch.Tensor:
+    """[.., 3, 3] rotation -> [.., 4] unit quaternion (wxyz, w >= 0)."""
+    m = rotation
+    m00, m11, m22 = m[..., 0, 0], m[..., 1, 1], m[..., 2, 2]
+    m01, m02, m10 = m[..., 0, 1], m[..., 0, 2], m[..., 1, 0]
+    m12, m20, m21 = m[..., 1, 2], m[..., 2, 0], m[..., 2, 1]
+    # Shepperd: take the largest of |w|, |x|, |y|, |z| from the diagonal and the
+    # other three from off-diagonal sums divided by it -- accurate for every
+    # rotation, unlike sqrt-ing each component (lossy near 0) or the trace-only
+    # formula (breaks near 180 degrees).
+    squares = torch.stack(
+        (
+            1 + m00 + m11 + m22,
+            1 + m00 - m11 - m22,
+            1 - m00 + m11 - m22,
+            1 - m00 - m11 + m22,
+        ),
+        dim=-1,
+    )
+    root = torch.sqrt(squares.clamp_min(1e-12))  # 2 * |component|
+    candidates = torch.stack(
+        (
+            torch.stack((root[..., 0] ** 2, m21 - m12, m02 - m20, m10 - m01), -1),
+            torch.stack((m21 - m12, root[..., 1] ** 2, m10 + m01, m02 + m20), -1),
+            torch.stack((m02 - m20, m10 + m01, root[..., 2] ** 2, m12 + m21), -1),
+            torch.stack((m10 - m01, m20 + m02, m21 + m12, root[..., 3] ** 2), -1),
+        ),
+        dim=-2,
+    ) / (2 * root[..., None])
+    best = squares.argmax(dim=-1)
+    q = torch.gather(
+        candidates, -2, best[..., None, None].expand(*best.shape, 1, 4)
+    ).squeeze(-2)
+    q = q / q.norm(dim=-1, keepdim=True)
+    return torch.where(q[..., :1] < 0, -q, q)
+
+
+def quat_to_matrix(q: torch.Tensor) -> torch.Tensor:
+    """[.., 4] quaternion (wxyz, any norm) -> [.., 3, 3] rotation."""
+    q = q / q.norm(dim=-1, keepdim=True)
+    w, x, y, z = q.unbind(-1)
+    return torch.stack(
+        (
+            1 - 2 * (y * y + z * z),
+            2 * (x * y - w * z),
+            2 * (x * z + w * y),
+            2 * (x * y + w * z),
+            1 - 2 * (x * x + z * z),
+            2 * (y * z - w * x),
+            2 * (x * z - w * y),
+            2 * (y * z + w * x),
+            1 - 2 * (x * x + y * y),
+        ),
+        dim=-1,
+    ).reshape(*q.shape[:-1], 3, 3)
+
+
+def quat_multiply(a: torch.Tensor, b: torch.Tensor) -> torch.Tensor:
+    """Hamilton product a * b (wxyz): rotate by b, then by a."""
+    aw, ax, ay, az = a.unbind(-1)
+    bw, bx, by, bz = b.unbind(-1)
+    return torch.stack(
+        (
+            aw * bw - ax * bx - ay * by - az * bz,
+            aw * bx + ax * bw + ay * bz - az * by,
+            aw * by - ax * bz + ay * bw + az * bx,
+            aw * bz + ax * by - ay * bx + az * bw,
+        ),
+        dim=-1,
+    )
+
+
+def rotate_sh(
+    colors: torch.Tensor, rotation: torch.Tensor, degree: int
+) -> torch.Tensor:
+    """Rotate SH coefficients [B, N, K, 3] by per-sample rotations [B, 3, 3].
+
+    The rotated field along d is the original along R^T d, in the real SH basis of
+    3DGS / gsplat, any degree. Degree 0 is rotation invariant. Degree 1 is
+    C1 * (-y c1 + z c2 - x c3), a linear form a . d with a = (-c3, -c1, c2), so it
+    turns into R a. Degrees >= 2 use `sh_rotation_blocks`, built in fp64.
+    """
+    if colors.shape[-2] != (degree + 1) ** 2:
+        raise ValueError(
+            f"degree {degree} needs {(degree + 1) ** 2} SH coefficients, got "
+            f"{tuple(colors.shape)}"
+        )
+    rotated = colors.clone()
+    if degree < 1:
+        return rotated
+    c1, c2, c3 = colors[:, :, 1], colors[:, :, 2], colors[:, :, 3]
+    a = torch.stack((-c3, -c1, c2), dim=2)  # [B, N, xyz, rgb]
+    a = torch.einsum("bij,bnjc->bnic", rotation, a)
+    rotated[:, :, 1] = -a[:, :, 1]
+    rotated[:, :, 2] = a[:, :, 2]
+    rotated[:, :, 3] = -a[:, :, 0]
+    if degree < 2:
+        return rotated
+    blocks = sh_rotation_blocks(rotation.double(), degree)
+    for degree_l in range(2, degree + 1):
+        band = slice(degree_l**2, (degree_l + 1) ** 2)
+        block = blocks[degree_l].to(colors.dtype)
+        rotated[:, :, band] = torch.einsum("bij,bnjc->bnic", block, colors[:, :, band])
+    return rotated
+
+
+def sh_rotation_blocks(rotation: torch.Tensor, degree: int) -> list[torch.Tensor]:
+    """Real SH rotation matrices D_l [.., 2l+1, 2l+1] for l = 0..degree, in the
+    3DGS / gsplat basis: coefficients c_l of degree l become D_l @ c_l.
+
+    Ivanic & Ruedenberg's recursion (J. Phys. Chem. 100, 6342 (1996); errata 102,
+    9099 (1998)), exact for any degree: D_l from D_{l-1} and D_1, which is R in the
+    (y, z, x) order of the real basis without the Condon-Shortley phase. gsplat's
+    basis carries that phase, (-1)^m on order m, so each D_l is conjugated by it.
+    """
+    order = [1, 2, 0]  # m = -1, 0, 1 <-> y, z, x
+    r1 = rotation[..., order, :][..., :, order]
+    blocks = [torch.ones_like(rotation[..., :1, :1]), r1]
+    for degree_l in range(2, degree + 1):
+        blocks.append(_ivanic_ruedenberg_step(r1, blocks[-1], degree_l))
+    out = []
+    for degree_l, block in enumerate(blocks[: degree + 1]):
+        m = torch.arange(-degree_l, degree_l + 1, device=rotation.device)
+        phase = (1 - 2 * (m % 2)).to(rotation.dtype)
+        out.append(block * phase[:, None] * phase[None, :])
+    return out
+
+
+def _ivanic_ruedenberg_step(
+    r1: torch.Tensor, prev: torch.Tensor, degree: int
+) -> torch.Tensor:
+    """D_l [.., 2l+1, 2l+1] from D_1 and D_{l-1}, indexed by order m in -l..l."""
+    lv = degree
+
+    def d1(i: int, j: int) -> torch.Tensor:
+        return r1[..., i + 1, j + 1]
+
+    def d_prev(a: int, b: int) -> torch.Tensor:
+        return prev[..., a + lv - 1, b + lv - 1]
+
+    def p(i: int, a: int, b: int) -> torch.Tensor:
+        if b == lv:
+            return d1(i, 1) * d_prev(a, lv - 1) - d1(i, -1) * d_prev(a, -lv + 1)
+        if b == -lv:
+            return d1(i, 1) * d_prev(a, -lv + 1) + d1(i, -1) * d_prev(a, lv - 1)
+        return d1(i, 0) * d_prev(a, b)
+
+    def u(m: int, n: int) -> torch.Tensor:
+        return p(0, m, n)
+
+    def v(m: int, n: int) -> torch.Tensor:
+        if m == 0:
+            return p(1, 1, n) + p(-1, -1, n)
+        if m > 0:
+            first = p(1, m - 1, n) * math.sqrt(1 + (m == 1))
+            return first if m == 1 else first - p(-1, -m + 1, n)
+        first = p(-1, -m - 1, n) * math.sqrt(1 + (m == -1))
+        return first if m == -1 else first + p(1, m + 1, n)
+
+    def w(m: int, n: int) -> torch.Tensor:
+        if m > 0:
+            return p(1, m + 1, n) + p(-1, -m - 1, n)
+        return p(1, m - 1, n) - p(-1, -m + 1, n)
+
+    rows = []
+    for m in range(-lv, lv + 1):
+        row = []
+        for n in range(-lv, lv + 1):
+            am = abs(m)
+            denom = 2 * lv * (2 * lv - 1) if abs(n) == lv else (lv + n) * (lv - n)
+            cu = math.sqrt((lv + m) * (lv - m) / denom)
+            cv = 0.5 * math.sqrt((1 + (m == 0)) * (lv + am - 1) * (lv + am) / denom)
+            cv = -cv if m == 0 else cv
+            cw = 0.0 if m == 0 else -0.5 * math.sqrt((lv - am - 1) * (lv - am) / denom)
+            entry = cv * v(m, n)
+            if cu:
+                entry = entry + cu * u(m, n)
+            if cw:
+                entry = entry + cw * w(m, n)
+            row.append(entry)
+        rows.append(torch.stack(row, dim=-1))
+    return torch.stack(rows, dim=-2)