File size: 5,118 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 111 112 113 114 115 116 117 118 119 | #!/usr/bin/env python3
"""End-to-end inference with DCVC-RT-selected patches on the release model.
Two modes:
* ``--asset_dir DIR`` use canvases already written by ``precompute_dcvc_rt.py``.
* ``--video PATH`` precompute on the fly (calls ``precompute_dcvc_rt.py``)
into ``--work_dir`` then runs inference.
Example::
python neural_codec/infer_dcvc_rt.py \
--video clip.mp4 \
--question "Describe what happens in this video." \
--max_pixels 150000 --qp 21
"""
from __future__ import annotations
import argparse
import os
import subprocess
import sys
from pathlib import Path
import torch
_HERE = os.path.dirname(os.path.abspath(__file__))
_MODEL_DEFAULT = os.path.dirname(_HERE) # Mage-VL-Ported/ (this file lives in neural_codec/)
sys.path.insert(0, _HERE)
from codec_loader import build_inputs_from_assets # noqa: E402
def _precompute_inline(args) -> str:
key = Path(args.video).stem
out_root = args.work_dir
asset_dir = os.path.join(out_root, "assets", key)
py = args.precompute_python or sys.executable
cmd = [
py, os.path.join(_HERE, "precompute_dcvc_rt.py"),
"--videos", args.video, "--out_root", out_root,
"--intra_tar", args.intra_tar, "--inter_tar", args.inter_tar,
"--target_canvas", str(args.target_canvas), "--seq_len_frames", str(args.seq_len_frames),
"--max_pixels", str(args.max_pixels), "--qp", str(args.qp),
"--dcvc_max_side", str(args.dcvc_max_side), "--cuda_idx", str(args.gpu),
"--num_workers", "1",
]
if args.canvas_token_side:
cmd += ["--canvas_token_side", str(args.canvas_token_side)]
print("[infer] precomputing:", " ".join(cmd), flush=True)
subprocess.run(cmd, check=True)
return asset_dir
def main():
ap = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
ap.add_argument("--model", default=_MODEL_DEFAULT)
src = ap.add_mutually_exclusive_group(required=True)
src.add_argument("--asset_dir", help="precomputed assets/<key> dir")
src.add_argument("--video", help="video to precompute + run")
ap.add_argument("--question", default="Describe this video in detail.")
ap.add_argument("--max_new_tokens", type=int, default=256)
ap.add_argument("--max_pixels", type=int, default=150000)
ap.add_argument("--gpu", type=int, default=0)
# inline-precompute knobs (only for --video)
ap.add_argument("--work_dir", default="/tmp/dcvc_rt_infer")
ap.add_argument("--precompute_python", default=None,
help="interpreter for the DCVC-RT precompute subprocess (needs the "
"'codec' env); default: the current interpreter (sys.executable)")
ap.add_argument("--intra_tar", default=os.path.join(_HERE, "dcvc_rt_intra.tar"))
ap.add_argument("--inter_tar", default=os.path.join(_HERE, "dcvc_rt_inter.tar"))
ap.add_argument("--target_canvas", type=int, default=32)
ap.add_argument("--seq_len_frames", type=int, default=64)
ap.add_argument("--canvas_token_side", type=int, default=None)
ap.add_argument("--qp", type=int, default=21)
ap.add_argument("--dcvc_max_side", type=int, default=512)
args = ap.parse_args()
asset_dir = args.asset_dir or _precompute_inline(args)
sys.path.insert(0, os.path.join(args.model, "processor")) # MageVLProcessor + siblings
from processing_magevl import MageVLProcessor
from transformers import AutoModelForCausalLM
device = torch.device(f"cuda:{args.gpu}" if torch.cuda.is_available() else "cpu")
print("[infer] loading model:", args.model, flush=True)
processor = MageVLProcessor.from_pretrained(args.model) # subdir-aware
model = AutoModelForCausalLM.from_pretrained(
os.path.join(args.model, "transformer"), trust_remote_code=True, torch_dtype="auto"
).to(device).eval()
# Build the chat-templated text with a video span for rewrite to replace.
messages = [{"role": "user", "content": [
{"type": "video"},
{"type": "text", "text": args.question},
]}]
text = processor.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
inputs = build_inputs_from_assets(processor, asset_dir, text,
max_pixels=args.max_pixels, device=device)
# match vision dtype to model
inputs["pixel_values"] = inputs["pixel_values"].to(model.dtype)
print("[infer] input_ids", tuple(inputs["input_ids"].shape),
"pixel_values", tuple(inputs["pixel_values"].shape),
"grid_thw", inputs["image_grid_thw"].tolist(),
"patch_positions", tuple(inputs["patch_positions"].shape), flush=True)
with torch.inference_mode():
gen = model.generate(**inputs, max_new_tokens=args.max_new_tokens, do_sample=False)
new_tokens = gen[0, inputs["input_ids"].shape[1]:]
answer = processor.tokenizer.decode(new_tokens, skip_special_tokens=True)
print("\n================ ANSWER ================\n" + answer.strip() + "\n")
if __name__ == "__main__":
main()
|