inference.py

#!/usr/bin/env python3
"""Run RefCaptioner with the MRVBench inference configuration."""

from __future__ import annotations

import argparse
import json
from pathlib import Path

import torch
from qwen_vl_utils import process_vision_info
from transformers import AutoProcessor, Qwen3VLForConditionalGeneration


INSTRUCTION = """You are a multi-reference video captioning model.

Input:
- Several reference images, each labeled as <Image_1>, <Image_2>, etc.
- Some reference images may be distractors that are not visible in the video.
- One reference video.

Task:
Write one fluent English video caption that describes the visible video content and locally binds each usable reference image tag to the visual phrase it grounds.

Output rules:
- Return only the caption: no markdown, bullets, JSON, explanations, or title.
- Write one natural English paragraph, usually 4 to 7 complete sentences and about 120 to 250 words.
- Cover the main video style or format, main subject, setting, referenced appearances or objects, action progression, and useful camera, lighting, color, or mood details.
- Keep the caption grounded in visible video evidence. Do not invent unseen names, relationships, causes, dialogue, or story details.
- Do not mention audio, music, speech, dialogue, transcript, voiceover, or sound. Mention visible subtitles, logos, or text only when visually important.

Reference binding rules:
- Use only the provided tags. Do not invent tags.
- Use a tag only when its reference image can be grounded to visible video content.
- Do not force every provided tag. Omit tags whose image is not visible in the video or cannot be confidently grounded.
- Place each used tag immediately after a concrete grounded phrase, such as "the woman <Image_1>" or "the red dress <Image_2>".
- Multiple tags may be stacked as one contiguous tag group only when they refer to the same concrete visual unit, such as the same person, animal, character, vehicle, room, landscape, outfit item, prop, action pose, lighting, mood, or visual style.
- Do not stack tags merely because the images are related. A person and their clothing, accessory, carried object, background, or action are usually different visual units and should be tagged on separate phrases.
- If several tags belong to the same visual unit, keep them together as a complete tag group whenever that unit is explicitly mentioned.
- If two tags need different phrases, do not stack them.
- Attach tags to explicit noun phrases, not pronouns such as "he", "she", "it", "they", or "them".
- Do not put all tags at the end of the caption.
- Do not write phrases like "from <Image_N>", "shown in <Image_N>", "as shown in <Image_N>", "same as <Image_N>", or "similar to <Image_N>"."""

FINAL_REQUEST = (
    "Now write the final caption. Use only visibly grounded provided tags exactly "
    "as tag tokens, attach them to concrete visual phrases, omit ungrounded "
    "distractor tags, and stack tags only when they refer to the same visual unit."
)


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser()
    parser.add_argument("--model", default="NOVAglow646/RefCaptioner")
    parser.add_argument("--video", type=Path, required=True)
    parser.add_argument("--images", type=Path, nargs="+", required=True)
    parser.add_argument("--config", type=Path, default=Path(__file__).with_name("inference_config.json"))
    return parser.parse_args()


def build_messages(video: Path, images: list[Path], cfg: dict) -> list[dict]:
    tags = [f"<Image_{index}>" for index in range(1, len(images) + 1)]
    content: list[dict] = [
        {"type": "text", "text": INSTRUCTION},
        {
            "type": "text",
            "text": (
                "\n\nCurrent sample starts here.\nReference image tags, in order: "
                + ", ".join(tags)
                + "\nEach following image is the visual reference for the tag immediately before it.\n\n"
            ),
        },
    ]
    for tag, image in zip(tags, images):
        content.extend(
            [
                {"type": "text", "text": f"{tag} reference image:\n"},
                {
                    "type": "image",
                    "image": str(image.resolve()),
                    "max_pixels": cfg["image_max_pixels"],
                },
                {"type": "text", "text": "\n\n"},
            ]
        )
    content.extend(
        [
            {"type": "text", "text": "Reference video:\n"},
            {
                "type": "video",
                "video": str(video.resolve()),
                "fps": cfg["video_fps"],
                "min_frames": cfg["video_min_frames"],
                "max_frames": cfg["video_max_frames"],
                "max_pixels": cfg["video_max_pixels"],
            },
            {"type": "text", "text": "\n\n" + FINAL_REQUEST},
        ]
    )
    return [{"role": "user", "content": content}]


def main() -> None:
    args = parse_args()
    if not args.video.is_file():
        raise FileNotFoundError(args.video)
    missing = [path for path in args.images if not path.is_file()]
    if missing:
        raise FileNotFoundError(f"Missing reference images: {missing}")

    cfg = json.loads(args.config.read_text(encoding="utf-8"))
    model = Qwen3VLForConditionalGeneration.from_pretrained(
        args.model,
        dtype=torch.bfloat16,
        device_map=cfg["device_map"],
    ).eval()
    processor = AutoProcessor.from_pretrained(args.model)
    messages = build_messages(args.video, args.images, cfg)
    text = processor.apply_chat_template(
        messages,
        tokenize=False,
        add_generation_prompt=True,
        enable_thinking=False,
    )
    image_inputs, video_inputs, video_kwargs = process_vision_info(
        messages, return_video_kwargs=True
    )
    if isinstance(video_kwargs, dict):
        video_kwargs = {
            key: value
            for key, value in video_kwargs.items()
            if not (isinstance(value, list) and not value)
        }
        if isinstance(video_kwargs.get("fps"), list) and video_kwargs["fps"]:
            video_kwargs["fps"] = video_kwargs["fps"][0]
    inputs = processor(
        text=[text],
        images=image_inputs,
        videos=video_inputs,
        padding=True,
        return_tensors="pt",
        **video_kwargs,
    )
    if inputs["input_ids"].shape[-1] > cfg["max_length"]:
        raise ValueError(
            f"Input has {inputs['input_ids'].shape[-1]} tokens; "
            f"limit is {cfg['max_length']}."
        )
    inputs = inputs.to(model.device)
    with torch.inference_mode():
        generated = model.generate(
            **inputs,
            max_new_tokens=cfg["max_new_tokens"],
            do_sample=False,
        )
    trimmed = [output[len(source) :] for source, output in zip(inputs.input_ids, generated)]
    caption = processor.batch_decode(
        trimmed,
        skip_special_tokens=True,
        clean_up_tokenization_spaces=False,
    )[0].strip()
    print(caption)


if __name__ == "__main__":
    main()