neural_codec/codec_loader.py

"""Load precomputed DCVC-RT canvases and turn them into model inputs.

This is the inference-side counterpart of ``precompute_dcvc_rt.py``. It reuses
the release codec helpers verbatim, so precomputed DCVC-RT assets go through the
*exact* same downstream the HEVC ``video_backend="codec"`` path uses.
"""

from __future__ import annotations

import importlib
import importlib.util
import json
import os
import sys
from pathlib import Path
from typing import Optional

import numpy as np
import torch
from PIL import Image

_REPO = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))  # Mage-VL-Ported/


def _load_release_codec_module():
    """Import ``codec_video_processing_magevl`` from the release dir."""
    try:
        return importlib.import_module("codec_video_processing_magevl")
    except Exception:
        path = os.path.join(_REPO, "processor", "codec_video_processing_magevl.py")
        spec = importlib.util.spec_from_file_location(
            "codec_video_processing_magevl", path
        )
        mod = importlib.util.module_from_spec(spec)
        sys.modules[spec.name] = mod  # required for dataclass + future annotations
        spec.loader.exec_module(mod)
        return mod


def load_precomputed(asset_dir: str) -> dict:
    """Read ``<asset_dir>`` written by precompute -> {images, src_positions, fps}.

    Mirrors ``codec_video_processing_magevl._load_codec_result``.
    """
    asset_dir = Path(asset_dir)
    with open(asset_dir / "meta.json", "r", encoding="utf-8") as f:
        meta = json.load(f)
    canvas_files = meta.get("canvas_files")
    if not canvas_files:
        canvas_files = sorted(p.name for p in asset_dir.glob("canvas_*.jpg"))
    images = [Image.open(asset_dir / n).convert("RGB") for n in canvas_files]
    src_positions = np.load(asset_dir / "src_patch_position.npy")
    return {
        "images": images,
        "src_positions": src_positions,
        "fps": float(meta.get("fps") or 30.0),
        "meta": meta,
    }


def build_inputs_from_assets(
    processor,
    asset_dir: str,
    text: str,
    max_pixels: Optional[int] = None,
    device: Optional[torch.device] = None,
) -> dict:
    """Build a model-ready input dict from precomputed DCVC-RT assets + a
    chat-templated ``text`` string (containing a ``<|vision_start|>...<|vision_end|>``
    video span). Returns tensors ready for ``model.generate``.
    """
    cm = _load_release_codec_module()
    payload = load_precomputed(asset_dir)

    if max_pixels is None:
        # Canvas budget lives in preprocessor_config.json's codec.dcvc (150000),
        # NOT the image_processor's global max_pixels (the full-frame budget, 4M).
        try:
            import codec_dcvc_config as _dc
            max_pixels = int(_dc.get("max_pixels"))
        except Exception:
            max_pixels = 150000

    imgs, src_positions, _ = cm.drop_padding_canvases(payload["images"], payload["src_positions"])
    if not imgs:
        raise RuntimeError(f"no usable canvases in {asset_dir}")

    image_data = cm.codec_image_processor_outputs(processor.image_processor, imgs, max_pixels=max_pixels)
    image_grid_thw = image_data["image_grid_thw"]
    patch_positions = cm.codec_positions_for_processor(
        src_positions, image_grid_thw, device=image_grid_thw.device
    )
    rewritten = cm.rewrite_text_with_codec_positions(
        text, patch_positions, fps=float(payload["fps"]), decimals=1
    )

    enc = processor.tokenizer(rewritten, return_tensors="pt")
    out = {
        "input_ids": enc["input_ids"],
        "attention_mask": enc["attention_mask"],
        "pixel_values": image_data["pixel_values"],
        "image_grid_thw": image_grid_thw,
        "patch_positions": patch_positions,
    }
    if device is not None:
        for k, v in out.items():
            if isinstance(v, torch.Tensor):
                out[k] = v.to(device)
    return out