Spaces:
Running on Zero
Running on Zero
| """Stage 2 of FineSplat: error-conditioned correction of the Stage-1 tokens. | |
| Stage 1 (2Xplat) predicts one token per 4x4 patch, decoded into 4 Gaussians. Each Stage-2 | |
| step renders the current Gaussians into the context views, builds a per-view error map, | |
| lifts it onto the tokens by projection and visibility pooling, runs PTv3 serialized-attention | |
| blocks over the token cloud, and predicts a token residual the frozen Stage-1 head decodes. | |
| python -m src.model.stage2 # self-check (needs a GPU) | |
| """ | |
| import contextlib | |
| import torch | |
| import torch.nn.functional as F | |
| import torch.utils.checkpoint | |
| from einops import rearrange | |
| from torch import nn | |
| from src.model.gaussians import GaussianRenderer | |
| from src.model.ptv3.ptv3 import Block, SerializedAttention | |
| from src.model.ptv3.structure import Point | |
| from src.model.sh_eval import _spherical_harmonics | |
| from src.utils.camera_utils import fxfycxcy_to_K | |
| EPS = 1e-6 | |
| # Cap on elements materialised by one sparse-gather chunk in ErrorLifter (~512 MB in fp32). | |
| _GATHER_CHUNK_ELEMS = 1 << 27 | |
| IMAGENET_MEAN = (0.485, 0.456, 0.406) | |
| IMAGENET_STD = (0.229, 0.224, 0.225) | |
| default_conf = { | |
| # "token": one point per patch token, residual on the token (decoded by the Stage-1 head). | |
| # "gaussian": one point per Gaussian, residual on its raw pre-activation parameters. | |
| "mode": "token", | |
| # Stage-2 steps per forward; each re-renders the context views from the previous step. | |
| "num_steps": 4, | |
| "network": { | |
| "width": 512, | |
| "depth": 4, | |
| "attn_channels": 256, # q/k/v width (num_heads x 16); == width gives plain PTv3 attention | |
| "num_heads": 16, | |
| "patch_size": 1024, # window = min(points per scene, patch_size) | |
| "order": ["z", "z-trans", "hilbert", "hilbert-trans"], | |
| "drop_path": 0.1, | |
| "grid_resolution": 384, # serialization + xCPE voxel grid (SplatFormer's value) | |
| }, | |
| "error": { | |
| "stride": 4, # maps at 56x56 for 224x224 inputs | |
| "width": 128, # per-view map = [error 128 | image 128] | |
| "lift_channels": 256, | |
| "occlusion_tol": 0.15, | |
| "reduce": "mean", # how visible views are pooled per token: mean | max | |
| "lift": "mean+home", # what reaches the token: mean+home | mean | home | |
| "use_rgb_diff": True, | |
| "use_image_embed": True, | |
| "backbone": "resnet18", # frozen error feature extractor: resnet18 | vgg16_lpips | |
| }, | |
| } | |
| def _sample(img: torch.Tensor, pix: torch.Tensor, size: int) -> torch.Tensor: | |
| """Bilinear sample img [N,C,h,w] at pixel coords pix [N,M,2] -> [N,M,C]; centres at i+0.5 (align_corners=False).""" | |
| grid = (pix / size * 2 - 1).unsqueeze(2) | |
| out = F.grid_sample(img, grid.to(img.dtype), mode="bilinear", padding_mode="border", align_corners=False) | |
| return out[..., 0].transpose(1, 2) # [N,M,C] | |
| def _sample_sparse(imgs: torch.Tensor, pix: torch.Tensor, img_idx: torch.Tensor, | |
| size: int) -> torch.Tensor: | |
| """Gather form of ``_sample``: imgs [N,C,h,w], pix [P,2] from image img_idx[p] -> [P,C].""" | |
| n, c, h, w = imgs.shape | |
| xs = (pix[:, 0] * (w / size) - 0.5).clamp(0, w - 1) | |
| ys = (pix[:, 1] * (h / size) - 0.5).clamp(0, h - 1) | |
| x0, y0 = xs.floor(), ys.floor() | |
| wx, wy = (xs - x0).to(imgs.dtype)[:, None], (ys - y0).to(imgs.dtype)[:, None] | |
| x0, y0 = x0.long(), y0.long() | |
| x1, y1 = (x0 + 1).clamp(max=w - 1), (y0 + 1).clamp(max=h - 1) | |
| flat = imgs.reshape(n, c, h * w) | |
| # advanced indices at dim 0 and 2 with a slice between them -> result is [P, C] | |
| top = flat[img_idx, :, y0 * w + x0] * (1 - wx) + flat[img_idx, :, y0 * w + x1] * wx | |
| bot = flat[img_idx, :, y1 * w + x0] * (1 - wx) + flat[img_idx, :, y1 * w + x1] * wx | |
| return top * (1 - wy) + bot * wy | |
| class ErrorMapEncoder(nn.Module): | |
| """Per-view error map at 1/stride resolution: [error `width` | image `img_width`] channels.""" | |
| def __init__(self, image_size: int, stride: int, width: int, | |
| use_rgb_diff: bool = True, use_image_embed: bool = True, | |
| backbone: str = "resnet18", use_feat_diff: bool = True, | |
| image_embed: str = "gt_rgb", mvp_dim: int = 256, | |
| img_width: int | None = None): | |
| super().__init__() | |
| assert backbone in ("resnet18", "vgg16_lpips"), backbone | |
| assert image_embed in ("gt_rgb", "mvp_s1"), image_embed | |
| self.backbone = backbone | |
| self.use_feat_diff = use_feat_diff | |
| self.image_embed = image_embed if use_image_embed else None | |
| self.use_err = use_feat_diff or use_rgb_diff | |
| assert self.use_err or use_image_embed, "no error branches and no image embed: empty map" | |
| if not use_feat_diff: | |
| self.stem = self.rest = self.layer2 = None | |
| self._stages = [] | |
| elif backbone == "resnet18": | |
| from torchvision.models import ResNet18_Weights, resnet18 | |
| # checkpoints store the frozen backbone under error.stem/rest/layer2 - do not rename | |
| r = resnet18(weights=ResNet18_Weights.IMAGENET1K_V1) | |
| self.stem = nn.Sequential(r.conv1) | |
| self.rest = nn.Sequential(r.bn1, r.relu, r.maxpool, r.layer1) | |
| self.layer2 = r.layer2 | |
| self._stages = [self.stem, self.rest, self.layer2] | |
| feat_channels = 64 + 64 + 128 | |
| else: | |
| from torchvision.models import VGG16_Weights, vgg16 | |
| vf = vgg16(weights=VGG16_Weights.IMAGENET1K_V1).features | |
| self.slice1, self.slice2, self.slice3 = vf[:4], vf[4:9], vf[9:16] | |
| self._stages = [self.slice1, self.slice2, self.slice3] | |
| feat_channels = 64 + 128 + 256 | |
| for m in self._stages: | |
| m.requires_grad_(False) | |
| self.register_buffer("mean", torch.tensor(IMAGENET_MEAN).view(1, 3, 1, 1), persistent=False) | |
| self.register_buffer("std", torch.tensor(IMAGENET_STD).view(1, 3, 1, 1), persistent=False) | |
| assert image_size % stride == 0, (image_size, stride) | |
| self.stride = stride | |
| self.grid_hw = (image_size // stride, image_size // stride) | |
| img_width = width if img_width is None else img_width | |
| self.channels = (width if self.use_err else 0) + (img_width if use_image_embed else 0) | |
| self.feat_proj = nn.Linear(feat_channels, width) if use_feat_diff else None | |
| self.rgb_proj = nn.Linear(3 * stride * stride, width) if use_rgb_diff else None | |
| self.err_norm = nn.LayerNorm(width) if self.use_err else None | |
| # "gt_rgb": pixel-unshuffled context image; "mvp_s1": Stage 1's own patch tokens for the view | |
| img_in = {"gt_rgb": 3 * stride * stride, "mvp_s1": mvp_dim}.get(self.image_embed) | |
| self.img_proj = ( | |
| nn.Sequential(nn.Linear(img_in, img_width), nn.LayerNorm(img_width)) | |
| if self.image_embed is not None else None | |
| ) | |
| def train(self, mode: bool = True): | |
| super().train(mode) | |
| for m in self._stages: | |
| m.eval() # frozen backbone: keep BatchNorm running stats fixed | |
| return self | |
| def features(self, img: torch.Tensor) -> torch.Tensor | None: | |
| """img [N,3,H,W] in [0,1] -> [N,C,gh,gw] (C = 256 resnet / 448 vgg).""" | |
| if not self.use_feat_diff: | |
| return None | |
| x = (img.float() - self.mean) / self.std | |
| feats = [] | |
| for m in self._stages: | |
| x = m(x) | |
| feats.append(x) | |
| if self.backbone == "vgg16_lpips": | |
| feats = [f / f.pow(2).sum(dim=1, keepdim=True).sqrt().clamp(min=1e-10) for f in feats] | |
| return torch.cat( | |
| [F.interpolate(f, size=self.grid_hw, mode="bilinear", align_corners=True) for f in feats], | |
| dim=1, | |
| ) | |
| def image_map(self, gt: torch.Tensor, mvp_feat: torch.Tensor | None = None): | |
| """The image half of the map, [N, gh, gw, img_width]; constant across steps, computed once per forward.""" | |
| if self.img_proj is None: | |
| return None | |
| if self.image_embed == "mvp_s1": | |
| if mvp_feat.shape[-2:] != self.grid_hw: | |
| mvp_feat = F.interpolate(mvp_feat, size=self.grid_hw, mode="bilinear", align_corners=True) | |
| return self.img_proj(mvp_feat.permute(0, 2, 3, 1)) | |
| return self.img_proj(F.pixel_unshuffle(gt, self.stride).permute(0, 2, 3, 1)) | |
| def forward(self, render: torch.Tensor, gt: torch.Tensor, | |
| gt_feat: torch.Tensor | None = None, | |
| mvp_feat: torch.Tensor | None = None, | |
| img: torch.Tensor | None = None) -> torch.Tensor: | |
| """render, gt [N,3,H,W] -> [N, width + img_width, gh, gw]; ``gt_feat``/``img`` may be precomputed.""" | |
| if not self.use_err: | |
| img = self.image_map(gt, mvp_feat) if img is None else img | |
| return img.permute(0, 3, 1, 2) | |
| err = None | |
| if self.use_feat_diff: | |
| if gt_feat is None: | |
| gt_feat = self.features(gt) | |
| err = self.feat_proj((self.features(render) - gt_feat).permute(0, 2, 3, 1)) | |
| if self.rgb_proj is not None: | |
| rgb = self.rgb_proj(F.pixel_unshuffle(render - gt, self.stride).permute(0, 2, 3, 1)) | |
| err = rgb if err is None else err + rgb | |
| err = self.err_norm(err) | |
| if self.img_proj is None: | |
| return err.permute(0, 3, 1, 2) | |
| if img is None: | |
| img = self.image_map(gt, mvp_feat) | |
| return torch.cat([err, img.to(err.dtype)], dim=-1).permute(0, 3, 1, 2) | |
| class _GlobalBlock(nn.Module): | |
| """Plain pre-norm transformer block with full self-attention.""" | |
| def __init__(self, dim: int, num_heads: int): | |
| super().__init__() | |
| self.norm1 = nn.LayerNorm(dim) | |
| self.attn = nn.MultiheadAttention(dim, num_heads, batch_first=True) | |
| self.norm2 = nn.LayerNorm(dim) | |
| self.mlp = nn.Sequential(nn.Linear(dim, 4 * dim), nn.GELU(), nn.Linear(4 * dim, dim)) | |
| def forward(self, x): | |
| h = self.norm1(x) | |
| x = x + self.attn(h, h, h, need_weights=False)[0] | |
| return x + self.mlp(self.norm2(x)) | |
| class ErrorGlobalAttn(nn.Module): | |
| """Optional cross-view attention over the per-view error maps, at 1/``down`` resolution (off by default).""" | |
| def __init__(self, channels: int, down: int = 4, attn_channels: int = 64, | |
| num_heads: int = 1, depth: int = 1): | |
| super().__init__() | |
| self.down = down | |
| self.proj0 = nn.Linear(channels, attn_channels) | |
| self.proj1 = nn.Linear(attn_channels * down ** 2, attn_channels) | |
| self.norm1 = nn.LayerNorm(attn_channels) | |
| self.blocks = nn.ModuleList([_GlobalBlock(attn_channels, num_heads) for _ in range(depth)]) | |
| self.proj2 = nn.Linear(attn_channels, attn_channels * down ** 2) | |
| self.norm2 = nn.LayerNorm(attn_channels * down ** 2) | |
| self.conv = nn.Conv2d(attn_channels, attn_channels, 3, 1, 1) | |
| self.proj3 = nn.Linear(attn_channels, channels) | |
| nn.init.zeros_(self.proj3.weight) | |
| nn.init.zeros_(self.proj3.bias) | |
| def forward(self, E: torch.Tensor) -> torch.Tensor: | |
| """E [B, V, C, h, w] -> same shape.""" | |
| b, v, _, h, w = E.shape | |
| x = rearrange(E, "b v c h w -> b (v h w) c") | |
| x = self.proj0(x) | |
| x = rearrange(x, "b (v h w) c -> (b v) c h w", v=v, h=h, w=w) | |
| x = F.pixel_unshuffle(x, self.down) | |
| x = rearrange(x, "(b v) c h w -> b (v h w) c", v=v) | |
| x = self.norm1(self.proj1(x)) | |
| for blk in self.blocks: | |
| x = blk(x) | |
| x = self.norm2(self.proj2(x)) | |
| x = rearrange(x, "b (v h w) c -> (b v) c h w", v=v, h=h // self.down, w=w // self.down) | |
| x = self.conv(F.pixel_shuffle(x, self.down)) | |
| x = rearrange(x, "(b v) c h w -> b (v h w) c", v=v) | |
| x = self.proj3(x) | |
| return E + rearrange(x, "b (v h w) c -> b v c h w", v=v, h=h, w=w) | |
| class ErrorLifter(nn.Module): | |
| """Per-view maps -> one feature per token: pooled over the views the token is visible in | |
| (home view always counts), concatenated with the map at the token's own patch.""" | |
| def __init__(self, in_channels: int, out_channels: int, occlusion_tol: float, | |
| reduce: str = "mean", lift: str = "mean+home"): | |
| super().__init__() | |
| assert reduce in ("mean", "max"), reduce | |
| assert lift in ("mean+home", "mean", "home"), lift | |
| self.tol = occlusion_tol | |
| self.reduce = reduce | |
| self.lift = lift | |
| n_in = 2 if lift == "mean+home" else 1 | |
| self.fuse = nn.Sequential(nn.Linear(n_in * in_channels, out_channels), nn.LayerNorm(out_channels)) | |
| def forward(self, E, depth, coord, i_c2w, i_fxfycxcy, patch: int, side: int, size: int): | |
| """E [B,V,C,gh,gw], depth [B,V,H,W] (render depth), coord [B,K,3] world points, | |
| i_c2w [B,V,4,4], i_fxfycxcy [B,V,4], tokens ordered (v, i, j) -> [B,K,C'], [B,K].""" | |
| B, V, C = E.shape[:3] | |
| K = coord.shape[1] | |
| dev = E.device | |
| home_view = torch.arange(K, device=dev) // (side * side) | |
| if self.lift == "home": | |
| cnt = torch.ones(B, K, 1, device=dev, dtype=E.dtype) | |
| pooled = None | |
| else: | |
| # project every token into every view; only in-frustum pairs are sampled below | |
| w2c = i_c2w.float().inverse() | |
| p_cam = (torch.einsum("bkj,bvij->bvki", coord.float(), w2c[:, :, :3, :3]) | |
| + w2c[:, :, None, :3, 3]) | |
| z = p_cam[..., 2] | |
| k_mat = fxfycxcy_to_K(i_fxfycxcy.reshape(B * V, 4)).reshape(B, V, 3, 3) | |
| uvw = torch.einsum("bvkj,bvij->bvki", p_cam, k_mat) | |
| pix = uvw[..., :2] / uvw[..., 2:3].clamp(min=EPS) | |
| in_frustum = (z > EPS) & (pix[..., 0] >= 0) & (pix[..., 0] < size) \ | |
| & (pix[..., 1] >= 0) & (pix[..., 1] < size) | |
| is_home = home_view[None, :] == torch.arange(V, device=dev)[:, None] | |
| b_i, v_i, k_i = (in_frustum | is_home).nonzero(as_tuple=True) | |
| img_i = b_i * V + v_i | |
| pix_c = pix[b_i, v_i, k_i] | |
| d_render = _sample_sparse( | |
| depth.reshape(B * V, 1, *depth.shape[-2:]).float(), pix_c, img_i, size)[:, 0] | |
| z_c = z[b_i, v_i, k_i] | |
| # visible: in the frustum and not behind the rendered surface; the home view always counts | |
| vis = (in_frustum[b_i, v_i, k_i] | |
| & ((d_render <= EPS) | (z_c <= d_render * (1 + self.tol)))) | is_home[v_i, k_i] | |
| sel = vis.nonzero(as_tuple=True)[0] | |
| flat_idx = b_i[sel] * K + k_i[sel] | |
| cnt = torch.zeros(B * K, 1, device=dev, dtype=E.dtype).index_add( | |
| 0, flat_idx, torch.ones(sel.numel(), 1, device=dev, dtype=E.dtype)).view(B, K, 1) | |
| # chunked: a single flat [P, C] buffer would dominate memory at high view counts | |
| e_flat = E.reshape(B * V, C, *E.shape[-2:]) | |
| acc = (torch.full((B * K, C), -torch.inf, device=dev, dtype=E.dtype) | |
| if self.reduce == "max" else torch.zeros(B * K, C, device=dev, dtype=E.dtype)) | |
| step = max(1, _GATHER_CHUNK_ELEMS // C) | |
| for s in range(0, sel.numel(), step): | |
| part, idx = sel[s : s + step], flat_idx[s : s + step] | |
| e_s = _sample_sparse(e_flat, pix_c[part], img_i[part], size) | |
| acc = (acc.scatter_reduce(0, idx[:, None].expand(-1, C), e_s, reduce="amax") | |
| if self.reduce == "max" else acc.index_add(0, idx, e_s)) | |
| pooled = acc.view(B, K, C) | |
| if self.reduce != "max": | |
| pooled = pooled / cnt.clamp(min=1) | |
| # home view: token (v, i, j) <-> patch centre ((j+0.5)p, (i+0.5)p) of view v | |
| ij = torch.stack(torch.meshgrid(torch.arange(side), torch.arange(side), indexing="ij"), -1) | |
| centres = ((ij.flip(-1).float() + 0.5) * patch).to(E.device) # [side, side, 2] as (x, y) | |
| centres = centres.view(1, side * side, 2).expand(B * V, -1, -1) | |
| home = _sample(rearrange(E, "b v c h w -> (b v) c h w"), centres, size) | |
| home = rearrange(home, "(b v) n c -> b (v n) c", b=B) | |
| assert home.shape[1] == K, (home.shape, K) | |
| feats = {"mean+home": lambda: torch.cat([pooled, home], dim=-1), | |
| "mean": lambda: pooled, | |
| "home": lambda: home}[self.lift]() | |
| return self.fuse(feats), cnt[..., 0] | |
| def _knn_indices(coord: torch.Tensor, offset: torch.Tensor, k: int) -> torch.Tensor: | |
| """[K,3] coords, scene-end offsets -> [K,k] neighbour indices (self included).""" | |
| out = torch.empty(coord.shape[0], k, dtype=torch.long, device=coord.device) | |
| start = 0 | |
| for end in offset.tolist(): | |
| pts = coord[start:end] | |
| for q0 in range(0, pts.shape[0], 8192): | |
| q = pts[q0:q0 + 8192] | |
| d = torch.cdist(q, pts) | |
| out[start + q0:start + q0 + q.shape[0]] = start + d.topk(k, largest=False).indices | |
| start = end | |
| return out | |
| def _knn_block_forward(blk, feat: torch.Tensor, idx: torch.Tensor) -> torch.Tensor: | |
| """One PTv3 block evaluated with gather attention over ``idx`` [N,k] neighbours instead of | |
| serialized windows; reuses the block's weights and mirrors Block.forward's residual layout.""" | |
| n, k = idx.shape | |
| a = blk.attn | |
| heads = a.num_heads | |
| shortcut = feat | |
| for m in blk.cpe._modules.values(): | |
| feat = m(feat) | |
| feat = shortcut + feat | |
| shortcut = feat | |
| x = feat | |
| for m in blk.norm1._modules.values(): | |
| x = m(x) | |
| qkv = a.qkv(x) | |
| cp = qkv.shape[-1] // 3 | |
| q, kk, v = qkv.split(cp, dim=-1) | |
| hd = cp // heads | |
| q = q.view(n, heads, hd) | |
| kn = kk[idx.reshape(-1)].view(n, k, heads, hd) | |
| vn = v[idx.reshape(-1)].view(n, k, heads, hd) | |
| attn = torch.einsum("nhd,nkhd->nhk", q * a.scale, kn) | |
| attn = attn.softmax(dim=-1) | |
| out = torch.einsum("nhk,nkhd->nhd", attn, vn).reshape(n, cp) | |
| out = a.proj(out) | |
| for m in blk.drop_path._modules.values(): | |
| out = m(out) | |
| feat = shortcut + out | |
| shortcut = feat | |
| x = feat | |
| for m in blk.norm2._modules.values(): | |
| x = m(x) | |
| for m in blk.mlp._modules.values(): | |
| x = m(x) | |
| for m in blk.drop_path._modules.values(): | |
| x = m(x) | |
| return shortcut + x | |
| def _grid_neighbors(mode: str, B: int, V: int, side: int, device, coord=None, c2w=None, | |
| fxfycxcy=None, depth=None, patch: int = 1, size: int = 1, | |
| tol: float = 0.15) -> torch.Tensor: | |
| """Image-space neighbour indices for tokens laid out (b, v, i, j) -> [B*V*side^2, n]. | |
| "inframe": 3x3 block in the token's own view (9). "xframe": 4-neighbourhood in its own view | |
| plus the 4-neighbourhood of its projection into the partner view v ^ 1 (10), falling back | |
| to self when the projection is out of frustum or fails the depth test. | |
| """ | |
| S2 = side * side | |
| K = V * S2 | |
| ar = torch.arange(side, device=device) | |
| ii, jj = torch.meshgrid(ar, ar, indexing="ij") | |
| ii, jj = ii.reshape(-1), jj.reshape(-1) | |
| base = (torch.arange(B, device=device)[:, None, None] * K | |
| + torch.arange(V, device=device)[None, :, None] * S2) | |
| self_idx = (base + torch.arange(S2, device=device)).reshape(-1) | |
| def offsets(ci, cj, offs): | |
| cols = [] | |
| for di, dj in offs: | |
| ni, nj = ci + di, cj + dj | |
| ok = (ni >= 0) & (ni < side) & (nj >= 0) & (nj < side) | |
| cols.append((ni.clamp(0, side - 1) * side + nj.clamp(0, side - 1), ok)) | |
| return cols | |
| if mode == "inframe": | |
| offs = [(di, dj) for di in (-1, 0, 1) for dj in (-1, 0, 1)] | |
| cols = [(base + c).reshape(-1) for c, _ in offsets(ii, jj, offs)] | |
| oks = [ok.expand(B, V, -1).reshape(-1) for _, ok in offsets(ii, jj, offs)] | |
| idx = torch.stack([torch.where(o, c, self_idx) for c, o in zip(cols, oks)], 1) | |
| return idx | |
| assert mode == "xframe", mode | |
| cross = [(0, 0), (-1, 0), (1, 0), (0, -1), (0, 1)] | |
| own = [torch.where(ok.expand(B, V, -1).reshape(-1), (base + c).reshape(-1), self_idx) | |
| for c, ok in offsets(ii, jj, cross)] | |
| # project every token into its partner view | |
| partner = torch.arange(V, device=device) ^ 1 | |
| pc2w = c2w[:, partner].float() | |
| w2c = pc2w.inverse() | |
| pts = coord.float().view(B, V, S2, 3) | |
| p_cam = torch.einsum("bvkj,bvij->bvki", pts, w2c[..., :3, :3]) + w2c[..., None, :3, 3] | |
| z = p_cam[..., 2] | |
| kmat = fxfycxcy_to_K(fxfycxcy[:, partner].reshape(B * V, 4).float()).reshape(B, V, 3, 3) | |
| uvw = torch.einsum("bvkj,bvij->bvki", p_cam, kmat) | |
| pix = uvw[..., :2] / uvw[..., 2:3].clamp(min=EPS) # [B,V,S2,2] (x, y) | |
| inside = (z > EPS) & (pix[..., 0] >= 0) & (pix[..., 0] < size) & (pix[..., 1] >= 0) & (pix[..., 1] < size) | |
| px = pix[..., 0].floor().long().clamp(0, size - 1) | |
| py = pix[..., 1].floor().long().clamp(0, size - 1) | |
| dpart = depth.float()[:, partner] | |
| d = torch.gather(dpart.reshape(B, V, -1), 2, py * dpart.shape[-1] + px) | |
| visible = inside & ((d <= EPS) | (z <= d * (1 + tol))) | |
| ci, cj = (py // patch).clamp(0, side - 1), (px // patch).clamp(0, side - 1) | |
| pbase = (torch.arange(B, device=device)[:, None, None] * K + partner[None, :, None] * S2) | |
| theirs = [] | |
| for di, dj in cross: | |
| ni, nj = ci + di, cj + dj | |
| ok = visible & (ni >= 0) & (ni < side) & (nj >= 0) & (nj < side) | |
| c = pbase + ni.clamp(0, side - 1) * side + nj.clamp(0, side - 1) | |
| theirs.append(torch.where(ok.reshape(-1), c.reshape(-1), self_idx)) | |
| return torch.stack(own + theirs, 1) | |
| def _block(channels, attn_channels, num_heads, patch_size, drop_path, order_index): | |
| """PTv3 block (xCPE -> window attention -> FFN, pre-norm) with q/k/v in `attn_channels`.""" | |
| blk = Block( | |
| channels=channels, num_heads=num_heads, patch_size=patch_size, mlp_ratio=4, qkv_bias=True, | |
| drop_path=drop_path, norm_layer=nn.LayerNorm, act_layer=nn.GELU, pre_norm=True, order_index=order_index, | |
| cpe_indice_key="stage0", enable_rpe=False, enable_flash=False, upcast_attention=False, upcast_softmax=False, | |
| ) | |
| if attn_channels != channels: | |
| attn = SerializedAttention( | |
| channels=attn_channels, num_heads=num_heads, patch_size=patch_size, qkv_bias=True, | |
| order_index=order_index, enable_rpe=False, enable_flash=False, upcast_attention=False, | |
| upcast_softmax=False, | |
| ) | |
| attn.qkv = nn.Linear(channels, 3 * attn_channels, bias=True) | |
| attn.proj = nn.Linear(attn_channels, channels) | |
| blk.attn = attn | |
| return blk | |
| class Stage2(nn.Module): | |
| """Token + error embeddings -> `depth` PTv3 blocks over the whole token cloud -> token residual.""" | |
| def __init__(self, conf, token_dim: int, error_channels: int, z_dim: int, | |
| aux_channels: int = 0, raw_input: bool = False): | |
| super().__init__() | |
| self.conf = conf | |
| self.z_dim = z_dim | |
| self.order = list(conf["order"]) | |
| W = conf["width"] | |
| # raw_input: raw Gaussian parameters enter the fuse normalised, unprojected; aux: per-point conditioning | |
| if raw_input: | |
| self.token_proj = nn.LayerNorm(token_dim) | |
| fuse_in = token_dim + error_channels + aux_channels | |
| else: | |
| self.token_proj = nn.Sequential(nn.LayerNorm(token_dim), nn.Linear(token_dim, W)) | |
| fuse_in = W + error_channels + aux_channels | |
| self.fuse = nn.Sequential(nn.Linear(fuse_in, W), nn.LayerNorm(W), nn.GELU()) | |
| dp = torch.linspace(0, conf["drop_path"], conf["depth"]).tolist() | |
| self.blocks = nn.ModuleList( | |
| [_block(W, conf["attn_channels"], conf["num_heads"], conf["patch_size"], dp[i], i % len(self.order)) | |
| for i in range(conf["depth"])] | |
| ) | |
| self.norm = nn.LayerNorm(W) | |
| # z_dim=None: no head, forward returns the normalised backbone features. | |
| # attn_mode: "serialized" (trained) or an inference-time gather swap: knn | inframe | xframe. | |
| self.attn_mode = conf.get("attn_mode", "serialized") | |
| self.knn_k = conf.get("knn_k", 8) | |
| self.head = nn.Linear(W, z_dim) if z_dim is not None else None | |
| if self.head is not None: | |
| nn.init.zeros_(self.head.weight) # zero-init -> step 0 reproduces the frozen model | |
| nn.init.zeros_(self.head.bias) | |
| self.window = 0 | |
| def _point(self, grid_coord, offset, feat): | |
| """Serialize `feat` under one grid. Rebuilt whenever the grid changes between blocks.""" | |
| grid = grid_coord.reshape(-1, 3).int() | |
| point = Point({ | |
| "coord": grid.float(), | |
| "grid_coord": grid, | |
| "grid_size": torch.ones(3, device=grid.device), | |
| "feat": feat, | |
| "offset": offset, | |
| }) | |
| point.serialization(order=self.order, shuffle_orders=self.training) | |
| point.sparsify() | |
| return point | |
| def forward(self, tokens, err, grids, aux=None): | |
| """tokens [B,K,D], err [B,K,E], grids [(grid_coord, offset)] cycled over by the blocks -> [B,K,z_dim].""" | |
| B, K = tokens.shape[:2] | |
| parts = [self.token_proj(tokens), err.to(tokens.dtype)] | |
| if aux is not None: | |
| parts.append(aux.to(tokens.dtype)) | |
| feat = self.fuse(torch.cat(parts, dim=-1)) | |
| feat = feat.reshape(B * K, -1) | |
| if self.attn_mode in ("knn", "inframe", "xframe"): | |
| if self.attn_mode == "knn": | |
| grid_coord, offset = grids[0] | |
| idx = _knn_indices(grid_coord.reshape(-1, 3).float(), offset, self.knn_k) | |
| else: | |
| idx = grids[0][1] | |
| for blk in self.blocks: | |
| feat = _knn_block_forward(blk, feat, idx) | |
| self.window = -idx.shape[1] # negative marks gather attention in the log line | |
| feat = self.norm(feat) | |
| if self.head is None: | |
| return feat.view(B, K, -1) | |
| return self.head(feat).view(B, K, self.z_dim) | |
| point, active = None, None | |
| for i, blk in enumerate(self.blocks): | |
| grid_coord, offset = grids[i % len(grids)] | |
| if active is not grid_coord: | |
| point = self._point(grid_coord, offset, feat) | |
| active = grid_coord | |
| else: | |
| point["feat"] = feat | |
| point = blk(point) | |
| feat = point.feat | |
| self.window = self.blocks[0].attn.patch_size | |
| feat = self.norm(feat) | |
| if self.head is None: | |
| return feat.view(B, K, -1) | |
| return self.head(feat).view(B, K, self.z_dim) | |
| class FineSplat(nn.Module): | |
| """FineSplat: Stage 1 (2Xplat) + Stage 2; ``forward`` returns the target-view render after every step.""" | |
| def __init__(self, model, conf: dict): | |
| super().__init__() | |
| self.e2e = conf.get("e2e", False) # true: train Stage 1 jointly; default: Stage 1 frozen | |
| self.grad_ckpt = conf.get("grad_ckpt", False) # recompute each step's Stage-2 forward in backward | |
| self.model = model | |
| if not self.e2e: | |
| self.model.requires_grad_(False) | |
| self.model.eval() | |
| self.conf = conf | |
| self.mode = conf.get("mode", "token") | |
| assert self.mode in ("token", "gaussian"), self.mode | |
| assert model.width == model.height, "square images assumed (side, home-patch lookup)" | |
| self.image_size = model.width | |
| expert = model.appearance_expert | |
| self.G = expert.gaussian_side ** 2 | |
| if self.mode == "token": | |
| self.patch = expert.patch_size | |
| else: | |
| # one point per Gaussian; the (v hh ph ww pw) unfold makes a view's Gaussians a | |
| # plain grid of sub-cells, so the same home-patch lookup works with the sub-cell stride | |
| self.patch = expert.gaussian_stride | |
| self.side = self.image_size // self.patch | |
| self.num_steps = conf.get("num_steps", 4) | |
| # Context views re-rendered with gradients per forward so the loss can score them | |
| # (int for every view count, or dict view_count -> quota); 0 disables. | |
| self.context_loss_views = conf.get("context_loss_views", 0) | |
| self.keep_step_gaussians = False | |
| err = conf["error"] | |
| self.error = ErrorMapEncoder( | |
| self.image_size, err["stride"], err["width"], | |
| use_rgb_diff=err.get("use_rgb_diff", True), | |
| use_image_embed=err.get("use_image_embed", True), | |
| backbone=err.get("backbone", "resnet18"), | |
| use_feat_diff=err.get("use_feat_diff", True), | |
| image_embed=err.get("image_embed", "gt_rgb"), | |
| mvp_dim=model.dim1, | |
| img_width=err.get("img_width"), | |
| ) | |
| self.lifter = ErrorLifter( | |
| self.error.channels, err["lift_channels"], err["occlusion_tol"], | |
| err["reduce"], err.get("lift", "mean+home"), | |
| ) | |
| self.error_global = ( | |
| ErrorGlobalAttn( | |
| self.error.channels, | |
| down=err.get("global_attn_down", 4), | |
| attn_channels=err.get("global_attn_channels", 64), | |
| num_heads=err.get("global_attn_heads", 1), | |
| depth=err.get("global_attn_blocks", 1), | |
| ) if err.get("global_attn", False) else None | |
| ) | |
| # gaussian_token_feat: per-Gaussian points carry (raw params | parent token | error); | |
| # gaussian_delta_z: same points, output is a token residual (G features concatenated in (ph, pw) order) | |
| self.gaussian_token_feat = self.mode == "gaussian" and conf.get("gaussian_token_feat", False) | |
| self.gaussian_delta_z = self.gaussian_token_feat and conf.get("gaussian_delta_z", False) | |
| point_dim = model.dim3 if self.mode == "token" else self.raw_dim(model) | |
| if self.gaussian_token_feat: | |
| self.token_broadcast = nn.Sequential( | |
| nn.LayerNorm(model.dim3), nn.Linear(model.dim3, conf.get("token_feat_channels", 512))) | |
| self.stage2 = Stage2( | |
| conf["network"], point_dim, err["lift_channels"], | |
| None if self.gaussian_delta_z else point_dim, | |
| aux_channels=conf.get("token_feat_channels", 512), raw_input=True) | |
| if self.gaussian_delta_z: | |
| self.token_head = nn.Linear(self.G * conf["network"]["width"], model.dim3) | |
| nn.init.zeros_(self.token_head.weight) | |
| nn.init.zeros_(self.token_head.bias) | |
| else: | |
| self.stage2 = Stage2(conf["network"], point_dim, err["lift_channels"], point_dim) | |
| self._logged = False | |
| def _to_gaussians(self, x: torch.Tensor, v: int) -> torch.Tensor: | |
| """Broadcast a per-token tensor (B, K, C) onto its Gaussians in the (v hh ph ww pw) unfold | |
| order: a token's Gaussians are not adjacent, so a plain repeat would be wrong.""" | |
| side = self.image_size // self.model.appearance_expert.patch_size | |
| g = self.model.appearance_expert.gaussian_side | |
| x = rearrange(x, "b (v hh ww) c -> b v hh ww c", v=v, hh=side, ww=side) | |
| return rearrange(x.unsqueeze(3).unsqueeze(5).expand(-1, -1, -1, g, -1, g, -1), | |
| "b v hh ph ww pw c -> b (v hh ph ww pw) c") | |
| def _gs_to_token(self, x: torch.Tensor, v: int) -> torch.Tensor: | |
| """Inverse grouping of ``_to_gaussians``: per-Gaussian (B, N, C) -> per-token | |
| (B, K, G*C), the G features concatenated in their fixed (ph, pw) slot order.""" | |
| side = self.image_size // self.model.appearance_expert.patch_size | |
| g = self.model.appearance_expert.gaussian_side | |
| return rearrange(x, "b (v hh ph ww pw) c -> b (v hh ww) (ph pw c)", | |
| v=v, hh=side, ww=side, ph=g, pw=g) | |
| def raw_dim(model) -> int: | |
| """Width of one Gaussian's raw parameter vector: depth + offset + SH + scale + quat + opacity.""" | |
| e = model.appearance_expert | |
| return 1 + 2 + e.color_dim + 3 + 4 + e.opacity_dim | |
| def train(self, mode: bool = True): | |
| super().train(mode) | |
| if not self.e2e: | |
| self.model.eval() | |
| return self | |
| def trainable_parameters(self): | |
| return [p for p in self.parameters() if p.requires_grad] | |
| def _render_views(self, gaussians, c2w, fxfycxcy, with_grad: bool, render_mode="RGB", | |
| size: int | None = None): | |
| """Render every view -> (B,V,3,H,W) rgb and, for "RGB+ED", (B,V,H,W) depth; ``size`` sets the raster size only.""" | |
| model = self.model | |
| h = w = self.image_size if size is None else size | |
| b, v = c2w.shape[:2] | |
| if with_grad: | |
| dirs = gaussians.xyz[:, None, :, :] - c2w[..., :3, 3][..., None, :] | |
| opacity = torch.broadcast_to( | |
| gaussians.opacity[..., None, :, :, :], (b, v, gaussians.opacity.shape[1], -1, 1), | |
| ) | |
| opacity = _spherical_harmonics(model.opacity_degree, dirs, opacity) | |
| out = GaussianRenderer.apply( | |
| gaussians.xyz, gaussians.feature, gaussians.scale, gaussians.rotation, opacity, | |
| c2w, fxfycxcy, w, h, model.sh_degree, model.near_plane, model.far_plane, | |
| ) | |
| return out.permute(0, 1, 4, 2, 3).contiguous(), None | |
| # "RGB+ED" is not exposed through GaussianRenderer.apply: rasterize one view at a time | |
| rows = [] | |
| with torch.no_grad(), torch.autocast(device_type="cuda", enabled=False): | |
| for ib in range(b): | |
| xyz, feature = gaussians.xyz[ib], gaussians.feature[ib] | |
| scale, rotation = gaussians.scale[ib], gaussians.rotation[ib] | |
| opac = gaussians.opacity[ib] | |
| outs = [] | |
| for i in range(v): | |
| d = xyz - c2w[ib, i:i + 1, :3, 3][None, ...] | |
| opacity_i = _spherical_harmonics(model.opacity_degree, d, opac[None, ...])[0] | |
| outs.append(GaussianRenderer.render( | |
| xyz, feature, scale, rotation, opacity_i, c2w[ib, i], fxfycxcy[ib, i], | |
| w, h, model.sh_degree, model.near_plane, model.far_plane, render_mode, | |
| )) | |
| rows.append(torch.cat(outs, dim=0)) | |
| out = torch.stack(rows, dim=0) | |
| rgb = out[..., :3].permute(0, 1, 4, 2, 3).contiguous() | |
| depth = out[..., 3] if out.shape[-1] == 4 else None | |
| return rgb, depth | |
| def forward(self, input_data: dict, target_data: dict) -> dict: | |
| model, expert = self.model, self.model.appearance_expert | |
| h = w = self.image_size | |
| base_ctx = (contextlib.nullcontext() if (self.e2e and self.training) | |
| else torch.no_grad()) | |
| with base_ctx: | |
| cameras = model.geometry_expert.predict_cameras(input_data, target_data) | |
| raymap_images, Ks, i_w2c = model._build_raymap_input( | |
| input_data, cameras.pred_i_fxfycxcy, cameras.pred_i_c2w, | |
| ) | |
| v = raymap_images.shape[1] | |
| # model.dynamic_grouping: Stage 1 sees a sampled view-group size while it is being | |
| # trained (end-to-end); frozen or at inference it keeps model.group_size. | |
| with expert.grouping(v, enabled=self.e2e and self.training): | |
| x, s1 = expert._stage1_tokenize_and_encode(raymap_images) | |
| x, s2 = expert._stage2_cross_view_and_downsample(x, i_w2c, Ks, h, w) | |
| s3 = expert._stage3_global_cross_view(x, i_w2c, Ks) | |
| tokens = expert.fuse_tokens(s1, s2, s3, v, h, w).float() | |
| raw = expert.raw_from_tokens(tokens, v, h, w) | |
| g0 = expert.gaussians_from_raw( | |
| raw, cameras.pred_i_fxfycxcy, cameras.pred_i_c2w, cameras.pred_t_c2w, h, w, | |
| ) | |
| render0 = None | |
| if self.e2e and self.training: | |
| with torch.autocast(device_type="cuda", enabled=False): | |
| render0, _ = self._render_views( | |
| g0, cameras.pred_t_c2w, cameras.pred_t_fxfycxcy, with_grad=True, | |
| ) | |
| B = tokens.shape[0] | |
| K = tokens.shape[1] if self.mode == "token" else raw.shape[1] | |
| assert K == v * self.side ** 2, (K, v, self.side) | |
| gt_ctx = rearrange(input_data["image"], "b v c h w -> (b v) c h w").float() | |
| gt_feat = self.error.features(gt_ctx) | |
| mvp_feat = None | |
| if self.error.image_embed == "mvp_s1": | |
| gh = self.image_size // expert.patch_size | |
| mvp_feat = rearrange(s1.detach(), "n (hh ww) d -> n d hh ww", hh=gh, ww=gh) | |
| img_map = self.error.image_map(gt_ctx, mvp_feat) | |
| points = tokens if self.mode == "token" else raw.float() | |
| aux =(self._to_gaussians(self.token_broadcast(tokens), v) | |
| if self.gaussian_token_feat and not self.gaussian_delta_z else None) | |
| g, renders, n_vis = g0, [], None | |
| gaussians_all = [] | |
| ctx_renders, ctx_indices = [], None | |
| k_ctx = (self.context_loss_views if isinstance(self.context_loss_views, int) | |
| else self.context_loss_views.get(v, 0)) | |
| if k_ctx and self.training: | |
| ctx_indices = torch.randperm(v, device=gt_ctx.device)[:min(k_ctx, v)] | |
| for _ in range(self.num_steps): | |
| with torch.no_grad(): | |
| ctx_rgb, ctx_depth = self._render_views( | |
| g, cameras.pred_i_c2w, cameras.pred_i_fxfycxcy, | |
| with_grad=False, render_mode="RGB+ED", | |
| ) | |
| if ctx_indices is not None: | |
| # separate differentiable render: the error-map render above must stay detached | |
| with torch.autocast(device_type="cuda", enabled=False): | |
| sub, _ = self._render_views( | |
| g, cameras.pred_i_c2w[:, ctx_indices], | |
| cameras.pred_i_fxfycxcy[:, ctx_indices], with_grad=True, | |
| ) | |
| ctx_renders.append(sub) | |
| E = self.error(rearrange(ctx_rgb, "b v c h w -> (b v) c h w").float(), gt_ctx, | |
| gt_feat, mvp_feat, img=img_map) | |
| E = rearrange(E, "(b v) c h w -> b v c h w", b=B) | |
| if self.error_global is not None: | |
| E = self.error_global(E) | |
| coord = self._token_coord(g.xyz, v) if self.mode == "token" else g.xyz | |
| err, n_vis = self.lifter( | |
| E, ctx_depth, coord, cameras.pred_i_c2w, cameras.pred_i_fxfycxcy, | |
| self.patch, self.side, self.image_size, | |
| ) | |
| if self.stage2.attn_mode in ("inframe", "xframe"): | |
| grids = [("idx", _grid_neighbors( | |
| self.stage2.attn_mode, B, v, self.side, coord.device, coord=coord, | |
| c2w=cameras.pred_i_c2w, fxfycxcy=cameras.pred_i_fxfycxcy, | |
| depth=ctx_depth.view(B, v, *ctx_depth.shape[-2:]), patch=self.patch, | |
| size=self.image_size, tol=self.conf["error"]["occlusion_tol"]))] | |
| else: | |
| grids = self._grids(coord, g, B, K) | |
| if self.gaussian_delta_z: | |
| aux = self._to_gaussians(self.token_broadcast(tokens), v) | |
| feat = self._run_stage2(points.detach(), err, grids, aux=aux) | |
| tokens = tokens + self.token_head(self._gs_to_token(feat, v)) | |
| elif self.gaussian_token_feat: | |
| points = points + self._run_stage2(points.detach(), err, grids, aux=aux) | |
| else: | |
| points = points + self._run_stage2(points, err, grids) | |
| # Decode under the ambient autocast (as inference does) so a zero residual | |
| # reproduces Stage 1 bit for bit. | |
| if self.mode == "token": | |
| g = expert.decode_from_tokens( | |
| points, cameras.pred_i_fxfycxcy, cameras.pred_i_c2w, | |
| cameras.pred_t_c2w, h, w, | |
| ) | |
| else: | |
| # residual on the pre-activation parameters: only the activations are re-run, | |
| # so every bound they impose still holds; delta-z re-decodes raw from the tokens | |
| if self.gaussian_delta_z: | |
| points = expert.raw_from_tokens(tokens, v, h, w).float() | |
| g = expert.gaussians_from_raw( | |
| points, cameras.pred_i_fxfycxcy, cameras.pred_i_c2w, | |
| cameras.pred_t_c2w, h, w, | |
| ) | |
| with torch.autocast(device_type="cuda", enabled=False): | |
| render, _ = self._render_views( | |
| g, cameras.pred_t_c2w, cameras.pred_t_fxfycxcy, with_grad=True, | |
| ) | |
| renders.append(render) | |
| if self.keep_step_gaussians: | |
| gaussians_all.append(g) | |
| # each step is supervised on its own render and starts from a detached state | |
| points = points.detach() | |
| if self.gaussian_delta_z: | |
| tokens = tokens.detach() | |
| if not self._logged: | |
| print(f"flat PTv3 ({self.mode}): {K} points/scene, attention window " | |
| f"{self.stage2.window}, {self.num_steps} step(s) (batch of {B}); " | |
| f"tokens visible in {n_vis.mean():.2f} views on average", flush=True) | |
| self._logged = True | |
| return {"render": renders[-1], "renders": renders, "gaussians": g, | |
| "gaussians_all": gaussians_all, "gaussians_base": g0, | |
| "render0": render0, | |
| "tokens_final": points, "tokens_base": tokens, | |
| "context_renders": ctx_renders, "context_indices": ctx_indices, | |
| "cameras": cameras, "error_views": n_vis} | |
| def _run_stage2(self, *args, **kwargs): | |
| if self.grad_ckpt and self.training and torch.is_grad_enabled(): | |
| return torch.utils.checkpoint.checkpoint(self.stage2, *args, use_reentrant=False, **kwargs) | |
| return self.stage2(*args, **kwargs) | |
| def _token_coord(self, xyz: torch.Tensor, v: int) -> torch.Tensor: | |
| """Mean position of each token's Gaussians, (B, N, 3) -> (B, K, 3), honouring the | |
| (v hh ph ww pw) unfold order (a plain (K, G) reshape would mix neighbouring tokens).""" | |
| side = self.image_size // self.model.appearance_expert.patch_size | |
| g = self.model.appearance_expert.gaussian_side | |
| xyz = rearrange(xyz, "b (v hh ph ww pw) c -> b (v hh ww) (ph pw) c", | |
| v=v, hh=side, ww=side, ph=g, pw=g) | |
| return xyz.mean(dim=2) | |
| def _grids(self, coord, g0, b: int, k: int): | |
| """World-space serialization grid: the Stage-1 Gaussians' bounding box split | |
| ``grid_resolution`` ways per axis, one point cloud per scene -> [(grid, offset)].""" | |
| lo = g0.xyz.amin(dim=-2, keepdim=True) | |
| span = (g0.xyz.amax(dim=-2, keepdim=True) - lo).clamp(min=EPS) | |
| res = self.conf["network"]["grid_resolution"] | |
| grid = (((coord - lo) / span) * res).floor().int().clamp(0, res - 1) | |
| offset = torch.arange(1, b + 1, device=coord.device) * k | |
| return [(grid, offset)] | |
| def render_baseline_targets(self, gaussians_base, cameras) -> torch.Tensor: | |
| """Render the frozen Gaussians at the target views, for side-by-side logging.""" | |
| rgb, _ = self._render_views( | |
| gaussians_base, cameras.pred_t_c2w, cameras.pred_t_fxfycxcy, with_grad=False, | |
| ) | |
| return rgb | |