github-actions[bot]
456 lines
20 KiB
Python
456 lines
20 KiB
Python
"""
|
|
Audio-visual media captioning using multimodal models.
|
|
This module provides captioning capabilities for videos with audio using:
|
|
- Qwen3-Omni via a local vLLM server (default)
|
|
- Gemini Flash 3.5 (cloud API)
|
|
Both produce a single combined English caption per video as a single
|
|
continuous paragraph of prose.
|
|
The Qwen3-Omni backend runs in a separately-launched vLLM server rather than
|
|
in-process, so vLLM's heavy CUDA dependencies stay out of this package. The
|
|
captioner talks to it over the OpenAI-compatible HTTP API.
|
|
Launch the server once (in an isolated environment) with:
|
|
.. code-block:: bash
|
|
uv run python scripts/serve_captioner.py
|
|
That helper picks BF16 vs FP8 dynamic quantization based on the GPU's free
|
|
memory and forwards everything else to ``vllm serve``. To check the recommended
|
|
command without running it, pass ``--print-cmd``.
|
|
To use Gemini instead, install ``google-genai`` and either set ``GEMINI_API_KEY``
|
|
(Gemini Developer API) or have Google Cloud credentials available (gcloud / an
|
|
attached service account), in which case it uses Vertex AI automatically.
|
|
"""
|
|
|
|
import json
|
|
import os
|
|
import re
|
|
import subprocess
|
|
import tempfile
|
|
from abc import ABC, abstractmethod
|
|
from enum import Enum
|
|
from pathlib import Path
|
|
from typing import ClassVar
|
|
|
|
DEFAULT_VIDEO_CAPTION_INSTRUCTION = """\
|
|
Analyze this video and produce a single detailed caption covering both its visual content and its audio. Be \
|
|
detailed enough that someone reading the caption could form an accurate mental picture of what happens on screen \
|
|
and what can be heard. Be exhaustive: include every meaningful detail you can see and hear, including small \
|
|
objects, textures, secondary movements, and minor background sounds.
|
|
|
|
Begin the caption directly with the action or visual detail; do not preface it with phrases like \
|
|
"The video opens with...", "The scene shows...", "We see...", or "There is...".
|
|
|
|
For every shot, include:
|
|
- The shot type and framing (extreme wide / wide / medium / medium close-up / close-up / extreme close-up) and any \
|
|
camera motion.
|
|
- Characters' clothing, appearance, posture, and movement (direction, speed, quality).
|
|
- The environment's materials, textures, lighting, and colors.
|
|
- All audio: spoken dialogue (quoted exactly in the original language), tone of voice, music (style, mood, \
|
|
volume changes), and environmental sounds. If a category is absent -- for example no music is playing, or no one is \
|
|
speaking -- state that explicitly. Do not invent specific instruments, music genres, moods, or ambient sounds \
|
|
that are not actually present.
|
|
- Any on-screen text (signs, titles, labels).
|
|
|
|
Describe only what is visible or audible. Do not infer emotions, intentions, or anything outside the segment. \
|
|
Refer to people descriptively (e.g., "the man in the blue jacket"). Narrate strictly in chronological order; if \
|
|
the video contains multiple shots, describe each one in turn.
|
|
|
|
Write everything as a single continuous paragraph of prose. Do not use section headers, bullet points, or labels \
|
|
like "Audio:" / "Visual:" / "Shot:". Integrate visual and audio details naturally within the same sentences.
|
|
|
|
Return a JSON object with exactly one key:
|
|
|
|
{"combined_caption_english": "<your caption here>"}"""
|
|
|
|
|
|
DEFAULT_IMAGE_CAPTION_INSTRUCTION = """\
|
|
Analyze this image and produce a single detailed caption of its visual content. Be detailed enough that \
|
|
someone reading the caption could form an accurate mental picture of the image. Be thorough: include every meaningful \
|
|
detail that is actually present, including small objects, textures, and background elements.
|
|
|
|
Begin the caption directly with the main subject or a visual detail; do not preface it with phrases like \
|
|
"The image shows...", "This is a photo of...", "We see...", or "There is...".
|
|
|
|
Include:
|
|
- The framing and composition (close-up / medium / wide / overhead, etc.) and the vantage point.
|
|
- The medium or style if distinctive (photograph, illustration, 3D render, painting).
|
|
- People's clothing, appearance, and posture, and what they are doing.
|
|
- The setting's materials, textures, lighting, and colors.
|
|
- Transcribe any visible text verbatim (signs, labels, titles, captions).
|
|
|
|
Describe only what is visible. Do not infer emotions or intentions, and do not describe sounds, motion, or \
|
|
events before or after the moment shown -- this is a single still image. When something is ambiguous, describe \
|
|
the visible cue (e.g., "warm low-angle light") rather than guessing the underlying fact (e.g., "sunrise"). \
|
|
Refer to people descriptively (e.g., "the man in the blue jacket").
|
|
|
|
Only describe what is present. Never state that something is absent or missing -- do not write phrases like \
|
|
"there is no text", "no people are present", or "no other objects". If a category such as people or text does \
|
|
not appear, simply leave it out.
|
|
|
|
Write everything as a single continuous paragraph of prose. Do not use section headers, bullet points, or \
|
|
labels.
|
|
|
|
Return a JSON object with exactly one key:
|
|
|
|
{"combined_caption_english": "<your caption here>"}"""
|
|
|
|
|
|
# Default model served by ``scripts/serve_captioner.py``. The captioner does not
|
|
# download or load this model itself -- it just sends requests to the vLLM
|
|
# server, which already has the model loaded.
|
|
DEFAULT_QWEN_MODEL = "Qwen/Qwen3-Omni-30B-A3B-Thinking"
|
|
DEFAULT_VLLM_BASE_URL = "http://127.0.0.1:8001/v1"
|
|
|
|
# Key the combined-caption prompt asks the model to return its caption under.
|
|
_CAPTION_JSON_KEY = "combined_caption_english"
|
|
|
|
|
|
class CaptionerType(str, Enum):
|
|
"""Enum for different types of media captioners."""
|
|
|
|
QWEN_OMNI = "qwen_omni" # Qwen3-Omni via local vLLM HTTP server
|
|
GEMINI_FLASH = "gemini_flash" # Gemini Flash 3.5 cloud API
|
|
|
|
|
|
def create_captioner(captioner_type: CaptionerType, **kwargs) -> "MediaCaptioningModel":
|
|
"""Factory function to create a media captioner."""
|
|
match captioner_type:
|
|
case CaptionerType.QWEN_OMNI:
|
|
return QwenOmniCaptioner(**kwargs)
|
|
case CaptionerType.GEMINI_FLASH:
|
|
return GeminiFlashCaptioner(**kwargs)
|
|
case _:
|
|
raise ValueError(f"Unsupported captioner type: {captioner_type}")
|
|
|
|
|
|
class MediaCaptioningModel(ABC):
|
|
"""Abstract base class for audio-visual media captioning models."""
|
|
|
|
instruction: str | None = None
|
|
|
|
@abstractmethod
|
|
def caption(self, path: str | Path, **kwargs) -> str:
|
|
"""Generate a caption for the given video or image."""
|
|
|
|
def _resolve_instruction(self, path: str | Path) -> str:
|
|
"""Return the custom instruction, or the image/video default for this input."""
|
|
if self.instruction is not None:
|
|
return self.instruction
|
|
return DEFAULT_IMAGE_CAPTION_INSTRUCTION if self._is_image_file(path) else DEFAULT_VIDEO_CAPTION_INSTRUCTION
|
|
|
|
@staticmethod
|
|
def _is_image_file(path: str | Path) -> bool:
|
|
return str(path).lower().endswith((".png", ".jpg", ".jpeg", ".heic", ".heif", ".webp"))
|
|
|
|
@staticmethod
|
|
def _is_video_file(path: str | Path) -> bool:
|
|
return str(path).lower().endswith((".mp4", ".avi", ".mov", ".mkv", ".webm"))
|
|
|
|
|
|
class QwenOmniCaptioner(MediaCaptioningModel):
|
|
"""Audio-visual captioning via a local vLLM server running Qwen3-Omni.
|
|
The vLLM server must already be running. See ``scripts/serve_captioner.py``
|
|
for a helper that launches one in an isolated environment (no impact on
|
|
this package's dependency tree).
|
|
The captioner uses the OpenAI-compatible chat completions API. It sends
|
|
a ``file://`` URL pointing at the local video, the default combined-caption
|
|
prompt, and parses the JSON-wrapped response.
|
|
"""
|
|
|
|
def __init__(
|
|
self,
|
|
base_url: str = DEFAULT_VLLM_BASE_URL,
|
|
model: str = DEFAULT_QWEN_MODEL,
|
|
api_key: str = "EMPTY",
|
|
instruction: str | None = None,
|
|
max_tokens: int = 4096,
|
|
enable_thinking: bool = False,
|
|
timeout_s: float = 600.0,
|
|
):
|
|
"""Initialize the Qwen3-Omni captioner.
|
|
Args:
|
|
base_url: Base URL of the vLLM OpenAI-compatible server (default
|
|
``http://127.0.0.1:8001/v1``).
|
|
model: Model identifier the server is serving. Must match the
|
|
server's ``--served-model-name`` (defaults to the HuggingFace
|
|
model ID).
|
|
api_key: Token sent in the ``Authorization`` header. vLLM accepts
|
|
any value by default.
|
|
instruction: Custom instruction prompt. If ``None``, uses the
|
|
default combined-caption prompt.
|
|
max_tokens: Maximum new tokens to generate per caption. 4096 leaves
|
|
comfortable headroom for both ``enable_thinking`` modes.
|
|
enable_thinking: Whether to let the Thinking model produce a
|
|
``<think>...</think>`` chain-of-thought before the caption.
|
|
Off by default: it makes captioning ~5x slower with little
|
|
quality benefit and occasionally introduces hallucinations
|
|
(e.g., inventing dialogue or background music).
|
|
timeout_s: Per-request HTTP timeout.
|
|
"""
|
|
from openai import OpenAI # noqa: PLC0415
|
|
|
|
self.model = model
|
|
self.instruction = instruction
|
|
self.max_tokens = max_tokens
|
|
self.enable_thinking = enable_thinking
|
|
self._client = OpenAI(base_url=base_url, api_key=api_key, timeout=timeout_s)
|
|
|
|
def caption(
|
|
self,
|
|
path: str | Path,
|
|
fps: int = 2,
|
|
) -> str:
|
|
"""Generate a caption for the given video or image.
|
|
Args:
|
|
path: Path to the video/image file to caption.
|
|
fps: Frames per second to sample from the video. Passed through to
|
|
vLLM's multimodal processor (``mm_processor_kwargs.fps``).
|
|
Default 2 is a typical choice for video MLLMs at this resolution.
|
|
Ignored for image inputs.
|
|
Returns:
|
|
The extracted caption string.
|
|
"""
|
|
path = Path(path)
|
|
is_image = self._is_image_file(path)
|
|
is_video = self._is_video_file(path)
|
|
if not (is_image or is_video):
|
|
raise ValueError(f"Unsupported media file: {path}")
|
|
|
|
instruction = self._resolve_instruction(path)
|
|
|
|
if is_image:
|
|
content = [
|
|
{"type": "image_url", "image_url": {"url": f"file://{path.resolve()}"}},
|
|
{"type": "text", "text": instruction},
|
|
]
|
|
return _parse_caption_response(self._chat(content)).strip()
|
|
|
|
return self._caption_video(path, instruction, fps)
|
|
|
|
def _chat(self, content: list[dict], mm_kwargs: dict | None = None) -> str:
|
|
"""Send one chat-completions request and return the raw response text."""
|
|
extra_body: dict = {
|
|
"repetition_penalty": 1.05,
|
|
"chat_template_kwargs": {"enable_thinking": self.enable_thinking},
|
|
}
|
|
if mm_kwargs:
|
|
extra_body["mm_processor_kwargs"] = mm_kwargs
|
|
response = self._client.chat.completions.create(
|
|
model=self.model,
|
|
messages=[{"role": "user", "content": content}],
|
|
max_tokens=self.max_tokens,
|
|
temperature=0.0,
|
|
extra_body=extra_body,
|
|
)
|
|
return response.choices[0].message.content or ""
|
|
|
|
def _caption_video(self, path: Path, instruction: str, fps: int) -> str:
|
|
"""Caption a video, sending its audio track as a separate modality.
|
|
vLLM does not extract a video's audio on its own (and its
|
|
``use_audio_in_video`` path is broken server-side), so we pull the audio
|
|
into a 16 kHz mono WAV and send it alongside the video -- otherwise the
|
|
model only sees frames and fabricates any spoken content.
|
|
"""
|
|
with tempfile.TemporaryDirectory(prefix="qwencap_") as tmp:
|
|
work = Path(tmp)
|
|
|
|
# Best-effort: ffmpeg fails (and we send video only) if there's no audio.
|
|
audio_url: str | None = None
|
|
try:
|
|
wav = work / "audio.wav"
|
|
_extract_audio_wav(path, wav)
|
|
audio_url = f"file://{wav.resolve()}"
|
|
except subprocess.CalledProcessError:
|
|
pass
|
|
|
|
def content(video: Path) -> list[dict]:
|
|
parts: list[dict] = [{"type": "video_url", "video_url": {"url": f"file://{video.resolve()}"}}]
|
|
if audio_url:
|
|
parts.append({"type": "audio_url", "audio_url": {"url": audio_url}})
|
|
parts.append({"type": "text", "text": instruction})
|
|
return parts
|
|
|
|
mm_kwargs = {"fps": fps}
|
|
try:
|
|
raw = self._chat(content(path), mm_kwargs)
|
|
except Exception as e:
|
|
# Raw / variable-frame-rate videos over-report their frame count, which
|
|
# breaks the server's frame sampler ("... frames from video"). Re-encode
|
|
# to a constant frame rate and retry once.
|
|
if "frames from video" not in str(e):
|
|
raise
|
|
cfr = work / "video_cfr.mp4"
|
|
_transcode_cfr(path, cfr)
|
|
raw = self._chat(content(cfr), mm_kwargs)
|
|
|
|
return _parse_caption_response(raw).strip()
|
|
|
|
|
|
class GeminiFlashCaptioner(MediaCaptioningModel):
|
|
"""Audio-visual captioning using Google's Gemini via the Google Gen AI SDK.
|
|
Uses the ``google-genai`` package (the current SDK; ``google-generativeai``
|
|
is deprecated). Auth is resolved automatically:
|
|
1. If an API key is given (``api_key`` argument, or ``GEMINI_API_KEY`` /
|
|
``GOOGLE_API_KEY`` in the environment) -> the Gemini Developer API (AI Studio).
|
|
2. Otherwise, if Google Cloud Application Default Credentials are available
|
|
(an attached service account or ``gcloud auth application-default login``)
|
|
-> Vertex AI. The project comes from ADC (or ``GOOGLE_CLOUD_PROJECT``) and
|
|
the location defaults to ``global`` (override with ``GOOGLE_CLOUD_LOCATION``).
|
|
This means it "just works" on a gcloud-authed GCP VM with no env vars.
|
|
If neither is available, a clear error explains how to authenticate.
|
|
Media is sent inline (``Part.from_bytes``), which works on both backends.
|
|
"""
|
|
|
|
MODEL_ID = "gemini-3.5-flash"
|
|
|
|
_MIME_TYPES: ClassVar[dict[str, str]] = {
|
|
".jpg": "image/jpeg",
|
|
".jpeg": "image/jpeg",
|
|
".png": "image/png",
|
|
".webp": "image/webp",
|
|
".heic": "image/heic",
|
|
".heif": "image/heif",
|
|
".mp4": "video/mp4",
|
|
".mov": "video/quicktime",
|
|
".avi": "video/x-msvideo",
|
|
".mkv": "video/x-matroska",
|
|
".webm": "video/webm",
|
|
}
|
|
|
|
def __init__(
|
|
self,
|
|
api_key: str | None = None,
|
|
instruction: str | None = None,
|
|
model: str | None = None,
|
|
):
|
|
"""Initialize the Gemini captioner.
|
|
Args:
|
|
api_key: Gemini Developer API key. If ``None``, falls back to
|
|
``GEMINI_API_KEY`` / ``GOOGLE_API_KEY``; if no key is set at all,
|
|
uses Vertex AI via Application Default Credentials.
|
|
instruction: Custom instruction prompt. If ``None``, uses the default
|
|
image or video prompt depending on the input.
|
|
model: Override the served model id (defaults to ``MODEL_ID``).
|
|
"""
|
|
self.instruction = instruction
|
|
self.model = model or self.MODEL_ID
|
|
self._client = self._make_client(api_key)
|
|
|
|
def caption(
|
|
self,
|
|
path: str | Path,
|
|
fps: int = 2, # noqa: ARG002 - kept for API compatibility
|
|
) -> str:
|
|
from google.genai import types # noqa: PLC0415
|
|
|
|
path = Path(path)
|
|
instruction = self._resolve_instruction(path)
|
|
media = types.Part.from_bytes(data=path.read_bytes(), mime_type=self._mime_type(path))
|
|
response = self._client.models.generate_content(
|
|
model=self.model,
|
|
contents=[media, instruction],
|
|
config=types.GenerateContentConfig(temperature=0.0),
|
|
)
|
|
|
|
# Gemini may also return JSON if it followed our prompt format.
|
|
return _parse_caption_response(response.text or "").strip()
|
|
|
|
@classmethod
|
|
def _mime_type(cls, path: Path) -> str:
|
|
try:
|
|
return cls._MIME_TYPES[path.suffix.lower()]
|
|
except KeyError:
|
|
raise ValueError(f"Unsupported media type for Gemini: {path.suffix}") from None
|
|
|
|
def _make_client(self, api_key: str | None): # noqa: ANN202 - genai.Client type is lazy-imported
|
|
from google import genai # noqa: PLC0415
|
|
|
|
# 1. API key (explicit arg or env) -> Gemini Developer API.
|
|
key = api_key or os.environ.get("GEMINI_API_KEY") or os.environ.get("GOOGLE_API_KEY")
|
|
if key:
|
|
return genai.Client(api_key=key)
|
|
|
|
# 2. No key -> Vertex AI via Application Default Credentials (gcloud / service account).
|
|
import google.auth # noqa: PLC0415
|
|
|
|
try:
|
|
_, adc_project = google.auth.default()
|
|
except Exception as e:
|
|
raise ValueError(
|
|
"No Gemini credentials found. Provide an API key (--api-key, or "
|
|
"GEMINI_API_KEY / GOOGLE_API_KEY), or set up Google Cloud credentials "
|
|
"for Vertex AI (e.g. `gcloud auth application-default login` or an "
|
|
"attached service account)."
|
|
) from e
|
|
|
|
project = os.environ.get("GOOGLE_CLOUD_PROJECT") or adc_project
|
|
location = os.environ.get("GOOGLE_CLOUD_LOCATION", "global")
|
|
return genai.Client(vertexai=True, project=project, location=location)
|
|
|
|
|
|
def _parse_caption_response(raw: str) -> str:
|
|
"""Extract the caption text from a model response.
|
|
Backend-agnostic: works for any model that follows the combined-caption
|
|
prompt. Handles the formats a model may produce:
|
|
- Plain caption text
|
|
- JSON ``{"combined_caption_english": "..."}``
|
|
- ``<think>...</think>`` chain-of-thought followed by either of the above
|
|
- Truncated JSON (when generation hits a token limit mid-string)
|
|
"""
|
|
text = re.sub(r"<think>[\s\S]*?</think>", "", raw).strip()
|
|
|
|
# Thinking models (e.g. Qwen3-Omni-*-Thinking) emit the reasoning trace
|
|
# without an opening ``<think>`` tag, because the chat template injects it
|
|
# for them -- so the response starts mid-thought and is terminated by a lone
|
|
# ``</think>`` before the real answer. Drop everything up to that closer.
|
|
if "</think>" in text:
|
|
text = text.rsplit("</think>", 1)[1].strip()
|
|
|
|
if not text:
|
|
return raw.strip()
|
|
|
|
try:
|
|
parsed = json.loads(text)
|
|
if isinstance(parsed, dict) and _CAPTION_JSON_KEY in parsed:
|
|
return parsed[_CAPTION_JSON_KEY]
|
|
except (json.JSONDecodeError, ValueError):
|
|
pass
|
|
|
|
match = re.search(rf"\{{[^{{}}]*\"{_CAPTION_JSON_KEY}\"[^{{}}]*\}}", text)
|
|
if match:
|
|
try:
|
|
parsed = json.loads(match.group())
|
|
if isinstance(parsed, dict) and _CAPTION_JSON_KEY in parsed:
|
|
return parsed[_CAPTION_JSON_KEY]
|
|
except (json.JSONDecodeError, ValueError):
|
|
pass
|
|
|
|
# Truncated JSON: extract the string value even if the closing quote/brace is missing.
|
|
match = re.search(rf'"{_CAPTION_JSON_KEY}"\s*:\s*"((?:[^"\\]|\\.)*)', text)
|
|
if match:
|
|
try:
|
|
return json.loads('"' + match.group(1) + '"')
|
|
except (json.JSONDecodeError, ValueError):
|
|
return match.group(1)
|
|
|
|
return text
|
|
|
|
|
|
def _run_ffmpeg(args: list[str]) -> None:
|
|
"""Run the ffmpeg binary bundled with ``imageio-ffmpeg`` (a dependency)."""
|
|
import imageio_ffmpeg # noqa: PLC0415
|
|
|
|
cmd = [imageio_ffmpeg.get_ffmpeg_exe(), "-y", "-loglevel", "error", *args]
|
|
subprocess.run(cmd, check=True, capture_output=True)
|
|
|
|
|
|
def _extract_audio_wav(src: Path, dest: Path) -> None:
|
|
"""Extract the audio track to a 16 kHz mono PCM WAV (matches pretraining).
|
|
Raises ``CalledProcessError`` when the video has no audio stream.
|
|
"""
|
|
_run_ffmpeg(["-i", str(src), "-vn", "-ac", "1", "-ar", "16000", "-c:a", "pcm_s16le", str(dest)])
|
|
|
|
|
|
def _transcode_cfr(src: Path, dest: Path) -> None:
|
|
"""Re-encode the video to a constant frame rate so the server's frame sampler can
|
|
read every requested index (raw / variable-frame-rate videos over-report frames)."""
|
|
_run_ffmpeg(["-i", str(src), "-fps_mode", "cfr", "-c:v", "libx264", "-pix_fmt", "yuv420p", "-an", str(dest)])
|