"""DCVC-RT engine for patch-selection bit-cost maps. This wraps the bundled DCVC-RT neural video codec (``neural_codec/DCVC/``) and adds a *bit-cost bitmap* output that DCVC-RT does not expose natively: for every latent spatial location we estimate how many bits the entropy model spends, summed over channels. That per-frame ``(H/16, W/16)`` map is the DCVC-RT analogue of the dev codec's ``bitmap_y`` and is the importance signal used to select which video patches to keep for VLM inference. Design notes ------------ * We *subclass* ``DMCI`` / ``DMC`` and add ``compute_bitmap`` rather than editing the DCVC repo. The parent ``compress_prior_2x/4x`` collapse channels for the entropy writer (``single_part_for_writing_*``), which loses the per-element information needed for bits, so we re-run the same prior loop and accumulate bits from each ``process_with_mask`` output before that collapse. * We deliberately **do not** call ``model.update()`` or the arithmetic coder: the bitmap only needs the analysis/synthesis transforms + ``process_with_mask``. This also avoids depending on the compiled ``MLCodec_extensions_cpp`` RANS coder (whose build in some envs mismatches this source tree). * Frames are fed as ``[0, 1]`` YCbCr (DCVC-RT convention; ``x_hat`` is clamped to ``[0, 1]``) — note this differs from the dev codec which used ``[-0.5, 0.5]``. Requires the ``codec`` conda env (torch + DCVC-RT importable). GPU recommended; the CUDA inference kernels fall back to PyTorch if unavailable (slower but numerically fine). """ from __future__ import annotations import math import os import sys from typing import Dict import numpy as np import torch # --- make DCVC-RT importable ------------------------------------------------- # The DCVC-RT source (MIT, github.com/microsoft/DCVC) is bundled next to this file # at ``neural_codec/DCVC/`` — no external checkout or env var needed. _DCVC_SRC_ROOT = os.path.join(os.path.dirname(os.path.abspath(__file__)), "DCVC") if not os.path.isdir(os.path.join(_DCVC_SRC_ROOT, "src")): raise RuntimeError( f"bundled DCVC-RT source not found at {_DCVC_SRC_ROOT!r} (expected a 'src/' " "directory); the neural_codec/DCVC/ folder looks incomplete." ) if _DCVC_SRC_ROOT not in sys.path: sys.path.insert(0, _DCVC_SRC_ROOT) from src.utils.common import get_state_dict, set_torch_env # noqa: E402 from src.models.image_model import DMCI # noqa: E402 from src.models.video_model import DMC # noqa: E402 from src.layers.cuda_inference import ( # noqa: E402 replicate_pad, round_and_to_int8, add_and_multiply, ) from src.utils.transforms import rgb2ycbcr # noqa: E402 _SCALE_MIN = 0.11 # GaussianEncoder.scale_min in DCVC-RT _SCALE_MAX = 16.0 # GaussianEncoder.scale_max in DCVC-RT _INV_SQRT2 = 1.0 / math.sqrt(2.0) def _gaussian_bits(y_q: torch.Tensor, s_hat: torch.Tensor) -> torch.Tensor: """Estimate per-element bits of a quantised residual under a zero-mean Gaussian. ``y_q`` is the (masked) integer residual and ``s_hat`` the (masked) predicted std. Masked-out positions have ``y_q == 0`` and ``s_hat == 0``; after clamping the std to ``scale_min`` they contribute ~0 bits, so summing over the complementary masks recovers each element's bits exactly once. """ q = y_q.float() s = s_hat.float().clamp_(_SCALE_MIN, _SCALE_MAX) inv = _INV_SQRT2 / s upper = 0.5 * torch.erf((q + 0.5) * inv) lower = 0.5 * torch.erf((q - 0.5) * inv) prob = (upper - lower).clamp_min_(1e-9) return -torch.log2(prob) class DMCIBitmap(DMCI): """Intra codec with a per-frame bit-cost bitmap.""" @torch.inference_mode() def compute_bitmap(self, x: torch.Tensor, qp: int) -> Dict[str, torch.Tensor]: curr_q_enc = self.q_scale_enc[qp:qp + 1, :, :, :] curr_q_dec = self.q_scale_dec[qp:qp + 1, :, :, :] y = self.enc(x, curr_q_enc) y_pad = self.pad_for_y(y) z = self.hyper_enc(y_pad) z_hat, _ = round_and_to_int8(z) params = self.y_prior_fusion(self.hyper_dec(z_hat)) _, _, yH, yW = y.shape params = params[:, :, :yH, :yW].contiguous() bitmap, y_hat = self._prior_4x_bits(y, params) x_hat = self.dec(y_hat, curr_q_dec).clamp_(0, 1) return {"bitmap": bitmap, "x_hat": x_hat} def _prior_4x_bits(self, y, common_params): """Mirror ``CompressionModel.compress_prior_4x`` but accumulate bits.""" q_enc, q_dec, scales, means = self.separate_prior(common_params, False) common_params = self.y_spatial_prior_reduction(common_params) B, C, H, W = y.size() mask_0, mask_1, mask_2, mask_3 = self.get_mask_4x(B, C, H, W, y.dtype, y.device) y = y * q_enc _, y_q_0, y_hat_0, s_hat_0 = self.process_with_mask(y, scales, means, mask_0) bits = _gaussian_bits(y_q_0, s_hat_0) y_hat_so_far = y_hat_0 params = torch.cat((y_hat_so_far, common_params), dim=1) scales, means = self.y_spatial_prior(self.y_spatial_prior_adaptor_1(params)).chunk(2, 1) _, y_q_1, y_hat_1, s_hat_1 = self.process_with_mask(y, scales, means, mask_1) bits = bits + _gaussian_bits(y_q_1, s_hat_1) y_hat_so_far = y_hat_so_far + y_hat_1 params = torch.cat((y_hat_so_far, common_params), dim=1) scales, means = self.y_spatial_prior(self.y_spatial_prior_adaptor_2(params)).chunk(2, 1) _, y_q_2, y_hat_2, s_hat_2 = self.process_with_mask(y, scales, means, mask_2) bits = bits + _gaussian_bits(y_q_2, s_hat_2) y_hat_so_far = y_hat_so_far + y_hat_2 params = torch.cat((y_hat_so_far, common_params), dim=1) scales, means = self.y_spatial_prior(self.y_spatial_prior_adaptor_3(params)).chunk(2, 1) _, y_q_3, y_hat_3, s_hat_3 = self.process_with_mask(y, scales, means, mask_3) bits = bits + _gaussian_bits(y_q_3, s_hat_3) y_hat = (y_hat_so_far + y_hat_3) * q_dec bitmap = bits.sum(dim=1) # (B, H, W) return bitmap, y_hat class DMCBitmap(DMC): """Inter (P-frame) codec with a per-frame bit-cost bitmap. ``compute_bitmap`` mirrors ``DMC.compress`` (minus the arithmetic coder) and, unlike the parent, does **not** push a reference frame by default so the caller controls DPB propagation (needed for the double-encode-on-reset trick). Pass ``add_ref=True`` for the encode whose reconstructed feature should propagate to the next frame. """ @torch.inference_mode() def compute_bitmap(self, x: torch.Tensor, qp: int, add_ref: bool = True) -> Dict[str, torch.Tensor]: q_encoder = self.q_encoder[qp:qp + 1, :, :, :] q_decoder = self.q_decoder[qp:qp + 1, :, :, :] q_feature = self.q_feature[qp:qp + 1, :, :, :] feature = self.apply_feature_adaptor() ctx, ctx_t = self.feature_extractor(feature, q_feature) y = self.encoder(x, ctx, q_encoder) hyper_inp = self.pad_for_y(y) z = self.hyper_encoder(hyper_inp) z_hat, _ = round_and_to_int8(z) params = self.res_prior_param_decoder(z_hat, ctx_t) bitmap, y_hat = self._prior_2x_bits(y, params) feature = self.decoder(y_hat, ctx, q_decoder) if add_ref: self.add_ref_frame(feature, None) return {"bitmap": bitmap, "feature": feature} def _prior_2x_bits(self, y, common_params): """Mirror ``CompressionModel.compress_prior_2x`` but accumulate bits.""" y, q_dec, scales, means = self.separate_prior_for_video_encoding(common_params, y) B, C, H, W = y.size() mask_0, mask_1 = self.get_mask_2x(B, C, H, W, y.dtype, y.device) _, y_q_0, y_hat_0, s_hat_0 = self.process_with_mask(y, scales, means, mask_0) bits = _gaussian_bits(y_q_0, s_hat_0) cat_params = torch.cat((y_hat_0, common_params), dim=1) scales, means = self.y_spatial_prior(cat_params).chunk(2, 1) _, y_q_1, y_hat_1, s_hat_1 = self.process_with_mask(y, scales, means, mask_1) bits = bits + _gaussian_bits(y_q_1, s_hat_1) y_hat = add_and_multiply(y_hat_0, y_hat_1, q_dec) bitmap = bits.sum(dim=1) # (B, H, W) return bitmap, y_hat # --------------------------------------------------------------------------- # # engine / frame loop # # --------------------------------------------------------------------------- # # Same feature-adaptor schedule DCVC-RT's test_video.py uses for P-frames. _INDEX_MAP = [0, 1, 0, 2, 0, 2, 0, 2] class DCVCRTEngine: """Loads DCVC-RT intra/inter nets and produces per-frame bit-cost bitmaps.""" def __init__( self, intra_ckpt: str, inter_ckpt: str, device: str = "cuda:0", half: bool = True, intra_period: int = -1, reset_interval: int = 32, ): set_torch_env() self.device = torch.device(device) self.half = bool(half) self.intra_period = int(intra_period) self.reset_interval = int(reset_interval) i_net = DMCIBitmap() i_net.load_state_dict(get_state_dict(intra_ckpt)) p_net = DMCBitmap() p_net.load_state_dict(get_state_dict(inter_ckpt)) self.i_net = i_net.to(self.device).eval() self.p_net = p_net.to(self.device).eval() if self.half: self.i_net.half() self.p_net.half() @property def _dtype(self): return torch.float16 if self.half else torch.float32 def _frame_to_input(self, rgb_hwc_uint8: np.ndarray, padding_b: int, padding_r: int) -> torch.Tensor: """RGB HxWx3 uint8 -> padded [0,1] YCbCr tensor (1,3,H',W').""" t = torch.from_numpy(rgb_hwc_uint8).to(self.device).permute(2, 0, 1).unsqueeze(0) t = t.float().div_(255.0) x = rgb2ycbcr(t).to(self._dtype) return replicate_pad(x, padding_b, padding_r) # --- stateful streaming API (memory-light for long videos) -------------- def reset_sequence(self, height: int, width: int, qp_i: int, qp_p: int) -> None: """Begin a new video sequence of the given (decoded) frame size.""" self._seq_qp_i = int(qp_i) self._seq_qp_p = int(qp_p) self._seq_last_qp = int(qp_i) self._seq_pad_r, self._seq_pad_b = DMCI.get_padding_size(int(height), int(width), 16) self.p_net.set_curr_poc(0) self.p_net.clear_dpb() @torch.inference_mode() def step(self, frame_idx: int, rgb: np.ndarray) -> torch.Tensor: """Encode one frame (in order from 0) and return its bit-cost bitmap ``(H/16, W/16)`` on GPU. Must be called sequentially; DPB state carries over. Call :meth:`reset_sequence` first. """ x_padded = self._frame_to_input(rgb, self._seq_pad_b, self._seq_pad_r) is_intra = (frame_idx == 0) or ( self.intra_period > 0 and frame_idx % self.intra_period == 0 ) if is_intra: out = self.i_net.compute_bitmap(x_padded, self._seq_qp_i) bitmap = out["bitmap"] self.p_net.clear_dpb() self.p_net.add_ref_frame(None, out["x_hat"]) self._seq_last_qp = self._seq_qp_i else: fa_idx = _INDEX_MAP[frame_idx % 8] curr_qp = self.p_net.shift_qp(self._seq_qp_p, fa_idx) is_reset = self.reset_interval > 0 and frame_idx % self.reset_interval == 1 if is_reset: # Clean bitmap from a non-reset encode (no DPB mutation) ... bitmap = self.p_net.compute_bitmap(x_padded, curr_qp, add_ref=False)["bitmap"] # ... then reset + re-encode to propagate the correct ref feature. self.p_net.prepare_feature_adaptor_i(self._seq_last_qp) self.p_net.compute_bitmap(x_padded, curr_qp, add_ref=True) else: bitmap = self.p_net.compute_bitmap(x_padded, curr_qp, add_ref=True)["bitmap"] self._seq_last_qp = curr_qp return bitmap[0]