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 = "\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))