neural_codec/dcvc_rt_engine.py

"""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]