src/model/encoder/encoder_codecsplat.py

import io
import time
import zipfile
import zlib
from dataclasses import dataclass, field
from typing import Literal, Optional, List

import numpy as np
import torch
from einops import rearrange, repeat
from torch import nn

from ...dataset.shims.patch_shim import apply_patch_shim
from ...dataset.types import BatchedExample, DataShim
from ...geometry.projection import sample_image_grid
from ..types import Gaussians
from .common.depth_guided_mv_refiner import DepthGuidedMVRefiner
from .common.gaussian_adapter import GaussianAdapter, GaussianAdapterCfg
from .common.gs_decoder import GaussianLatentDecoder, GaussianLatentDepthDecoder
from .encoder import Encoder
from .vae import coding
from .vae import VariableRateLossyVAE
from .vae.config import Config as VaeConfig

import torch.nn.functional as F
from numpy.lib.format import write_array

from .unimatch.mv_unimatch import MultiViewUniMatch
from .unimatch.dpt_head import DPTHead


@dataclass
class EncoderCodecSplatCfg:
    name: Literal["codecsplat"]
    num_depth_candidates: int
    num_surfaces: int
    gaussian_adapter: GaussianAdapterCfg
    downscale_factor: int
    shim_patch_size: int

    # mv_unimatch
    num_scales: int
    upsample_factor: int
    lowest_feature_resolution: int
    depth_unet_channels: int
    grid_sample_disable_cudnn: bool

    # color branch
    init_sh_input_img: bool
    gaussian_regressor_channels: int

    # loss config
    supervise_intermediate_depth: bool
    return_depth: bool

    # monodepth config
    monodepth_vit_type: str

    # multi-view matching
    local_mv_match: int

    # feature codec path
    enable_feature_codec: bool = False
    enable_depth_guided_mv_refiner: bool = False
    enable_post_vae_mv_refiner: bool = False
    codec_loss_weight: float = 1.0
    codec_depth_loss_weight: float = 0.1
    codec_lmb_range: List[float] = field(default_factory=lambda: [16.0, 1024.0])
    freeze_legacy_modules_for_codec: bool = False
    codec_estimated_feature_num_bytes: bool = False

    # Test-only ablation: entropy-code the stage-1 latent feature/depth tensors
    # with plain deflate instead of using the learned VAE codec.
    deflate_feature_ablation: bool = False
    deflate_feature_dtype: Literal["float16", "float32"] = "float16"
    deflate_feature_level: int = 9


class EncoderCodecSplat(Encoder[EncoderCodecSplatCfg]):
    def __init__(self, cfg: EncoderCodecSplatCfg) -> None:
        super().__init__(cfg)
        if cfg.enable_feature_codec:
            assert not cfg.init_sh_input_img, (
                "Feature codec mode requires init_sh_input_img=False because "
                "decoder-side reconstruction cannot access original input images."
            )
        if cfg.deflate_feature_ablation:
            assert not cfg.enable_feature_codec, (
                "deflate_feature_ablation is a non-learned baseline and should be "
                "run with enable_feature_codec=False."
            )

        self.depth_predictor = MultiViewUniMatch(
            num_scales=cfg.num_scales,
            upsample_factor=cfg.upsample_factor,
            lowest_feature_resolution=cfg.lowest_feature_resolution,
            num_depth_candidates=cfg.num_depth_candidates,
            vit_type=cfg.monodepth_vit_type,
            unet_channels=cfg.depth_unet_channels,
            grid_sample_disable_cudnn=cfg.grid_sample_disable_cudnn,
        )

        self.mv_refiner = None
        self.post_mv_refiner = None
        self.feature_codec = None
        self.gs_decoder = None
        self.codec_depth_head = None

        # upsample features to the original resolution
        model_configs = {
            'vits': {'in_channels': 384, 'features': 64, 'out_channels': [48, 96, 192, 384]},
            'vitb': {'in_channels': 768, 'features': 96, 'out_channels': [96, 192, 384, 768]},
            'vitl': {'in_channels': 1024, 'features': 128, 'out_channels': [128, 256, 512, 1024]},
        }

        self.feature_upsampler = DPTHead(**model_configs[cfg.monodepth_vit_type],
                                        downsample_factor=cfg.upsample_factor,
                                        return_feature=True,
                                        num_scales=cfg.num_scales,
                                        )
        feature_upsampler_channels = model_configs[cfg.monodepth_vit_type]["features"]

        # gaussians adapter
        self.gaussian_adapter = GaussianAdapter(cfg.gaussian_adapter)

        # If a pre-refiner is enabled, let it handle explicit depth conditioning and keep
        # the gaussian_regressor input focused on image / match / feature cues.
        self.gaussian_regressor_uses_depth = not cfg.enable_depth_guided_mv_refiner
        gaussian_regressor_in_ch = 3 + 1 + feature_upsampler_channels
        if self.gaussian_regressor_uses_depth:
            gaussian_regressor_in_ch += 1
        channels = self.cfg.gaussian_regressor_channels

        self.gaussian_regressor = nn.Sequential(
            nn.Conv2d(gaussian_regressor_in_ch, channels, 3, 1, 1),
            nn.GELU(),
            nn.Conv2d(channels, channels, 3, 1, 1),
        )

        # predict gaussian parameters: scale, q, sh, offset, opacity
        num_gaussian_parameters = self.gaussian_adapter.d_in + 2 + 1
        self.num_gaussian_parameters = num_gaussian_parameters
        self.gaussian_regressor_out_channels = channels

        if self.cfg.enable_depth_guided_mv_refiner:
            self.mv_refiner = DepthGuidedMVRefiner(
                feature_channels=channels,
                hidden_channels=channels,
                concat_ref_depth=True,
            )
        if self.cfg.enable_feature_codec:
            codec_lmb_range = tuple(float(x) for x in self.cfg.codec_lmb_range)
            self.feature_codec = VariableRateLossyVAE(
                channels,
                VaeConfig(
                    lmb_range=codec_lmb_range,
                    lmb=float(codec_lmb_range[-1]),
                ),
            )
            self.codec_depth_head = GaussianLatentDepthDecoder(
                latent_channels=channels,
                hidden_channels=channels,
            )
            if self.cfg.enable_post_vae_mv_refiner:
                self.post_mv_refiner = DepthGuidedMVRefiner(
                    feature_channels=channels,
                    hidden_channels=channels,
                )
        self.gs_decoder = GaussianLatentDecoder(
            latent_channels=channels,
            hidden_channels=channels,
            num_gaussian_parameters=num_gaussian_parameters,
            init_sh_input_img=self.cfg.init_sh_input_img,
        )
        if self.cfg.enable_feature_codec:
            if self.cfg.freeze_legacy_modules_for_codec:
                self.freeze_legacy_modules_for_codec()

    def freeze_legacy_modules_for_codec(self) -> None:
        for param in self.parameters():
            param.requires_grad_(False)

        modules_to_train = [
            self.gaussian_regressor,
            self.mv_refiner,
            self.feature_codec,
            self.codec_depth_head,
            self.post_mv_refiner,
            self.gs_decoder,
        ]

        for module in modules_to_train:
            if module is None:
                continue
            for param in module.parameters():
                param.requires_grad_(True)

    def forward(
        self,
        context: dict,
        global_step: int,
        visualization_dump: Optional[dict] = None,
        scene_names: Optional[list] = None,
        return_bitstream: bool = False,
    ):
        if self.training:
            return self.normal_forward(
                context,
                global_step,
                visualization_dump=visualization_dump,
                scene_names=scene_names,
            )
        if self.cfg.deflate_feature_ablation:
            return self.deflate_feature_forward(
                context,
                global_step,
                visualization_dump=visualization_dump,
                scene_names=scene_names,
            )
        if self.feature_codec is not None:
            return self.codec_forward(
                context,
                global_step,
                visualization_dump=visualization_dump,
                scene_names=scene_names,
                return_bitstream=return_bitstream,
            )
        return self.normal_forward(
            context,
            global_step,
            visualization_dump=visualization_dump,
            scene_names=scene_names,
        )

    def compute_view_neighbor_indices(self, extrinsics: torch.Tensor):
        v = extrinsics.shape[1]
        if v <= 3:
            return None

        with torch.no_grad():
            xyzs = extrinsics[:, :, :3, -1].detach().float()
            cameras_dist_matrix = torch.cdist(xyzs, xyzs, p=2)
            cameras_dist_index = torch.argsort(cameras_dist_matrix)
            return cameras_dist_index[:, :, :(self.cfg.local_mv_match + 1)]

    def encode_camera_stream(self, context: dict) -> bytes:
        _, v, _, h, w = context["image"].shape
        camera_arrays = {
            "image_shape": np.array([h, w], dtype=np.int32),
            "num_context_views": np.array([v], dtype=np.int32),
            "intrinsics": context["intrinsics"][0].detach().cpu().float().numpy(),
            "extrinsics": context["extrinsics"][0].detach().cpu().float().numpy(),
            "near": context["near"][0].detach().cpu().float().numpy(),
            "far": context["far"][0].detach().cpu().float().numpy(),
        }
        camera_buffer = io.BytesIO()
        with zipfile.ZipFile(
            camera_buffer,
            mode="w",
            compression=zipfile.ZIP_DEFLATED,
            compresslevel=9,
        ) as camera_zip:
            for name, array in camera_arrays.items():
                array_buffer = io.BytesIO()
                write_array(array_buffer, np.asanyarray(array), allow_pickle=False)
                camera_zip.writestr(f"{name}.npy", array_buffer.getvalue())
        camera_stream = camera_buffer.getvalue()
        camera_buffer.close()
        return camera_stream

    @torch.no_grad()
    def compress(self, context: dict, features: torch.Tensor) -> tuple[bytes, dict]:
        assert self.feature_codec is not None
        b, v, _, h, w = context["image"].shape
        assert b == 1, "Scene bitstream export currently supports batch size 1."
        assert features.shape[0] == v, (
            f"Expected one feature tensor per view, got {features.shape[0]=}, {v=}."
        )
        assert v > 1, "Scene feature stream packing expects at least two context views."

        if features.is_cuda:
            torch.cuda.synchronize(features.device)
        start = time.perf_counter()

        camera_stream = self.encode_camera_stream(context)
        feature_streams = [
            self.feature_codec.compress(features[i : i + 1])
            for i in range(features.shape[0])
        ]
        feature_stream = coding.concat_bytes_list(feature_streams)

        camera_num_bytes = len(camera_stream)
        assert camera_num_bytes < (1 << 24), (
            f"Camera stream is too large for the 3-byte length header: {camera_num_bytes}"
        )
        bytestream = (
            camera_num_bytes.to_bytes(3, byteorder="big", signed=False)
            + camera_stream
            + feature_stream
        )

        if features.is_cuda:
            torch.cuda.synchronize(features.device)
        compress_seconds = time.perf_counter() - start

        return bytestream, {
            "codec_compress_seconds": compress_seconds,
            "codec_camera_num_bytes": float(camera_num_bytes),
            "codec_feature_num_bytes": float(len(feature_stream)),
            "codec_num_bytes": float(len(bytestream)),
        }

    @torch.no_grad()
    def run_deflate_feature_ablation(
        self,
        context: dict,
        features: torch.Tensor,
        depth: torch.Tensor,
    ) -> tuple[torch.Tensor, torch.Tensor, dict]:
        assert not self.training, "deflate_feature_ablation is intended for evaluation."

        device = features.device
        if features.is_cuda:
            torch.cuda.synchronize(device)
        compress_start = time.perf_counter()

        dtype = self.cfg.deflate_feature_dtype
        np_dtype = np.float16 if dtype == "float16" else np.float32
        level = int(self.cfg.deflate_feature_level)

        def tensor_to_npy_bytes(tensor: torch.Tensor) -> bytes:
            tensor_np = tensor.detach().cpu().to(torch.float32).numpy().astype(
                np_dtype,
                copy=False,
            )
            tensor_buffer = io.BytesIO()
            write_array(tensor_buffer, np.asanyarray(tensor_np), allow_pickle=False)
            stream = tensor_buffer.getvalue()
            tensor_buffer.close()
            return stream

        camera_stream = self.encode_camera_stream(context)
        feature_raw_stream = tensor_to_npy_bytes(features)
        depth_raw_stream = tensor_to_npy_bytes(depth)
        feature_stream = zlib.compress(feature_raw_stream, level=level)
        depth_stream = zlib.compress(depth_raw_stream, level=level)

        if features.is_cuda:
            torch.cuda.synchronize(device)
        compress_seconds = time.perf_counter() - compress_start

        if features.is_cuda:
            torch.cuda.synchronize(device)
        decompress_start = time.perf_counter()

        features_hat = torch.from_numpy(
            np.load(io.BytesIO(zlib.decompress(feature_stream)), allow_pickle=False)
        ).to(device=device, dtype=features.dtype)
        depth_hat = torch.from_numpy(
            np.load(io.BytesIO(zlib.decompress(depth_stream)), allow_pickle=False)
        ).to(device=device, dtype=depth.dtype)

        if features.is_cuda:
            torch.cuda.synchronize(device)
        decompress_seconds = time.perf_counter() - decompress_start

        header_num_bytes = 24  # three uint64 lengths for camera/feature/depth streams
        codec_num_bytes = (
            header_num_bytes
            + len(camera_stream)
            + len(feature_stream)
            + len(depth_stream)
        )
        return features_hat, depth_hat, {
            "codec_compress_seconds": compress_seconds,
            "codec_decompress_seconds": decompress_seconds,
            "codec_camera_num_bytes": float(len(camera_stream)),
            "codec_feature_num_bytes": float(len(feature_stream) + len(depth_stream)),
            "codec_latent_feature_num_bytes": float(len(feature_stream)),
            "codec_depth_num_bytes": float(len(depth_stream)),
            "codec_feature_raw_num_bytes": float(len(feature_raw_stream)),
            "codec_depth_raw_num_bytes": float(len(depth_raw_stream)),
            "codec_num_bytes": float(codec_num_bytes),
        }

    @torch.no_grad()
    def decompress(
        self,
        bytestream: bytes,
        visualization_dump: Optional[dict] = None,
    ) -> tuple[Gaussians, Optional[torch.Tensor], dict, dict]:
        assert self.feature_codec is not None
        param = next(self.parameters())
        device = param.device

        if param.is_cuda:
            torch.cuda.synchronize(device)
        start = time.perf_counter()

        camera_num_bytes = int.from_bytes(
            bytestream[:3], byteorder="big", signed=False
        )
        camera_stream = bytestream[3 : 3 + camera_num_bytes]
        feature_stream = bytestream[3 + camera_num_bytes :]

        with np.load(io.BytesIO(camera_stream)) as camera_data:
            image_shape = tuple(int(x) for x in camera_data["image_shape"])
            num_context_views = int(camera_data["num_context_views"][0])
            intrinsics_np = camera_data["intrinsics"]
            extrinsics_np = camera_data["extrinsics"]
            near_np = camera_data["near"]
            far_np = camera_data["far"]

        feature_streams = coding.split_bytes_list(feature_stream, num_context_views)
        features_hat = torch.cat(
            [self.feature_codec.decompress(stream) for stream in feature_streams],
            dim=0,
        )

        decoded_context = {}
        decoded_context["image_shape"] = image_shape
        decoded_context["intrinsics"] = torch.from_numpy(intrinsics_np).to(
            device=device,
            dtype=torch.float32,
        ).unsqueeze(0)
        decoded_context["extrinsics"] = torch.from_numpy(extrinsics_np).to(
            device=device,
            dtype=torch.float32,
        ).unsqueeze(0)
        decoded_context["near"] = torch.from_numpy(near_np).to(
            device=device,
            dtype=torch.float32,
        ).unsqueeze(0)
        decoded_context["far"] = torch.from_numpy(far_np).to(
            device=device,
            dtype=torch.float32,
        ).unsqueeze(0)

        neighbor_indices = self.compute_view_neighbor_indices(
            decoded_context["extrinsics"]
        )
        decode_state = {
            "device": device,
            "b": 1,
            "v": num_context_views,
            "h": image_shape[0],
            "w": image_shape[1],
            "refiner_source_indices": (
                neighbor_indices[..., 1:] if neighbor_indices is not None else None
            ),
        }
        gaussians_raw, depth, _, _ = self.decode_codec_features(
            features_hat,
            decoded_context,
            decode_state,
            codec_depth_targets=None,
        )
        decode_state["depth_preds"] = [depth]
        gaussians, depths = self.adapt_gaussians(
            decoded_context,
            decoded_context,
            gaussians_raw,
            depth,
            decode_state,
            visualization_dump,
        )

        if param.is_cuda:
            torch.cuda.synchronize(device)
        decompress_seconds = time.perf_counter() - start

        return gaussians, depths, decoded_context, {
            "codec_decompress_seconds": decompress_seconds,
        }

    def encode_context_features(self, context: dict):
        device = context["image"].device
        b, v, _, h, w = context["image"].shape

        cameras_dist_index = self.compute_view_neighbor_indices(context["extrinsics"])
        refiner_source_indices = (
            cameras_dist_index[..., 1:] if cameras_dist_index is not None else None
        )

        results_dict = self.depth_predictor(
            context["image"],
            attn_splits_list=[2],
            min_depth=1. / context["far"],
            max_depth=1. / context["near"],
            intrinsics=context["intrinsics"],
            extrinsics=context["extrinsics"],
            nn_matrix=cameras_dist_index,
            supervise_intermediate_depth=self.cfg.supervise_intermediate_depth,
        )

        depth_preds = results_dict['depth_preds']
        depth = depth_preds[-1]

        features = self.feature_upsampler(
            results_dict["features_mono_intermediate"],
            cnn_features=results_dict["features_cnn_all_scales"][::-1],
            mv_features=results_dict["features_mv"][0]
            if self.cfg.num_scales == 1
            else results_dict["features_mv"][::-1],
        )

        match_prob = results_dict['match_probs'][-1]
        match_prob = torch.max(match_prob, dim=1, keepdim=True)[0]
        match_prob = F.interpolate(match_prob, size=depth.shape[-2:], mode='nearest')

        regressor_inputs = [
            rearrange(context["image"], "b v c h w -> (b v) c h w"),
            match_prob,
            features,
        ]
        if self.mv_refiner is None:
            regressor_inputs.insert(1, rearrange(depth, "b v h w -> (b v) () h w"))

        out = self.gaussian_regressor(torch.cat(regressor_inputs, dim=1))

        if self.mv_refiner is not None:
            out = rearrange(out, "(b v) c h w -> b v c h w", b=b, v=v)
            out = self.mv_refiner(
                out,
                depth,
                context["intrinsics"],
                context["extrinsics"],
                source_view_indices=refiner_source_indices,
            )
            out = rearrange(out, "b v c h w -> (b v) c h w")

        return {
            "device": device,
            "b": b,
            "v": v,
            "h": h,
            "w": w,
            "depth_preds": depth_preds,
            "depth": depth,
            "teacher_depth": depth,
            "features": features,
            "match_prob": match_prob,
            "out": out,
            "refiner_source_indices": refiner_source_indices,
        }

    def decode_codec_features(
        self,
        out_hat: torch.Tensor,
        decode_context: dict,
        state: dict,
        codec_depth_targets: Optional[torch.Tensor] = None,
    ):
        device = state["device"]
        b = state["b"]
        v = state["v"]

        near_flat = rearrange(decode_context["near"], "b v -> (b v) () () ()")
        far_flat = rearrange(decode_context["far"], "b v -> (b v) () () ()")
        (
            codec_depth_flat,
            _codec_inv_depth_flat,
            codec_norm_inv_depth_flat,
        ) = self.codec_depth_head(
            out_hat,
            near_flat,
            far_flat,
            return_inverse=True,
            return_normalized_inverse=True,
        )
        codec_depth = rearrange(codec_depth_flat, "(b v) () h w -> b v h w", b=b, v=v)
        codec_depth_loss_pred = None
        codec_depth_loss_target = None

        if codec_depth_targets is not None:
            codec_norm_inv_depth = rearrange(
                codec_norm_inv_depth_flat, "(b v) () h w -> b v h w", b=b, v=v
            )
            with torch.autocast(device_type=device.type, enabled=False):
                near = rearrange(decode_context["near"].float(), "b v -> b v () ()")
                far = rearrange(decode_context["far"].float(), "b v -> b v () ()")
                inv_min = far.clamp_min(torch.finfo(torch.float32).eps).reciprocal()
                inv_max = near.clamp_min(torch.finfo(torch.float32).eps).reciprocal()
                inv_range = (inv_max - inv_min).clamp_min(torch.finfo(torch.float32).eps)
                teacher_inv_depth = codec_depth_targets.float().clamp_min(
                    torch.finfo(torch.float32).eps
                ).reciprocal()
                teacher_inv_depth = teacher_inv_depth.clamp(min=inv_min, max=inv_max)
                codec_depth_loss_target = ((teacher_inv_depth - inv_min) / inv_range).clamp(0.0, 1.0)
                codec_depth_loss_pred = codec_norm_inv_depth

        if self.post_mv_refiner is not None:
            out_hat = rearrange(out_hat, "(b v) c h w -> b v c h w", b=b, v=v)
            out_hat = self.post_mv_refiner(
                out_hat,
                codec_depth,
                decode_context["intrinsics"],
                decode_context["extrinsics"],
                source_view_indices=state["refiner_source_indices"],
            )
            out_hat = rearrange(out_hat, "b v c h w -> (b v) c h w")

        depth_flat = rearrange(codec_depth, "b v h w -> (b v) () h w")
        assert self.gs_decoder is not None
        gaussians = self.gs_decoder(out_hat, depth_flat)
        return gaussians, codec_depth, codec_depth_loss_pred, codec_depth_loss_target

    def adapt_gaussians(
        self,
        context: dict,
        adapter_context: dict,
        gaussians: torch.Tensor,
        depth: torch.Tensor,
        state: dict,
        visualization_dump: Optional[dict],
    ):
        device = state["device"]
        b = state["b"]
        v = state["v"]
        h = state["h"]
        w = state["w"]
        depth_preds = state["depth_preds"]

        gaussians = rearrange(gaussians, "(b v) c h w -> b v c h w", b=b, v=v)
        depths = rearrange(depth, "b v h w -> b v (h w) () ()")
        raw_gaussians = rearrange(gaussians, "b v c h w -> b v (h w) c")

        if self.cfg.supervise_intermediate_depth and len(depth_preds) > 1:
            num_depths = len(depth_preds)
            intermediate_depths = torch.cat(depth_preds[:(num_depths - 1)], dim=0)
            intermediate_depths = rearrange(
                intermediate_depths, "b v h w -> b v (h w) () ()"
            )
            depths = torch.cat((intermediate_depths, depths), dim=0)
            raw_gaussians = torch.cat([raw_gaussians] * num_depths, dim=0)
            b *= num_depths

        with torch.autocast(device_type=device.type, enabled=False):
            raw_gaussians = raw_gaussians.float()
            depths = depths.float()
            opacities = raw_gaussians[..., :1].sigmoid().unsqueeze(-1)
            raw_gaussians = raw_gaussians[..., 1:]

            xy_ray, _ = sample_image_grid((h, w), device)
            xy_ray = rearrange(xy_ray, "h w xy -> (h w) () xy")
            gaussians = rearrange(
                raw_gaussians,
                "... (srf c) -> ... srf c",
                srf=self.cfg.num_surfaces,
            )
            offset_xy = gaussians[..., :2].sigmoid()
            pixel_size = 1 / torch.tensor((w, h), dtype=torch.float32, device=device)
            xy_ray = xy_ray + (offset_xy - 0.5) * pixel_size

        sh_input_images = context.get("image")
        if self.cfg.init_sh_input_img:
            assert sh_input_images is not None, (
                "init_sh_input_img=True requires input images, but decoded "
                "bitstreams intentionally do not contain original RGB images."
            )
        if self.cfg.supervise_intermediate_depth and len(depth_preds) > 1:
            context_extrinsics = torch.cat(
                [adapter_context["extrinsics"]] * len(depth_preds), dim=0
            )
            context_intrinsics = torch.cat(
                [adapter_context["intrinsics"]] * len(depth_preds), dim=0
            )
            gaussians = self.gaussian_adapter.forward(
                rearrange(context_extrinsics, "b v i j -> b v () () () i j"),
                rearrange(context_intrinsics, "b v i j -> b v () () () i j"),
                rearrange(xy_ray, "b v r srf xy -> b v r srf () xy"),
                depths,
                opacities,
                rearrange(gaussians[..., 2:], "b v r srf c -> b v r srf () c"),
                (h, w),
                input_images=sh_input_images.repeat(len(depth_preds), 1, 1, 1, 1)
                if self.cfg.init_sh_input_img
                else None,
            )
        else:
            gaussians = self.gaussian_adapter.forward(
                rearrange(adapter_context["extrinsics"], "b v i j -> b v () () () i j"),
                rearrange(adapter_context["intrinsics"], "b v i j -> b v () () () i j"),
                rearrange(xy_ray, "b v r srf xy -> b v r srf () xy"),
                depths,
                opacities,
                rearrange(gaussians[..., 2:], "b v r srf c -> b v r srf () c"),
                (h, w),
                input_images=sh_input_images if self.cfg.init_sh_input_img else None,
            )

        if visualization_dump is not None:
            visualization_dump["depth"] = rearrange(
                depths, "b v (h w) srf s -> b v h w srf s", h=h, w=w
            )
            visualization_dump["scales"] = rearrange(
                gaussians.scales, "b v r srf spp xyz -> b (v r srf spp) xyz"
            )
            visualization_dump["rotations"] = rearrange(
                gaussians.rotations, "b v r srf spp xyzw -> b (v r srf spp) xyzw"
            )

        gaussians = Gaussians(
            means=rearrange(gaussians.means, "b v r srf spp xyz -> b (v r srf spp) xyz"),
            covariances=rearrange(
                gaussians.covariances, "b v r srf spp i j -> b (v r srf spp) i j"
            ),
            harmonics=rearrange(
                gaussians.harmonics, "b v r srf spp c d_sh -> b (v r srf spp) c d_sh"
            ),
            opacities=rearrange(gaussians.opacities, "b v r srf spp -> b (v r srf spp)"),
            scales=rearrange(gaussians.scales, "b v r srf spp xyz -> b (v r srf spp) xyz"),
            rotations=rearrange(
                gaussians.rotations, "b v r srf spp xyzw -> b (v r srf spp) xyzw"
            ),
            rotations_unnorm=rearrange(
                gaussians.rotations_unnorm, "b v r srf spp xyzw -> b (v r srf spp) xyzw"
            ),
        )
        if self.cfg.return_depth:
            depths = rearrange(
                depths, "b v (h w) srf s -> b v h w srf s", h=h, w=w
            ).squeeze(-1).squeeze(-1)
        return gaussians, depths

    @torch.no_grad()
    def codec_forward(
        self,
        context: dict,
        global_step: int,
        visualization_dump: Optional[dict] = None,
        scene_names: Optional[list] = None,
        return_bitstream: bool = False,
    ):
        b = context["image"].shape[0]
        assert b == 1, "Codec bitstream evaluation currently supports batch size 1."
        assert self.feature_codec is not None
        assert not self.training

        input_device = context["image"].device
        if context["image"].is_cuda:
            torch.cuda.synchronize(input_device)
        encode_context_start = time.perf_counter()
        state = self.encode_context_features(context)
        if context["image"].is_cuda:
            torch.cuda.synchronize(input_device)
        encode_context_seconds = time.perf_counter() - encode_context_start
        assert "depth_only_result" not in state, "Codec bitstream evaluation needs feature latents."

        estimated_feature_num_bytes = None
        if self.cfg.codec_estimated_feature_num_bytes:
            estimated_feature_bpp, _, _ = self.feature_codec.train_forward(
                state["out"],
                lmb=self.feature_codec.cfg.lmb,
            )
            feature_height, feature_width = state["out"].shape[-2:]
            estimated_feature_num_bytes = (
                estimated_feature_bpp.float().sum()
                * float(feature_height * feature_width)
                / 8.0
            )

        bytestream, compress_stats = self.compress(context, state["out"])
        gaussians, depths, codec_decode_context, decompress_stats = self.decompress(
            bytestream,
            visualization_dump=visualization_dump,
        )
        assert codec_decode_context["image_shape"] == (state["h"], state["w"]), (
            f"Decoded image_shape={codec_decode_context['image_shape']} "
            f"does not match context image_shape={(state['h'], state['w'])}."
        )
        assert codec_decode_context["intrinsics"].shape[1] == state["v"], (
            "Decoded camera stream context view count does not match "
            f"the current context: {codec_decode_context['intrinsics'].shape[1]} vs {state['v']}."
        )

        device = state["device"]
        codec_num_bytes = torch.tensor(
            compress_stats["codec_num_bytes"], device=device, dtype=torch.float32
        )
        total_pixels_per_scene = state["out"].shape[-2] * state["out"].shape[-1]
        result = {
            "gaussians": gaussians,
            "codec_num_bytes": codec_num_bytes,
            "codec_bpp": codec_num_bytes * 8.0 / float(total_pixels_per_scene),
            "codec_encode_context_seconds": torch.tensor(
                encode_context_seconds, device=device, dtype=torch.float32
            ),
            "codec_compress_seconds": torch.tensor(
                compress_stats["codec_compress_seconds"], device=device, dtype=torch.float32
            ),
            "codec_decompress_seconds": torch.tensor(
                decompress_stats["codec_decompress_seconds"], device=device, dtype=torch.float32
            ),
            "codec_camera_num_bytes": torch.tensor(
                compress_stats["codec_camera_num_bytes"], device=device, dtype=torch.float32
            ),
            "codec_feature_num_bytes": torch.tensor(
                compress_stats["codec_feature_num_bytes"], device=device, dtype=torch.float32
            ),
        }
        if self.cfg.return_depth:
            result["depths"] = depths
        if estimated_feature_num_bytes is not None:
            result["codec_estimated_feature_num_bytes"] = (
                estimated_feature_num_bytes
            )
        if return_bitstream:
            result["bitstream"] = bytestream
        return result

    @torch.no_grad()
    def deflate_feature_forward(
        self,
        context: dict,
        global_step: int,
        visualization_dump: Optional[dict] = None,
        scene_names: Optional[list] = None,
    ):
        assert not self.training
        assert self.feature_codec is None
        assert context["image"].shape[0] == 1, (
            "Deflate feature evaluation currently supports batch size 1."
        )

        state = self.encode_context_features(context)
        assert "depth_only_result" not in state, (
            "Deflate feature evaluation needs Gaussian-generation features."
        )

        device = state["device"]
        out, depth, deflate_stats = self.run_deflate_feature_ablation(
            context,
            state["out"],
            state["depth"],
        )
        state["depth"] = depth
        state["depth_preds"] = [*state["depth_preds"][:-1], depth]

        depth_flat = rearrange(depth, "b v h w -> (b v) () h w")
        assert self.gs_decoder is not None
        gaussians_raw = self.gs_decoder(out, depth_flat)
        gaussians, depths = self.adapt_gaussians(
            context,
            context,
            gaussians_raw,
            depth,
            state,
            visualization_dump,
        )

        codec_num_bytes = torch.tensor(
            deflate_stats["codec_num_bytes"],
            device=device,
            dtype=torch.float32,
        )
        total_pixels_per_scene = out.shape[-2] * out.shape[-1]
        result = {
            "gaussians": gaussians,
            "codec_num_bytes": codec_num_bytes,
            "codec_bpp": codec_num_bytes * 8.0 / float(total_pixels_per_scene),
        }
        result.update(
            {
                name: torch.tensor(value, device=device, dtype=torch.float32)
                for name, value in deflate_stats.items()
                if name != "codec_num_bytes"
            }
        )
        if self.cfg.return_depth:
            result["depths"] = depths
        return result

    def normal_forward(
        self,
        context: dict,
        global_step: int,
        visualization_dump: Optional[dict] = None,
        scene_names: Optional[list] = None,
    ):
        state = self.encode_context_features(context)
        b = state["b"]
        v = state["v"]
        out = state["out"]
        depth = state["depth"]

        codec_loss = None
        codec_lmb = None
        codec_depth_targets = None
        codec_depth_loss_pred = None
        codec_depth_loss_target = None

        if self.feature_codec is not None:
            assert self.training, (
                "Learned codec evaluation must use the complete bitstream path."
            )
            codec_depth_targets = state["teacher_depth"].detach()
            scene_lmb = self.feature_codec.sample_lmb(b, out.device)
            view_lmb = repeat(scene_lmb, "b -> (b v)", v=v)
            codec_loss, out_hat, _ = self.feature_codec(out, lmb=view_lmb)
            codec_loss = rearrange(codec_loss, "(b v) -> b v", b=b, v=v)
            codec_loss = codec_loss.sum(dim=1)
            codec_lmb = scene_lmb

            gaussians_raw, depth, codec_depth_loss_pred, codec_depth_loss_target = (
                self.decode_codec_features(
                    out_hat,
                    context,
                    state,
                    codec_depth_targets=codec_depth_targets,
                )
            )
            adapter_context = context
        else:
            depth_flat = rearrange(depth, "b v h w -> (b v) () h w")
            assert self.gs_decoder is not None
            gaussians_raw = self.gs_decoder(out, depth_flat)
            adapter_context = context

        gaussians, depths = self.adapt_gaussians(
            context,
            adapter_context,
            gaussians_raw,
            depth,
            state,
            visualization_dump,
        )
        if self.cfg.return_depth:
            result = {"gaussians": gaussians, "depths": depths}
            if codec_depth_targets is not None:
                result["codec_depth_targets"] = codec_depth_targets
            if codec_depth_loss_pred is not None:
                result["codec_depth_loss_pred"] = codec_depth_loss_pred
            if codec_depth_loss_target is not None:
                result["codec_depth_loss_target"] = codec_depth_loss_target
            if codec_loss is not None:
                return result, codec_loss, codec_lmb
            return result

        if codec_loss is not None:
            return gaussians, codec_loss, codec_lmb
        return gaussians

    def get_data_shim(self) -> DataShim:
        def data_shim(batch: BatchedExample) -> BatchedExample:
            batch = apply_patch_shim(
                batch,
                patch_size=self.cfg.shim_patch_size
                * self.cfg.downscale_factor,
            )

            return batch

        return data_shim