vllm_caption_runtime.py

from __future__ import annotations

import subprocess
import sys
from pathlib import Path
from typing import Any

import cv2
import numpy as np
from transformers import AutoProcessor
from vllm import LLM, SamplingParams

CAPTION_LENGTH_LABELS = ("very small", "small", "medium", "large", "very large")
CAPTION_SETTING_FIELD_CHOICES = {
    "vulgarity": ("none", "low", "medium", "high"),
    "uncertainty": ("none", "low", "medium", "high"),
    "character_names": ("none", "ambiguous", "single", "multiple"),
    "fluff": ("none", "low", "medium", "high"),
    "speculation": ("none", "low", "medium", "high"),
    "temporal_detail": ("static", "low", "medium", "high"),
    "visual_specificity": ("generic", "moderate", "detailed", "excessive"),
    "camera_detail": ("none", "low", "medium", "high"),
    "caption_style": ("plain", "verbose", "ornate", "robotic"),
}


def bool_from_value(value: Any, field_name: str) -> bool:
    if isinstance(value, bool):
        return value
    if isinstance(value, str):
        lowered = value.strip().lower()
        if lowered in {"1", "true", "yes", "y", "on"}:
            return True
        if lowered in {"0", "false", "no", "n", "off"}:
            return False
    raise ValueError(f"{field_name} must be boolean-like, got {value!r}")


def format_caption_settings_prompt(settings: dict[str, Any]) -> str:
    caption_length = str(settings["caption_length"]).strip().lower()
    if caption_length not in CAPTION_LENGTH_LABELS:
        raise ValueError(f"caption_length must be one of {CAPTION_LENGTH_LABELS}, got {caption_length!r}")

    include_watermark_info = bool_from_value(
        settings["include_watermark_info"], "include_watermark_info"
    )
    has_repetition = bool_from_value(settings["has_repetition"], "has_repetition")
    has_thinking = bool_from_value(settings["has_thinking"], "has_thinking")

    normalized: dict[str, str] = {}
    for field_name, choices in CAPTION_SETTING_FIELD_CHOICES.items():
        value = str(settings[field_name]).strip().lower()
        if value not in choices:
            raise ValueError(f"{field_name} must be one of {choices}, got {value!r}")
        normalized[field_name] = value

    watermark_instruction = (
        "Include watermark info." if include_watermark_info else "Do not include watermark info."
    )
    thinking_instruction = (
        "Output thought JSON before the final caption."
        if has_thinking
        else "Do not output thought JSON; output only the caption."
    )
    setting_text = (
        f"vulgarity={normalized['vulgarity']}; "
        f"uncertainty={normalized['uncertainty']}; "
        f"character_names={normalized['character_names']}; "
        f"fluff={normalized['fluff']}; "
        f"has_repetition={str(has_repetition).lower()}; "
        f"has_thinking={str(has_thinking).lower()}; "
        f"speculation={normalized['speculation']}; "
        f"temporal_detail={normalized['temporal_detail']}; "
        f"visual_specificity={normalized['visual_specificity']}; "
        f"camera_detail={normalized['camera_detail']}; "
        f"caption_style={normalized['caption_style']}"
    )
    return (
        f"Write a {caption_length} caption for this clip using both the visuals and the audio. "
        f"{watermark_instruction} {thinking_instruction} Match these caption settings: {setting_text}."
    )


def load_video_frames(path: str | Path, num_frames: int) -> tuple[np.ndarray, dict[str, Any]]:
    video_path = str(path)
    cap = cv2.VideoCapture(video_path)
    if not cap.isOpened():
        raise RuntimeError(f"Could not open video: {video_path}")

    total_frames = int(cap.get(cv2.CAP_PROP_FRAME_COUNT) or 0)
    fps = float(cap.get(cv2.CAP_PROP_FPS) or 24.0)
    if total_frames <= 0:
        cap.release()
        raise RuntimeError(f"Could not determine frame count: {video_path}")

    indices = np.linspace(0, max(0, total_frames - 1), num_frames).round().astype(int)
    frames = []
    for idx in indices:
        cap.set(cv2.CAP_PROP_POS_FRAMES, int(idx))
        ok, frame_bgr = cap.read()
        if ok:
            frames.append(cv2.cvtColor(frame_bgr, cv2.COLOR_BGR2RGB))
    cap.release()

    if not frames:
        raise RuntimeError(f"Could not read frames: {video_path}")

    metadata = {
        "fps": fps,
        "duration": total_frames / fps if fps > 0 else 0.0,
        "total_num_frames": total_frames,
        "frames_indices": [int(x) for x in indices[: len(frames)]],
        "video_backend": "opencv",
        "do_sample_frames": False,
    }
    return np.stack(frames, axis=0), metadata


def load_audio(path: str | Path, sampling_rate: int) -> tuple[np.ndarray, int]:
    video_path = str(path)
    cmd = [
        "ffmpeg",
        "-hide_banner",
        "-loglevel",
        "error",
        "-i",
        video_path,
        "-vn",
        "-ac",
        "1",
        "-ar",
        str(sampling_rate),
        "-f",
        "f32le",
        "-",
    ]
    result = subprocess.run(cmd, check=False, capture_output=True)
    if result.returncode == 0 and result.stdout:
        return np.frombuffer(result.stdout, dtype=np.float32), sampling_rate

    frames, metadata = load_video_frames(video_path, 2)
    del frames
    duration = max(1.0, float(metadata.get("duration") or 1.0))
    return np.zeros(max(1, int(duration * sampling_rate)), dtype=np.float32), sampling_rate


def build_prompt(
    *,
    processor: AutoProcessor,
    model_dir: str | Path,
    video_path: str | Path,
    prompt_override: str,
    prompt_settings: dict[str, Any],
) -> str:
    del model_dir
    prompt_text = prompt_override.strip() or format_caption_settings_prompt(prompt_settings)
    messages = [
        {
            "role": "user",
            "content": [
                {"type": "video", "path": str(video_path)},
                {"type": "text", "text": prompt_text},
                {"type": "audio", "path": "/tmp/polished_model_v3_audio.wav"},
            ],
        },
        {"role": "assistant", "content": [{"type": "text", "text": ""}]},
    ]
    prompt = processor.apply_chat_template(messages, tokenize=False)
    assistant_tail = "<turn|>\n"
    if prompt.endswith(assistant_tail):
        prompt = prompt[: -len(assistant_tail)]
    return prompt


def build_vllm_request(
    *,
    processor: AutoProcessor,
    model_dir: str | Path,
    video_path: str | Path,
    num_frames: int,
    sampling_rate: int,
    prompt_override: str,
    prompt_settings: dict[str, Any],
) -> dict[str, Any]:
    video, video_metadata = load_video_frames(video_path, num_frames)
    audio, sr = load_audio(video_path, sampling_rate)
    prompt = build_prompt(
        processor=processor,
        model_dir=model_dir,
        video_path=video_path,
        prompt_override=prompt_override,
        prompt_settings=prompt_settings,
    )
    return {
        "prompt": prompt,
        "multi_modal_data": {
            "video": [(video, video_metadata)],
            "audio": (audio, sr),
        },
    }


def load_processor(model_dir: str | Path) -> AutoProcessor:
    return AutoProcessor.from_pretrained(str(model_dir), local_files_only=True)


def load_llm(
    *,
    model_dir: str | Path,
    max_model_len: int,
    max_num_seqs: int,
    gpu_memory_utilization: float,
    dtype: str,
    quantization: str | None,
    enforce_eager: bool,
    enable_prefix_caching: bool,
    trust_remote_code: bool,
) -> LLM:
    kwargs: dict[str, Any] = {
        "model": str(model_dir),
        "tokenizer": str(model_dir),
        "max_model_len": max_model_len,
        "max_num_seqs": max_num_seqs,
        "gpu_memory_utilization": gpu_memory_utilization,
        "limit_mm_per_prompt": {"video": 1, "audio": 1, "image": 0},
        "enforce_eager": enforce_eager,
        "enable_prefix_caching": enable_prefix_caching,
        "trust_remote_code": trust_remote_code,
        "dtype": dtype,
    }
    if quantization:
        kwargs["quantization"] = quantization
    return LLM(**kwargs)


def sampling_params(
    *,
    temperature: float,
    max_tokens: int,
    top_p: float,
    repetition_penalty: float,
) -> SamplingParams:
    kwargs: dict[str, Any] = {
        "temperature": temperature,
        "max_tokens": max_tokens,
        "repetition_penalty": repetition_penalty,
    }
    if temperature > 0:
        kwargs["top_p"] = top_p
    return SamplingParams(**kwargs)


def ensure_local_vllm_source(script_dir: Path) -> None:
    local_vllm = script_dir / "vllm"
    if local_vllm.is_dir():
        sys.path.insert(0, str(local_vllm))