File size: 3,899 Bytes
12acbba | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 | """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
|