llava_onevision2.py · tduncan/Nvidia-Optimization-Stack-VLLM

@tduncan / Nvidia-Optimization-Stack-VLLM

Code Commits Issues Pulls Cloud Agent CI/CD

main

py 2,315 lines 91 KB

Raw

1	# SPDX-License-Identifier: Apache-2.0
2	# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3	"""Inference-only LLaVA-OneVision-2 (OV2) model for vLLM.
4
5	Architecture notes:
6
7	* LLM backbone is plain Qwen3-8B with 1-D position_ids (no M-RoPE).
8	* Vision tower removes the CLS token (no class_embedding/class_pos_emb).
9	* Vision RoPE is 3-D (T:H:W) with a 4:6:6 head_dim split and uses
10	``patch_positions`` instead of grid_thw to compute per-token freqs.
11	* ``rotate_half`` is interleaved (``(::2, 1::2)``) rather than
12	split-half.
13	* Vision attention uses windowed ``cu_seqlens`` (``frame_windows_size``
14	in T-dim); two backends implemented (SDPA + flash_attn varlen).
15	* ``patch_positions: [total_patches, 3]`` is plumbed as a first-class
16	MM kwarg alongside ``pixel_values`` / ``image_grid_thw``.
17	* Video frame-backend and codec-backend both alias to the image path
18	inside the HF processor, so the model implements a single visual
19	code path.
20	"""
21
22	from __future__ import annotations
23
24	import hashlib
25	import importlib
26	from collections.abc import Callable, Iterable, Mapping, Sequence, Sized
27	from functools import lru_cache
28	from typing import (
29	Annotated,
30	Any,
31	Literal,
32	SupportsIndex,
33	SupportsInt,
34	TypedDict,
35	)
36
37	import numpy as np
38	import regex as re
39	import torch
40	import torch.nn as nn
41	import torch.nn.functional as F
42	from PIL import Image
43	from transformers import AutoProcessor, AutoTokenizer, BatchFeature
44	from transformers.dynamic_module_utils import (
45	get_class_from_dynamic_module,
46	resolve_trust_remote_code,
47	)
48	from transformers.models.qwen2_vl import Qwen2VLImageProcessor
49	from transformers.models.qwen2_vl.image_processing_qwen2_vl import smart_resize
50	from typing_extensions import Buffer
51
52	from vllm.compilation.decorators import (
53	should_torch_compile_mm_encoder,
54	support_torch_compile,
55	)
56	from vllm.config import VllmConfig
57	from vllm.config.multimodal import MultiModalDummyOptions
58	from vllm.distributed import parallel_state
59	from vllm.distributed import utils as dist_utils
60	from vllm.inputs import ModalityData, MultiModalDataDict
61	from vllm.logger import init_logger
62	from vllm.model_executor.layers.attention import MMEncoderAttention
63	from vllm.model_executor.layers.linear import (
64	ColumnParallelLinear,
65	QKVParallelLinear,
66	RowParallelLinear,
67	)
68	from vllm.model_executor.layers.quantization import QuantizationConfig
69	from vllm.model_executor.models.interfaces import (
70	MultiModalEmbeddings,
71	SupportsMultiModal,
72	SupportsPP,
73	)
74	from vllm.model_executor.models.module_mapping import MultiModelKeys
75	from vllm.model_executor.models.utils import (
76	AutoWeightsLoader,
77	WeightsMapper,
78	init_vllm_registered_model,
79	maybe_prefix,
80	)
81	from vllm.model_executor.models.utils import (
82	_merge_multimodal_embeddings as merge_multimodal_embeddings,
83	)
84	from vllm.model_executor.models.vision import get_vit_attn_backend
85	from vllm.multimodal import MULTIMODAL_REGISTRY
86	from vllm.multimodal.inputs import (
87	ImageItem,
88	MultiModalFieldConfig,
89	MultiModalKwargsItems,
90	)
91	from vllm.multimodal.parse import (
92	DictEmbeddingItems,
93	ImageSize,
94	ModalityDataItems,
95	MultiModalDataItems,
96	MultiModalDataParser,
97	)
98	from vllm.multimodal.processing import (
99	BaseMultiModalProcessor,
100	BaseProcessingInfo,
101	PromptReplacement,
102	PromptUpdate,
103	PromptUpdateDetails,
104	cached_encode,
105	)
106	from vllm.multimodal.processing.dummy_inputs import BaseDummyInputsBuilder
107	from vllm.multimodal.video import (
108	VIDEO_LOADER_REGISTRY,
109	VideoBackend,
110	VideoSourceMetadata,
111	VideoTargetMetadata,
112	)
113	from vllm.sequence import IntermediateTensors
114	from vllm.transformers_utils.processor import _merge_mm_kwargs
115	from vllm.transformers_utils.repo_utils import get_hf_file_to_dict
116	from vllm.transformers_utils.utils import convert_model_repo_to_path
117	from vllm.utils.tensor_schema import TensorSchema, TensorShape
118
119	logger = init_logger(__name__)
120
121
122	@lru_cache
123	def _load_ov2_processor(
124	model: str,
125	revision: str \| None,
126	trust_remote_code: bool,
127	**kwargs: Any,
128	):
129	# OV2's trust_remote_code processor is a bare class (not a ProcessorMixin),
130	# so the shared type-checked get_hf_processor rejects it. We also can't use
131	# AutoProcessor.from_pretrained: OV2's remote from_pretrained drops
132	# trust_remote_code before building its nested tokenizer, which makes that
133	# nested load fall back to an interactive stdin prompt that hangs in
134	# non-interactive CI. Instead, assemble the processor here with
135	# trust_remote_code threaded through every component explicitly.
136	path = convert_model_repo_to_path(model)
137	revision = revision or "main"
138
139	resolve_trust_remote_code(
140	trust_remote_code,
141	model,
142	has_local_code=False,
143	has_remote_code=True,
144	)
145
146	processor_cls = get_class_from_dynamic_module(
147	"processing_llava_onevision2.LlavaOnevision2Processor",
148	path,
149	revision=revision,
150	)
151	video_processor_cls = get_class_from_dynamic_module(
152	"video_processing_llava_onevision2.LlavaOnevision2VideoProcessor",
153	path,
154	revision=revision,
155	)
156
157	# Slow Qwen2VLImageProcessor mirrors the remote processor (the Fast variant
158	# has normalization rounding differences that change pixel_values).
159	image_processor = Qwen2VLImageProcessor.from_pretrained(
160	path, revision=revision, **kwargs
161	)
162	tokenizer = AutoTokenizer.from_pretrained(
163	path, revision=revision, trust_remote_code=trust_remote_code, **kwargs
164	)
165	video_processor = video_processor_cls(
166	image_processor=image_processor,
167	min_pixels=getattr(image_processor, "min_pixels", 256 * 28 * 28),
168	max_pixels=getattr(image_processor, "max_pixels", 1605632),
169	patch_size=getattr(image_processor, "patch_size", 14),
170	spatial_merge_size=getattr(image_processor, "merge_size", 2),
171	)
172
173	# Codec defaults live under the "codec" key of preprocessor_config.json,
174	# which Qwen2VLImageProcessor does not preserve; read them directly so the
175	# codec video backend keeps its configured defaults.
176	codec_config: dict = {}
177	try:
178	preprocessor_config = get_hf_file_to_dict(
179	"preprocessor_config.json", path, revision
180	)
181	codec_config = (preprocessor_config or {}).get("codec", {}) or {}
182	except Exception:
183	logger.debug("OV2: no codec defaults found in preprocessor_config.json")
184
185	return processor_cls(
186	image_processor=image_processor,
187	tokenizer=tokenizer,
188	video_processor=video_processor,
189	codec_config=codec_config,
190	)
191
192
193	# Upper bound on frames used when profiling the worst-case video item, mirroring
194	# Qwen2-VL. The real frame count is decided by the HF VideoProcessor at apply
195	# time; this only sizes the memory-profiling estimate.
196	_MAX_FRAMES_PER_VIDEO = 14
197
198
199	def _pack_timestamps(per_video: list[list[float]]) -> torch.Tensor:
200	if not per_video:
201	return torch.empty((0, 0), dtype=torch.float32)
202	t_max = max((len(ts) for ts in per_video), default=0)
203	padded = torch.zeros((len(per_video), t_max), dtype=torch.float32)
204	for i, ts in enumerate(per_video):
205	padded[i, : len(ts)] = torch.tensor(ts, dtype=torch.float32)
206	return padded
207
208
209	def _validate_video_source(path: str, model_config) -> str:
210	"""Confine a codec video path to ``--allowed-local-media-path``.
211
212	The codec backend keeps the raw path string alive past vLLM's
213	``MultiModalDataParser`` and hands it to the trust-remote-code codec
214	module, which opens it directly via ``cv2.VideoCapture`` / ffmpeg. That
215	bypasses both ``MediaConnector``'s access controls and its redirect
216	handling (``VLLM_MEDIA_URL_ALLOW_REDIRECTS``), so we restrict the codec
217	backend to local files only: remote ``http(s)`` / ``data`` URLs are
218	rejected here and must instead go through the frame backend (a registered
219	``VIDEO_LOADER_REGISTRY`` loader), which rides vLLM's connector and its
220	domain/redirect gates.
221
222	Returns the resolved absolute path so the codec module opens exactly the
223	file that was validated, closing the validate-vs-open (symlink-retarget)
224	window. Mirrors the confinement in ``MediaConnector._load_file_url``.
225	"""
226	from pathlib import Path
227	from urllib.request import url2pathname
228
229	from urllib3.util import parse_url
230
231	allowed_local = getattr(model_config, "allowed_local_media_path", "") or ""
232
233	parsed = parse_url(str(path))
234	scheme = (parsed.scheme or "").lower()
235
236	if scheme in ("http", "https", "data"):
237	raise ValueError(
238	f"The codec video backend does not support remote {scheme!r} URLs: "
239	f"its trust-remote-code decoder fetches them outside vLLM's domain "
240	f"and redirect controls. Use a local file path, or the frame "
241	f"backend for remote videos."
242	)
243	if scheme not in ("", "file"):
244	raise ValueError(
245	f"Unsupported codec video URL scheme {scheme!r}; only local file "
246	f"paths or file:// URLs are supported."
247	)
248
249	# Local file access is opt-in: require --allowed-local-media-path and
250	# confine the resolved path to that directory (connector.py:253-271).
251	if not allowed_local:
252	raise ValueError(
253	"Local video file access is disabled. Set "
254	"--allowed-local-media-path to enable reading local videos."
255	)
256	if scheme == "file":
257	# Decode percent-encoding (mirrors MediaConnector._load_file_url),
258	# including the netloc so file://host/path is handled identically.
259	local = Path(url2pathname((parsed.netloc or "") + (parsed.path or "")))
260	else:
261	local = Path(str(path))
262	# Require an absolute path: resolving a relative path against an ambiguous
263	# CWD before the confinement check is brittle/unsafe.
264	if not local.is_absolute():
265	raise ValueError(
266	f"Local video path {str(path)!r} must be absolute; "
267	f"relative paths are not supported."
268	)
269	allowed_root = Path(allowed_local).resolve()
270	resolved = local.resolve()
271	if resolved != allowed_root and allowed_root not in resolved.parents:
272	raise ValueError(
273	f"Video path {str(path)!r} is outside the allowed local media "
274	f"directory {allowed_local!r}."
275	)
276	return str(resolved)
277
278
279	def _validate_video_sources(paths, model_config) -> list[str]:
280	# Return the resolved paths so the codec module opens exactly what was
281	# validated (no validate-vs-open / symlink-retarget differential).
282	return [_validate_video_source(path, model_config) for path in paths]
283
284
285	# Design note: the two video backends take deliberately different paths.
286	#
287	# * frame backend: a normal vLLM VIDEO_LOADER_REGISTRY loader
288	# (``LlavaOnevision2VideoBackend`` below). Decoding to RGB
289	# ``(frames, metadata)`` is exactly what loaders are for, so frame sampling
290	# participates in the standard decode-stage pipeline.
291	#
292	# * codec backend: NOT a loader. OV2's codec path needs the video path string
293	# to survive into ``_apply_hf_processor_main``, where the HF processor builds the
294	# codec canvas + smart_resize + patchify
295	# (pixel_values/image_grid_thw/patch_positions). That transform is
296	# path-level and inseparable; it cannot be reconstructed from pre-decoded
297	# RGB frames, so it must run at the processor stage rather than the decode
298	# stage. The small marker/parser machinery below keeps the path alive for
299	# codec; ``_validate_video_source`` then confines it to a local file (codec
300	# does not support remote URLs -- see its docstring).
301	_CODEC_VIDEO_MARKER = "ov2_codec_video"
302
303
304	def prepare_codec_video_input(video_path: str) -> tuple:
305	"""Wrap a video path for vLLM's MultiModalDataParser + OV2 codec backend.
306
307	Returns ``(dummy_ndarray, metadata)`` where the ndarray satisfies the
308	parser's 4-D shape check and the metadata carries the actual path to
309	our ``_apply_hf_processor_main``. Use as::
310
311	multi_modal_data = {"video": prepare_codec_video_input("foo.mp4")}
312
313	The dummy ndarray bytes encode a hash of ``video_path`` so distinct codec
314	videos get distinct mm_hashes: the parser drops the metadata dict before
315	hashing (only the ndarray reaches MultiModalHasher), so without this
316	variance every video after the first would collide and skip the encoder.
317	"""
318	path_str = str(video_path)
319	digest = hashlib.blake2b(path_str.encode("utf-8"), digest_size=16).digest()
320	dummy = np.frombuffer(digest, dtype=np.uint8).reshape(1, 1, 16, 1)
321	dummy = np.broadcast_to(dummy, (1, 1, 16, 3)).copy()
322	return (dummy, {_CODEC_VIDEO_MARKER: str(video_path)})
323
324
325	def _extract_codec_video_paths(videos: Any) -> list[str] \| None:
326	# vLLM's parser yields list-of-(ndarray, metadata-dict) for tuple inputs.
327	# We accept either that shape or a single raw tuple (pre-parser cases).
328	def _path_from(item):
329	if (
330	isinstance(item, tuple)
331	and len(item) == 2
332	and isinstance(item[1], dict)
333	and _CODEC_VIDEO_MARKER in item[1]
334	):
335	return item[1][_CODEC_VIDEO_MARKER]
336	return None
337
338	if isinstance(videos, list):
339	paths: list[str] = []
340	for item in videos:
341	p = _path_from(item)
342	if p is None:
343	return None
344	paths.append(p)
345	return paths if paths else None
346	p = _path_from(videos)
347	return [p] if p is not None else None
348
349
350	_CODEC_FPS_CACHE: dict[str, float] = {}
351
352
353	def _codec_fps_for(video_path: str, hf_processor) -> float:
354	if video_path in _CODEC_FPS_CACHE:
355	return _CODEC_FPS_CACHE[video_path]
356	# The codec module is shipped inside the HF transformers_modules package
357	# for OV2, so an absolute import does not resolve. Locate it relative to
358	# the processor module (which lives in the same package).
359	proc_module_name = type(hf_processor).__module__
360	pkg = proc_module_name.rsplit(".", 1)[0] if "." in proc_module_name else ""
361	codec_mod = importlib.import_module(
362	f"{pkg}.codec_video_processing_llava_onevision2"
363	if pkg
364	else "codec_video_processing_llava_onevision2"
365	)
366	CodecConfig = codec_mod.CodecConfig
367	process_codec_video = codec_mod.process_codec_video
368	cfg_defaults = dict(getattr(hf_processor, "_codec_config_defaults", {}))
369	cfg = CodecConfig(**cfg_defaults)
370	payload = process_codec_video(video_path, cfg)
371	fps = float(payload["fps"])
372	_CODEC_FPS_CACHE[video_path] = fps
373	return fps
374
375
376	def _codec_timestamp_runs(
377	patch_positions: torch.Tensor,
378	fps: float,
379	spatial_merge_size: int,
380	) -> list[tuple[float, int]]:
381	# Mirrors HF's _timestamp_runs (codec_video_processing_llava_onevision2.py)
382	# exactly: same column, same merge factor, same negative-t skip, same
383	# zero-token-count skip. Keeping the logic local avoids importing the
384	# private helper from transformers_modules.
385	t_values = patch_positions[:, 0]
386	unique_t, counts = torch.unique_consecutive(t_values, return_counts=True)
387	merge_factor = int(spatial_merge_size) ** 2
388	runs: list[tuple[float, int]] = []
389	for t_val, count in zip(unique_t.tolist(), counts.tolist()):
390	if int(t_val) < 0:
391	continue
392	token_count = int(count) // merge_factor
393	if token_count <= 0:
394	continue
395	runs.append((float(t_val) / float(fps), token_count))
396	return runs
397
398
399	def _create_field_factory(
400	spatial_merge_size: int,
401	) -> Callable[[Mapping[str, torch.Tensor]], Mapping[str, MultiModalFieldConfig]]:
402	"""Build the per-batch field-config callback.
403
404	OV2-specific: also exposes ``patch_positions`` as a flat-from-sizes
405	field, sized by the total per-image patch count (THW). The merger and
406	the 3-D RoPE both consume it.
407	"""
408
409	def _field_config(hf_inputs: Mapping[str, torch.Tensor]):
410	image_grid_thw = hf_inputs.get("image_grid_thw", torch.empty((0, 3)))
411	image_pixel_grid_sizes = image_grid_thw.prod(-1)
412	image_embed_grid_sizes = (
413	image_pixel_grid_sizes // spatial_merge_size // spatial_merge_size
414	)
415
416	video_grid_thw = hf_inputs.get("video_grid_thw", torch.empty((0, 3)))
417	# OV2 emits one grid_thw row per frame, so vLLM's per-video sharding
418	# requires explicit frame counts. video_patch_sizes sums H*W over the
419	# frames that belong to each video; video_grid_thw uses the frame
420	# count directly (one row per frame).
421	video_num_frames = hf_inputs.get(
422	"video_num_frames", torch.empty((0,), dtype=torch.long)
423	)
424	if video_num_frames.numel() > 0:
425	per_row_patches = video_grid_thw.prod(-1)
426	offsets = torch.cumsum(
427	torch.cat([torch.zeros(1, dtype=torch.long), video_num_frames[:-1]]), 0
428	).tolist()
429	video_patch_sizes = torch.tensor(
430	[
431	int(per_row_patches[int(s) : int(s) + int(n)].sum())
432	for s, n in zip(offsets, video_num_frames.tolist())
433	],
434	dtype=torch.long,
435	)
436	else:
437	video_patch_sizes = torch.empty((0,), dtype=torch.long)
438
439	return dict(
440	pixel_values=MultiModalFieldConfig.flat_from_sizes(
441	"image", image_pixel_grid_sizes
442	),
443	image_embeds=MultiModalFieldConfig.flat_from_sizes(
444	"image", image_embed_grid_sizes
445	),
446	image_grid_thw=MultiModalFieldConfig.batched("image", keep_on_cpu=True),
447	# OV2 first-class MM kwarg: per-patch (t,h,w)
448	# positions required by the 3-D vision RoPE.
449	patch_positions=MultiModalFieldConfig.flat_from_sizes(
450	"image", image_pixel_grid_sizes
451	),
452	pixel_values_videos=MultiModalFieldConfig.flat_from_sizes(
453	"video", video_patch_sizes
454	),
455	video_grid_thw=MultiModalFieldConfig.flat_from_sizes(
456	"video", video_num_frames, keep_on_cpu=True
457	),
458	patch_positions_videos=MultiModalFieldConfig.flat_from_sizes(
459	"video", video_patch_sizes
460	),
461	video_num_frames=MultiModalFieldConfig.batched("video", keep_on_cpu=True),
462	frame_timestamps=MultiModalFieldConfig.batched("video", keep_on_cpu=True),
463	# Per-video flag: 0 = frame-sampling backend, 1 = codec backend.
464	# Drives codec-aware ``\n`` insertion in ``get_video_replacement``.
465	video_is_codec=MultiModalFieldConfig.batched("video", keep_on_cpu=True),
466	# Codec backend: per-video source-frame fps. Needed at
467	# replacement time to convert patch_positions t-indices into
468	# the timestamp tags HF writes (``<sec seconds>``).
469	codec_fps=MultiModalFieldConfig.batched("video", keep_on_cpu=True),
470	)
471
472	return _field_config
473
474
475	# ---------------------------------------------------------------------------
476	# Frame backend helpers
477	# ---------------------------------------------------------------------------
478	# The default video pathway materialises each video as a series of PIL frames
479	# (decoded + sampled by the registered ``LlavaOnevision2VideoBackend``) and
480	# feeds them to the HF processor through the image branch (per-frame timestamp
481	# marker + ``<\|image_pad\|>``). This mirrors the validated lmms-eval
482	# ``vllm_hf_chat`` adapter, and empirically scores higher on Video-MME than
483	# OV2's native VideoProcessor frame extractor.
484	_DEFAULT_TIMESTAMP_DECIMALS = 1
485	_DEFAULT_FPS = 1.0
486	_DEFAULT_MAX_FRAMES = 32
487	# OV2 vision tower has spatial_merge_size=2 -> temporal frame count must be
488	# even. The hf-chat reference pads by repeating the last frame; same here.
489	_TEMPORAL_MERGE_SIZE = 2
490
491	# Token sequence emitted by the OV2 vLLM dummy inputs builder for each video
492	# item. The frame backend expands each marker into per-frame image markers
493	# (timestamp + image_pad block).
494	_VIDEO_MARKER = "<\|vision_start\|><\|video_pad\|><\|vision_end\|>"
495	_IMAGE_MARKER = "<\|vision_start\|><\|image_pad\|><\|vision_end\|>"
496
497
498	def _frame_video_to_pil_and_timestamps(
499	item: Any,
500	) -> tuple[list[Image.Image], list[float]]:
501	"""Convert a ``(frames_ndarray, metadata)`` video item into
502	``(pil_frames, timestamps_seconds)``.
503
504	Both real ``video_url`` inputs (decoded + sampled by the registered
505	``LlavaOnevision2VideoBackend``) and dummy profiling videos arrive here as a
506	``(frames, metadata)`` tuple because the data parser runs with
507	``video_needs_metadata=True``. ``frames`` is a ``(T, H, W, C)`` uint8 array;
508	``metadata`` carries ``frames_indices`` and the source ``fps``.
509
510	Timestamps follow the qwen_vl_utils policy: ``frame_index / original_fps``.
511	The frame count is padded up to ``_TEMPORAL_MERGE_SIZE`` (repeating the last
512	frame) because OV2's vision tower merges frames temporally in pairs.
513	"""
514	if not (isinstance(item, (tuple, list)) and len(item) == 2):
515	raise ValueError(
516	"LlavaOnevision2 frame backend expects each video as a "
517	f"(frames_ndarray, metadata) tuple; got {type(item).__name__}. "
518	"Pass videos via `video_url` so the registered backend can decode "
519	"and sample them."
520	)
521	frames, metadata = item
522	if isinstance(frames, torch.Tensor):
523	frames_np = frames.cpu().numpy()
524	else:
525	frames_np = np.asarray(frames)
526
527	pil_frames = [Image.fromarray(f.astype(np.uint8)) for f in frames_np]
528
529	indices = metadata.get("frames_indices") if isinstance(metadata, Mapping) else None
530	if indices is None:
531	indices = list(range(len(pil_frames)))
532	elif not isinstance(indices, list):
533	indices = list(indices)
534	# Keep indices aligned with the actual frame count.
535	if len(indices) != len(pil_frames):
536	if len(indices) > len(pil_frames):
537	indices = indices[: len(pil_frames)]
538	else:
539	indices = list(indices) + [indices[-1] if indices else 0] * (
540	len(pil_frames) - len(indices)
541	)
542
543	fps = _DEFAULT_FPS
544	if isinstance(metadata, Mapping) and metadata.get("fps"):
545	fps = float(metadata["fps"])
546	if fps <= 0:
547	fps = _DEFAULT_FPS
548
549	# OV2 vision tower: temporal merge=2 -> frame count must be even.
550	if len(pil_frames) % _TEMPORAL_MERGE_SIZE != 0:
551	pad = _TEMPORAL_MERGE_SIZE - len(pil_frames) % _TEMPORAL_MERGE_SIZE
552	pil_frames = pil_frames + [pil_frames[-1]] * pad
553	indices = indices + [indices[-1]] * pad
554
555	timestamps = [idx / fps for idx in indices]
556	return pil_frames, timestamps
557
558
559	def _expand_video_markers_in_prompt(
560	prompt: str,
561	per_video_timestamps: list[list[float]],
562	*,
563	timestamp_decimals: int,
564	) -> str:
565	"""Replace each ``<\|vision_start\|><\|video_pad\|><\|vision_end\|>`` with a
566	sequence of ``<{t:.Nf} seconds><\|vision_start\|><\|image_pad\|><\|vision_end\|>``
567	blocks -- one per frame -- matching ``vllm_hf_chat._build_prompt``.
568
569	Replacement is positional: the i-th marker consumes
570	``per_video_timestamps[i]``.
571	"""
572	parts: list[str] = []
573	cursor = 0
574	idx = 0
575	pattern = re.escape(_VIDEO_MARKER)
576	for m in re.finditer(pattern, prompt):
577	parts.append(prompt[cursor : m.start()])
578	if idx >= len(per_video_timestamps):
579	raise ValueError(
580	f"Prompt has more video markers than supplied timestamp "
581	f"groups ({len(per_video_timestamps)})"
582	)
583	timestamps = per_video_timestamps[idx]
584	expanded = "".join(
585	f"<{t:.{timestamp_decimals}f} seconds>{_IMAGE_MARKER}" for t in timestamps
586	)
587	parts.append(expanded)
588	cursor = m.end()
589	idx += 1
590	parts.append(prompt[cursor:])
591	if idx != len(per_video_timestamps):
592	raise ValueError(
593	f"Prompt has {idx} video markers but {len(per_video_timestamps)} "
594	f"timestamp groups were supplied"
595	)
596	return "".join(parts)
597
598
599	# ---------------------------------------------------------------------------
600	# Registered video loader backend (frame sampling)
601	# ---------------------------------------------------------------------------
602	# OV2 frame-sampling policy, expressed as a vLLM ``VideoBackend`` so videos
603	# entering through the standard ``video_url`` -> MediaConnector -> VideoMediaIO
604	# path are decoded + sampled inside vLLM (no qwen_vl_utils dependency, and the
605	# connector's SSRF / local-file gates apply automatically).
606	#
607	# Parity note: ``compute_frames_index_to_sample`` replicates
608	# ``qwen_vl_utils.smart_nframes`` (the policy the OV2 hf-chat recipe validated)
609	# -- frame count and index selection are byte-identical to qwen. The one
610	# residual difference is the source frame count: qwen decodes via decord
611	# whereas vLLM uses OpenCV/PyAV, whose ``CAP_PROP_FRAME_COUNT`` (a duration x
612	# fps estimate) can differ by +/-1 frame on some containers, shifting sampled
613	# indices by one frame on those files. Downstream metrics stay within noise.
614	_OV2_FRAME_FACTOR = 2
615	_OV2_FPS_MIN_FRAMES = 4
616
617
618	def _ceil_by_factor(n: float, factor: int) -> int:
619	import math as _math
620
621	return _math.ceil(n / factor) * factor
622
623
624	def _floor_by_factor(n: float, factor: int) -> int:
625	import math as _math
626
627	return _math.floor(n / factor) * factor
628
629
630	def _ov2_smart_nframes(
631	total_frames: int,
632	video_fps: float,
633	*,
634	fps: float,
635	min_frames: int,
636	max_frames: int,
637	) -> int:
638	"""Replicate ``qwen_vl_utils.smart_nframes`` (fps branch).
639
640	Returns an even frame count in ``[min_frames, min(max_frames, total)]``.
641	"""
642	min_frames = _ceil_by_factor(min_frames, _OV2_FRAME_FACTOR)
643	max_frames = _floor_by_factor(max_frames, _OV2_FRAME_FACTOR)
644	nframes = total_frames / video_fps * fps if video_fps > 0 else total_frames
645	nframes = min(min(max(nframes, min_frames), max_frames), total_frames)
646	nframes = _floor_by_factor(nframes, _OV2_FRAME_FACTOR)
647	return max(int(nframes), _OV2_FRAME_FACTOR)
648
649
650	@VIDEO_LOADER_REGISTRY.register(
651	"llava_onevision2",
652	video_processor="LlavaOnevision2VideoProcessor",
653	)
654	class LlavaOnevision2VideoBackend(VideoBackend):
655	"""Frame-sampling backend for LLaVA-OneVision-2.
656
657	Selected automatically for OV2 via the ``video_processor`` binding
658	(``video_processor_type == "LlavaOnevision2VideoProcessor"`` in the model's
659	``video_preprocessor_config.json``). Decoding uses the inherited OpenCV /
660	PyAV codecs; only the sampling index policy is overridden to match qwen.
661	"""
662
663	_sampling_suffix = "_llava_onevision2"
664
665	# OV2 hf-chat reference sampling constants (mirror the validated adapter).
666	_FPS = _DEFAULT_FPS
667	_MAX_FRAMES = _DEFAULT_MAX_FRAMES
668	_MIN_FRAMES = _OV2_FPS_MIN_FRAMES
669
670	@classmethod
671	def compute_frames_index_to_sample(
672	cls,
673	source: VideoSourceMetadata,
674	target: VideoTargetMetadata,
675	**kwargs,
676	) -> list[int]:
677	total = int(source.total_frames_num)
678	if total <= 0:
679	return []
680	video_fps = float(source.original_fps)
681	# Honor caller-provided sampling targets (via ``--media-io-kwargs`` →
682	# ``VideoTargetMetadata``) so benchmarks can override the conservative
683	# defaults (e.g. VSI-Bench needs max_frames=128). Fall back to the OV2
684	# hf-chat reference constants when the target leaves a field unset
685	# (sentinel ``<= 0``).
686	target_fps = float(target.fps) if target.fps > 0 else cls._FPS
687	target_max_frames = (
688	int(target.num_frames) if target.num_frames > 0 else cls._MAX_FRAMES
689	)
690	n = _ov2_smart_nframes(
691	total,
692	video_fps,
693	fps=target_fps,
694	min_frames=cls._MIN_FRAMES,
695	max_frames=target_max_frames,
696	)
697	# qwen uses linspace().round() (NOT the floor cast used by the base
698	# uniform backend), so replicate the rounding exactly.
699	idx = np.linspace(0, total - 1, n).round().astype(int).tolist()
700	# smart_nframes floors to FRAME_FACTOR so ``n`` is even; guard anyway
701	# since OV2's vision tower (temporal merge = 2) requires even counts.
702	if len(idx) % _OV2_FRAME_FACTOR != 0:
703	idx.append(idx[-1])
704	return idx
705
706
707	class LlavaOnevision2ImagePixelInputs(TensorSchema):
708	type: Literal["pixel_values"]
709
710	pixel_values: Annotated[torch.Tensor, TensorShape("np", "cps")]
711	image_grid_thw: Annotated[torch.Tensor, TensorShape("ni", 3)]
712	patch_positions: Annotated[torch.Tensor, TensorShape("np", 3)]
713
714
715	class LlavaOnevision2ImageEmbeddingInputs(TensorSchema):
716	type: Literal["image_embeds"]
717
718	image_embeds: Annotated[torch.Tensor, TensorShape("nf", "hs")]
719	image_grid_thw: Annotated[torch.Tensor, TensorShape("ni", 3)]
720
721
722	class LlavaOnevision2VideoPixelInputs(TensorSchema):
723	type: Literal["pixel_values_videos"]
724
725	pixel_values_videos: Annotated[torch.Tensor, TensorShape("np", "cps")]
726	video_grid_thw: Annotated[torch.Tensor, TensorShape("nf", 3)]
727	patch_positions_videos: Annotated[torch.Tensor, TensorShape("np", 3)]
728	video_num_frames: Annotated[torch.Tensor, TensorShape("nv")]
729
730
731	LlavaOnevision2ImageInputs = (
732	LlavaOnevision2ImagePixelInputs \| LlavaOnevision2ImageEmbeddingInputs
733	)
734
735
736	class LlavaOnevision2InputsByModality(TypedDict, total=False):
737	images: LlavaOnevision2ImageInputs \| None
738	videos: LlavaOnevision2VideoPixelInputs \| None
739
740
741	class LlavaOnevision2VisionRotaryEmbedding(nn.Module):
742	"""3-D rotary frequency constructor with 4:6:6 (T:H:W) split.
743
744	Mirrors ``VisionRotaryEmbedding`` in the HF reference
745	(``modeling_llava_onevision2.py`` L79-L210). The three ``inv_freq_*``
746	buffers are non-persistent — they are not in the checkpoint and must
747	be reconstructed at module init time (which we do here).
748
749	Public entry points used by the vision tower:
750	* ``forward_from_positions(patch_positions)`` — per-patch (t,h,w)
751	positions → per-token freqs [N, half].
752	"""
753
754	def __init__(self, head_dim: int, theta: float = 10000.0) -> None:
755	super().__init__()
756	assert head_dim % 2 == 0, "head_dim must be even"
757	assert head_dim % 16 == 0, "head_dim must be divisible by 16 (4:6:6)"
758	half = head_dim // 2
759	assert half % 16 == 0, "head_dim//2 must be divisible by 16"
760
761	self.head_dim = head_dim
762	self.half = half
763	self.base = float(theta)
764
765	unit = half // 16
766	self.t_size = 4 * unit
767	self.h_size = 6 * unit
768	self.w_size = 6 * unit
769	assert self.t_size + self.h_size + self.w_size == half
770
771	self.register_buffer(
772	"inv_freq_t",
773	1.0
774	/ (
775	self.base
776	** (torch.arange(self.t_size, dtype=torch.float32) / self.t_size)
777	),
778	persistent=False,
779	)
780	self.register_buffer(
781	"inv_freq_h",
782	1.0
783	/ (
784	self.base
785	** (torch.arange(self.h_size, dtype=torch.float32) / self.h_size)
786	),
787	persistent=False,
788	)
789	self.register_buffer(
790	"inv_freq_w",
791	1.0
792	/ (
793	self.base
794	** (torch.arange(self.w_size, dtype=torch.float32) / self.w_size)
795	),
796	persistent=False,
797	)
798
799	def forward_from_positions(self, patch_positions: torch.Tensor) -> torch.Tensor:
800	"""[N, 3] (t,h,w) int → [N, half] float frequencies."""
801	device = patch_positions.device
802	inv_t = self.inv_freq_t.to(device=device)
803	inv_h = self.inv_freq_h.to(device=device)
804	inv_w = self.inv_freq_w.to(device=device)
805
806	t_pos = patch_positions[:, 0].float()
807	h_pos = patch_positions[:, 1].float()
808	w_pos = patch_positions[:, 2].float()
809
810	ft = torch.outer(t_pos, inv_t)
811	fh = torch.outer(h_pos, inv_h)
812	fw = torch.outer(w_pos, inv_w)
813	return torch.cat([ft, fh, fw], dim=-1)
814
815
816	class LlavaOnevision2VisionEmbeddings(nn.Module):
817	def __init__(
818	self, patch_size: int = 14, in_channels: int = 3, embed_dim: int = 1024
819	) -> None:
820	super().__init__()
821	self.patch_size = patch_size
822	self.in_channels = in_channels
823	self.embed_dim = embed_dim
824	self.patch_embedding = nn.Conv2d(
825	in_channels,
826	embed_dim,
827	kernel_size=(patch_size, patch_size),
828	stride=(patch_size, patch_size),
829	bias=False,
830	)
831
832	def forward(self, x: torch.Tensor) -> torch.Tensor:
833	x = x.view(-1, self.in_channels, self.patch_size, self.patch_size)
834	x = self.patch_embedding(x).view(-1, self.embed_dim)
835	return x
836
837
838	class LlavaOnevision2VisionMLP(nn.Module):
839	def __init__(
840	self,
841	in_features: int,
842	hidden_features: int,
843	bias: bool = True,
844	quant_config: QuantizationConfig \| None = None,
845	prefix: str = "",
846	use_data_parallel: bool = False,
847	) -> None:
848	super().__init__()
849	if quant_config is not None:
850	raise RuntimeError("LLaVAOneVision2 does not support quantization")
851	self.fc1 = ColumnParallelLinear(
852	in_features,
853	hidden_features,
854	bias=bias,
855	quant_config=quant_config,
856	prefix=f"{prefix}.fc1",
857	disable_tp=use_data_parallel,
858	)
859	self.fc2 = RowParallelLinear(
860	hidden_features,
861	in_features,
862	bias=bias,
863	quant_config=quant_config,
864	prefix=f"{prefix}.fc2",
865	disable_tp=use_data_parallel,
866	)
867	self.act_fn = F.gelu
868
869	def forward(self, x: torch.Tensor) -> torch.Tensor:
870	x, _ = self.fc1(x)
871	x = self.act_fn(x)
872	x, _ = self.fc2(x)
873	return x
874
875
876	class LlavaOnevision2VisionAttn(nn.Module):
877	"""Vision self-attention with windowed cu_seqlens.
878
879	The HF checkpoint ships a fused qkv linear (``self_attn.qkv``), so
880	we load directly into ``QKVParallelLinear`` with no stacked_params
881	mapping. (Compare OV1.5, whose checkpoint had separate q/k/v.)
882	"""
883
884	def __init__(
885	self,
886	embed_dim: int,
887	num_heads: int,
888	projection_size: int,
889	quant_config: QuantizationConfig \| None = None,
890	prefix: str = "",
891	use_data_parallel: bool = False,
892	) -> None:
893	super().__init__()
894	if quant_config is not None:
895	raise RuntimeError("LLaVAOneVision2 does not support quantization")
896
897	self.tp_size = (
898	1
899	if use_data_parallel
900	else parallel_state.get_tensor_model_parallel_world_size()
901	)
902	self.tp_rank = parallel_state.get_tensor_model_parallel_rank()
903	self.num_heads = num_heads
904	self.hidden_size_per_attn_head = dist_utils.divide(projection_size, num_heads)
905	self.num_attn_heads_per_partition = dist_utils.divide(num_heads, self.tp_size)
906
907	self.qkv = QKVParallelLinear(
908	hidden_size=embed_dim,
909	head_size=self.hidden_size_per_attn_head,
910	total_num_heads=num_heads,
911	total_num_kv_heads=num_heads,
912	bias=True,
913	quant_config=quant_config,
914	prefix=f"{prefix}.qkv",
915	disable_tp=use_data_parallel,
916	)
917
918	self.proj = RowParallelLinear(
919	input_size=projection_size,
920	output_size=embed_dim,
921	quant_config=quant_config,
922	prefix=f"{prefix}.proj",
923	disable_tp=use_data_parallel,
924	)
925
926	self.attn = MMEncoderAttention(
927	num_heads=self.num_attn_heads_per_partition,
928	head_size=self.hidden_size_per_attn_head,
929	scale=self.hidden_size_per_attn_head**-0.5,
930	prefix=f"{prefix}.attn",
931	)
932
933	@staticmethod
934	def _rotate_half_interleaved(x: torch.Tensor) -> torch.Tensor:
935	"""OV2-specific interleaved rotate_half.
936
937	Pairs adjacent dims: (x[::2], x[1::2]) -> (-x[1::2], x[::2]).
938	NOT compatible with the split-half rotate used in OV1.5/LLaMA.
939	"""
940	x_even = x[..., 0::2]
941	x_odd = x[..., 1::2]
942	out = torch.stack((-x_odd, x_even), dim=-1)
943	return out.flatten(-2)
944
945	def _apply_rotary_pos_embed(
946	self, t: torch.Tensor, freqs: torch.Tensor
947	) -> torch.Tensor:
948	# freqs is [seq_len, half]; cat([f,f],-1) pair-repeat layout
949	# matches the interleaved rotate above.
950	orig_dtype = t.dtype
951	t = t.float()
952	emb = torch.cat((freqs, freqs), dim=-1)
953	cos = emb.cos().unsqueeze(-2).float()
954	sin = emb.sin().unsqueeze(-2).float()
955	t = (t * cos) + (self._rotate_half_interleaved(t) * sin)
956	return t.to(orig_dtype)
957
958	def forward(
959	self,
960	x: torch.Tensor,
961	cu_seqlens: torch.Tensor,
962	rotary_pos_emb: torch.Tensor,
963	max_seqlen: torch.Tensor \| None = None,
964	sequence_lengths: torch.Tensor \| None = None,
965	) -> torch.Tensor:
966	x, _ = self.qkv(x)
967	seq_len = x.shape[0]
968	# QKVParallelLinear packs q/k/v along the last dim as
969	# [q_heads, k_heads, v_heads]; view splits the three sections, each
970	# holding this partition's heads (no all-gather: MMEncoderAttention
971	# runs per-partition and ``proj`` reduces across TP ranks).
972	qkv = x.view(
973	seq_len,
974	3,
975	self.num_attn_heads_per_partition,
976	self.hidden_size_per_attn_head,
977	)
978	q, k, v = qkv.unbind(1)
979
980	if rotary_pos_emb is not None:
981	# OV2 uses interleaved RoPE on the raw freqs, applied outside
982	# MMEncoderAttention (which is rotary-agnostic).
983	q = self._apply_rotary_pos_embed(q, rotary_pos_emb)
984	k = self._apply_rotary_pos_embed(k, rotary_pos_emb)
985
986	# Add a leading batch dim (b=1) for MMEncoderAttention, which expects
987	# (batch, seq_len, num_heads, head_size) and windows via cu_seqlens.
988	output = self.attn(
989	query=q.unsqueeze(0),
990	key=k.unsqueeze(0),
991	value=v.unsqueeze(0),
992	cu_seqlens=cu_seqlens,
993	max_seqlen=max_seqlen,
994	sequence_lengths=sequence_lengths,
995	)
996	output = output.reshape(seq_len, -1)
997
998	output, _ = self.proj(output)
999	return output
1000
1001
1002	@support_torch_compile(
1003	dynamic_arg_dims={
1004	"x": 0,
1005	"cu_seqlens": 0,
1006	"rotary_pos_emb": 0,
1007	"sequence_lengths": 0,
1008	},
1009	enable_if=should_torch_compile_mm_encoder,
1010	is_encoder=True,
1011	)
1012	class LlavaOnevision2VisionTowerBlock(nn.Module):
1013	def __init__(
1014	self,
1015	dim: int,
1016	num_heads: int,
1017	mlp_hidden_dim: int,
1018	norm_eps: float = 1e-6,
1019	quant_config: QuantizationConfig \| None = None,
1020	prefix: str = "",
1021	use_data_parallel: bool = False,
1022	) -> None:
1023	super().__init__()
1024	self.layer_norm1 = nn.LayerNorm(dim, eps=norm_eps)
1025	self.layer_norm2 = nn.LayerNorm(dim, eps=norm_eps)
1026	self.self_attn = LlavaOnevision2VisionAttn(
1027	embed_dim=dim,
1028	num_heads=num_heads,
1029	projection_size=dim,
1030	quant_config=quant_config,
1031	prefix=f"{prefix}.self_attn",
1032	use_data_parallel=use_data_parallel,
1033	)
1034	self.mlp = LlavaOnevision2VisionMLP(
1035	dim,
1036	mlp_hidden_dim,
1037	bias=True,
1038	quant_config=quant_config,
1039	prefix=f"{prefix}.mlp",
1040	use_data_parallel=use_data_parallel,
1041	)
1042
1043	def forward(
1044	self,
1045	x: torch.Tensor,
1046	cu_seqlens: torch.Tensor,
1047	rotary_pos_emb: torch.Tensor,
1048	max_seqlen: torch.Tensor \| None = None,
1049	sequence_lengths: torch.Tensor \| None = None,
1050	) -> torch.Tensor:
1051	x = x + self.self_attn(
1052	self.layer_norm1(x),
1053	cu_seqlens=cu_seqlens,
1054	rotary_pos_emb=rotary_pos_emb,
1055	max_seqlen=max_seqlen,
1056	sequence_lengths=sequence_lengths,
1057	)
1058	x = x + self.mlp(self.layer_norm2(x))
1059	return x
1060
1061
1062	class LlavaOnevision2PatchMerger(nn.Module):
1063	def __init__(
1064	self,
1065	d_model: int,
1066	context_dim: int,
1067	spatial_merge_size: int = 2,
1068	norm_eps: float = 1e-6,
1069	quant_config: QuantizationConfig \| None = None,
1070	prefix: str = "",
1071	use_data_parallel: bool = False,
1072	) -> None:
1073	super().__init__()
1074	self.hidden_size = context_dim * (spatial_merge_size**2)
1075	self.ln_q = nn.LayerNorm(context_dim, eps=norm_eps)
1076	self.mlp = nn.ModuleList(
1077	[
1078	ColumnParallelLinear(
1079	self.hidden_size,
1080	self.hidden_size,
1081	bias=True,
1082	quant_config=quant_config,
1083	prefix=f"{prefix}.mlp.0",
1084	disable_tp=use_data_parallel,
1085	),
1086	nn.GELU(),
1087	RowParallelLinear(
1088	self.hidden_size,
1089	d_model,
1090	bias=True,
1091	quant_config=quant_config,
1092	prefix=f"{prefix}.mlp.2",
1093	disable_tp=use_data_parallel,
1094	),
1095	]
1096	)
1097
1098	def forward(
1099	self, x: torch.Tensor, patch_positions: torch.Tensor \| None = None
1100	) -> torch.Tensor:
1101	# patch_positions accepted for API symmetry with the HF impl,
1102	# but unused: pixel_values already arrive in spatial-merge block
1103	# order (Qwen2VLImageProcessor handles that across image / video-
1104	# frames / video-codec backends). See codec_video_processing for
1105	# the codec backend's ``codec_positions_for_processor`` call.
1106	del patch_positions
1107	x = self.ln_q(x)
1108	x = x.view(-1, self.hidden_size)
1109	fc1, act, fc2 = self.mlp
1110	x, _ = fc1(x)
1111	x = act(x)
1112	x, _ = fc2(x)
1113	return x
1114
1115
1116	class LlavaOnevision2VisionTower(nn.Module):
1117	"""OV2 vision tower (no CLS token, 3-D RoPE, windowed attention).
1118
1119	Module attribute names mirror HF checkpoint names verbatim so the
1120	WeightsMapper only needs prefix rewrites (no substring rules, which would
1121	otherwise collide with the Qwen3 text-path ``self_attn`` modules):
1122	visual.embeddings.patch_embedding
1123	visual.layernorm_pre
1124	visual.encoder.layers.{i}.self_attn.{qkv,proj}
1125	visual.encoder.layers.{i}.layer_norm{1,2}
1126	visual.encoder.layers.{i}.mlp.fc{1,2}
1127	visual.merger.{ln_q, mlp.{0,2}}
1128	visual.rotary_pos_emb (non-persistent inv_freq buffers)
1129	"""
1130
1131	def __init__(
1132	self,
1133	vision_config,
1134	text_hidden_size: int,
1135	norm_eps: float = 1e-6,
1136	quant_config: QuantizationConfig \| None = None,
1137	prefix: str = "",
1138	use_data_parallel: bool = False,
1139	) -> None:
1140	super().__init__()
1141	if quant_config is not None:
1142	raise RuntimeError("LLaVAOneVision2 does not support quantization")
1143
1144	patch_size = vision_config.patch_size
1145	spatial_merge_size = vision_config.spatial_merge_size
1146	in_channels = getattr(vision_config, "num_channels", 3)
1147	hidden_size = vision_config.hidden_size
1148	embed_dim = hidden_size
1149	depth = vision_config.num_hidden_layers
1150	num_heads = vision_config.num_attention_heads
1151	mlp_hidden_dim = vision_config.intermediate_size
1152	frame_windows_size = getattr(vision_config, "frame_windows_size", 4)
1153	rope_theta = getattr(vision_config, "rope_theta", 10000.0)
1154
1155	self.spatial_merge_size = spatial_merge_size
1156	self.frame_windows_size = int(frame_windows_size)
1157	self.num_heads = num_heads
1158	self.embed_dim = embed_dim
1159	self.head_dim = embed_dim // num_heads
1160	self.use_data_parallel = use_data_parallel
1161	self.tp_size = (
1162	1
1163	if use_data_parallel
1164	else parallel_state.get_tensor_model_parallel_world_size()
1165	)
1166
1167	self.embeddings = LlavaOnevision2VisionEmbeddings(
1168	patch_size=patch_size, in_channels=in_channels, embed_dim=embed_dim
1169	)
1170	self.layernorm_pre = nn.LayerNorm(embed_dim, eps=norm_eps)
1171
1172	self.rotary_pos_emb = LlavaOnevision2VisionRotaryEmbedding(
1173	self.head_dim, theta=rope_theta
1174	)
1175
1176	self.encoder = nn.Module()
1177	self.encoder.layers = nn.ModuleList(
1178	[
1179	LlavaOnevision2VisionTowerBlock(
1180	dim=embed_dim,
1181	num_heads=num_heads,
1182	mlp_hidden_dim=mlp_hidden_dim,
1183	norm_eps=norm_eps,
1184	quant_config=quant_config,
1185	prefix=f"{prefix}.encoder.layers.{i}",
1186	use_data_parallel=use_data_parallel,
1187	)
1188	for i in range(depth)
1189	]
1190	)
1191
1192	self.merger = LlavaOnevision2PatchMerger(
1193	d_model=text_hidden_size,
1194	context_dim=embed_dim,
1195	spatial_merge_size=spatial_merge_size,
1196	norm_eps=norm_eps,
1197	quant_config=quant_config,
1198	prefix=f"{prefix}.merger",
1199	use_data_parallel=use_data_parallel,
1200	)
1201
1202	# Vision attention backend; mirrors the one MMEncoderAttention picks
1203	# internally so the cu_seqlens / max_seqlen metadata computed below
1204	# matches the kernel actually used.
1205	self.attn_backend = get_vit_attn_backend(
1206	head_size=self.head_dim, dtype=torch.get_default_dtype()
1207	)
1208
1209	@property
1210	def dtype(self) -> torch.dtype:
1211	return self.embeddings.patch_embedding.weight.dtype
1212
1213	@property
1214	def device(self) -> torch.device:
1215	return self.embeddings.patch_embedding.weight.device
1216
1217	def _build_window_cu_seqlens(self, grid_thw: torch.Tensor) -> np.ndarray:
1218	"""Build cu_seqlens that chunk each sample's T-axis into windows of
1219	``frame_windows_size`` frames.
1220
1221	Returns an int32 ``np.ndarray`` of shape [num_windows+1] (the
1222	canonical prefix-sum format). Backend-specific transforms are applied
1223	afterwards via ``MMEncoderAttention.maybe_recompute_cu_seqlens``.
1224	"""
1225	win = self.frame_windows_size
1226	chunk_lengths: list[int] = []
1227	for row in grid_thw.tolist():
1228	t, h, w = int(row[0]), int(row[1]), int(row[2])
1229	per_frame = h * w
1230	t_remaining = t
1231	while t_remaining > 0:
1232	this_t = min(win, t_remaining)
1233	chunk_lengths.append(this_t * per_frame)
1234	t_remaining -= this_t
1235	cu = np.concatenate(
1236	[
1237	np.zeros(1, dtype=np.int32),
1238	np.array(chunk_lengths, dtype=np.int32).cumsum(dtype=np.int32),
1239	]
1240	)
1241	return cu
1242
1243	def forward(
1244	self,
1245	pixel_values: torch.Tensor,
1246	grid_thw: torch.Tensor,
1247	patch_positions: torch.Tensor,
1248	) -> torch.Tensor:
1249	x = pixel_values.to(device=self.device, dtype=self.dtype)
1250	x = self.embeddings(x)
1251	x = self.layernorm_pre(x)
1252
1253	rotary_pos_emb = self.rotary_pos_emb.forward_from_positions(
1254	patch_positions.to(self.device)
1255	)
1256
1257	# Build window cu_seqlens, then derive backend-specific attention
1258	# metadata (passthrough for FA/SDPA; transformed for FlashInfer).
1259	cu_seqlens_np = self._build_window_cu_seqlens(grid_thw)
1260	sequence_lengths = MMEncoderAttention.maybe_compute_seq_lens(
1261	self.attn_backend, cu_seqlens_np, self.device
1262	)
1263	max_seqlen = torch.tensor(
1264	MMEncoderAttention.compute_max_seqlen(self.attn_backend, cu_seqlens_np),
1265	dtype=torch.int32,
1266	)
1267	cu_seqlens = MMEncoderAttention.maybe_recompute_cu_seqlens(
1268	self.attn_backend,
1269	cu_seqlens_np,
1270	self.embed_dim,
1271	self.tp_size,
1272	self.device,
1273	)
1274
1275	for blk in self.encoder.layers:
1276	x = blk(
1277	x,
1278	cu_seqlens=cu_seqlens,
1279	rotary_pos_emb=rotary_pos_emb,
1280	max_seqlen=max_seqlen,
1281	sequence_lengths=sequence_lengths,
1282	)
1283
1284	return self.merger(x, patch_positions=patch_positions)
1285
1286
1287	class LlavaOnevision2ProcessingInfo(BaseProcessingInfo):
1288	def get_hf_config(self):
1289	return self.ctx.get_hf_config()
1290
1291	def get_data_parser(self) -> MultiModalDataParser:
1292	# ``video_needs_metadata=True`` makes the parser preserve both the
1293	# ``(frames, metadata)`` tuples from the frame backend and the
1294	# ``(dummy, {marker: path})`` tuples from prepare_codec_video_input;
1295	# both are dispatched by metadata content in ``_apply_hf_processor_main``.
1296	return LlavaOnevision2MultiModalDataParser(
1297	self.get_hf_config().vision_config.spatial_merge_size,
1298	video_needs_metadata=True,
1299	allow_missing_mm_embeddings=self.allow_missing_mm_embeddings,
1300	)
1301
1302	def get_hf_processor(self, **kwargs: object):
1303	# OV2's trust_remote_code processor is a bare class (not a
1304	# ProcessorMixin), so the shared get_hf_processor cannot load it; load
1305	# via AutoProcessor here (as other trust_remote_code models do).
1306	model_config = self.ctx.model_config
1307	# ``_merge_mm_kwargs`` restricts caller ``mm_processor_kwargs`` to known
1308	# processor args and wraps values as hashable for the lru_cache.
1309	merged = _merge_mm_kwargs(model_config, AutoProcessor, **kwargs)
1310	return _load_ov2_processor(
1311	model_config.model,
1312	model_config.revision,
1313	model_config.trust_remote_code,
1314	**merged,
1315	)
1316
1317	def get_image_processor(self, **kwargs: object) -> Qwen2VLImageProcessor:
1318	return self.get_hf_processor(**kwargs).image_processor
1319
1320	def get_mm_max_tokens_per_item(
1321	self,
1322	seq_len: int,
1323	mm_counts: Mapping[str, int],
1324	) -> Mapping[str, int]:
1325	return {
1326	"image": self.get_max_image_tokens(),
1327	"video": self.get_max_video_tokens(seq_len, mm_counts),
1328	}
1329
1330	def _get_vision_info(
1331	self,
1332	*,
1333	image_width: int,
1334	image_height: int,
1335	num_frames: int = 1,
1336	do_resize: bool = True,
1337	image_processor: Qwen2VLImageProcessor \| None,
1338	) -> tuple[ImageSize, int]:
1339	if image_processor is None:
1340	image_processor = self.get_image_processor()
1341	hf_config = self.get_hf_config()
1342	vision_config = hf_config.vision_config
1343	patch_size = vision_config.patch_size
1344	merge_size = vision_config.spatial_merge_size
1345	temporal_patch_size = getattr(vision_config, "temporal_patch_size", 1)
1346	if do_resize:
1347	min_pixels = getattr(image_processor, "min_pixels", None)
1348	max_pixels = getattr(image_processor, "max_pixels", None)
1349	if min_pixels is None or max_pixels is None:
1350	size = image_processor.size
1351	min_pixels = (
1352	getattr(size, "shortest_edge", None) or size["shortest_edge"]
1353	)
1354	max_pixels = getattr(size, "longest_edge", None) or size["longest_edge"]
1355	rh, rw = smart_resize(
1356	height=image_height,
1357	width=image_width,
1358	factor=patch_size * merge_size,
1359	min_pixels=min_pixels,
1360	max_pixels=max_pixels,
1361	)
1362	preprocessed = ImageSize(width=rw, height=rh)
1363	else:
1364	preprocessed = ImageSize(width=image_width, height=image_height)
1365	padded_frames = num_frames + (-num_frames % temporal_patch_size)
1366	grid_t = max(padded_frames // temporal_patch_size, 1)
1367	grid_h = preprocessed.height // patch_size
1368	grid_w = preprocessed.width // patch_size
1369	num_patches = grid_t * grid_h * grid_w
1370	return preprocessed, num_patches // (merge_size**2)
1371
1372	def get_num_image_tokens(
1373	self,
1374	*,
1375	image_width: int,
1376	image_height: int,
1377	image_processor: Qwen2VLImageProcessor \| None,
1378	) -> int:
1379	_, n = self._get_vision_info(
1380	image_width=image_width,
1381	image_height=image_height,
1382	image_processor=image_processor,
1383	)
1384	return n
1385
1386	def get_image_size_with_most_features(self) -> ImageSize:
1387	sz, _ = self._get_vision_info(
1388	image_width=1800, image_height=1800, image_processor=None
1389	)
1390	return sz
1391
1392	def get_max_image_tokens(self) -> int:
1393	w, h = self.get_image_size_with_most_features()
1394	return self.get_num_image_tokens(
1395	image_width=w, image_height=h, image_processor=None
1396	)
1397
1398	def get_num_video_tokens(
1399	self,
1400	*,
1401	image_width: int,
1402	image_height: int,
1403	num_frames: int,
1404	image_processor: Qwen2VLImageProcessor \| None = None,
1405	) -> int:
1406	_, n = self._get_vision_info(
1407	image_width=image_width,
1408	image_height=image_height,
1409	num_frames=num_frames,
1410	image_processor=image_processor,
1411	)
1412	return n
1413
1414	def _get_max_video_frames(self, max_tokens: int, start_num_frames: int = 1) -> int:
1415	w, h = self.get_image_size_with_most_features()
1416	num_frames = start_num_frames
1417	while True:
1418	next_num_frames = num_frames + 1
1419	next_max_tokens = self.get_num_video_tokens(
1420	image_width=w,
1421	image_height=h,
1422	num_frames=next_num_frames,
1423	)
1424	if next_max_tokens > max_tokens:
1425	break
1426	num_frames = next_num_frames
1427	return num_frames
1428
1429	def get_num_frames_with_most_features(
1430	self,
1431	seq_len: int,
1432	mm_counts: Mapping[str, int],
1433	max_frames_per_video: int = _MAX_FRAMES_PER_VIDEO,
1434	) -> int:
1435	max_videos = mm_counts.get("video", 0)
1436	max_total_frames = self._get_max_video_frames(seq_len)
1437	max_frames_per_video = min(
1438	max_total_frames // max(max_videos, 1), max_frames_per_video
1439	)
1440	return max(max_frames_per_video, 1)
1441
1442	def get_max_video_tokens(
1443	self,
1444	seq_len: int,
1445	mm_counts: Mapping[str, int],
1446	) -> int:
1447	w, h = self.get_image_size_with_most_features()
1448	return self.get_num_video_tokens(
1449	image_width=w,
1450	image_height=h,
1451	num_frames=self.get_num_frames_with_most_features(seq_len, mm_counts),
1452	)
1453
1454	def get_supported_mm_limits(self) -> Mapping[str, int \| None]:
1455	return {"image": None, "video": None}
1456
1457
1458	class LlavaOnevision2DummyInputsBuilder(
1459	BaseDummyInputsBuilder[LlavaOnevision2ProcessingInfo]
1460	):
1461	def get_dummy_text(self, mm_counts: Mapping[str, int]) -> str:
1462	n_img = mm_counts.get("image", 0)
1463	n_vid = mm_counts.get("video", 0)
1464	return (
1465	"<\|vision_start\|><\|image_pad\|><\|vision_end\|>" * n_img
1466	+ "<\|vision_start\|><\|video_pad\|><\|vision_end\|>" * n_vid
1467	)
1468
1469	def get_dummy_mm_data(
1470	self,
1471	seq_len: int,
1472	mm_counts: Mapping[str, int],
1473	mm_options: MultiModalDummyOptions \| None = None,
1474	) -> MultiModalDataDict:
1475	n_img = mm_counts.get("image", 0)
1476	n_vid = mm_counts.get("video", 0)
1477	w, h = self.info.get_image_size_with_most_features()
1478	out: dict[str, object] = {}
1479	if n_img:
1480	out["image"] = self._get_dummy_images(width=w, height=h, num_images=n_img)
1481	if n_vid:
1482	# 4 frames per dummy video keeps profiling cheap; the real frame
1483	# count is decided by the HF VideoProcessor at apply time.
1484	out["video"] = self._get_dummy_videos(
1485	width=w, height=h, num_frames=4, num_videos=n_vid
1486	)
1487	return out
1488
1489	def _get_dummy_videos(
1490	self,
1491	*,
1492	width: int,
1493	height: int,
1494	num_frames: int,
1495	num_videos: int,
1496	overrides=None,
1497	):
1498	# ``video_needs_metadata=True`` (see ProcessingInfo.get_data_parser)
1499	# makes the parser require a metadata dict on every video item, so the
1500	# dummy profiling videos must carry one too. ``do_sample_frames=False``
1501	# plus ``frames_indices=range(T)`` tells the frame path to consume the
1502	# pre-built frames verbatim (no resampling) -- mirrors GLM-4V.
1503	# OV2's vision tower (temporal merge = 2) needs an even frame count.
1504	num_frames = max(num_frames, _OV2_FRAME_FACTOR)
1505	if num_frames % _OV2_FRAME_FACTOR != 0:
1506	num_frames += 1
1507	videos = super()._get_dummy_videos(
1508	width=width,
1509	height=height,
1510	num_frames=num_frames,
1511	num_videos=num_videos,
1512	overrides=overrides,
1513	)
1514	video_items = []
1515	for video in videos:
1516	t = video.shape[0]
1517	metadata = {
1518	"fps": 1.0,
1519	"duration": float(t),
1520	"total_num_frames": int(t),
1521	"frames_indices": list(range(t)),
1522	"video_backend": "llava_onevision2",
1523	"do_sample_frames": False,
1524	}
1525	video_items.append((video, metadata))
1526	return video_items
1527
1528
1529	class LlavaOnevision2MultiModalDataParser(MultiModalDataParser):
1530	# The patch grid is what sizes the placeholder range.
1531	embedding_fields = {
1532	"image": {"image_embeds": "values", "image_grid_thw": "metadata"},
1533	}
1534
1535	def __init__(self, spatial_merge_size: int, args, *kwargs):
1536	self._spatial_merge_size = spatial_merge_size
1537	super().__init__(args, *kwargs)
1538
1539	def _parse_image_data(
1540	self,
1541	data: dict[str, torch.Tensor] \| ModalityData[ImageItem],
1542	) -> ModalityDataItems[Any, Any] \| None:
1543	if isinstance(data, dict):
1544	required, optional = self.embedding_field_sets("image")
1545	return DictEmbeddingItems(
1546	data,
1547	modality="image",
1548	required_fields=required,
1549	optional_fields=optional,
1550	fields_factory=_create_field_factory(self._spatial_merge_size),
1551	)
1552	return super()._parse_image_data(data)
1553
1554
1555	class LlavaOnevision2MultiModalProcessor(
1556	BaseMultiModalProcessor[LlavaOnevision2ProcessingInfo]
1557	):
1558	def _get_hf_mm_text(self, mm_counts: Mapping[str, int]) -> str:
1559	return self.dummy_inputs.get_dummy_text(mm_counts)
1560
1561	def _apply_hf_processor_main(
1562	self,
1563	mm_items: MultiModalDataItems,
1564	hf_kwargs: Mapping[str, object],
1565	) -> BatchFeature:
1566	hf_data, hf_kwargs, passthrough_data = self._get_hf_mm_inputs(
1567	mm_items, hf_kwargs
1568	)
1569
1570	if not hf_data:
1571	return self._finalize_hf_mm_data(hf_data, hf_kwargs, passthrough_data)
1572
1573	prompt_text = hf_data.pop("text")
1574	assert isinstance(prompt_text, str)
1575
1576	hf_processor = self.info.get_hf_processor(**hf_kwargs)
1577	merged_kwargs = self.info.ctx.get_merged_mm_kwargs(hf_kwargs)
1578	merged_kwargs.setdefault("return_tensors", "pt")
1579	call_kwargs = {
1580	k: v
1581	for k, v in merged_kwargs.items()
1582	if k
1583	in {
1584	"return_tensors",
1585	"padding",
1586	"num_frames",
1587	"max_frames",
1588	"target_fps",
1589	"video_backend",
1590	"max_pixels",
1591	"codec_config",
1592	}
1593	}
1594
1595	# Explicit None + length checks: ``mm_data[...]`` may be a list, numpy
1596	# array, or tensor, and ``and <array>`` would raise on the ambiguous
1597	# truth value of a multi-element array.
1598	_videos = hf_data.get("videos")
1599	if _videos is not None and not isinstance(_videos, Sized):
1600	raise TypeError(f"object of type '{type(_videos).__name__}' has no len()")
1601	videos_present = _videos is not None and len(_videos) > 0
1602
1603	codec_video_paths = (
1604	_extract_codec_video_paths(hf_data["videos"]) if videos_present else None
1605	)
1606	is_codec_marker = codec_video_paths is not None
1607	# Fallback: caller passed video_backend=codec via mm_processor_kwargs
1608	# without wrapping paths through prepare_codec_video_input (e.g.
1609	# lmms-eval's chat/vllm.py ov2_path_video=True). Recover the path
1610	# strings directly from mm_data["videos"] so the codec rename
1611	# branch still fires and video-modality fields get populated.
1612	is_codec_kwarg = (
1613	not is_codec_marker
1614	and videos_present
1615	and call_kwargs.get("video_backend") == "codec"
1616	)
1617	if is_codec_kwarg:
1618	raw = hf_data["videos"]
1619	if isinstance(raw, str):
1620	codec_video_paths = [raw]
1621	elif isinstance(raw, (list, tuple)) and all(
1622	isinstance(x, str) for x in raw
1623	):
1624	codec_video_paths = list(raw)
1625	else:
1626	# Non-string payload (PIL/ndarray/etc.) - codec backend
1627	# cannot consume pre-decoded frames; fall through to frame
1628	# path.
1629	is_codec_kwarg = False
1630	is_codec = is_codec_marker or is_codec_kwarg
1631
1632	if is_codec:
1633	# Confine codec paths to --allowed-local-media-path (local-only;
1634	# SSRF / local-file-read protection) and use the resolved paths
1635	# downstream so the codec module opens exactly the file that was
1636	# validated (no symlink-retarget / validate-vs-open gap).
1637	codec_video_paths = _validate_video_sources(
1638	codec_video_paths, self.info.ctx.model_config
1639	)
1640	# Codec backend: HF processor consumes the path string directly and
1641	# performs decode + canvas-packing internally. The dummy ndarray
1642	# we attached during prepare_codec_video_input is discarded here.
1643	hf_data["videos"] = (
1644	codec_video_paths
1645	if len(codec_video_paths) > 1
1646	else codec_video_paths[0]
1647	)
1648	# Route through ``ctx.call_hf_processor`` so float-tensor
1649	# dtype postprocessing is applied automatically; inject
1650	# ``video_backend="codec"`` via mm_kwargs so the wrapped processor
1651	# dispatches to its codec branch.
1652	output = self.info.ctx.call_hf_processor(
1653	self.info.get_hf_processor({hf_kwargs, "video_backend": "codec"}),
1654	dict(text=prompt_text, **hf_data),
1655	{**hf_kwargs, "video_backend": "codec"},
1656	)
1657	data = dict(output)
1658	processed_data = BatchFeature(
1659	self._rename_codec_outputs_to_video(
1660	data, codec_video_paths, hf_processor
1661	)
1662	)
1663	return self._finalize_hf_mm_data(
1664	hf_data, hf_kwargs, passthrough_data, processed_data
1665	)
1666
1667	# ---- Frame backend (registered LlavaOnevision2VideoBackend) ------
1668	# Every non-codec video reaches here as a ``(frames_ndarray, metadata)``
1669	# tuple (``video_needs_metadata=True``): the connector decoded + sampled
1670	# it through ``LlavaOnevision2VideoBackend`` for real ``video_url``
1671	# inputs, or the dummy-inputs builder attached synthetic metadata during
1672	# profiling. We materialise the frames as PIL images + per-frame
1673	# timestamp markers and feed them through the HF processor's image
1674	# branch (smart_resize + patchify), then re-tag the image-series outputs
1675	# as video-series so vLLM's ``<\|video_pad\|>`` placeholder replacement
1676	# finds them. Sampling parity with the original qwen_vl_utils policy is
1677	# provided by the backend's ``compute_frames_index_to_sample``; SSRF /
1678	# local-file gating is enforced by the connector before decoding.
1679	if videos_present:
1680	timestamp_decimals = hf_kwargs.get(
1681	"timestamp_decimals", _DEFAULT_TIMESTAMP_DECIMALS
1682	)
1683	if not isinstance(
1684	timestamp_decimals, (str, Buffer, SupportsInt, SupportsIndex)
1685	):
1686	raise TypeError(
1687	"int() argument must be a string, a bytes-like object "
1688	f"or a real number, not '{type(timestamp_decimals).__name__}'"
1689	)
1690	timestamp_decimals = int(timestamp_decimals)
1691
1692	per_video_frames: list[list[Image.Image]] = []
1693	per_video_timestamps: list[list[float]] = []
1694	videos = hf_data["videos"]
1695	if not isinstance(videos, Iterable):
1696	raise TypeError(f"Unsupported videos input: {type(videos)}")
1697	for item in videos:
1698	pil_frames, timestamps = _frame_video_to_pil_and_timestamps(item)
1699	per_video_frames.append(pil_frames)
1700	per_video_timestamps.append(timestamps)
1701
1702	# Rewrite the prompt so each video marker becomes a sequence of
1703	# ``<{t} seconds><\|vision_start\|><\|image_pad\|><\|vision_end\|>``
1704	# blocks (matches the OV2 hf-chat reference exactly).
1705	new_prompt = _expand_video_markers_in_prompt(
1706	prompt_text,
1707	per_video_timestamps,
1708	timestamp_decimals=timestamp_decimals,
1709	)
1710
1711	# Build the merged ``images`` list the wrapped HF processor will
1712	# consume. The processor binds the merged list positionally to the
1713	# ``<\|image_pad\|>`` slots in prompt order, so it must follow the
1714	# interleaved marker order of the prompt (not a fixed "videos first"
1715	# order) -- otherwise mixed image+video requests bind frames to the
1716	# wrong placeholder. ``row_is_video`` labels each grid row so outputs
1717	# can be split back into per-modality keys below.
1718	merged_mm_data = dict(hf_data)
1719	existing_images = merged_mm_data.pop("images", None)
1720	caller_images: list[Image.Image] = []
1721	if existing_images:
1722	if isinstance(existing_images, list):
1723	caller_images.extend(existing_images)
1724	else:
1725	caller_images.append(existing_images)
1726
1727	marker_pattern = re.compile(
1728	"(?P<image>" + re.escape(_IMAGE_MARKER) + ")"
1729	"\|(?P<video>" + re.escape(_VIDEO_MARKER) + ")"
1730	)
1731	flat_frames: list[Image.Image] = []
1732	row_is_video: list[bool] = []
1733	vid_idx = 0
1734	img_idx = 0
1735	for marker in marker_pattern.finditer(prompt_text):
1736	if marker.lastgroup == "video":
1737	frames = per_video_frames[vid_idx]
1738	vid_idx += 1
1739	flat_frames.extend(frames)
1740	row_is_video.extend([True] * len(frames))
1741	else:
1742	flat_frames.append(caller_images[img_idx])
1743	img_idx += 1
1744	row_is_video.append(False)
1745	# Fallback: if the prompt carried no parseable markers (defensive --
1746	# should not happen for well-formed inputs), preserve the legacy
1747	# "video frames first, then caller images" ordering.
1748	if vid_idx == 0 and img_idx == 0:
1749	for frames in per_video_frames:
1750	flat_frames.extend(frames)
1751	row_is_video.extend([True] * len(frames))
1752	flat_frames.extend(caller_images)
1753	row_is_video.extend([False] * len(caller_images))
1754
1755	merged_mm_data.pop("videos", None)
1756	merged_mm_data["images"] = flat_frames
1757
1758	# Route through ``ctx.call_hf_processor`` (applies float-tensor
1759	# dtype postprocessing automatically). The wrapped processor's image
1760	# branch ignores video/codec-only kwargs and does not forward extra
1761	# **kwargs to the image processor, so passing the full merged kwarg
1762	# set here is a no-op beyond return_tensors/padding.
1763	output = self.info.ctx.call_hf_processor(
1764	self.info.get_hf_processor(**hf_kwargs),
1765	dict(text=new_prompt, **merged_mm_data),
1766	hf_kwargs,
1767	)
1768	data = dict(output)
1769
1770	# Split the image-series processor outputs back into video rows and
1771	# genuine-image rows (one grid row per frame/image, in flat_frames
1772	# order). Video rows become video-series keys; caller-image rows
1773	# stay under image-series keys. Without this split a mixed
1774	# image+video request silently drops the image modality.
1775	grid = data.get("image_grid_thw")
1776	has_caller_images = any(not v for v in row_is_video)
1777	if grid is not None and has_caller_images:
1778	pixel_values = data.pop("pixel_values")
1779	patch_positions = data.pop("patch_positions")
1780
1781	per_row_patches = grid.prod(-1).tolist()
1782	row_offsets = [0]
1783	for p in per_row_patches:
1784	row_offsets.append(row_offsets[-1] + int(p))
1785
1786	def _gather_rows(tensor, rows):
1787	if not rows:
1788	return tensor[:0]
1789	return torch.cat(
1790	[tensor[row_offsets[i] : row_offsets[i + 1]] for i in rows],
1791	dim=0,
1792	)
1793
1794	video_rows = [i for i, v in enumerate(row_is_video) if v]
1795	image_rows = [i for i, v in enumerate(row_is_video) if not v]
1796
1797	data["pixel_values_videos"] = _gather_rows(pixel_values, video_rows)
1798	data["patch_positions_videos"] = _gather_rows(
1799	patch_positions, video_rows
1800	)
1801	data["video_grid_thw"] = grid[video_rows]
1802
1803	data["pixel_values"] = _gather_rows(pixel_values, image_rows)
1804	data["patch_positions"] = _gather_rows(patch_positions, image_rows)
1805	data["image_grid_thw"] = grid[image_rows]
1806	else:
1807	# Video-only (no caller images): re-tag every image-series
1808	# output as video-series.
1809	if "pixel_values" in data:
1810	data["pixel_values_videos"] = data.pop("pixel_values")
1811	if "image_grid_thw" in data:
1812	data["video_grid_thw"] = data.pop("image_grid_thw")
1813	if "patch_positions" in data:
1814	data["patch_positions_videos"] = data.pop("patch_positions")
1815	data["frame_timestamps"] = _pack_timestamps(per_video_timestamps)
1816	data["video_num_frames"] = torch.tensor(
1817	[len(ts) for ts in per_video_timestamps], dtype=torch.long
1818	)
1819	data["video_is_codec"] = torch.zeros(
1820	(len(per_video_timestamps),), dtype=torch.long
1821	)
1822
1823	processed_data = BatchFeature(data)
1824	return self._finalize_hf_mm_data(
1825	hf_data, hf_kwargs, passthrough_data, processed_data
1826	)
1827
1828	# ---- Image-only / text-only call --------------------------------
1829	# No videos present: delegate to ``ctx.call_hf_processor``, which
1830	# runs the wrapped processor over the (possibly empty) image set and
1831	# applies float-tensor dtype postprocessing automatically.
1832	processed_data = self.info.ctx.call_hf_processor(
1833	self.info.get_hf_processor(**hf_kwargs),
1834	dict(text=prompt_text, **hf_data),
1835	hf_kwargs,
1836	)
1837	return self._finalize_hf_mm_data(
1838	hf_data, hf_kwargs, passthrough_data, processed_data
1839	)
1840
1841	def _rename_codec_outputs_to_video(
1842	self,
1843	data: dict,
1844	codec_video_paths: list[str],
1845	hf_processor,
1846	) -> dict:
1847	# Codec branch emits image-series keys; vLLM expects video-series.
1848	# Timestamps are NOT synthesized here: the codec format packs multiple
1849	# source-frame timestamps into one canvas, so a per-canvas list is
1850	# ill-defined. We ship per-video fps and let get_video_replacement
1851	# reconstruct (timestamp, token_count) runs from patch_positions_videos.
1852	data["pixel_values_videos"] = data.pop("pixel_values")
1853	video_grid_thw = data.pop("image_grid_thw")
1854	data["video_grid_thw"] = video_grid_thw
1855	patch_positions = data.pop("patch_positions")
1856	data["patch_positions_videos"] = patch_positions
1857
1858	# Read per-video fps from processor output (order matches
1859	# codec_video_paths); fall back to _codec_fps_for only if the processor
1860	# didn't populate codec_fps (older model snapshots).
1861	processor_fps = data.pop("codec_fps", None)
1862	per_video_canvas_counts: list[int] = []
1863	per_video_fps: list[float] = []
1864	grid_offset = 0
1865	for idx, video_path in enumerate(codec_video_paths):
1866	num_canvases = self._count_canvases_for_video(
1867	patch_positions, video_grid_thw, 0, grid_offset
1868	)
1869	per_video_canvas_counts.append(num_canvases)
1870	if processor_fps is not None and idx < len(processor_fps):
1871	per_video_fps.append(float(processor_fps[idx]))
1872	else:
1873	per_video_fps.append(_codec_fps_for(video_path, hf_processor))
1874	grid_offset += num_canvases
1875
1876	data["video_num_frames"] = torch.tensor(
1877	per_video_canvas_counts, dtype=torch.long
1878	)
1879	data["video_is_codec"] = torch.ones(
1880	(len(per_video_canvas_counts),), dtype=torch.long
1881	)
1882	# Stored as scaled int64 (fps * 1e9 rounded) because vLLM casts
1883	# all float MM fields to model dtype (bfloat16), which rounds
1884	# 29.97003 fps to 30.0 and shifts every timestamp tag by ~0.1s.
1885	# Integer storage round-trips losslessly through the field-config
1886	# pipeline; we divide back to float at replacement time.
1887	data["codec_fps"] = torch.tensor(
1888	[int(round(f * 1_000_000_000)) for f in per_video_fps], dtype=torch.int64
1889	)
1890	# frame_timestamps is required by the field-config schema; emit an
1891	# empty per-video list (frames-backend populates it, codec ignores it).
1892	data["frame_timestamps"] = _pack_timestamps([[] for _ in codec_video_paths])
1893	return data
1894
1895	def _count_canvases_for_video(
1896	self,
1897	patch_positions: torch.Tensor,
1898	video_grid_thw: torch.Tensor,
1899	canvas_offset: int,
1900	grid_offset: int,
1901	) -> int:
1902	# Single-video case: every remaining grid row belongs to this video.
1903	# HF merges per-canvas rows into one [N,H,W] row per video, so the
1904	# canvas count lives in the t-dim, not shape[0].
1905	return int(video_grid_thw[grid_offset:, 0].sum().item())
1906
1907	def _get_prompt_updates(
1908	self,
1909	mm_items: MultiModalDataItems,
1910	hf_processor_mm_kwargs: Mapping[str, Any],
1911	out_mm_kwargs: MultiModalKwargsItems,
1912	) -> Sequence[PromptUpdate]:
1913	image_processor = self.info.get_image_processor(**hf_processor_mm_kwargs)
1914	tokenizer = self.info.get_tokenizer()
1915	vocab = tokenizer.get_vocab()
1916	image_pad_id = vocab["<\|image_pad\|>"]
1917	video_pad_id = vocab["<\|video_pad\|>"]
1918	vision_start_id = vocab["<\|vision_start\|>"]
1919	vision_end_id = vocab["<\|vision_end\|>"]
1920	newline_ids = cached_encode(tokenizer, "\n", add_special_tokens=False)
1921	merge_length = image_processor.merge_size**2
1922	decimals = int(hf_processor_mm_kwargs.get("timestamp_decimals", 1))
1923
1924	def get_image_replacement(item_idx: int):
1925	out_item = out_mm_kwargs["image"][item_idx]
1926	grid_thw = out_item["image_grid_thw"].data
1927	assert isinstance(grid_thw, torch.Tensor)
1928	n = int(grid_thw.prod(-1).sum()) // merge_length
1929	return [image_pad_id] * n
1930
1931	def get_video_replacement(item_idx: int):
1932	out_item = out_mm_kwargs["video"][item_idx]
1933	grid_thw = out_item["video_grid_thw"].data
1934	assert isinstance(grid_thw, torch.Tensor)
1935	is_codec_field = out_item.get("video_is_codec")
1936	if is_codec_field is not None:
1937	is_codec_data = is_codec_field.data
1938	assert isinstance(is_codec_data, torch.Tensor)
1939	is_codec = bool(int(is_codec_data.item()))
1940	else:
1941	is_codec = False
1942	tokens: list[int] = []
1943
1944	if is_codec:
1945	# Codec packs multiple source-frame timestamps into one canvas,
1946	# so timestamps are per-run (consecutive patches sharing the
1947	# same source-frame t), not per-canvas. Mirrors HF's
1948	# rewrite_text_with_codec_positions: group patch_positions by
1949	# run, emit ``<sec seconds><\|vision_start\|><pad*N><\|vision_end\|>\n``.
1950	patch_positions = out_item["patch_positions_videos"].data
1951	fps_t = out_item["codec_fps"].data
1952	assert isinstance(patch_positions, torch.Tensor)
1953	assert isinstance(fps_t, torch.Tensor)
1954	fps = float(int(fps_t.item())) / 1_000_000_000.0
1955	runs = _codec_timestamp_runs(
1956	patch_positions, fps, image_processor.merge_size
1957	)
1958	for sec, token_count in runs:
1959	tag = f"<{sec:.{decimals}f} seconds>"
1960	tag_ids = tokenizer.encode(tag, add_special_tokens=False)
1961	tokens.extend(tag_ids)
1962	tokens.append(vision_start_id)
1963	tokens.extend([image_pad_id] * token_count)
1964	tokens.append(vision_end_id)
1965	tokens.extend(newline_ids)
1966	else:
1967	timestamps = out_item["frame_timestamps"].data
1968	assert isinstance(timestamps, torch.Tensor)
1969	T_total = int(grid_thw.shape[0])
1970	for t in range(T_total):
1971	sec = float(timestamps[t].item())
1972	tag = f"<{sec:.{decimals}f} seconds>"
1973	tag_ids = tokenizer.encode(tag, add_special_tokens=False)
1974	n_per_frame = int(grid_thw[t].prod()) // merge_length
1975	tokens.extend(tag_ids)
1976	tokens.append(vision_start_id)
1977	tokens.extend([image_pad_id] * n_per_frame)
1978	tokens.append(vision_end_id)
1979	# Replacement mixes timestamp/marker tokens with image_pad
1980	# placeholders; only image_pad positions carry vision embeddings,
1981	# so a partial-embed mask is emitted via select_token_id.
1982	return PromptUpdateDetails.select_token_id(tokens, image_pad_id)
1983
1984	return [
1985	PromptReplacement(
1986	modality="image",
1987	target=[image_pad_id],
1988	replacement=get_image_replacement,
1989	),
1990	PromptReplacement(
1991	modality="video",
1992	target=[vision_start_id, video_pad_id, vision_end_id],
1993	replacement=get_video_replacement,
1994	),
1995	]
1996
1997	def _get_mm_fields_config(
1998	self,
1999	hf_inputs: BatchFeature,
2000	hf_processor_mm_kwargs: Mapping[str, object],
2001	) -> Mapping[str, MultiModalFieldConfig]:
2002	return _create_field_factory(
2003	self.info.get_hf_config().vision_config.spatial_merge_size,
2004	)(hf_inputs)
2005
2006
2007	@MULTIMODAL_REGISTRY.register_processor(
2008	LlavaOnevision2MultiModalProcessor,
2009	info=LlavaOnevision2ProcessingInfo,
2010	dummy_inputs=LlavaOnevision2DummyInputsBuilder,
2011	)
2012	class LlavaOnevision2ForConditionalGeneration(
2013	nn.Module, SupportsMultiModal, SupportsPP
2014	):
2015	"""vLLM-side OV2 top-level model.
2016
2017	Weight name rewriting (HF checkpoint → vLLM module tree):
2018	Prefix rewrites only (longest match wins). Vision tower attribute names
2019	mirror HF names verbatim, so no substring rules are needed. Substring
2020	rules would otherwise collide with the Qwen3 text-path ``self_attn``
2021	modules and break the language-model loader.
2022	``model.language_model.`` → ``language_model.model.``
2023	``model.visual.`` → ``visual.``
2024	``lm_head.`` → ``language_model.lm_head.``
2025	``model.`` (fallback) → ``language_model.model.``
2026	"""
2027
2028	hf_to_vllm_mapper = WeightsMapper(
2029	orig_to_new_prefix={
2030	"model.language_model.": "language_model.model.",
2031	"model.visual.": "visual.",
2032	"lm_head.": "language_model.lm_head.",
2033	"model.": "language_model.model.",
2034	}
2035	)
2036
2037	@classmethod
2038	def get_placeholder_str(cls, modality: str, i: int) -> str \| None:
2039	if modality.startswith("image"):
2040	return "<\|vision_start\|><\|image_pad\|><\|vision_end\|>"
2041	if modality.startswith("video"):
2042	return "<\|vision_start\|><\|video_pad\|><\|vision_end\|>"
2043	raise ValueError("Only image or video modality is supported")
2044
2045	def __init__(self, *, vllm_config: VllmConfig, prefix: str = "") -> None:
2046	super().__init__()
2047	config = vllm_config.model_config.hf_config
2048	multimodal_config = vllm_config.model_config.multimodal_config
2049
2050	self.config = config
2051	self.multimodal_config = multimodal_config
2052
2053	# Build the vision tower under the tower marker so it is shared by the
2054	# image, video-frame and video-codec backends (all routed through
2055	# ``self.visual``). When both modalities are disabled via
2056	# ``--limit-mm-per-prompt`` the marker turns the tower into a skipped
2057	# placeholder whose weights are dropped automatically by the loader.
2058	with self._mark_tower_model(vllm_config, {"image", "video"}):
2059	self.visual = LlavaOnevision2VisionTower(
2060	config.vision_config,
2061	text_hidden_size=config.text_config.hidden_size,
2062	norm_eps=getattr(config.vision_config, "layer_norm_eps", 1e-6),
2063	quant_config=None,
2064	prefix=maybe_prefix(prefix, "visual"),
2065	)
2066
2067	# OV2 LLM is plain Qwen3 -- 1-D positions, no M-RoPE. The wrapper
2068	# LlavaOnevision2Config keeps text/vision configs nested (unlike OV1.5
2069	# which promoted text fields), so explicitly hand text_config to the
2070	# Qwen3 init path or qwen3.py will hit AttributeError on vocab_size.
2071	self.language_model = init_vllm_registered_model(
2072	vllm_config=vllm_config,
2073	hf_config=config.text_config,
2074	prefix=maybe_prefix(prefix, "language_model"),
2075	architectures=["Qwen3ForCausalLM"],
2076	)
2077
2078	self.make_empty_intermediate_tensors = (
2079	self.language_model.make_empty_intermediate_tensors
2080	)
2081
2082	def _validate_and_reshape_mm_tensor(
2083	self, mm_input: object, name: str
2084	) -> torch.Tensor:
2085	if not isinstance(mm_input, (torch.Tensor, list)):
2086	raise ValueError(f"Incorrect type of {name}: {type(mm_input)}")
2087	if isinstance(mm_input, torch.Tensor):
2088	if mm_input.ndim == 2:
2089	return mm_input
2090	if mm_input.ndim != 3:
2091	raise ValueError(
2092	f"{name} must be 2D or batched-3D, got shape={mm_input.shape}"
2093	)
2094	# Flatten the leading batch dim into the patch dim: (b, n, d) ->
2095	# (b*n, d), avoiding a Python list of row tensors.
2096	return mm_input.flatten(0, 1)
2097	return torch.concat(mm_input)
2098
2099	def _parse_and_validate_image_input(
2100	self, **kwargs: object
2101	) -> LlavaOnevision2ImageInputs \| None:
2102	pixel_values = kwargs.pop("pixel_values", None)
2103	image_embeds = kwargs.pop("image_embeds", None)
2104	image_grid_thw = kwargs.pop("image_grid_thw", None)
2105	patch_positions = kwargs.pop("patch_positions", None)
2106
2107	if pixel_values is None and image_embeds is None:
2108	return None
2109
2110	if pixel_values is not None:
2111	pixel_values = self._validate_and_reshape_mm_tensor(
2112	pixel_values, "image pixel values"
2113	)
2114	image_grid_thw = self._validate_and_reshape_mm_tensor(
2115	image_grid_thw, "image grid_thw"
2116	)
2117	if patch_positions is None:
2118	raise ValueError(
2119	"OV2 requires patch_positions alongside pixel_values; "
2120	"ensure the HF processor produces it (image, video-"
2121	"frames, and video-codec backends all do)."
2122	)
2123	patch_positions = self._validate_and_reshape_mm_tensor(
2124	patch_positions, "image patch_positions"
2125	)
2126	return LlavaOnevision2ImagePixelInputs(
2127	type="pixel_values",
2128	pixel_values=pixel_values,
2129	image_grid_thw=image_grid_thw,
2130	patch_positions=patch_positions,
2131	)
2132
2133	# image_embeds path
2134	image_embeds = self._validate_and_reshape_mm_tensor(
2135	image_embeds, "image embeds"
2136	)
2137	image_grid_thw = self._validate_and_reshape_mm_tensor(
2138	image_grid_thw, "image grid_thw"
2139	)
2140	return LlavaOnevision2ImageEmbeddingInputs(
2141	type="image_embeds",
2142	image_embeds=image_embeds,
2143	image_grid_thw=image_grid_thw,
2144	)
2145
2146	def _process_image_input(
2147	self, image_input: LlavaOnevision2ImageInputs
2148	) -> tuple[torch.Tensor, ...]:
2149	grid_thw = image_input["image_grid_thw"]
2150	assert grid_thw.ndim == 2
2151
2152	if isinstance(image_input, LlavaOnevision2ImageEmbeddingInputs):
2153	image_embeds = image_input["image_embeds"]
2154	else:
2155	image_embeds = self.visual(
2156	image_input["pixel_values"],
2157	grid_thw=grid_thw,
2158	patch_positions=image_input["patch_positions"],
2159	)
2160
2161	merge_size = self.visual.spatial_merge_size
2162	sizes = grid_thw.prod(-1) // merge_size // merge_size
2163	return image_embeds.split(sizes.tolist())
2164
2165	def _parse_and_validate_video_input(
2166	self, **kwargs: object
2167	) -> LlavaOnevision2VideoPixelInputs \| None:
2168	pixel_values_videos = kwargs.pop("pixel_values_videos", None)
2169	video_grid_thw = kwargs.pop("video_grid_thw", None)
2170	patch_positions_videos = kwargs.pop("patch_positions_videos", None)
2171	video_num_frames = kwargs.pop("video_num_frames", None)
2172	kwargs.pop("frame_timestamps", None)
2173	kwargs.pop("video_is_codec", None)
2174
2175	if pixel_values_videos is None:
2176	return None
2177
2178	pixel_values_videos = self._validate_and_reshape_mm_tensor(
2179	pixel_values_videos, "video pixel values"
2180	)
2181	video_grid_thw = self._validate_and_reshape_mm_tensor(
2182	video_grid_thw, "video grid_thw"
2183	)
2184	if patch_positions_videos is None:
2185	raise ValueError(
2186	"OV2 requires patch_positions_videos alongside pixel_values_videos."
2187	)
2188	patch_positions_videos = self._validate_and_reshape_mm_tensor(
2189	patch_positions_videos, "video patch_positions"
2190	)
2191	if video_num_frames is None:
2192	raise ValueError(
2193	"OV2 requires video_num_frames alongside pixel_values_videos."
2194	)
2195	if isinstance(video_num_frames, list):
2196	video_num_frames = torch.cat([v.flatten() for v in video_num_frames])
2197	elif isinstance(video_num_frames, torch.Tensor):
2198	video_num_frames = video_num_frames.flatten()
2199	else:
2200	raise TypeError(
2201	"video_num_frames must be a tensor or a list of tensors, "
2202	f"got {type(video_num_frames)}"
2203	)
2204	return LlavaOnevision2VideoPixelInputs(
2205	type="pixel_values_videos",
2206	pixel_values_videos=pixel_values_videos,
2207	video_grid_thw=video_grid_thw,
2208	patch_positions_videos=patch_positions_videos,
2209	video_num_frames=video_num_frames,
2210	)
2211
2212	def _process_video_input(
2213	self,
2214	video_input: LlavaOnevision2VideoPixelInputs,
2215	) -> tuple[torch.Tensor, ...]:
2216	# OV2 encodes each video frame independently through the same
2217	# vision stack as still images, then splits the resulting token
2218	# stream into per-video chunks using video_num_frames.
2219	grid_thw = video_input["video_grid_thw"]
2220	assert grid_thw.ndim == 2
2221	video_embeds = self.visual(
2222	video_input["pixel_values_videos"],
2223	grid_thw=grid_thw,
2224	patch_positions=video_input["patch_positions_videos"],
2225	)
2226	merge_size = self.visual.spatial_merge_size
2227	per_frame_tokens = grid_thw.prod(-1) // merge_size // merge_size
2228	# Aggregate per-frame token counts into per-video token counts.
2229	num_frames = video_input["video_num_frames"].tolist()
2230	sizes: list[int] = []
2231	cursor = 0
2232	for n in num_frames:
2233	n = int(n)
2234	sizes.append(int(per_frame_tokens[cursor : cursor + n].sum()))
2235	cursor += n
2236	return video_embeds.split(sizes)
2237
2238	def _parse_and_validate_multimodal_inputs(
2239	self, **kwargs: object
2240	) -> LlavaOnevision2InputsByModality:
2241	modalities = LlavaOnevision2InputsByModality()
2242	for key in kwargs:
2243	if key in ("pixel_values", "image_embeds") and "images" not in modalities:
2244	modalities["images"] = self._parse_and_validate_image_input(**kwargs)
2245	if key == "pixel_values_videos" and "videos" not in modalities:
2246	modalities["videos"] = self._parse_and_validate_video_input(**kwargs)
2247	return modalities
2248
2249	def get_language_model(self) -> torch.nn.Module:
2250	return self.language_model
2251
2252	def embed_multimodal(self, **kwargs: object) -> MultiModalEmbeddings:
2253	modalities = self._parse_and_validate_multimodal_inputs(**kwargs)
2254	if not modalities:
2255	return []
2256	multimodal_embeddings: tuple[torch.Tensor, ...] = ()
2257	for modality in modalities:
2258	if modality == "images":
2259	image_input = modalities["images"]
2260	assert image_input is not None
2261	multimodal_embeddings += self._process_image_input(image_input)
2262	elif modality == "videos":
2263	video_input = modalities["videos"]
2264	assert video_input is not None
2265	multimodal_embeddings += self._process_video_input(video_input)
2266	return multimodal_embeddings
2267
2268	def get_input_embeddings(
2269	self,
2270	input_ids: torch.Tensor,
2271	multimodal_embeddings: MultiModalEmbeddings \| None = None,
2272	) -> torch.Tensor:
2273	inputs_embeds = self.language_model.get_input_embeddings(input_ids)
2274	if multimodal_embeddings is not None and len(multimodal_embeddings) != 0:
2275	inputs_embeds = merge_multimodal_embeddings(
2276	inputs_embeds,
2277	multimodal_embeddings,
2278	input_ids == self.config.image_token_id,
2279	)
2280	return inputs_embeds
2281
2282	def forward(
2283	self,
2284	input_ids: torch.Tensor,
2285	positions: torch.Tensor,
2286	intermediate_tensors: IntermediateTensors \| None = None,
2287	inputs_embeds: torch.Tensor \| None = None,
2288	**kwargs: object,
2289	) -> torch.Tensor \| IntermediateTensors:
2290	# V1 precomputes multimodal ``inputs_embeds`` via ``embed_multimodal`` /
2291	# ``get_input_embeddings``, so the model forward only threads them into
2292	# the language model (image + video share the same embedding merge).
2293	if intermediate_tensors is not None:
2294	inputs_embeds = None
2295
2296	return self.language_model.model(
2297	input_ids=input_ids,
2298	positions=positions,
2299	intermediate_tensors=intermediate_tensors,
2300	inputs_embeds=inputs_embeds,
2301	)
2302
2303	def compute_logits(self, hidden_states: torch.Tensor):
2304	return self.language_model.compute_logits(hidden_states)
2305
2306	def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
2307	loader = AutoWeightsLoader(self)
2308	return loader.load_weights(weights, mapper=self.hf_to_vllm_mapper)
2309
2310	def get_mm_mapping(self) -> MultiModelKeys:
2311	return MultiModelKeys.from_string_field(
2312	language_model="language_model",
2313	connector="visual.merger.",
2314	tower_model="visual.",
2315	)