@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / 7e6f123810

Fork 0

A high-throughput and memory-efficient inference and serving engine for LLMs

Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git

Synced Nov 10, 04:00 UTC

Add Molmo2 multimodal model support (#30997)

Signed-off-by: sanghol <sanghol@allenai.org> Signed-off-by: Isotr0py <mozf@mail2.sysu.edu.cn> Co-authored-by: Isotr0py <mozf@mail2.sysu.edu.cn>

sangho.lee committed Jan 13, 2026 at 23:33 UTC 7e6f12381092cd0b457ad10c57fa48ca73c415a7

11 files changed +3215 -1

docs/models/supported_models.md

@@ -698,6 +698,7 @@ These models primarily accept the [`LLM.generate`](./generative_models.md#llmgen

698 | `MiniMaxVL01ForConditionalGeneration` | MiniMax-VL | T + IE+ | `MiniMaxAI/MiniMax-VL-01`, etc. | | ✅︎ |

699 | `Mistral3ForConditionalGeneration` | Mistral3 (HF Transformers) | T + I+ | `mistralai/Mistral-Small-3.1-24B-Instruct-2503`, etc. | ✅︎ | ✅︎ |

700 | `MolmoForCausalLM` | Molmo | T + I+ | `allenai/Molmo-7B-D-0924`, `allenai/Molmo-7B-O-0924`, etc. | ✅︎ | ✅︎ |

701 +| `Molmo2ForConditionalGeneration` | Molmo2 | T + I+ / V | `allenai/Molmo2-4B`, `allenai/Molmo2-8B`, `allenai/Molmo2-O-7B` | ✅︎ | ✅︎ |

702 | `NVLM_D_Model` | NVLM-D 1.0 | T + I+ | `nvidia/NVLM-D-72B`, etc. | | ✅︎ |

703 | `OpenCUAForConditionalGeneration` | OpenCUA-7B | T + IE+ | `xlangai/OpenCUA-7B` | ✅︎ | ✅︎ |

704 | `Ovis` | Ovis2, Ovis1.6 | T + I+ | `AIDC-AI/Ovis2-1B`, `AIDC-AI/Ovis1.6-Llama3.2-3B`, etc. | | ✅︎ |

examples/offline_inference/vision_language.py

+32

@@ -1227,6 +1227,36 @@ def run_molmo(questions: list[str], modality: str) -> ModelRequestData:

1227 )

1228

1229

1230 +# Molmo2

1231 +def run_molmo2(questions: list[str], modality: str) -> ModelRequestData:

1232 + model_name = "allenai/Molmo2-8B"

1233 +

1234 + engine_args = EngineArgs(

1235 + model=model_name,

1236 + trust_remote_code=True,

1237 + dtype="bfloat16",

1238 + limit_mm_per_prompt={modality: 1},

1239 + max_num_batched_tokens=36864,

1240 + )

1241 +

1242 + if modality == "image":

1243 + placeholder = "<|image|>"

1244 + elif modality == "video":

1245 + placeholder = "<|video|>"

1246 + else:

1247 + raise ValueError(f"Unsupported modality for molmo2: {modality}")

1248 +

1249 + prompts = [

1250 + f"{placeholder}<|im_start|>user\n{question}<|im_end|>\n<|im_start|>assistant\n"

1251 + for question in questions

1252 + ]

1253 +

1254 + return ModelRequestData(

1255 + engine_args=engine_args,

1256 + prompts=prompts,

1257 + )

1258 +

1259 +

1260 # Nemontron_VL

1261 def run_nemotron_vl(questions: list[str], modality: str) -> ModelRequestData:

1262 model_name = "nvidia/Llama-3.1-Nemotron-Nano-VL-8B-V1"

@@ -1920,6 +1950,7 @@ model_example_map = {

1950 "minimax_vl_01": run_minimax_vl_01,

1951 "mistral3": run_mistral3,

1952 "molmo": run_molmo,

1953 + "molmo2": run_molmo2,

1954 "nemotron_vl": run_nemotron_vl,

1955 "NVLM_D": run_nvlm_d,

1956 "ovis": run_ovis,

@@ -1949,6 +1980,7 @@ MODELS_NEED_VIDEO_METADATA = [

1980 "glm4_1v",

1981 "glm4_5v",

1982 "glm4_5v_fp8",

1983 + "molmo2",

1984 "qwen3_vl",

1985 "qwen3_vl_moe",

1986 ]

examples/offline_inference/vision_language_multi_image.py

+38

@@ -1301,6 +1301,43 @@ def load_glm4_5v_fp8(question: str, image_urls: list[str]) -> ModelRequestData:

1301 )

1302

1303

1304 +def load_molmo2(question: str, image_urls: list[str]) -> ModelRequestData:

1305 + model_name = "allenai/Molmo2-8B"

1306 +

1307 + engine_args = EngineArgs(

1308 + model=model_name,

1309 + trust_remote_code=True,

1310 + dtype="bfloat16",

1311 + limit_mm_per_prompt={"image": len(image_urls)},

1312 + max_num_batched_tokens=36864,

1313 + )

1314 +

1315 + placeholders = [{"type": "image", "image": url} for url in image_urls]

1316 + messages = [

1317 + {

1318 + "role": "user",

1319 + "content": [

1320 + *placeholders,

1321 + {"type": "text", "text": question},

1322 + ],

1323 + },

1324 + ]

1325 +

1326 + processor = AutoProcessor.from_pretrained(model_name)

1327 +

1328 + prompt = processor.apply_chat_template(

1329 + messages, tokenize=False, add_generation_prompt=True

1330 + )

1331 +

1332 + image_data = [fetch_image(url) for url in image_urls]

1333 +

1334 + return ModelRequestData(

1335 + engine_args=engine_args,

1336 + prompt=prompt,

1337 + image_data=image_data,

1338 + )

1339 +

1340 +

1341 model_example_map = {

1342 "aria": load_aria,

1343 "aya_vision": load_aya_vision,

@@ -1323,6 +1360,7 @@ model_example_map = {

1360 "llava-next": load_llava_next,

1361 "llava-onevision": load_llava_onevision,

1362 "mistral3": load_mistral3,

1363 + "molmo2": load_molmo2,

1364 "NVLM_D": load_nvlm_d,

1365 "ovis": load_ovis,

1366 "ovis2_5": load_ovis2_5,

tests/models/multimodal/processing/test_common.py

@@ -123,6 +123,7 @@ MM_DATA_PATCHES = {

123 "glm4v": glm4_1v_patch_mm_data,

124 "glm4v_moe": glm4_1v_patch_mm_data,

125 "glmasr": glmasr_patch_mm_data,

126 + "molmo2": qwen3_vl_patch_mm_data,

127 "qwen3_vl": qwen3_vl_patch_mm_data,

128 "qwen3_vl_moe": qwen3_vl_patch_mm_data,

129 }

tests/models/registry.py

+13

@@ -92,6 +92,11 @@ class _HfExamplesInfo:

92 length that is too large to fit into memory in CI.

93 """

95 + max_num_batched_tokens: int | None = None

96 + """

97 + The maximum number of tokens to be processed in a single batch.

98 + """

99 +

100 revision: str | None = None

101 """

102 The specific revision (commit hash, tag, or branch) to use for the model.

@@ -817,6 +822,14 @@ _MULTIMODAL_EXAMPLE_MODELS = {

822 extras={"olmo": "allenai/Molmo-7B-O-0924"},

823 trust_remote_code=True,

824 ),

825 + "Molmo2ForConditionalGeneration": _HfExamplesInfo(

826 + "allenai/Molmo2-8B",

827 + extras={"olmo": "allenai/Molmo2-O-7B"},

828 + min_transformers_version="4.51",

829 + trust_remote_code=True,

830 + # required by current PrefixLM implementation

831 + max_num_batched_tokens=31872,

832 + ),

833 "NVLM_D": _HfExamplesInfo("nvidia/NVLM-D-72B", trust_remote_code=True),

834 "Llama_Nemotron_Nano_VL": _HfExamplesInfo(

835 "nvidia/Llama-3.1-Nemotron-Nano-VL-8B-V1",

tests/models/test_initialization.py

@@ -140,6 +140,7 @@ def can_initialize(

140 else None,

141 trust_remote_code=model_info.trust_remote_code,

142 max_model_len=model_info.max_model_len,

143 + max_num_batched_tokens=model_info.max_num_batched_tokens,

144 # these tests seem to produce leftover memory

145 gpu_memory_utilization=0.80,

146 load_format="dummy",

vllm/config/model.py

@@ -1127,6 +1127,7 @@ class ModelConfig:

1127 """Whether to use bidirectional attention for mm positions."""

1128 MM_PREFIX_LM_MODELS = (

1129 "gemma3",

1130 + "molmo2",

1131 "paligemma",

1132 )

1133 if not hasattr(self.hf_config, "model_type"):

vllm/model_executor/models/molmo2.py new

+2793

@@ -0,0 +1,2793 @@

1 +# SPDX-License-Identifier: Apache-2.0

2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project

3 +import math

4 +from collections.abc import Iterable, Mapping, Sequence

5 +from dataclasses import dataclass, fields

6 +from functools import cached_property, partial

7 +from itertools import islice

8 +from typing import Annotated, Any

9 +

10 +import numpy as np

11 +import torch

12 +import torch.nn as nn

13 +import torch.nn.functional as F

14 +from PIL import ImageOps

15 +from PIL.Image import Image

16 +from transformers import (

17 + BatchFeature,

18 + PretrainedConfig,

19 + ProcessorMixin,

20 + TensorType,

21 +)

22 +from transformers.image_utils import ImageInput

23 +from transformers.tokenization_utils_base import TextInput

24 +from transformers.video_utils import VideoInput, VideoMetadata

25 +

26 +from vllm.attention.layer import Attention

27 +from vllm.compilation.decorators import support_torch_compile

28 +from vllm.config import CacheConfig, VllmConfig

29 +from vllm.config.multimodal import BaseDummyOptions, VideoDummyOptions

30 +from vllm.distributed import (

31 + get_pp_group,

32 + get_tensor_model_parallel_rank,

33 + get_tensor_model_parallel_world_size,

34 + split_tensor_along_last_dim,

35 + tensor_model_parallel_all_gather,

36 +)

37 +from vllm.logger import init_logger

38 +from vllm.model_executor.layers.activation import MulAndSilu, SiluAndMul, get_act_fn

39 +from vllm.model_executor.layers.attention.mm_encoder_attention import MMEncoderAttention

40 +from vllm.model_executor.layers.layernorm import RMSNorm

41 +from vllm.model_executor.layers.linear import (

42 + ColumnParallelLinear,

43 + MergedColumnParallelLinear,

44 + QKVParallelLinear,

45 + RowParallelLinear,

46 +)

47 +from vllm.model_executor.layers.logits_processor import LogitsProcessor

48 +from vllm.model_executor.layers.quantization import QuantizationConfig

49 +from vllm.model_executor.layers.rotary_embedding import get_rope

50 +from vllm.model_executor.layers.vocab_parallel_embedding import (

51 + ParallelLMHead,

52 + VocabParallelEmbedding,

53 +)

54 +from vllm.model_executor.model_loader.weight_utils import default_weight_loader

55 +from vllm.model_executor.models.module_mapping import MultiModelKeys

56 +from vllm.multimodal import MULTIMODAL_REGISTRY

57 +from vllm.multimodal.inputs import (

58 + MultiModalDataDict,

59 + MultiModalFieldConfig,

60 + MultiModalKwargsItems,

61 + VideoItem,

62 +)

63 +from vllm.multimodal.parse import (

64 + ImageProcessorItems,

65 + ImageSize,

66 + MultiModalDataItems,

67 + MultiModalDataParser,

68 +)

69 +from vllm.multimodal.processing import (

70 + BaseMultiModalProcessor,

71 + BaseProcessingInfo,

72 + PromptReplacement,

73 + PromptUpdate,

74 + PromptUpdateDetails,

75 +)

76 +from vllm.multimodal.profiling import BaseDummyInputsBuilder

77 +from vllm.sequence import IntermediateTensors

78 +from vllm.utils.math_utils import round_down

79 +from vllm.utils.tensor_schema import TensorSchema, TensorShape

80 +

81 +from .interfaces import (

82 + MultiModalEmbeddings,

83 + SupportsLoRA,

84 + SupportsMultiModal,

85 + SupportsPP,

86 + SupportsQuant,

87 +)

88 +from .utils import (

89 + AutoWeightsLoader,

90 + WeightsMapper,

91 + _merge_multimodal_embeddings,

92 + extract_layer_index,

93 + is_pp_missing_parameter,

94 + make_empty_intermediate_tensors_factory,

95 + make_layers,

96 + maybe_prefix,

97 +)

98 +

99 +logger = init_logger(__name__)

100 +

101 +

102 +# Special tokens. These should be present in any tokenizer we use

103 +# because the preprocessor relies on them.

104 +IMAGE_PROMPT = "<|image|>"

105 +VIDEO_PROMPT = "<|video|>"

106 +_MAX_VIDEO_FPS = 8

107 +

108 +

109 +class Molmo2ImageInputs(TensorSchema):

110 + """

111 + Dimensions:

112 + - nc: The total number of crops (dynamic)

113 + - np: The total number of patches per crop

114 + - cps: Number of channels * patch_size * patch_size

115 + - npp: Number of pooled patches (dynamic)

116 + - pp: pooling_size * pooling_size

117 + - ni: Number of images

118 + - nt: Number of image tokens (dynamic)

119 + """

120 +

121 + pixel_values: Annotated[torch.Tensor, TensorShape("nc", "np", "cps")]

122 +

123 + token_pooling: Annotated[torch.Tensor, TensorShape("npp", "pp")]

124 + """

125 + An index tensor that maps image features to their corresponding

126 + patch tokens before pooling.

127 + """

128 +

129 + num_pooled_patches: Annotated[torch.Tensor, TensorShape("ni")]

130 +

131 + image_tokens: Annotated[torch.BoolTensor, TensorShape("nt")]

132 +

133 + num_image_tokens: Annotated[torch.Tensor, TensorShape("ni")]

134 +

135 +

136 +class Molmo2VideoInputs(TensorSchema):

137 + """

138 + Dimensions:

139 + - nc: The total number of frames (dynamic)

140 + - np: The total number of patches per frame

141 + - cps: Number of channels * patch_size * patch_size

142 + - npp: Number of pooled patches (dynamic)

143 + - pp: pooling_size * pooling_size

144 + - nv: Number of videos

145 + - nt: Number of video tokens (dynamic)

146 + """

147 +

148 + pixel_values_videos: Annotated[torch.Tensor, TensorShape("nc", "np", "cps")]

149 +

150 + token_pooling: Annotated[torch.Tensor, TensorShape("npp", "pp")]

151 + """

152 + An index tensor that maps image features to their corresponding

153 + patch tokens before pooling.

154 + """

155 +

156 + num_pooled_patches: Annotated[torch.Tensor, TensorShape("nv")]

157 +

158 + video_tokens: Annotated[torch.BoolTensor, TensorShape("nt")]

159 +

160 + num_video_tokens: Annotated[torch.Tensor, TensorShape("nv")]

161 +

162 +

163 +@dataclass

164 +class VitConfig:

165 + """Config for a vision transformer"""

166 +

167 + hidden_size: int = 1152

168 + intermediate_size: int = 4304

169 + num_hidden_layers: int = 27

170 + num_attention_heads: int = 16

171 + num_key_value_heads: int = 16

172 + head_dim: int = 72

173 + hidden_act: str = "gelu_pytorch_tanh"

174 + layer_norm_eps: float = 1e-6

175 + image_default_input_size: tuple[int, int] = (378, 378)

176 + image_patch_size: int = 14

177 + image_num_pos: int = 577

178 +

179 + def __post_init__(self):

180 + self.image_default_input_size = tuple(self.image_default_input_size) # type: ignore[assignment]

181 +

182 + @property

183 + def image_num_patch(self):

184 + h, w = self.image_default_input_size

185 + return h // self.image_patch_size, w // self.image_patch_size

186 +

187 +

188 +@dataclass

189 +class AdapterConfig:

190 + """Config for a vit-llm adapter"""

191 +

192 + vit_layers: tuple[int, int] = (-3, -9)

193 + pooling_attention_mask: bool = False

194 + hidden_size: int = 1152

195 + num_attention_heads: int = 16

196 + num_key_value_heads: int = 16

197 + head_dim: int = 72

198 + hidden_act: str = "silu"

199 + intermediate_size: int = 18944

200 + text_hidden_size: int = 3584

201 +

202 +

203 +@dataclass

204 +class TextConfig:

205 + """Configuration for a text model transformer"""

206 +

207 + hidden_size: int = 3584

208 + """

209 + The hidden size of the model.

210 + """

211 +

212 + num_attention_heads: int = 28

213 + """

214 + The number of self-attention heads.

215 + """

216 +

217 + num_key_value_heads: int = 4

218 + """

219 + The number of heads to use for keys and values.

220 + """

221 +

222 + head_dim: int = 128

223 + """

224 + The head dimensionality for the attention mechanism.

225 + """

226 +

227 + vocab_size: int = 152064

228 + """Vocabulary size of the model."""

229 +

230 + additional_vocab_size: int = 128

231 + """Number of additional tokens to have the input embeddings for"""

232 +

233 + qkv_bias: bool = True

234 + """

235 + Do QKV projection a bias

236 + """

237 +

238 + num_hidden_layers: int = 48

239 + """

240 + The number of layers/blocks.

241 + """

242 +

243 + intermediate_size: int = 18944

244 + """

245 + The hidden size for the MLP.

246 + """

247 +

248 + hidden_act: str = "silu"

249 + """

250 + The activation function to use within the MLP layers.

251 + """

252 +

253 + max_position_embeddings: int = 4096

254 + """

255 + Max positional embeddings to use in RoPE cache

256 + """

257 +

258 + rope_theta: float = 1000000.0

259 + """

260 + RoPE theta parameter.

261 + """

262 +

263 + use_qk_norm: bool = False

264 + """

265 + Apply layer norm to the keys and queries within the attention mechanism.

266 + This can help stabilize training.

267 + """

268 +

269 + qk_norm_type: str = "olmo"

270 + """

271 + The type of layer norm to use for the keys and queries.

272 + Can be "olmo" or "qwen3".

273 + """

274 +

275 + layer_norm_eps: float = 1e-6

276 + """

277 + epsilon for layer norms

278 + """

279 +

280 + norm_after: bool = False

281 + """Apply layer norm before and after the attention and MLP blocks."""

282 +

283 + rope_scaling_layers: tuple[int, ...] | None = None

284 + """

285 + RoPE scaling layers.

286 + """

287 +

288 +

289 +class ViTMLP(nn.Module):

290 + """MLP used in Vision Transformer."""

291 +

292 + def __init__(

293 + self,

294 + dim: int,

295 + hidden_dim: int,

296 + hidden_act: str,

297 + quant_config: QuantizationConfig | None = None,

298 + ) -> None:

299 + super().__init__()

300 + self.w1 = ColumnParallelLinear(

301 + dim,

302 + hidden_dim,

303 + bias=True,

304 + quant_config=quant_config,

305 + )

306 + # Activation function.

307 + self.act = get_act_fn(hidden_act)

308 + self.w2 = RowParallelLinear(

309 + hidden_dim,

310 + dim,

311 + bias=True,

312 + quant_config=quant_config,

313 + )

314 +

315 + def forward(self, x: torch.Tensor) -> torch.Tensor:

316 + x, _ = self.w1(x)

317 + x = self.act(x)

318 + x, _ = self.w2(x)

319 + return x

320 +

321 +

322 +class ViTMultiHeadDotProductAttention(nn.Module):

323 + """Multi-head attention used in Vision Transformer."""

324 +

325 + def __init__(

326 + self,

327 + hidden_size: int,

328 + num_heads: int,

329 + num_key_value_heads: int,

330 + head_dim: int,

331 + use_bias: bool = True,

332 + quant_config: QuantizationConfig | None = None,

333 + prefix: str = "",

334 + ) -> None:

335 + super().__init__()

336 +

337 + self.hidden_size = hidden_size

338 + self.total_num_heads = num_heads

339 + tp_size = get_tensor_model_parallel_world_size()

340 +

341 + assert self.hidden_size % self.total_num_heads == 0

342 + assert self.total_num_heads % tp_size == 0

343 +

344 + self.num_heads = self.total_num_heads // tp_size

345 + self.head_dim = head_dim

346 +

347 + assert self.head_dim == self.hidden_size // self.total_num_heads

348 +

349 + self.total_num_kv_heads = num_key_value_heads

350 + if self.total_num_kv_heads >= tp_size:

351 + assert self.total_num_kv_heads % tp_size == 0

352 + else:

353 + assert tp_size % self.total_num_kv_heads == 0

354 +

355 + self.num_kv_heads = max(1, self.total_num_kv_heads // tp_size)

356 +

357 + self.q_size = self.num_heads * self.head_dim

358 + self.kv_size = self.num_kv_heads * self.head_dim

359 +

360 + self.merged_qkv = QKVParallelLinear(

361 + self.hidden_size,

362 + self.head_dim,

363 + self.total_num_heads,

364 + self.total_num_kv_heads,

365 + bias=use_bias,

366 + quant_config=quant_config,

367 + )

368 + self.wo = RowParallelLinear(

369 + self.total_num_heads * self.head_dim,

370 + self.hidden_size,

371 + bias=use_bias,

372 + quant_config=quant_config,

373 + )

374 + self.scale = self.head_dim**-0.5

375 + self.attn = MMEncoderAttention(

376 + self.num_heads,

377 + self.head_dim,

378 + self.scale,

379 + num_kv_heads=self.num_kv_heads,

380 + prefix=f"{prefix}.attn",

381 + )

382 +

383 + def forward(self, inputs: torch.Tensor) -> torch.Tensor:

384 + qkv, _ = self.merged_qkv(inputs)

385 + xq, xk, xv = qkv.split([self.q_size, self.kv_size, self.kv_size], dim=-1)

386 +

387 + output = self.attn(xq, xk, xv)

388 +

389 + output, _ = self.wo(output)

390 +

391 + return output

392 +

393 +

394 +class Molmo2VisionBlock(nn.Module):

395 + """Residual attention block used in Vision Transformer."""

396 +

397 + def __init__(

398 + self,

399 + config: VitConfig,

400 + quant_config: QuantizationConfig | None = None,

401 + prefix: str = "",

402 + ) -> None:

403 + super().__init__()

404 + self.attention = ViTMultiHeadDotProductAttention(

405 + hidden_size=config.hidden_size,

406 + num_heads=config.num_attention_heads,

407 + num_key_value_heads=config.num_key_value_heads,

408 + head_dim=config.head_dim,

409 + quant_config=quant_config,

410 + prefix=f"{prefix}.attention",

411 + )

412 + self.feed_forward = ViTMLP(

413 + dim=config.hidden_size,

414 + hidden_dim=config.intermediate_size,

415 + hidden_act=config.hidden_act,

416 + quant_config=quant_config,

417 + )

418 + self.attention_norm = nn.LayerNorm(

419 + config.hidden_size,

420 + eps=config.layer_norm_eps,

421 + )

422 + self.ffn_norm = nn.LayerNorm(

423 + config.hidden_size,

424 + eps=config.layer_norm_eps,

425 + )

426 +

427 + def forward(self, x: torch.Tensor) -> torch.Tensor:

428 + x = x + self.attention(self.attention_norm(x))

429 + x = x + self.feed_forward(self.ffn_norm(x))

430 + return x

431 +

432 +

433 +class Molmo2VisionBlockCollection(nn.Module):

434 + """Collection of residual attention blocks used in Vision Transformer."""

435 +

436 + def __init__(

437 + self,

438 + config: VitConfig,

439 + quant_config: QuantizationConfig | None = None,

440 + prefix: str = "",

441 + ) -> None:

442 + super().__init__()

443 + self.resblocks = nn.ModuleList(

444 + [

445 + Molmo2VisionBlock(

446 + config,

447 + quant_config,

448 + prefix=f"{prefix}.resblocks.{layer_idx}",

449 + )

450 + for layer_idx in range(config.num_hidden_layers)

451 + ]

452 + )

453 +

454 + def forward(self, x: torch.Tensor) -> list[torch.Tensor]:

455 + hidden_states = []

456 + for r in self.resblocks:

457 + x = r(x)

458 + hidden_states.append(x)

459 + return hidden_states

460 +

461 +

462 +class Molmo2VisionTransformer(nn.Module):

463 + """Vision Transformer used in Vision Backbone."""

464 +

465 + def __init__(

466 + self,

467 + config: VitConfig,

468 + quant_config: QuantizationConfig | None = None,

469 + prefix: str = "",

470 + ) -> None:

471 + super().__init__()

472 + scale = config.hidden_size**-0.5

473 + self.num_prefix_tokens: int = 0 # no class embeddings

474 + self.patch_num = config.image_num_patch

475 + self.positional_embedding = nn.Parameter(

476 + torch.randn(config.image_num_pos, config.hidden_size) * scale,

477 + )

478 + image_patch_size = config.image_patch_size

479 + self.patch_embedding = nn.Linear(

480 + image_patch_size * image_patch_size * 3,

481 + config.hidden_size,

482 + bias=True,

483 + )

484 + self.transformer = Molmo2VisionBlockCollection(

485 + config,

486 + quant_config,

487 + prefix=f"{prefix}.transformer",

488 + )

489 +

490 + def add_pos_emb(self, x: torch.Tensor, patch_num: int) -> torch.Tensor:

491 + pos_emb = self.positional_embedding

492 +

493 + pos_emb = pos_emb.reshape(

494 + (

495 + int(math.sqrt(pos_emb.shape[0])),

496 + int(math.sqrt(pos_emb.shape[0])),

497 + pos_emb.shape[1],

498 + )

499 + )

500 +

501 + (patch_num_0, patch_num_1) = patch_num

502 +

503 + if pos_emb.shape[0] != patch_num_0 or pos_emb.shape[1] != patch_num_1:

504 + # from https://github.com/facebookresearch/mae/blob/main/util/pos_embed.py

505 + pos_emb = pos_emb.unsqueeze(0).permute(0, 3, 1, 2)

506 + pos_emb = F.interpolate(

507 + pos_emb,

508 + size=(patch_num_0, patch_num_1),

509 + mode="bicubic",

510 + align_corners=False,

511 + antialias=True,

512 + )

513 + pos_emb = pos_emb.permute(0, 2, 3, 1).squeeze(0)

514 +

515 + pos_emb = pos_emb.reshape(-1, pos_emb.shape[-1])

516 + x = x + pos_emb[None, :, :].to(x.dtype)

517 + return x

518 +

519 + def forward(

520 + self,

521 + x: torch.Tensor,

522 + patch_num: int | None = None,

523 + ) -> list[torch.Tensor]:

524 + """

525 + : param x: (batch_size, num_patch, n_pixels)

526 + """

527 + if patch_num is None:

528 + patch_num = self.patch_num

529 +

530 + x = self.patch_embedding(x)

531 +

532 + x = self.add_pos_emb(x, patch_num)

533 +

534 + hidden_states = self.transformer(x)

535 + return hidden_states

536 +

537 +

538 +class ImagePoolingAttention(nn.Module):

539 + """Multi-head attention used for image pooling"""

540 +

541 + def __init__(

542 + self,

543 + input_dim: int,

544 + hidden_size: int,

545 + num_heads: int,

546 + num_key_value_heads: int,

547 + head_dim: int,

548 + use_bias: bool = True,

549 + use_pytorch_sdpa: bool = False,

550 + quant_config: QuantizationConfig | None = None,

551 + ) -> None:

552 + super().__init__()

553 +

554 + self.input_dim = input_dim

555 + self.hidden_size = hidden_size

556 + self.total_num_heads = num_heads

557 + tp_size = get_tensor_model_parallel_world_size()

558 +

559 + assert self.hidden_size % self.total_num_heads == 0

560 + assert self.total_num_heads % tp_size == 0

561 +

562 + self.num_heads = self.total_num_heads // tp_size

563 + self.head_dim = head_dim

564 +

565 + assert self.head_dim == self.hidden_size // self.total_num_heads

566 +

567 + self.total_num_kv_heads = num_key_value_heads

568 + if self.total_num_kv_heads >= tp_size:

569 + assert self.total_num_kv_heads % tp_size == 0

570 + else:

571 + assert tp_size % self.total_num_kv_heads == 0

572 +

573 + self.num_kv_heads = max(1, self.total_num_kv_heads // tp_size)

574 +

575 + self.kv_size = self.num_kv_heads * self.head_dim

576 +

577 + self.q_proj = ColumnParallelLinear(

578 + self.input_dim,

579 + self.total_num_heads * self.head_dim,

580 + bias=use_bias,

581 + quant_config=quant_config,

582 + )

583 + self.merged_kv = MergedColumnParallelLinear(

584 + self.input_dim,

585 + [self.total_num_kv_heads * self.head_dim] * 2,

586 + bias=use_bias,

587 + quant_config=quant_config,

588 + )

589 + self.o_proj = RowParallelLinear(

590 + self.total_num_heads * self.head_dim,

591 + self.hidden_size,

592 + bias=use_bias,

593 + quant_config=quant_config,

594 + )

595 + self.scale = self.head_dim**-0.5

596 + self.use_pytorch_sdpa = use_pytorch_sdpa

597 + if use_pytorch_sdpa:

598 + self.attn = None

599 + else:

600 + self.attn = MMEncoderAttention(

601 + self.num_heads,

602 + self.head_dim,

603 + self.scale,

604 + num_kv_heads=self.num_kv_heads,

605 + )

606 +

607 + def forward_sdpa(

608 + self,

609 + query: torch.Tensor,

610 + key: torch.Tensor,

611 + value: torch.Tensor,

612 + attn_mask: torch.Tensor | None = None,

613 + ) -> torch.Tensor:

614 + bsz, q_len, _ = query.size()

615 + kv_len = key.size(1)

616 +

617 + query = query.view(bsz, q_len, self.num_heads, self.head_dim)

618 + key = key.view(bsz, kv_len, self.num_kv_heads, self.head_dim)

619 + value = value.view(bsz, kv_len, self.num_kv_heads, self.head_dim)

620 +

621 + if self.num_heads != self.num_kv_heads:

622 + key = torch.repeat_interleave(

623 + key,

624 + self.num_heads // self.num_kv_heads,

625 + dim=2,

626 + )

627 + value = torch.repeat_interleave(

628 + value,

629 + self.num_heads // self.num_kv_heads,

630 + dim=2,

631 + )

632 +

633 + query, key, value = (x.transpose(1, 2) for x in (query, key, value))

634 +

635 + out = F.scaled_dot_product_attention(

636 + query,

637 + key,

638 + value,

639 + attn_mask=attn_mask,

640 + is_causal=False,

641 + ).transpose(1, 2)

642 +

643 + return out.reshape(bsz, q_len, -1)

644 +

645 + def forward(

646 + self,

647 + inputs_q: torch.Tensor,

648 + inputs_kv: torch.Tensor,

649 + attn_mask: torch.Tensor | None = None,

650 + ) -> torch.Tensor:

651 + xq, _ = self.q_proj(inputs_q)

652 + kv, _ = self.merged_kv(inputs_kv)

653 + xk, xv = kv.split([self.kv_size, self.kv_size], dim=-1)

654 +

655 + if self.use_pytorch_sdpa:

656 + output = self.forward_sdpa(xq, xk, xv, attn_mask)

657 + else:

658 + output = self.attn(xq, xk, xv)

659 +

660 + output, _ = self.o_proj(output)

661 +

662 + return output

663 +

664 +

665 +class ImageProjectorMLP(nn.Module):

666 + """MLP used for the image projector"""

667 +

668 + def __init__(

669 + self,

670 + input_dim: int,

671 + hidden_dim: int,

672 + output_dim: int,

673 + hidden_act: str,

674 + quant_config: QuantizationConfig | None = None,

675 + ) -> None:

676 + super().__init__()

677 +

678 + self.merged_linear = MergedColumnParallelLinear(

679 + input_dim,

680 + [hidden_dim] * 2,

681 + bias=False,

682 + quant_config=quant_config,

683 + )

684 + # Activation function.

685 + assert hidden_act == "silu"

686 + self.act_fn = SiluAndMul()

687 +

688 + # Feed-forward output projection.

689 + self.down_proj = RowParallelLinear(

690 + hidden_dim,

691 + output_dim,

692 + bias=False,

693 + quant_config=quant_config,

694 + )

695 +

696 + def forward(self, x: torch.Tensor) -> torch.Tensor:

697 + x, _ = self.merged_linear(x)

698 + x = self.act_fn(x)

699 + x, _ = self.down_proj(x)

700 + return x

701 +

702 +

703 +class Molmo2VisionBackbone(nn.Module, SupportsQuant):

704 + packed_modules_mapping = {

705 + "merged_qkv": ["wq", "wk", "wv"], # vision backbone

706 + "merged_kv": ["k_proj", "v_proj"], # image_pooling_2d

707 + "merged_linear": ["gate_proj", "up_proj"],

708 + }

709 +

710 + def __init__(

711 + self,

712 + vit_config: VitConfig,

713 + adapter_config: AdapterConfig,

714 + quant_config: QuantizationConfig | None = None,

715 + prefix: str = "",

716 + ) -> None:

717 + super().__init__()

718 + self.vit_config = vit_config

719 + self.adapter_config = adapter_config

720 +

721 + self.vit_layers = []

722 + for layer in adapter_config.vit_layers:

723 + if layer >= 0:

724 + self.vit_layers.append(layer)

725 + else:

726 + self.vit_layers.append(layer + vit_config.num_hidden_layers)

727 +

728 + last_layer_needed = max(self.vit_layers) + 1

729 + if last_layer_needed < vit_config.num_hidden_layers:

730 + vit_config.num_hidden_layers = last_layer_needed

731 + self.image_vit = Molmo2VisionTransformer(

732 + vit_config,

733 + quant_config,

734 + prefix=f"{prefix}.image_vit",

735 + )

736 +

737 + self.num_prefix_tokens: int = self.image_vit.num_prefix_tokens

738 +

739 + pool_dim = vit_config.hidden_size * len(adapter_config.vit_layers)

740 + self.image_pooling_2d = ImagePoolingAttention(

741 + input_dim=pool_dim,

742 + hidden_size=adapter_config.hidden_size,

743 + num_heads=adapter_config.num_attention_heads,

744 + num_key_value_heads=adapter_config.num_key_value_heads,

745 + head_dim=adapter_config.head_dim,

746 + use_pytorch_sdpa=adapter_config.pooling_attention_mask,

747 + quant_config=quant_config,

748 + )

749 + self.image_projector = ImageProjectorMLP(

750 + input_dim=adapter_config.hidden_size,

751 + hidden_dim=adapter_config.intermediate_size,

752 + output_dim=adapter_config.text_hidden_size,

753 + hidden_act=adapter_config.hidden_act,

754 + quant_config=quant_config,

755 + )

756 +

757 + @property

758 + def dtype(self) -> torch.dtype:

759 + return self.image_vit.patch_embedding.weight.dtype

760 +

761 + @property

762 + def device(self) -> torch.device:

763 + return self.image_vit.patch_embedding.weight.device

764 +

765 + def encode_image(self, images: torch.Tensor) -> torch.Tensor:

766 + """

767 + : param images: (batch_size, num_crops, num_patch, n_pixels)

768 + """

769 + B, T, N, D = images.shape

770 + images = images.view(B * T, N, D)

771 + image_features = self.image_vit(images)

772 +

773 + features = []

774 + for layer in self.vit_layers:

775 + features.append(image_features[layer])

776 + image_features = torch.cat(features, dim=-1)

777 +

778 + if self.num_prefix_tokens > 0:

779 + image_features = image_features[:, 1:]

780 + image_features = image_features.view(B, T, N, -1)

781 + return image_features

782 +

783 + def forward(

784 + self,

785 + images: torch.Tensor,

786 + token_pooling: torch.Tensor,

787 + ) -> torch.Tensor:

788 + # image_features shape:

789 + # (batch_size, num_crops(=num_image), num_patch, nximage_emb_dim)

790 + batch_size, num_image = images.shape[:2]

791 + images = images.to(device=self.device, dtype=self.dtype)

792 + image_features = self.encode_image(images)

793 +

794 + dim = image_features.shape[-1]

795 + valid = token_pooling >= 0

796 + valid_token = torch.any(valid, -1)

797 +

798 + # Use `token_pooling` to arange the features for image pooling

799 + batch_idx = torch.arange(

800 + token_pooling.shape[0],

801 + dtype=torch.long,

802 + device=token_pooling.device,

803 + )

804 + batch_idx = torch.tile(

805 + batch_idx.view(batch_size, 1, 1),

806 + [1, token_pooling.shape[1], token_pooling.shape[2]],

807 + )

808 +

809 + # Now [batch, num_features, num_pooled_patches, dim]

810 + to_pool = image_features.reshape(batch_size, -1, dim)[

811 + batch_idx, torch.clip(token_pooling, 0)

812 + ]

813 + to_pool = to_pool * valid.to(self.dtype)[:, :, :, None]

814 + to_pool = to_pool.reshape([-1, token_pooling.shape[-1], dim])

815 + if self.adapter_config.pooling_attention_mask:

816 + attn_mask = valid.reshape([-1, 1, 1, valid.shape[-1]])

817 + denom = valid.view(-1, to_pool.shape[-2]).float().sum(-1)

818 + denom = torch.where(denom == 0, 1, denom)

819 + query = to_pool.sum(-2, keepdim=True) / denom[:, None, None].to(

820 + to_pool.dtype

821 + )

822 + else:

823 + attn_mask = None

824 + query = to_pool.mean(-2, keepdim=True)

825 +

826 + pooled_features = self.image_pooling_2d(query, to_pool, attn_mask=attn_mask)

827 + pooled_features = pooled_features.reshape(

828 + [batch_size, -1, pooled_features.shape[-1]]

829 + )

830 +

831 + # MLP layer to map the feature.

832 + pooled_features = self.image_projector(pooled_features)

833 + return pooled_features.view(-1, pooled_features.shape[-1])[

834 + valid_token.flatten()

835 + ]

836 +

837 + def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:

838 + stacked_params_mapping = [

839 + # (param_name, shard_name, shard_id)

840 + ("merged_qkv", "wq", "q"),

841 + ("merged_qkv", "wk", "k"),

842 + ("merged_qkv", "wv", "v"),

843 + ("merged_kv", "k_proj", 0),

844 + ("merged_kv", "v_proj", 1),

845 + ("merged_linear", "gate_proj", 0),

846 + ("merged_linear", "up_proj", 1),

847 + ]

848 + params_dict = dict(self.named_parameters())

849 + loaded_params: set[str] = set()

850 +

851 + for name, loaded_weight in weights:

852 + for param_name, weight_name, shard_id in stacked_params_mapping:

853 + if weight_name not in name:

854 + continue

855 + name = name.replace(weight_name, param_name)

856 + # Skip loading extra bias for GPTQ models.

857 + if name.endswith(".bias") and name not in params_dict:

858 + continue

859 + if is_pp_missing_parameter(name, self):

860 + continue

861 + param = params_dict[name]

862 + weight_loader = param.weight_loader

863 + weight_loader(param, loaded_weight, shard_id)

864 + break

865 + else:

866 + if name.endswith(".bias") and name not in params_dict:

867 + continue

868 + if is_pp_missing_parameter(name, self):

869 + continue

870 + param = params_dict[name]

871 + weight_loader = getattr(param, "weight_loader", default_weight_loader)

872 + weight_loader(param, loaded_weight)

873 + loaded_params.add(name)

874 + return loaded_params

875 +

876 +

877 +class Molmo2Attention(nn.Module):

878 + """Molmo2's LLM Attention."""

879 +

880 + def __init__(

881 + self,

882 + config: TextConfig,

883 + rope_parameters: dict[str, Any],

884 + cache_config: CacheConfig | None = None,

885 + quant_config: QuantizationConfig | None = None,

886 + prefix: str = "",

887 + ) -> None:

888 + super().__init__()

889 + self.hidden_size = config.hidden_size

890 + self.tp_size = get_tensor_model_parallel_world_size()

891 + self.total_num_heads = config.num_attention_heads

892 +

893 + assert self.hidden_size % self.total_num_heads == 0

894 + assert self.total_num_heads % self.tp_size == 0

895 +

896 + self.num_heads = self.total_num_heads // self.tp_size

897 + self.total_num_kv_heads = config.num_key_value_heads

898 + if self.total_num_kv_heads >= self.tp_size:

899 + assert self.total_num_kv_heads % self.tp_size == 0

900 + else:

901 + assert self.tp_size % self.total_num_kv_heads == 0

902 + self.num_kv_heads = max(1, self.total_num_kv_heads // self.tp_size)

903 + self.head_dim = config.head_dim

904 +

905 + self.q_size = self.num_heads * self.head_dim

906 + self.kv_size = self.num_kv_heads * self.head_dim

907 + self.max_position_embeddings = config.max_position_embeddings

908 + self.rope_theta = config.rope_theta

909 +

910 + # Attention input projection. Projects x -> (q, k, v)

911 + self.qkv_proj = QKVParallelLinear(

912 + self.hidden_size,

913 + self.head_dim,

914 + self.total_num_heads,

915 + self.total_num_kv_heads,

916 + bias=config.qkv_bias,

917 + quant_config=quant_config,

918 + )

919 +

920 + self.tp_rank: int | None = None

921 + self.k_norm: nn.Module | None = None

922 + self.q_norm: nn.Module | None = None

923 + self.qk_norm_type: str | None = None

924 + if config.use_qk_norm:

925 + k_norm_size = (

926 + self.head_dim

927 + if config.qk_norm_type == "qwen3"

928 + else self.total_num_kv_heads * self.head_dim

929 + )

930 + self.tp_rank = get_tensor_model_parallel_rank()

931 + self.k_norm = RMSNorm(k_norm_size, eps=config.layer_norm_eps)

932 + q_norm_size = (

933 + self.head_dim

934 + if config.qk_norm_type == "qwen3"

935 + else self.total_num_heads * self.head_dim

936 + )

937 + self.q_norm = RMSNorm(q_norm_size, eps=config.layer_norm_eps)

938 + self.qk_norm_type = config.qk_norm_type

939 + # Rotary embeddings. Rope scaling is only applied on full attention layers.

940 + layer_idx = extract_layer_index(prefix)

941 + if (

942 + config.rope_scaling_layers is not None

943 + and layer_idx not in config.rope_scaling_layers

944 + ):

945 + rope_theta = rope_parameters["rope_theta"]

946 + rope_parameters = {"rope_type": "default", "rope_theta": rope_theta}

947 + self.rotary_emb = get_rope(

948 + self.head_dim,

949 + max_position=self.max_position_embeddings,

950 + rope_parameters=rope_parameters,

951 + )

952 + self.scaling = self.head_dim**-0.5

953 + self.attn = Attention(

954 + self.num_heads,

955 + self.head_dim,

956 + self.scaling,

957 + num_kv_heads=self.num_kv_heads,

958 + cache_config=cache_config,

959 + quant_config=quant_config,

960 + prefix=f"{prefix}.attn",

961 + )

962 +

963 + # Attention output projection.

964 + self.o_proj = RowParallelLinear(

965 + self.total_num_heads * self.head_dim,

966 + self.hidden_size,

967 + bias=False,

968 + quant_config=quant_config,

969 + )

970 +

971 + def _apply_qk_norm(

972 + self,

973 + q: torch.Tensor,

974 + k: torch.Tensor,

975 + ) -> tuple[torch.Tensor, torch.Tensor]:

976 + if self.tp_size > 1:

977 + q = tensor_model_parallel_all_gather(q.contiguous())

978 + k = tensor_model_parallel_all_gather(k.contiguous())

979 + q = self.q_norm(q)

980 + k = self.k_norm(k)

981 + if self.tp_size > 1:

982 + splitter = partial(split_tensor_along_last_dim, num_partitions=self.tp_size)

983 + q = splitter(q)[self.tp_rank]

984 + k = splitter(k)[self.tp_rank]

985 + return q, k

986 +

987 + def forward(

988 + self,

989 + positions: torch.Tensor,

990 + hidden_states: torch.Tensor,

991 + **kwargs: object,

992 + ) -> torch.Tensor:

993 + qkv, _ = self.qkv_proj(hidden_states)

994 + q, k, v = qkv.split([self.q_size, self.kv_size, self.kv_size], dim=-1)

995 + if (

996 + self.q_norm is not None

997 + and self.k_norm is not None

998 + and self.qk_norm_type == "olmo"

999 + ):

1000 + q, k = self._apply_qk_norm(q, k)

1001 + elif self.q_norm is not None and self.k_norm is not None:

1002 + q_by_head = q.view(

1003 + *q.shape[:-1],

1004 + q.shape[-1] // self.head_dim,

1005 + self.head_dim,

1006 + )

1007 + q_by_head = self.q_norm(q_by_head)

1008 + q = q_by_head.view(q.shape)

1009 + k_by_head = k.view(

1010 + *k.shape[:-1],

1011 + k.shape[-1] // self.head_dim,

1012 + self.head_dim,

1013 + )

1014 + k_by_head = self.k_norm(k_by_head)

1015 + k = k_by_head.view(k.shape)

1016 + q, k = self.rotary_emb(positions, q, k)

1017 + attn_output = self.attn(q, k, v)

1018 +

1019 + output, _ = self.o_proj(attn_output)

1020 + return output

1021 +

1022 +

1023 +class LanguageModelMLP(nn.Module):

1024 + """Molmo2's LLM mlp."""

1025 +

1026 + def __init__(

1027 + self,

1028 + input_dim: int,

1029 + intermediate_size: int,

1030 + hidden_act: str,

1031 + quant_config: QuantizationConfig | None = None,

1032 + ) -> None:

1033 + super().__init__()

1034 +

1035 + self.up_gate_proj = MergedColumnParallelLinear(

1036 + input_dim,

1037 + [intermediate_size] * 2,

1038 + bias=False,

1039 + quant_config=quant_config,

1040 + )

1041 + # Activation function.

1042 + assert hidden_act == "silu"

1043 + self.act_fn = MulAndSilu()

1044 + # Feed-forward output projection.

1045 + self.down_proj = RowParallelLinear(

1046 + intermediate_size,

1047 + input_dim,

1048 + bias=False,

1049 + quant_config=quant_config,

1050 + )

1051 +

1052 + def forward(

1053 + self,

1054 + x: torch.Tensor,

1055 + ) -> torch.Tensor:

1056 + up_gate, _ = self.up_gate_proj(x)

1057 + x = self.act_fn(up_gate)

1058 + x, _ = self.down_proj(x)

1059 + return x

1060 +

1061 +

1062 +class Molmo2DecoderLayer(nn.Module):

1063 + def __init__(

1064 + self,

1065 + config: TextConfig,

1066 + rope_parameters: dict[str, Any],

1067 + cache_config: CacheConfig | None = None,

1068 + quant_config: QuantizationConfig | None = None,

1069 + prefix: str = "",

1070 + ) -> None:

1071 + super().__init__()

1072 + # Attention block.

1073 + self.self_attn = Molmo2Attention(

1074 + config,

1075 + rope_parameters,

1076 + cache_config,

1077 + quant_config,

1078 + prefix=f"{prefix}.self_attn",

1079 + )

1080 +

1081 + # MLP block.

1082 + self.mlp = LanguageModelMLP(

1083 + config.hidden_size,

1084 + config.intermediate_size,

1085 + config.hidden_act,

1086 + quant_config,

1087 + )

1088 +

1089 + # LayerNorm

1090 + self.input_layernorm = RMSNorm(config.hidden_size, eps=config.layer_norm_eps)

1091 + self.post_attention_layernorm = RMSNorm(

1092 + config.hidden_size,

1093 + eps=config.layer_norm_eps,

1094 + )

1095 +

1096 + def forward(

1097 + self,

1098 + positions: torch.Tensor,

1099 + hidden_states: torch.Tensor,

1100 + residual: torch.Tensor | None,

1101 + **kwargs: object,

1102 + ) -> tuple[torch.Tensor, tuple[torch.Tensor, torch.Tensor] | None]:

1103 + # Self Attention

1104 + if residual is None:

1105 + residual = hidden_states

1106 + hidden_states = self.input_layernorm(hidden_states)

1107 + else:

1108 + hidden_states, residual = self.input_layernorm(hidden_states, residual)

1109 + hidden_states = self.self_attn(

1110 + positions=positions,

1111 + hidden_states=hidden_states,

1112 + **kwargs,

1113 + )

1114 +

1115 + hidden_states, residual = self.post_attention_layernorm(hidden_states, residual)

1116 + hidden_states = self.mlp(hidden_states)

1117 + return hidden_states, residual

1118 +

1119 +

1120 +class Molmo2DecoderNormAfterLayer(Molmo2DecoderLayer):

1121 + def forward(

1122 + self,

1123 + positions: torch.Tensor,

1124 + hidden_states: torch.Tensor,

1125 + residual: torch.Tensor | None,

1126 + **kwargs: object,

1127 + ) -> tuple[torch.Tensor, tuple[torch.Tensor, torch.Tensor] | None]:

1128 + # Self Attention

1129 + residual = hidden_states

1130 + hidden_states = self.self_attn(

1131 + positions=positions,

1132 + hidden_states=hidden_states,

1133 + **kwargs,

1134 + )

1135 +

1136 + hidden_states = self.input_layernorm(hidden_states)

1137 + hidden_states = hidden_states + residual

1138 + residual = hidden_states

1139 +

1140 + hidden_states = self.mlp(hidden_states)

1141 + hidden_states = self.post_attention_layernorm(hidden_states)

1142 + hidden_states = hidden_states + residual

1143 + residual = None

1144 + return hidden_states, residual

1145 +

1146 +

1147 +@support_torch_compile

1148 +class Molmo2TextModel(nn.Module, SupportsQuant):

1149 + def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):

1150 + super().__init__()

1151 +

1152 + config = vllm_config.model_config.hf_config

1153 + cache_config = vllm_config.cache_config

1154 + quant_config = vllm_config.quant_config

1155 +

1156 + self.config = config

1157 +

1158 + if hasattr(config, "text_config"):

1159 + hf_text_config = config.text_config

1160 + else:

1161 + hf_text_config = config.llm_config

1162 +

1163 + kwargs = {}

1164 + for field in fields(TextConfig):

1165 + kwargs[field.name] = getattr(hf_text_config, field.name)

1166 + text_config = TextConfig(**kwargs)

1167 +

1168 + self.embedding_size = text_config.vocab_size

1169 + self.embedding_size += text_config.additional_vocab_size or 0

1170 + self.embed_tokens = VocabParallelEmbedding(

1171 + self.embedding_size,

1172 + text_config.hidden_size,

1173 + quant_config=quant_config,

1174 + )

1175 +

1176 + decoder_layer = (

1177 + Molmo2DecoderNormAfterLayer

1178 + if text_config.norm_after

1179 + else Molmo2DecoderLayer

1180 + )

1181 + self.start_layer, self.end_layer, self.layers = make_layers(

1182 + text_config.num_hidden_layers,

1183 + lambda prefix: decoder_layer(

1184 + text_config,

1185 + hf_text_config.rope_parameters,

1186 + cache_config=cache_config,

1187 + quant_config=quant_config,

1188 + prefix=prefix,

1189 + ),

1190 + prefix=f"{prefix}.layers",

1191 + )

1192 +

1193 + self.norm = RMSNorm(text_config.hidden_size, eps=text_config.layer_norm_eps)

1194 +

1195 + self.make_empty_intermediate_tensors = make_empty_intermediate_tensors_factory(

1196 + ["hidden_states", "residual"],

1197 + text_config.hidden_size,

1198 + )

1199 +

1200 + def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor:

1201 + return self.embed_tokens(input_ids)

1202 +

1203 + def forward(

1204 + self,

1205 + input_ids: torch.Tensor,

1206 + positions: torch.Tensor,

1207 + intermediate_tensors: IntermediateTensors | None = None,

1208 + inputs_embeds: torch.Tensor | None = None,

1209 + **kwargs: object,

1210 + ) -> torch.Tensor:

1211 + if get_pp_group().is_first_rank:

1212 + if inputs_embeds is not None:

1213 + hidden_states = inputs_embeds

1214 + else:

1215 + hidden_states = self.embed_tokens(input_ids)

1216 + residual = None

1217 + else:

1218 + assert intermediate_tensors is not None

1219 + hidden_states = intermediate_tensors["hidden_states"]

1220 + residual = intermediate_tensors["residual"]

1221 +

1222 + # Apply blocks one-by-one.

1223 + for layer in islice(self.layers, self.start_layer, self.end_layer):

1224 + hidden_states, residual = layer(

1225 + positions,

1226 + hidden_states,

1227 + residual,

1228 + **kwargs,

1229 + )

1230 + if not get_pp_group().is_last_rank:

1231 + return IntermediateTensors(

1232 + {"hidden_states": hidden_states, "residual": residual}

1233 + )

1234 + if residual is not None:

1235 + hidden_states, _ = self.norm(hidden_states, residual)

1236 + else:

1237 + hidden_states = self.norm(hidden_states)

1238 + return hidden_states

1239 +

1240 + def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:

1241 + params_dict = dict(self.named_parameters())

1242 + loaded_params: set[str] = set()

1243 +

1244 + for name, loaded_weight in weights:

1245 + if name.endswith(".bias") and name not in params_dict:

1246 + continue

1247 + if is_pp_missing_parameter(name, self):

1248 + continue

1249 +

1250 + param = params_dict[name]

1251 + weight_loader = getattr(param, "weight_loader", default_weight_loader)

1252 + weight_loader(param, loaded_weight)

1253 + loaded_params.add(name)

1254 + return loaded_params

1255 +

1256 +

1257 +def get_patches_grid_size(

1258 + *,

1259 + image_h: int,

1260 + image_w: int,

1261 + patch_size: int,

1262 + pool_h: int,

1263 + pool_w: int,

1264 +) -> tuple[int, int]:

1265 + patch_h = image_h // patch_size

1266 + patch_w = image_w // patch_size

1267 + h_pad = round_down(patch_h + pool_h - 1, pool_h) - patch_h

1268 + w_pad = round_down(patch_w + pool_w - 1, pool_w) - patch_w

1269 + nrows = (patch_h + h_pad) // pool_h

1270 + ncols = (patch_w + w_pad) // pool_w

1271 +

1272 + return nrows, ncols

1273 +

1274 +

1275 +def get_candidate_tilings(max_num: int) -> list[tuple[int, int]]:

1276 + tilings = [

1277 + (i, j)

1278 + for i in range(1, max_num + 1)

1279 + for j in range(1, max_num + 1)

1280 + if i * j <= max_num

1281 + ]

1282 + return sorted(tilings, key=lambda x: (x[0] * x[1], x[0]))

1283 +

1284 +

1285 +def select_tiling(

1286 + *,

1287 + height: int,

1288 + width: int,

1289 + patch_size: int,

1290 + max_num_patches: int,

1291 +):

1292 + tilings = get_candidate_tilings(max_num_patches)

1293 + candidate_tilings = np.array(tilings, dtype=np.int32)

1294 + candidate_resolutions = candidate_tilings * patch_size

1295 +

1296 + original_size = np.array([height, width], dtype=np.float32)

1297 + required_scale_d = candidate_resolutions.astype(np.float32) / original_size

1298 + required_scale = required_scale_d.min(axis=-1, keepdims=True)

1299 +

1300 + if (required_scale < 1).all():

1301 + ix = required_scale.argmax()

1302 + else:

1303 + ix = np.where(required_scale < 1.0, 10e9, required_scale).argmin()

1304 +

1305 + return candidate_tilings[ix]

1306 +

1307 +

1308 +def get_image_size(image: ImageInput) -> ImageSize:

1309 + if isinstance(image, Image):

1310 + return ImageSize(*image.size)

1311 + elif isinstance(image, (np.ndarray, torch.Tensor)):

1312 + assert image.ndim == 3

1313 + h, w, c = image.shape

1314 + assert c in [1, 3]

1315 + return ImageSize(w, h)

1316 + else:

1317 + raise ValueError(f"Unknown image type: {type(image)}")

1318 +

1319 +

1320 +def exif_tranpose(

1321 + images: ImageInput | None,

1322 +) -> ImageInput | None:

1323 + if images is None:

1324 + return None

1325 + if images is not None and isinstance(images, (list, tuple)):

1326 + images = [

1327 + exif_tranpose(img) if isinstance(img, Image) else img for img in images

1328 + ]

1329 + elif images is not None and isinstance(images, Image):

1330 + images = ImageOps.exif_transpose(images)

1331 + return images

1332 +

1333 +

1334 +def build_flat_image_bool_length(

1335 + image_grids: torch.LongTensor,

1336 + image_patch_id: int,

1337 + low_res_image_start_id: int,

1338 + image_start_id: int,

1339 + image_col_id: int,

1340 + image_end_id: int,

1341 +) -> tuple[torch.LongTensor, torch.LongTensor]:

1342 + device = image_grids.device

1343 + B = image_grids.shape[0]

1344 +

1345 + resized_h = image_grids[:, 0]

1346 + resized_w = image_grids[:, 1]

1347 + h = image_grids[:, 2]

1348 + w = image_grids[:, 3]

1349 +

1350 + lengths = resized_h * resized_w + h * (w + 1) + 4 # [B]

1351 + total_len = int(lengths.sum().item())

1352 +

1353 + flat = torch.empty(total_len, dtype=torch.long, device=device)

1354 +

1355 + offset = 0

1356 + for i in range(B):

1357 + resized_h_i, resized_w_i, h_i, w_i = image_grids[i].tolist()

1358 + L_i = int(lengths[i].item())

1359 +

1360 + num_low_res_patches = resized_h_i * resized_w_i

1361 +

1362 + idx = offset

1363 +

1364 + flat[idx] = low_res_image_start_id

1365 + idx += 1

1366 +

1367 + if num_low_res_patches > 0:

1368 + flat[idx : idx + num_low_res_patches] = image_patch_id

1369 + idx += num_low_res_patches

1370 +

1371 + flat[idx] = image_end_id

1372 + idx += 1

1373 +

1374 + flat[idx] = image_start_id

1375 + idx += 1

1376 +

1377 + block_len = w_i + 1

1378 + if block_len > 0 and h_i > 0:

1379 + line = torch.empty(block_len, dtype=torch.long, device=device)

1380 + if w_i > 0:

1381 + line[:w_i] = image_patch_id

1382 + line[w_i] = image_col_id

1383 +

1384 + block = line.repeat(h_i)

1385 + flat[idx : idx + h_i * block_len] = block

1386 + idx += h_i * block_len

1387 +

1388 + flat[idx] = image_end_id

1389 + idx += 1

1390 +

1391 + assert idx - offset == L_i

1392 +

1393 + offset += L_i

1394 +

1395 + return flat, lengths

1396 +

1397 +

1398 +def build_flat_video_bool_length(

1399 + video_grids: torch.LongTensor,

1400 + image_patch_id: int,

1401 + frame_start_id: int,

1402 + frame_end_id: int,

1403 +) -> tuple[torch.LongTensor, torch.LongTensor]:

1404 + device = video_grids.device

1405 + B = video_grids.shape[0]

1406 +

1407 + t = video_grids[:, 0]

1408 + resized_h = video_grids[:, 1]

1409 + resized_w = video_grids[:, 2]

1410 +

1411 + P = resized_h * resized_w

1412 + block_len = P + 2

1413 + lengths = t * block_len

1414 +

1415 + total_len = int(lengths.sum().item())

1416 + flat = torch.empty(total_len, dtype=torch.long, device=device)

1417 +

1418 + offset = 0

1419 + for i in range(B):

1420 + ti = int(t[i].item())

1421 + Pi = int(P[i].item())

1422 + Li = int(lengths[i].item())

1423 +

1424 + block = torch.empty(Pi + 2, dtype=torch.long, device=device)

1425 + block[0] = frame_start_id

1426 + if Pi > 0:

1427 + block[1 : 1 + Pi] = image_patch_id

1428 + block[-1] = frame_end_id

1429 +

1430 + seq = block.repeat(ti)

1431 +

1432 + flat[offset : offset + Li] = seq

1433 + offset += Li

1434 +

1435 + return flat, lengths

1436 +

1437 +

1438 +class Molmo2ProcessorWrapper:

1439 + """

1440 + Wraps :class:`Molmo2Processor` so that it can be called directly.

1441 + """

1442 +

1443 + def __init__(self, processor: ProcessorMixin, hf_config: PretrainedConfig):

1444 + super().__init__()

1445 +

1446 + self.processor = processor

1447 + self.hf_config = hf_config

1448 +

1449 + @cached_property

1450 + def vocab(self) -> dict[str, int]:

1451 + return self.processor.tokenizer.vocab # type: ignore

1452 +

1453 + @cached_property

1454 + def max_crops(self) -> int:

1455 + image_processor = self.processor.image_processor # type: ignore

1456 +

1457 + max_crops = image_processor.max_crops

1458 + assert isinstance(max_crops, int)

1459 +

1460 + return max_crops

1461 +

1462 + @cached_property

1463 + def image_pooling_h(self) -> int:

1464 + image_processor = self.processor.image_processor # type: ignore

1465 +

1466 + image_pooling_h = image_processor.pooling_size[0]

1467 + assert isinstance(image_pooling_h, int)

1468 +

1469 + return image_pooling_h

1470 +

1471 + @cached_property

1472 + def image_pooling_w(self) -> int:

1473 + image_processor = self.processor.image_processor # type: ignore

1474 +

1475 + image_pooling_w = image_processor.pooling_size[1]

1476 + assert isinstance(image_pooling_w, int)

1477 +

1478 + return image_pooling_w

1479 +

1480 + @cached_property

1481 + def video_pooling_h(self) -> int:

1482 + video_processor = self.processor.video_processor # type: ignore

1483 +

1484 + video_pooling_h = video_processor.pooling_size[0]

1485 + assert isinstance(video_pooling_h, int)

1486 +

1487 + return video_pooling_h

1488 +

1489 + @cached_property

1490 + def video_pooling_w(self) -> int:

1491 + video_processor = self.processor.video_processor # type: ignore

1492 +

1493 + video_pooling_w = video_processor.pooling_size[1]

1494 + assert isinstance(video_pooling_w, int)

1495 +

1496 + return video_pooling_w

1497 +

1498 + @cached_property

1499 + def base_image_input_size(self) -> tuple[int, int]:

1500 + if getattr(self.processor, "image_processor", None) is not None:

1501 + processor = self.processor.image_processor # type: ignore

1502 + else:

1503 + processor = self.processor.video_processor # type: ignore

1504 +

1505 + base_image_input_size = (processor.size["height"], processor.size["width"])

1506 +

1507 + return base_image_input_size

1508 +

1509 + @cached_property

1510 + def image_patch_size(self) -> int:

1511 + if getattr(self.processor, "image_processor", None) is not None:

1512 + processor = self.processor.image_processor # type: ignore

1513 + else:

1514 + processor = self.processor.video_processor # type: ignore

1515 +

1516 + image_patch_size = processor.patch_size

1517 + assert isinstance(image_patch_size, int)

1518 +

1519 + return image_patch_size

1520 +

1521 + @cached_property

1522 + def overlap_margins(self) -> tuple[int, int]:

1523 + image_processor = self.processor.image_processor # type: ignore

1524 +

1525 + left_margin, right_margin = image_processor.overlap_margins

1526 + assert isinstance(left_margin, int)

1527 + assert isinstance(right_margin, int)

1528 +

1529 + return left_margin, right_margin

1530 +

1531 + @cached_property

1532 + def bos_token(self) -> str:

1533 + return self.processor.tokenizer.bos_token or self.processor.tokenizer.eos_token

1534 +

1535 + @cached_property

1536 + def image_patch_id(self) -> int:

1537 + return self.hf_config.image_patch_id

1538 +

1539 + @cached_property

1540 + def im_col_id(self) -> int:

1541 + return self.hf_config.image_col_id

1542 +

1543 + @cached_property

1544 + def im_start_id(self) -> int:

1545 + return self.hf_config.image_start_token_id

1546 +

1547 + @cached_property

1548 + def im_end_id(self) -> int:

1549 + return self.hf_config.image_end_token_id

1550 +

1551 + @cached_property

1552 + def low_res_im_start_id(self) -> int:

1553 + return self.hf_config.low_res_image_start_token_id

1554 +

1555 + @cached_property

1556 + def frame_start_id(self) -> int:

1557 + return self.hf_config.frame_start_token_id

1558 +

1559 + @cached_property

1560 + def frame_end_id(self) -> int:

1561 + return self.hf_config.frame_end_token_id

1562 +

1563 + @cached_property

1564 + def im_low_res_id(self) -> int:

1565 + return self.hf_config.image_low_res_id

1566 +

1567 + @cached_property

1568 + def image_placeholder_id(self) -> int:

1569 + return self.vocab[IMAGE_PROMPT]

1570 +

1571 + @cached_property

1572 + def video_placeholder_id(self) -> int:

1573 + return self.vocab[VIDEO_PROMPT]

1574 +

1575 + @cached_property

1576 + def image_token_ids(self) -> list[int]:

1577 + return [

1578 + self.image_patch_id,

1579 + self.im_col_id,

1580 + self.im_start_id,

1581 + self.low_res_im_start_id,

1582 + self.frame_start_id,

1583 + self.im_end_id,

1584 + self.frame_end_id,

1585 + self.im_low_res_id,

1586 + ]

1587 +

1588 + def select_tiling(

1589 + self,

1590 + *,

1591 + image_height: int,

1592 + image_width: int,

1593 + ) -> tuple[int, int]:

1594 + max_crops = self.max_crops

1595 + left_margin, right_margin = self.overlap_margins

1596 + base_image_input_size = self.base_image_input_size

1597 + base_image_input_d = self.image_patch_size

1598 +

1599 + total_margin_pixels = base_image_input_d * (right_margin + left_margin)

1600 + crop_patches = base_image_input_size[0] // base_image_input_d

1601 + crop_window_patches = crop_patches - (right_margin + left_margin)

1602 + crop_window_size = crop_window_patches * base_image_input_d

1603 + tiling_h, tiling_w = select_tiling(

1604 + height=image_height - total_margin_pixels,

1605 + width=image_width - total_margin_pixels,

1606 + patch_size=crop_window_size,

1607 + max_num_patches=max_crops,

1608 + )

1609 +

1610 + return tiling_h, tiling_w

1611 +

1612 + def get_base_grid_size(self, is_video: bool) -> tuple[int, int]:

1613 + base_image_input_size = self.base_image_input_size

1614 +

1615 + return get_patches_grid_size(

1616 + image_h=base_image_input_size[0],

1617 + image_w=base_image_input_size[1],

1618 + patch_size=self.image_patch_size,

1619 + pool_h=self.video_pooling_h if is_video else self.image_pooling_h,

1620 + pool_w=self.video_pooling_w if is_video else self.image_pooling_w,

1621 + )

1622 +

1623 + def get_patches_grid_size(

1624 + self,

1625 + *,

1626 + image_height: int,

1627 + image_width: int,

1628 + ) -> tuple[int, int]:

1629 + left_margin, right_margin = self.overlap_margins

1630 + base_image_input_size = self.base_image_input_size

1631 + base_image_input_d = self.image_patch_size

1632 +

1633 + total_margin_pixels = base_image_input_d * (right_margin + left_margin)

1634 + crop_patches = base_image_input_size[0] // base_image_input_d

1635 + crop_window_patches = crop_patches - (right_margin + left_margin)

1636 + crop_window_size = crop_window_patches * base_image_input_d

1637 +

1638 + tiling_h, tiling_w = self.select_tiling(

1639 + image_height=image_height,

1640 + image_width=image_width,

1641 + )

1642 +

1643 + h, w = [

1644 + tiling_h * crop_window_size + total_margin_pixels,

1645 + tiling_w * crop_window_size + total_margin_pixels,

1646 + ]

1647 + nrows, ncols = get_patches_grid_size(

1648 + image_h=h,

1649 + image_w=w,

1650 + patch_size=base_image_input_d,

1651 + pool_h=self.image_pooling_h,

1652 + pool_w=self.image_pooling_w,

1653 + )

1654 +

1655 + return nrows, ncols

1656 +

1657 + def __call__(

1658 + self,

1659 + text: TextInput | list[TextInput] | None = None,

1660 + images: ImageInput | None = None,

1661 + videos: VideoInput | None = None,

1662 + return_tensors: str | TensorType = None,

1663 + **kwargs: object,

1664 + ) -> BatchFeature:

1665 + inputs = [text]

1666 + images = exif_tranpose(images)

1667 + if getattr(self.processor, "image_processor", None) is not None:

1668 + inputs.append(images)

1669 + if getattr(self.processor, "video_processor", None) is not None:

1670 + inputs.append(videos)

1671 + outputs = self.processor( # type: ignore

1672 + *inputs,

1673 + return_tensors=return_tensors,

1674 + **kwargs,

1675 + )

1676 +

1677 + # revert insert bos token

1678 + if outputs["input_ids"][0, 0] == self.vocab[self.bos_token]:

1679 + outputs["input_ids"] = outputs["input_ids"][:, 1:]

1680 +

1681 + if images is None:

1682 + images = []

1683 + if not isinstance(images, list):

1684 + images = [images]

1685 +

1686 + if videos is None:

1687 + videos = []

1688 + if not isinstance(videos, list):

1689 + videos = [videos]

1690 +

1691 + assert len(videos) in {0, 1}, "At most one video is supported for Molmo2"

1692 +

1693 + _attention_mask: torch.Tensor = outputs.pop("attention_mask")

1694 + _token_type_ids: torch.Tensor = outputs.pop("token_type_ids", None)

1695 +

1696 + if len(images) > 0:

1697 + # For each image: tiling_h * tiling_w + global view

1698 + num_crops = []

1699 + for image in images:

1700 + image_size = get_image_size(image)

1701 + tiling = self.select_tiling(

1702 + image_height=image_size.height,

1703 + image_width=image_size.width,

1704 + )

1705 + num_crops.append(np.prod(tiling) + 1)

1706 +

1707 + assert sum(num_crops) == len(outputs["pixel_values"])

1708 + assert sum(num_crops) == outputs["image_num_crops"].sum().item()

1709 + image_grids: torch.Tensor = outputs.pop("image_grids")

1710 + image_num_pooled_patches: torch.Tensor = image_grids[:, :2].prod(

1711 + dim=1

1712 + ) + image_grids[:, 2:].prod(dim=1)

1713 + outputs["image_num_pooled_patches"] = image_num_pooled_patches

1714 + n_patches = outputs["pixel_values"].shape[1]

1715 + outputs["image_num_patches"] = outputs["image_num_crops"] * n_patches

1716 + image_tokens, num_image_tokens = build_flat_image_bool_length(

1717 + image_grids,

1718 + self.image_patch_id,

1719 + self.low_res_im_start_id,

1720 + self.im_start_id,

1721 + self.im_col_id,

1722 + self.im_end_id,

1723 + )

1724 + outputs["image_tokens"] = image_tokens

1725 + outputs["num_image_tokens"] = num_image_tokens

1726 +

1727 + if len(videos) > 0:

1728 + video_grids: torch.Tensor = outputs.pop("video_grids")

1729 + assert video_grids[:, 0].sum() == len(outputs["pixel_values_videos"])

1730 + outputs["video_num_crops"] = video_grids[:, 0]

1731 + outputs["video_num_pooled_patches"] = video_grids.prod(dim=1)

1732 + n_patches = outputs["pixel_values_videos"].shape[1]

1733 + outputs["video_num_patches"] = outputs["video_num_crops"] * n_patches

1734 + video_tokens, num_video_tokens = build_flat_video_bool_length(

1735 + video_grids,

1736 + self.image_patch_id,

1737 + self.frame_start_id,

1738 + self.frame_end_id,

1739 + )

1740 + outputs["video_tokens"] = video_tokens

1741 + outputs["num_video_tokens"] = num_video_tokens

1742 +

1743 + return BatchFeature(outputs)

1744 +

1745 +

1746 +def get_candidate_target_fps(

1747 + video_fps: int | float,

1748 + sampling_fps: int | float,

1749 + max_fps: int | float = _MAX_VIDEO_FPS,

1750 +) -> list[float]:

1751 + """

1752 + Return the subset of `video_fps` factors that remain multiples

1753 + of `sampling_fps`.

1754 +

1755 + Examples:

1756 + >>> get_candidate_target_fps(video_fps=6, sampling_fps=2)

1757 + [2, 6]

1758 + >>> get_candidate_target_fps(video_fps=5, sampling_fps=1)

1759 + [1, 5]

1760 + >>> get_candidate_target_fps(video_fps=2, sampling_fps=2)

1761 + [2]

1762 + >>> get_candidate_target_fps(video_fps=5, sampling_fps=2)

1763 + Traceback (most recent call last):

1764 + ...

1765 + ValueError: sampling_fps=2 must divide video_fps=5 to produce

1766 + consistent frame steps.

1767 + """

1768 + video_fps = int(video_fps)

1769 + sampling_fps = int(sampling_fps)

1770 + max_fps = int(max_fps)

1771 +

1772 + if sampling_fps is None:

1773 + raise ValueError("sampling_fps must be provided")

1774 + if video_fps <= 0 or sampling_fps <= 0:

1775 + raise ValueError(

1776 + "video_fps and sampling_fps must be positive "

1777 + f"(got {video_fps}, {sampling_fps})"

1778 + )

1779 + if video_fps % sampling_fps != 0:

1780 + raise ValueError(

1781 + f"sampling_fps={sampling_fps} must divide video_fps={video_fps}."

1782 + )

1783 +

1784 + candidates = []

1785 + for candidate in range(sampling_fps, video_fps + 1, sampling_fps):

1786 + if candidate > max_fps:

1787 + break

1788 + if video_fps % candidate == 0:

1789 + candidates.append(float(candidate))

1790 +

1791 + return candidates

1792 +

1793 +

1794 +def get_target_fps(

1795 + video_fps: float,

1796 + max_frames: int,

1797 + total_frames: int,

1798 + frame_sample_mode: str,

1799 + candidate_target_fps: list[float],

1800 +) -> float | None:

1801 + """

1802 + Get the target fps that best spans the video and has the most frames sampled

1803 + """

1804 + num_frames_sampled = 0

1805 + selected_target_fps = None

1806 + for target_fps in candidate_target_fps:

1807 + step_size = max(int(video_fps / target_fps), 1)

1808 + num_frames_sampled_at_fps = int(total_frames / step_size)

1809 + if num_frames_sampled == 0:

1810 + if (

1811 + "uniform" in frame_sample_mode

1812 + and num_frames_sampled_at_fps > max_frames

1813 + ):

1814 + break

1815 + selected_target_fps = target_fps

1816 + num_frames_sampled = num_frames_sampled_at_fps

1817 +

1818 + else:

1819 + # the candidate sampling fps increases so frame count can't decrease

1820 + assert num_frames_sampled <= num_frames_sampled_at_fps

1821 + if num_frames_sampled_at_fps > max_frames:

1822 + # choose the sampling fps that spans the video

1823 + continue

1824 +

1825 + elif num_frames_sampled_at_fps > num_frames_sampled:

1826 + # both are less than max_frames; choose the one with higher

1827 + # density of frames sampled

1828 + selected_target_fps = target_fps

1829 + num_frames_sampled = num_frames_sampled_at_fps

1830 + return selected_target_fps

1831 +

1832 +

1833 +def get_frame_times_and_chosen_fps(

1834 + selected_target_fps, total_frames, max_frames, video_fps

1835 +):

1836 + if selected_target_fps is None:

1837 + frame_indices = np.linspace(

1838 + 0, total_frames, max_frames, endpoint=False, dtype=int

1839 + )

1840 + else:

1841 + step_size = max(int(video_fps / selected_target_fps), 1)

1842 + frame_indices = np.arange(0, total_frames, step_size)

1843 + if len(frame_indices) > max_frames:

1844 + frame_indices = frame_indices[:max_frames]

1845 + return selected_target_fps, frame_indices

1846 +

1847 +

1848 +class Molmo2ProcessingInfo(BaseProcessingInfo):

1849 + def get_hf_processor(self, **kwargs: object) -> Molmo2ProcessorWrapper:

1850 + processor = self.ctx.get_hf_processor(**kwargs)

1851 + hf_config = self.ctx.get_hf_config()

1852 + return Molmo2ProcessorWrapper(processor, hf_config)

1853 +

1854 + def get_supported_mm_limits(self) -> Mapping[str, int | None]:

1855 + return {"image": None, "video": 1}

1856 +

1857 + def get_num_image_tokens(

1858 + self,

1859 + *,

1860 + image_height: int,

1861 + image_width: int,

1862 + processor: Molmo2ProcessorWrapper | None = None,

1863 + ) -> int:

1864 + if processor is None:

1865 + processor = self.get_hf_processor()

1866 +

1867 + hf_processor = processor.processor # type: ignore

1868 +

1869 + resize_nrows, resize_cols = processor.get_base_grid_size(is_video=False)

1870 + # start/end tokens + image patch token + col tokens

1871 + if hf_processor.use_single_crop_col_tokens is not None:

1872 + use_col_tokens = hf_processor.use_single_crop_col_tokens

1873 + else:

1874 + use_col_tokens = hf_processor.image_use_col_tokens

1875 + extra = 2 + resize_nrows * (resize_cols + int(use_col_tokens))

1876 + overlap_nrows, overlap_ncols = processor.get_patches_grid_size(

1877 + image_height=image_height,

1878 + image_width=image_width,

1879 + )

1880 + joint = 2 + overlap_nrows * (

1881 + overlap_ncols + int(hf_processor.image_use_col_tokens)

1882 + )

1883 +

1884 + return extra + joint

1885 +

1886 + def get_num_video_tokens(

1887 + self,

1888 + *,

1889 + num_frames: int,

1890 + processor: Molmo2ProcessorWrapper | None = None,

1891 + ) -> int:

1892 + if processor is None:

1893 + processor = self.get_hf_processor()

1894 +

1895 + resize_nrows, resize_cols = processor.get_base_grid_size(is_video=True)

1896 + # start/end tokens

1897 + extra = 2 + resize_nrows * (

1898 + resize_cols + int(processor.processor.video_use_col_tokens)

1899 + )

1900 + return num_frames * extra

1901 +

1902 + def get_image_size_with_most_features(self) -> ImageSize:

1903 + processor = self.get_hf_processor()

1904 +

1905 + left_margin, right_margin = processor.overlap_margins

1906 + base_image_input_size = processor.base_image_input_size

1907 + base_image_input_d = processor.image_patch_size

1908 +

1909 + total_margin_pixels = base_image_input_d * (right_margin + left_margin)

1910 + crop_patches = base_image_input_size[0] // base_image_input_d

1911 + crop_window_patches = crop_patches - (right_margin + left_margin)

1912 + crop_window_size = crop_window_patches * base_image_input_d

1913 +

1914 + tilings = get_candidate_tilings(processor.max_crops)

1915 + largest_feature_size, largest_feature_pinpoint = 0, None

1916 +

1917 + for hr, wr in tilings:

1918 + height = hr * crop_window_size + total_margin_pixels

1919 + width = wr * crop_window_size + total_margin_pixels

1920 +

1921 + feat_size = self.get_num_image_tokens(

1922 + image_height=height, image_width=width, processor=processor

1923 + )

1924 + if feat_size > largest_feature_size:

1925 + largest_feature_size = feat_size

1926 + largest_feature_pinpoint = ImageSize(width=width, height=height)

1927 +

1928 + if largest_feature_size == 0 or largest_feature_pinpoint is None:

1929 + raise ValueError("Cannot have a largest feature size of 0!")

1930 +

1931 + return largest_feature_pinpoint

1932 +

1933 + def _get_max_video_frames(self, max_tokens: int) -> int:

1934 + num_tokens_per_frame = self.get_num_video_tokens(num_frames=1)

1935 + max_frames = max_tokens // num_tokens_per_frame

1936 + return max(max_frames, 1)

1937 +

1938 + def get_num_frames_with_most_features(

1939 + self,

1940 + seq_len: int,

1941 + mm_counts: Mapping[str, int],

1942 + ) -> int:

1943 + video_processor = self.get_hf_processor().processor.video_processor

1944 + num_frames = video_processor.num_frames

1945 + max_videos = mm_counts.get("video", 0)

1946 + max_total_frames = self._get_max_video_frames(seq_len)

1947 + max_frames_per_video = min(

1948 + max_total_frames // max(max_videos, 1),

1949 + num_frames,

1950 + )

1951 + return max(max_frames_per_video, 1)

1952 +

1953 + def _sample_frames(

1954 + self,

1955 + total_num_frames: int,

1956 + video_fps: float,

1957 + duration: float,

1958 + frame_sample_mode: str,

1959 + num_frames: int,

1960 + max_fps: int,

1961 + sampling_fps: int,

1962 + ) -> np.ndarray:

1963 + if frame_sample_mode == "uniform_last_frame" and max_fps is not None:

1964 + if total_num_frames <= 2:

1965 + indices = np.arange(total_num_frames).astype(int)

1966 + elif duration > (num_frames - 1) / max_fps: # -1 to include the last frame

1967 + # uniform fallback

1968 + indices = np.linspace(

1969 + 0,

1970 + total_num_frames - 1,

1971 + num=min(num_frames, total_num_frames),

1972 + endpoint=True,

1973 + ).astype(int)

1974 + else:

1975 + float_indices = np.arange(

1976 + 0.0,

1977 + stop=total_num_frames - 1,

1978 + step=float(video_fps / max_fps),

1979 + )

1980 + if np.round(float_indices[-1]) != total_num_frames - 1:

1981 + float_indices = np.concatenate(

1982 + [float_indices, [total_num_frames - 1]], axis=0

1983 + )

1984 + indices = np.round(float_indices).astype(int)

1985 + assert indices[-1] < total_num_frames

1986 + assert len(float_indices) <= num_frames

1987 + elif frame_sample_mode == "uniform_last_frame":

1988 + indices = np.linspace(

1989 + 0,

1990 + total_num_frames - 1,

1991 + num=min(num_frames, total_num_frames),

1992 + endpoint=True,

1993 + ).astype(int)

1994 + elif frame_sample_mode == "fps":

1995 + candidate_target_fps = get_candidate_target_fps(video_fps, sampling_fps)

1996 + selected_target_fps = get_target_fps(

1997 + video_fps,

1998 + num_frames,

1999 + total_num_frames,

2000 + frame_sample_mode,

2001 + candidate_target_fps,

2002 + )

2003 + _, indices = get_frame_times_and_chosen_fps(

2004 + selected_target_fps,

2005 + total_num_frames,

2006 + num_frames,

2007 + video_fps,

2008 + )

2009 + else:

2010 + raise NotImplementedError(frame_sample_mode)

2011 +

2012 + return indices

2013 +

2014 + def _get_video_second_idx(

2015 + self,

2016 + metadata: dict[str, Any],

2017 + do_sample_frames: bool | None = None,

2018 + ) -> list[float]:

2019 + video_processor = self.get_hf_processor().processor.video_processor

2020 + # metadata["fps"] refers to the true fps of the input video.

2021 + video_fps = metadata["fps"]

2022 + frames_indices = metadata.get("frames_indices")

2023 + if do_sample_frames is None:

2024 + do_sample_frames = metadata.get("do_sample_frames", False)

2025 +

2026 + if do_sample_frames:

2027 + # Frame-based sampling is applied in HF video processor

2028 + total_num_frames = metadata["total_num_frames"]

2029 + duration = total_num_frames / video_fps

2030 + frame_sample_mode = video_processor.frame_sample_mode

2031 + num_frames = video_processor.num_frames

2032 + max_fps = video_processor.max_fps

2033 + sampling_fps = video_processor.sampling_fps

2034 + frames_indices = self._sample_frames(

2035 + total_num_frames,

2036 + video_fps,

2037 + duration,

2038 + frame_sample_mode,

2039 + num_frames,

2040 + max_fps,

2041 + sampling_fps,

2042 + )

2043 + else:

2044 + # Time-based sampling is done in vllm molmo2 video loader or molmo_utils

2045 + assert frames_indices is not None

2046 + timestamps = [frame_idx / video_fps for frame_idx in frames_indices]

2047 + return timestamps

2048 +

2049 +

2050 +class Molmo2DummyInputsBuilder(BaseDummyInputsBuilder[Molmo2ProcessingInfo]):

2051 + def get_dummy_text(self, mm_counts: Mapping[str, int]) -> str:

2052 + num_images = mm_counts.get("image", 0)

2053 + num_videos = mm_counts.get("video", 0)

2054 +

2055 + image_placeholder_token = IMAGE_PROMPT

2056 + video_placeholder_token = VIDEO_PROMPT

2057 +

2058 + if num_images == 1:

2059 + image_string = image_placeholder_token

2060 + else:

2061 + image_string = "".join(

2062 + [f"Image {i + 1}" + image_placeholder_token for i in range(num_images)]

2063 + )

2064 +

2065 + return image_string + video_placeholder_token * num_videos

2066 +

2067 + def get_dummy_mm_data(

2068 + self,

2069 + seq_len: int,

2070 + mm_counts: Mapping[str, int],

2071 + mm_options: Mapping[str, BaseDummyOptions] | None = None,

2072 + ) -> MultiModalDataDict:

2073 + num_images = mm_counts.get("image", 0)

2074 + num_videos = mm_counts.get("video", 0)

2075 +

2076 + dummy_images = []

2077 + dummy_videos = []

2078 +

2079 + if num_images > 0:

2080 + target_width, target_height = self.info.get_image_size_with_most_features()

2081 +

2082 + image_overrides = mm_options.get("image") if mm_options else None

2083 +

2084 + dummy_images = self._get_dummy_images(

2085 + width=target_width,

2086 + height=target_height,

2087 + num_images=num_images,

2088 + overrides=image_overrides,

2089 + )

2090 +

2091 + if num_videos > 0:

2092 + processor = self.info.get_hf_processor()

2093 + base_image_input_size = processor.base_image_input_size

2094 + target_num_frames = self.info.get_num_frames_with_most_features(

2095 + seq_len, mm_counts

2096 + )

2097 +

2098 + video_overrides = mm_options.get("video") if mm_options else None

2099 +

2100 + if video_overrides:

2101 + assert isinstance(video_overrides, VideoDummyOptions)

2102 + num_frames_override = video_overrides.num_frames

2103 + if num_frames_override:

2104 + if num_frames_override > target_num_frames:

2105 + logger.warning(

2106 + "video.num_frames override (%d) exceeds model's "

2107 + "maximum number of frames (%d), will be ignored",

2108 + num_frames_override,

2109 + target_num_frames,

2110 + )

2111 + if num_frames_override < 2:

2112 + logger.warning(

2113 + "video.num_frames override (%d) cannot be less "

2114 + "than 2, will be ignored",

2115 + num_frames_override,

2116 + )

2117 + target_num_frames = min(target_num_frames, num_frames_override)

2118 +

2119 + dummy_videos = self._get_dummy_videos(

2120 + width=base_image_input_size[1],

2121 + height=base_image_input_size[0],

2122 + num_frames=target_num_frames,

2123 + num_videos=num_videos,

2124 + )

2125 +

2126 + return {

2127 + "image": dummy_images,

2128 + "video": dummy_videos,

2129 + }

2130 +

2131 + def _get_dummy_videos(

2132 + self,

2133 + *,

2134 + width: int,

2135 + height: int,

2136 + num_frames: int,

2137 + num_videos: int,

2138 + ) -> list[VideoItem]:

2139 + video = np.full((num_frames, height, width, 3), 255, dtype=np.uint8)

2140 + video_items = []

2141 + for i in range(num_videos):

2142 + video_metadata = {

2143 + "fps": 2.0,

2144 + "duration": num_frames / 2.0,

2145 + "total_num_frames": num_frames,

2146 + "frames_indices": list(range(num_frames)),

2147 + "video_backend": "decord",

2148 + "do_sample_frames": False,

2149 + "height": height,

2150 + "width": width,

2151 + }

2152 + video_item = (video.copy(), video_metadata)

2153 + video_items.append(video_item)

2154 + return video_items

2155 +

2156 +

2157 +class Molmo2MultiModalProcessor(BaseMultiModalProcessor[Molmo2ProcessingInfo]):

2158 + def _apply_hf_processor_tokens_only(

2159 + self,

2160 + prompt_tokens: list[int],

2161 + ) -> list[int]:

2162 + processor = self.info.get_hf_processor()

2163 + tokenizer = processor.processor.tokenizer

2164 + bos_token_id = tokenizer.bos_token_id or tokenizer.eos_token_id

2165 +

2166 + if len(prompt_tokens) > 0 and prompt_tokens[0] != bos_token_id:

2167 + # Prepend the bos token to the prompt tokens

2168 + prompt_tokens = [bos_token_id] + prompt_tokens

2169 +

2170 + return prompt_tokens

2171 +

2172 + def _get_data_parser(self) -> MultiModalDataParser:

2173 + return MultiModalDataParser(video_needs_metadata=True)

2174 +

2175 + def _call_hf_processor(

2176 + self,

2177 + prompt: str,

2178 + mm_data: Mapping[str, object],

2179 + mm_kwargs: Mapping[str, object],

2180 + tok_kwargs: Mapping[str, object],

2181 + ) -> BatchFeature:

2182 + mm_data = dict(mm_data)

2183 + processor = self.info.get_hf_processor(**mm_kwargs)

2184 +

2185 + if videos := mm_data.pop("videos", []):

2186 + pixel_values_videos_lst = []

2187 + video_token_pooling_lst = []

2188 + video_num_crops_lst = []

2189 + video_num_pooled_patches_lst = []

2190 + video_num_patches_lst = []

2191 + video_tokens_lst = []

2192 + num_video_tokens_lst = []

2193 +

2194 + for item in videos:

2195 + video_array, metadata = item

2196 +

2197 + # NOTE: metadata.frames_indices indicates

2198 + # the sampled frames indices of pre-sampled videos, which is

2199 + # used to calculate the timestamps. Make sure that

2200 + # do_sample_frames in mm_kwargs is false for presampled videos.

2201 +

2202 + # NOTE: a copy of mm_kwargs is created to update do_sample_frames,

2203 + # otherwise mm_hash for the object will be incorrect.

2204 + video_mm_kwargs = dict(**mm_kwargs)

2205 + if "do_sample_frames" not in video_mm_kwargs:

2206 + # molmo_utils already has "do_sample_frames" in

2207 + # mm_kwargs, don't overwrite it.

2208 + video_mm_kwargs["do_sample_frames"] = metadata.get(

2209 + "do_sample_frames", False

2210 + )

2211 +

2212 + metadata = VideoMetadata(

2213 + **{k: metadata[k] for k in metadata if k != "do_sample_frames"}

2214 + )

2215 +

2216 + video_mm_data = dict()

2217 + video_mm_data["videos"] = [[video_array]]

2218 + video_mm_data["video_metadata"] = [[metadata]]

2219 +

2220 + video_outputs = super()._call_hf_processor(

2221 + prompt=VIDEO_PROMPT,

2222 + mm_data=video_mm_data,

2223 + mm_kwargs=video_mm_kwargs,

2224 + tok_kwargs=tok_kwargs,

2225 + )

2226 + input_ids = video_outputs.pop("input_ids")

2227 + video_string = processor.processor.tokenizer.batch_decode(input_ids)[0]

2228 + prompt = prompt.replace(

2229 + VIDEO_PROMPT,

2230 + video_string,

2231 + 1,

2232 + )

2233 +

2234 + pixel_values_videos_lst.append(video_outputs["pixel_values_videos"])

2235 + video_token_pooling_lst.append(video_outputs["video_token_pooling"])

2236 + video_num_crops_lst.append(video_outputs["video_num_crops"])

2237 + video_num_pooled_patches_lst.append(

2238 + video_outputs["video_num_pooled_patches"]

2239 + )

2240 + video_num_patches_lst.append(video_outputs["video_num_patches"])

2241 + video_tokens_lst.append(video_outputs["video_tokens"])

2242 + num_video_tokens_lst.append(video_outputs["num_video_tokens"])

2243 +

2244 + video_outputs = dict(

2245 + pixel_values_videos=torch.cat(pixel_values_videos_lst),

2246 + video_token_pooling=torch.cat(video_token_pooling_lst),

2247 + video_num_crops=torch.cat(video_num_crops_lst),

2248 + video_num_pooled_patches=torch.cat(video_num_pooled_patches_lst),

2249 + video_num_patches=torch.cat(video_num_patches_lst),

2250 + video_tokens=torch.cat(video_tokens_lst),

2251 + num_video_tokens=torch.cat(num_video_tokens_lst),

2252 + )

2253 + else:

2254 + video_outputs = dict()

2255 +

2256 + processed_outputs = super()._call_hf_processor(

2257 + prompt=prompt,

2258 + mm_data=mm_data,

2259 + mm_kwargs=mm_kwargs,

2260 + tok_kwargs=tok_kwargs,

2261 + )

2262 +

2263 + bos_token_id = processor.vocab[processor.bos_token]

2264 + input_ids = processed_outputs["input_ids"]

2265 + # add bos token back to prompt start

2266 + if input_ids.numel() > 0 and input_ids[0, 0] != bos_token_id:

2267 + bos_token_id_tensor = torch.tensor(

2268 + [[bos_token_id]], device=input_ids.device, dtype=input_ids.dtype

2269 + )

2270 + processed_outputs["input_ids"] = torch.concat(

2271 + [bos_token_id_tensor, input_ids], dim=1

2272 + )

2273 + combined_outputs = dict(

2274 + processed_outputs,

2275 + **video_outputs,

2276 + )

2277 + return BatchFeature(combined_outputs)

2278 +

2279 + def _get_mm_fields_config(

2280 + self,

2281 + hf_inputs: BatchFeature,

2282 + hf_processor_mm_kwargs: Mapping[str, object],

2283 + ) -> Mapping[str, MultiModalFieldConfig]:

2284 + image_num_crops = hf_inputs.get("image_num_crops", torch.empty(0))

2285 + image_num_pooled_patches = hf_inputs.get(

2286 + "image_num_pooled_patches", torch.empty(0)

2287 + )

2288 + video_num_crops = hf_inputs.get("video_num_crops", torch.empty(0))

2289 + video_num_pooled_patches = hf_inputs.get(

2290 + "video_num_pooled_patches", torch.empty(0)

2291 + )

2292 + num_image_tokens = hf_inputs.get("num_image_tokens", torch.empty(0))

2293 + num_video_tokens = hf_inputs.get("num_video_tokens", torch.empty(0))

2294 +

2295 + return dict(

2296 + pixel_values=MultiModalFieldConfig.flat_from_sizes(

2297 + "image", image_num_crops

2298 + ),

2299 + image_token_pooling=MultiModalFieldConfig.flat_from_sizes(

2300 + "image", image_num_pooled_patches

2301 + ),

2302 + image_num_crops=MultiModalFieldConfig.batched("image"),

2303 + image_num_pooled_patches=MultiModalFieldConfig.batched("image"),

2304 + image_num_patches=MultiModalFieldConfig.batched("image"),

2305 + image_tokens=MultiModalFieldConfig.flat_from_sizes(

2306 + "image", num_image_tokens

2307 + ),

2308 + num_image_tokens=MultiModalFieldConfig.batched("image"),

2309 + pixel_values_videos=MultiModalFieldConfig.flat_from_sizes(

2310 + "video", video_num_crops

2311 + ),

2312 + video_token_pooling=MultiModalFieldConfig.flat_from_sizes(

2313 + "video", video_num_pooled_patches

2314 + ),

2315 + video_num_crops=MultiModalFieldConfig.batched("video"),

2316 + video_num_pooled_patches=MultiModalFieldConfig.batched("video"),

2317 + video_num_patches=MultiModalFieldConfig.batched("video"),

2318 + video_tokens=MultiModalFieldConfig.flat_from_sizes(

2319 + "video", num_video_tokens

2320 + ),

2321 + num_video_tokens=MultiModalFieldConfig.batched("video"),

2322 + )

2323 +

2324 + def _get_prompt_updates(

2325 + self,

2326 + mm_items: MultiModalDataItems,

2327 + hf_processor_mm_kwargs: Mapping[str, object],

2328 + out_mm_kwargs: MultiModalKwargsItems,

2329 + ) -> Sequence[PromptUpdate]:

2330 + processor = self.info.get_hf_processor(**hf_processor_mm_kwargs)

2331 + img_patch_id = processor.image_patch_id

2332 + img_col_id = processor.im_col_id

2333 + img_start_id = processor.im_start_id

2334 + img_end_id = processor.im_end_id

2335 + image_use_col_tokens = processor.processor.image_use_col_tokens

2336 + use_single_crop_col_tokens = processor.processor.use_single_crop_col_tokens

2337 + use_single_crop_start_token = processor.processor.use_single_crop_start_token

2338 + video_use_col_tokens = processor.processor.video_use_col_tokens

2339 + use_frame_special_tokens = processor.processor.use_frame_special_tokens

2340 +

2341 + def get_image_replacement_molmo2(item_idx: int) -> list[int]:

2342 + images = mm_items.get_items("image", ImageProcessorItems)

2343 + image = images.get(item_idx)

2344 + image = exif_tranpose(image)

2345 +

2346 + resize_nrows, resize_cols = processor.get_base_grid_size(is_video=False)

2347 + if use_single_crop_col_tokens is not None:

2348 + use_col_tokens = use_single_crop_col_tokens

2349 + else:

2350 + use_col_tokens = image_use_col_tokens

2351 + if use_single_crop_start_token:

2352 + start_id = processor.low_res_im_start_id

2353 + else:

2354 + start_id = img_start_id

2355 + extra_row = [img_patch_id] * resize_cols + [img_col_id] * int(

2356 + use_col_tokens

2357 + )

2358 + extra_joint = [start_id] + extra_row * resize_nrows + [img_end_id]

2359 +

2360 + image_size = get_image_size(image)

2361 +

2362 + nrows, ncols = processor.get_patches_grid_size(

2363 + image_height=image_size.height,

2364 + image_width=image_size.width,

2365 + )

2366 +

2367 + joint_row = [img_patch_id] * ncols + [img_col_id] * int(

2368 + image_use_col_tokens

2369 + )

2370 + joint = [img_start_id] + joint_row * nrows + [img_end_id]

2371 + img_token_ids = extra_joint + joint

2372 +

2373 + return PromptUpdateDetails.select_token_ids(

2374 + img_token_ids,

2375 + processor.image_token_ids,

2376 + )

2377 +

2378 + def get_video_replacement_molmo2(item_idx: int) -> list[int]:

2379 + video, metadata = mm_items["video"][item_idx]

2380 + do_sample_frames = hf_processor_mm_kwargs.get("do_sample_frames")

2381 +

2382 + timestamps = self.info._get_video_second_idx(metadata, do_sample_frames)

2383 + nrows, ncols = processor.get_base_grid_size(is_video=True)

2384 +

2385 + if use_frame_special_tokens:

2386 + start_id = processor.frame_start_id

2387 + end_id = processor.frame_end_id

2388 + else:

2389 + start_id = img_start_id

2390 + end_id = img_end_id

2391 +

2392 + img_token_ids = []

2393 +

2394 + for frame_idx, frame_time in enumerate(timestamps):

2395 + prev_space = " " if frame_idx > 0 else ""

2396 + frame_prefix = (

2397 + prev_space + f"{frame_time:.1f} "

2398 + ) # explicit whitespace before/after image tokens

2399 +

2400 + img_token_ids += processor.processor.tokenizer.encode(

2401 + frame_prefix,

2402 + add_special_tokens=False,

2403 + )

2404 +

2405 + joint_row = [img_patch_id] * ncols + [img_col_id] * int(

2406 + video_use_col_tokens

2407 + )

2408 + joint = [start_id] + nrows * joint_row + [end_id]

2409 + img_token_ids += joint

2410 +

2411 + return PromptUpdateDetails.select_token_ids(

2412 + img_token_ids,

2413 + processor.image_token_ids,

2414 + )

2415 +

2416 + return [

2417 + PromptReplacement(

2418 + modality=modality,

2419 + target=[target],

2420 + replacement=replacement_fn,

2421 + )

2422 + for modality, target, replacement_fn in zip(

2423 + ["image", "video"],

2424 + [processor.image_placeholder_id, processor.video_placeholder_id],

2425 + [get_image_replacement_molmo2, get_video_replacement_molmo2],

2426 + )

2427 + ]

2428 +

2429 +

2430 +@MULTIMODAL_REGISTRY.register_processor(

2431 + Molmo2MultiModalProcessor,

2432 + info=Molmo2ProcessingInfo,

2433 + dummy_inputs=Molmo2DummyInputsBuilder,

2434 +)

2435 +class Molmo2ForConditionalGeneration(

2436 + nn.Module, SupportsMultiModal, SupportsPP, SupportsLoRA, SupportsQuant

2437 +):

2438 + hf_to_vllm_mapper = WeightsMapper(

2439 + orig_to_new_substr={

2440 + # vision backbone mapping

2441 + "image_pooling_2d.wq.": "image_pooling_2d.q_proj.",

2442 + "image_pooling_2d.wk.": "image_pooling_2d.k_proj.",

2443 + "image_pooling_2d.wv.": "image_pooling_2d.v_proj.",

2444 + "image_pooling_2d.wo.": "image_pooling_2d.o_proj.",

2445 + "image_projector.w1.": "image_projector.gate_proj.",

2446 + "image_projector.w3.": "image_projector.up_proj.",

2447 + "image_projector.w2.": "image_projector.down_proj.",

2448 + # language backbone mapping

2449 + "att_proj": "qkv_proj",

2450 + "attn_out": "o_proj",

2451 + "q_norm": "q_norm",

2452 + "k_norm": "k_norm",

2453 + "ff_proj": "up_gate_proj",

2454 + "ff_out": "down_proj",

2455 + "attn_norm": "input_layernorm",

2456 + "ff_norm": "post_attention_layernorm",

2457 + },

2458 + orig_to_new_prefix={

2459 + # vision backbone mapping

2460 + "model.vision_backbone.": "vision_backbone.",

2461 + # language backbone mapping

2462 + "model.transformer.blocks.": "model.layers.",

2463 + "model.transformer.ln_f.": "model.norm.",

2464 + },

2465 + )

2466 +

2467 + packed_modules_mapping = {

2468 + "qkv_proj": ["qkv_proj"],

2469 + "up_gate_proj": ["up_gate_proj"], # language model

2470 + "merged_qkv": ["wq", "wk", "wv"], # vision backbone

2471 + "merged_kv": ["k_proj", "v_proj"], # image_pooling_2d

2472 + "merged_linear": ["gate_proj", "up_proj"], # image_projector

2473 + }

2474 +

2475 + @classmethod

2476 + def get_placeholder_str(cls, modality: str, i: int) -> str | None:

2477 + if modality.startswith("image"):

2478 + return IMAGE_PROMPT

2479 + if modality.startswith("video"):

2480 + return VIDEO_PROMPT

2481 +

2482 + raise ValueError("Only image or video modality is supported")

2483 +

2484 + def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):

2485 + super().__init__()

2486 + config = vllm_config.model_config.hf_config

2487 + quant_config = vllm_config.quant_config

2488 + multimodal_config = vllm_config.model_config.multimodal_config

2489 + self.config = config

2490 + self.multimodal_config = multimodal_config

2491 +

2492 + kwargs = {}

2493 + for field in fields(VitConfig):

2494 + kwargs[field.name] = getattr(config.vit_config, field.name)

2495 + vit_config = VitConfig(**kwargs)

2496 +

2497 + kwargs = {}

2498 + for field in fields(AdapterConfig):

2499 + kwargs[field.name] = getattr(config.adapter_config, field.name)

2500 + adapter_config = AdapterConfig(**kwargs)

2501 +

2502 + self.vision_backbone = Molmo2VisionBackbone(

2503 + vit_config,

2504 + adapter_config,

2505 + quant_config,

2506 + prefix=maybe_prefix(prefix, "vision_backbone"),

2507 + )

2508 + self.model = Molmo2TextModel(

2509 + vllm_config=vllm_config,

2510 + prefix=maybe_prefix(prefix, "model"),

2511 + )

2512 +

2513 + self.img_patch_id = config.image_patch_id

2514 +

2515 + if hasattr(config, "text_config"):

2516 + hf_text_config = config.text_config

2517 + else:

2518 + hf_text_config = config.llm_config

2519 +

2520 + self.lm_head = ParallelLMHead(

2521 + hf_text_config.vocab_size,

2522 + hf_text_config.hidden_size,

2523 + quant_config=quant_config,

2524 + )

2525 + self.logits_processor = LogitsProcessor(hf_text_config.vocab_size)

2526 +

2527 + self.make_empty_intermediate_tensors = (

2528 + self.model.make_empty_intermediate_tensors

2529 + )

2530 +

2531 + @property

2532 + def dtype(self):

2533 + return next(self.parameters()).dtype

2534 +

2535 + def _parse_and_validate_image_input(

2536 + self,

2537 + **kwargs: object,

2538 + ) -> Molmo2ImageInputs | None:

2539 + pixel_values = kwargs.pop("pixel_values", None)

2540 + if pixel_values is None:

2541 + return None

2542 +

2543 + token_pooling = kwargs.pop("image_token_pooling", None)

2544 + num_pooled_patches = kwargs.pop("image_num_pooled_patches", None)

2545 + num_patches = kwargs.pop("image_num_patches", None)

2546 + image_tokens = kwargs.pop("image_tokens", None)

2547 + num_image_tokens = kwargs.pop("num_image_tokens", None)

2548 +

2549 + accum_patches = [0] + num_patches.cumsum(dim=0)[:-1].tolist()

2550 + patch_offset = 0

2551 + new_token_pooling = token_pooling.clone()

2552 + for i, n in enumerate(num_pooled_patches):

2553 + cur_slice = token_pooling[patch_offset : patch_offset + n]

2554 + index_offset = int(accum_patches[i])

2555 + new_token_pooling[patch_offset : patch_offset + n] = torch.where(

2556 + cur_slice >= 0,

2557 + cur_slice + index_offset,

2558 + cur_slice,

2559 + )

2560 + patch_offset += n

2561 +

2562 + return Molmo2ImageInputs(

2563 + pixel_values=pixel_values,

2564 + token_pooling=new_token_pooling,

2565 + num_pooled_patches=num_pooled_patches,

2566 + image_tokens=image_tokens,

2567 + num_image_tokens=num_image_tokens,

2568 + )

2569 +

2570 + def _parse_and_validate_video_input(

2571 + self,

2572 + **kwargs: object,

2573 + ) -> Molmo2VideoInputs | None:

2574 + pixel_values_videos = kwargs.pop("pixel_values_videos", None)

2575 + if pixel_values_videos is None:

2576 + return None

2577 +

2578 + token_pooling = kwargs.pop("video_token_pooling", None)

2579 + num_pooled_patches = kwargs.pop("video_num_pooled_patches", None)

2580 + num_patches = kwargs.pop("video_num_patches", None)

2581 + video_tokens = kwargs.pop("video_tokens", None)

2582 + num_video_tokens = kwargs.pop("num_video_tokens", None)

2583 +

2584 + accum_patches = [0] + num_patches.cumsum(dim=0)[:-1].tolist()

2585 + patch_offset = 0

2586 + new_token_pooling = token_pooling.clone()

2587 + for i, n in enumerate(num_pooled_patches):

2588 + cur_slice = token_pooling[patch_offset : patch_offset + n]

2589 + index_offset = int(accum_patches[i])

2590 + new_token_pooling[patch_offset : patch_offset + n] = torch.where(

2591 + cur_slice >= 0,

2592 + cur_slice + index_offset,

2593 + cur_slice,

2594 + )

2595 + patch_offset += n

2596 +

2597 + return Molmo2VideoInputs(

2598 + pixel_values_videos=pixel_values_videos,

2599 + token_pooling=new_token_pooling,

2600 + num_pooled_patches=num_pooled_patches,

2601 + video_tokens=video_tokens,

2602 + num_video_tokens=num_video_tokens,

2603 + )

2604 +

2605 + def _parse_and_validate_multimodal_inputs(self, **kwargs: object) -> dict:

2606 + modalities = {}

2607 +

2608 + for input_key in kwargs:

2609 + if input_key in ("pixel_values",) and "images" not in modalities:

2610 + modalities["images"] = self._parse_and_validate_image_input(**kwargs)

2611 + if input_key in ("pixel_values_videos",) and "videos" not in modalities:

2612 + modalities["videos"] = self._parse_and_validate_video_input(**kwargs)

2613 + return modalities

2614 +

2615 + def _process_image_input(

2616 + self,

2617 + image_input: Molmo2ImageInputs,

2618 + ) -> tuple[torch.Tensor, ...]:

2619 + pixel_values = image_input["pixel_values"]

2620 + token_pooling = image_input["token_pooling"]

2621 + num_pooled_patches = image_input["num_pooled_patches"]

2622 + image_tokens = image_input["image_tokens"]

2623 + num_image_tokens = image_input["num_image_tokens"]

2624 +

2625 + image_features_flat = self.vision_backbone(

2626 + images=pixel_values.unsqueeze(0),

2627 + token_pooling=token_pooling.unsqueeze(0),

2628 + )

2629 +

2630 + assert len(image_features_flat) == num_pooled_patches.sum()

2631 + image_features_list = image_features_flat.split(

2632 + num_pooled_patches.tolist(), dim=0

2633 + )

2634 + image_tokens_list = image_tokens.split(num_image_tokens.tolist(), dim=0)

2635 + out = []

2636 + for image_features_i, image_tokens_i in zip(

2637 + image_features_list, image_tokens_list

2638 + ):

2639 + out_features = self.get_language_model().embed_input_ids(image_tokens_i)

2640 + is_image_patch = image_tokens_i == self.img_patch_id

2641 + out_features[is_image_patch] = image_features_i

2642 + out.append(out_features)

2643 + return tuple(out)

2644 +

2645 + def _process_video_input(

2646 + self,

2647 + video_input: Molmo2VideoInputs,

2648 + ) -> tuple[torch.Tensor, ...]:

2649 + pixel_values_videos = video_input["pixel_values_videos"]

2650 + token_pooling = video_input["token_pooling"]

2651 + num_pooled_patches = video_input["num_pooled_patches"]

2652 + video_tokens = video_input["video_tokens"]

2653 + num_video_tokens = video_input["num_video_tokens"]

2654 +

2655 + image_features_flat = self.vision_backbone(

2656 + images=pixel_values_videos.unsqueeze(0),

2657 + token_pooling=token_pooling.unsqueeze(0),

2658 + )

2659 +

2660 + assert len(image_features_flat) == num_pooled_patches.sum()

2661 + image_features_list = image_features_flat.split(

2662 + num_pooled_patches.tolist(), dim=0

2663 + )

2664 + video_tokens_list = video_tokens.split(num_video_tokens.tolist(), dim=0)

2665 + out = []

2666 + for image_features_i, video_tokens_i in zip(

2667 + image_features_list, video_tokens_list

2668 + ):

2669 + out_features = self.get_language_model().embed_input_ids(video_tokens_i)

2670 + is_image_patch = video_tokens_i == self.img_patch_id

2671 + out_features[is_image_patch] = image_features_i

2672 + out.append(out_features)

2673 + return tuple(out)

2674 +

2675 + def get_language_model(self) -> torch.nn.Module:

2676 + return self.model

2677 +

2678 + def embed_multimodal(self, **kwargs: object) -> MultiModalEmbeddings | None:

2679 + modalities = self._parse_and_validate_multimodal_inputs(**kwargs)

2680 + if not modalities:

2681 + return []

2682 +

2683 + multimodal_embeddings: tuple[torch.Tensor, ...] = ()

2684 +

2685 + for modality in modalities:

2686 + if modality == "images":

2687 + image_input = modalities["images"]

2688 + image_embeddings = self._process_image_input(image_input)

2689 + multimodal_embeddings += image_embeddings

2690 + if modality == "videos":

2691 + video_input = modalities["videos"]

2692 + video_embeddings = self._process_video_input(video_input)

2693 + multimodal_embeddings += video_embeddings

2694 +

2695 + return multimodal_embeddings

2696 +

2697 + def embed_input_ids(

2698 + self,

2699 + input_ids: torch.Tensor,

2700 + multimodal_embeddings: MultiModalEmbeddings | None = None,

2701 + *,

2702 + is_multimodal: torch.Tensor | None = None,

2703 + handle_oov_mm_token: bool = False,

2704 + ) -> torch.Tensor:

2705 + inputs_embeds = self._embed_text_input_ids(

2706 + input_ids,

2707 + self.get_language_model().embed_input_ids,

2708 + is_multimodal=is_multimodal,

2709 + handle_oov_mm_token=handle_oov_mm_token,

2710 + )

2711 +

2712 + if multimodal_embeddings is None or len(multimodal_embeddings) == 0:

2713 + return inputs_embeds

2714 +

2715 + if is_multimodal is None:

2716 + raise ValueError(

2717 + "`embed_input_ids` now requires `is_multimodal` arg, "

2718 + "please update your model runner according to "

2719 + "https://github.com/vllm-project/vllm/pull/16229."

2720 + )

2721 +

2722 + inputs_embeds = _merge_multimodal_embeddings(

2723 + inputs_embeds=inputs_embeds,

2724 + multimodal_embeddings=multimodal_embeddings,

2725 + is_multimodal=is_multimodal,

2726 + )

2727 + return inputs_embeds

2728 +

2729 + def forward(

2730 + self,

2731 + input_ids: torch.LongTensor,

2732 + positions: torch.LongTensor,

2733 + intermediate_tensors: IntermediateTensors | None = None,

2734 + inputs_embeds: torch.Tensor | None = None,

2735 + **kwargs: object,

2736 + ) -> torch.Tensor:

2737 + if intermediate_tensors is not None:

2738 + inputs_embeds = None

2739 +

2740 + hidden_states = self.model(

2741 + input_ids,

2742 + positions,

2743 + intermediate_tensors,

2744 + inputs_embeds=inputs_embeds,

2745 + **kwargs,

2746 + )

2747 +

2748 + return hidden_states

2749 +

2750 + def compute_logits(self, hidden_states: torch.Tensor) -> torch.Tensor:

2751 + logits = self.logits_processor(self.lm_head, hidden_states)

2752 + return logits

2753 +

2754 + def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]):

2755 + loader = AutoWeightsLoader(self)

2756 + weights = _get_weights_with_merged_embedding(weights)

2757 + return loader.load_weights(weights, mapper=self.hf_to_vllm_mapper)

2758 +

2759 + def get_mm_mapping(self) -> MultiModelKeys:

2760 + """

2761 + Get the module prefix in multimodal models

2762 + """

2763 + return MultiModelKeys.from_string_field(

2764 + language_model="model",

2765 + connector="vision_backbone.image_projector",

2766 + tower_model="vision_backbone",

2767 + )

2768 +

2769 +

2770 +def _get_weights_with_merged_embedding(

2771 + weights: Iterable[tuple[str, torch.Tensor]],

2772 +) -> Iterable[tuple[str, torch.Tensor]]:

2773 + embedding_weights = {}

2774 + for name, weight in weights:

2775 + if "wte.embedding" in name:

2776 + embedding_weights["embedding"] = weight

2777 + elif "wte.new_embedding" in name:

2778 + embedding_weights["new_embedding"] = weight

2779 + else:

2780 + yield (name, weight)

2781 + # this is compatible with most of quantization,

2782 + # because they won't quantize embed_tokens

2783 + if "embedding" not in embedding_weights or "new_embedding" not in embedding_weights:

2784 + raise ValueError(

2785 + "Checkpoint is missing 'wte.embedding' or "

2786 + "'wte.new_embedding' weights required for Molmo2."

2787 + )

2788 +

2789 + embedding_weights = torch.cat(

2790 + [embedding_weights["embedding"], embedding_weights["new_embedding"]],

2791 + dim=0,

2792 + )

2793 + yield ("model.embed_tokens.weight", embedding_weights)

vllm/model_executor/models/registry.py

@@ -384,6 +384,7 @@ _MULTIMODAL_MODELS = {

384 "Mistral3ForConditionalGeneration",

385 ),

386 "MolmoForCausalLM": ("molmo", "MolmoForCausalLM"),

387 + "Molmo2ForConditionalGeneration": ("molmo2", "Molmo2ForConditionalGeneration"),

388 "NVLM_D": ("nvlm_d", "NVLM_D_Model"),

389 "Ovis": ("ovis", "Ovis"),

390 "Ovis2_5": ("ovis2_5", "Ovis2_5"),

vllm/multimodal/processing.py

+15

@@ -386,6 +386,21 @@ class PromptUpdateDetails(Generic[_S]):

386

387 return PromptUpdateDetails(full=seq, is_embed=is_embed)

388

389 + @staticmethod

390 + def select_token_ids(

391 + seq: _S,

392 + embed_token_ids: list[int],

393 + ) -> "PromptUpdateDetails[_S]":

394 + def is_embed(tokenizer: TokenizerLike | None, full: PromptSeq) -> torch.Tensor:

395 + token_ids = _seq2tokens(tokenizer, full)

396 +

397 + return torch.isin(

398 + torch.tensor(token_ids),

399 + torch.tensor(embed_token_ids),

400 + )

401 +

402 + return PromptUpdateDetails(full=seq, is_embed=is_embed)

403 +

404

405 PromptUpdateInfo: TypeAlias = PromptSeq | PromptUpdateDetails

406 """

vllm/multimodal/video.py

+319 -1

@@ -6,7 +6,7 @@ from abc import abstractmethod

6 from functools import partial

7 from io import BytesIO

8 from pathlib import Path

9 -from typing import TYPE_CHECKING, Any

9 +from typing import TYPE_CHECKING, Any, cast

11 import numpy as np

12 import numpy.typing as npt

@@ -439,6 +439,324 @@ class OpenCVDynamicVideoBackend(OpenCVVideoBackend):

439 return frames, metadata

440

441

442 +@VIDEO_LOADER_REGISTRY.register("molmo2")

443 +class Molmo2VideoBackend(VideoLoader):

444 + def get_cv2_video_api(self):

445 + import cv2.videoio_registry as vr

446 +

447 + api_pref = None

448 + for backend in vr.getStreamBufferedBackends():

449 + if not vr.hasBackend(backend):

450 + continue

451 + if not vr.isBackendBuiltIn(backend):

452 + _, abi, api = vr.getStreamBufferedBackendPluginVersion(backend)

453 + if abi < 1 or (abi == 1 and api < 2):

454 + continue

455 + api_pref = backend

456 + break

457 + return api_pref

458 +

459 + @classmethod

460 + def get_candidate_target_fps(

461 + cls,

462 + video_fps: float,

463 + sampling_fps: float,

464 + max_fps: float = 8.0,

465 + ) -> list[float]:

466 + """

467 + Return the subset of `video_fps` factors that remain multiples

468 + of `sampling_fps`.

469 +

470 + Examples:

471 + >>> get_candidate_target_fps(video_fps=6, sampling_fps=2)

472 + [2, 6]

473 + >>> get_candidate_target_fps(video_fps=5, sampling_fps=1)

474 + [1, 5]

475 + >>> get_candidate_target_fps(video_fps=2, sampling_fps=2)

476 + [2]

477 + >>> get_candidate_target_fps(video_fps=5, sampling_fps=2)

478 + Traceback (most recent call last):

479 + ...

480 + ValueError: sampling_fps=2 must divide video_fps=5 to produce

481 + consistent frame steps.

482 + """

483 + video_fps = int(video_fps)

484 + sampling_fps = int(sampling_fps)

485 + max_fps = int(max_fps)

486 +

487 + if sampling_fps is None:

488 + raise ValueError("sampling_fps must be provided")

489 + if video_fps <= 0 or sampling_fps <= 0:

490 + raise ValueError(

491 + "video_fps and sampling_fps must be positive "

492 + f"(got {video_fps}, {sampling_fps})"

493 + )

494 + if video_fps % sampling_fps != 0:

495 + raise ValueError(

496 + f"sampling_fps={sampling_fps} must divide video_fps={video_fps}."

497 + )

498 +

499 + candidates = []

500 + for candidate in range(sampling_fps, video_fps + 1, sampling_fps):

501 + if candidate > max_fps:

502 + break

503 + if video_fps % candidate == 0:

504 + candidates.append(float(candidate))

505 +

506 + return candidates

507 +

508 + @classmethod

509 + def get_target_fps(

510 + cls,

511 + video_fps: float,

512 + max_frames: int,

513 + total_frames: int,

514 + frame_sample_mode: str,

515 + candidate_target_fps: list[float],

516 + ) -> float | None:

517 + """

518 + Get the target fps that best spans the videoand has the most frames sampled

519 + """

520 + num_frames_sampled = 0

521 + selected_target_fps = None

522 + for target_fps in candidate_target_fps:

523 + step_size = max(int(video_fps / target_fps), 1)

524 + num_frames_sampled_at_fps = int(total_frames / step_size)

525 + if num_frames_sampled == 0:

526 + if (

527 + "uniform" in frame_sample_mode

528 + and num_frames_sampled_at_fps > max_frames

529 + ):

530 + break

531 + selected_target_fps = target_fps

532 + num_frames_sampled = num_frames_sampled_at_fps

533 +

534 + else:

535 + # the candidate sampling fps increases so frame count can't decrease

536 + assert num_frames_sampled <= num_frames_sampled_at_fps

537 + if num_frames_sampled_at_fps > max_frames:

538 + # choose the sampling fps that spans the video

539 + continue

540 +

541 + elif num_frames_sampled_at_fps > num_frames_sampled:

542 + # both are less than max_frames; choose the one with higher

543 + # density of frames sampled

544 + selected_target_fps = target_fps

545 + num_frames_sampled = num_frames_sampled_at_fps

546 + return selected_target_fps

547 +

548 + @classmethod

549 + def get_frame_times_and_chosen_fps(

550 + cls,

551 + selected_target_fps: float | None,

552 + total_frames: int,

553 + max_frames: int,

554 + video_fps: float,

555 + ) -> tuple[float | None, npt.NDArray]:

556 + if selected_target_fps is None:

557 + frame_indices = np.linspace(

558 + 0, total_frames, max_frames, endpoint=False, dtype=int

559 + )

560 + else:

561 + step_size = max(int(video_fps / selected_target_fps), 1)

562 + frame_indices = np.arange(0, total_frames, step_size)

563 + if len(frame_indices) > max_frames:

564 + frame_indices = frame_indices[:max_frames]

565 + return selected_target_fps, frame_indices

566 +

567 + @classmethod

568 + def sample_times(

569 + cls,

570 + duration: float,

571 + max_frames: int,

572 + frame_sample_mode: str,

573 + max_fps: int | None,

574 + candidate_target_fps: list[float] | None = None,

575 + **kwargs,

576 + ) -> npt.NDArray:

577 + if frame_sample_mode == "fps":

578 + assert candidate_target_fps is not None

579 + # Try larger and larger FPSs until we hit one that can't span the video

580 + sampling_fps = candidate_target_fps[0]

581 + for candidate_fps in candidate_target_fps[1:]:

582 + if max_frames / candidate_fps < duration:

583 + break

584 + sampling_fps = candidate_fps

585 + times = np.arange(0, max_frames) / sampling_fps

586 + times = times[times < duration]

587 + return times

588 + elif frame_sample_mode == "uniform_last_frame":

589 + if max_fps is not None:

590 + max_duration = (

591 + max_frames - 1

592 + ) / max_fps # -1 to include the last frame

593 + if max_duration < duration:

594 + times = np.linspace(

595 + 0, duration, num=max_frames, endpoint=True, dtype=np.float64

596 + )

597 + else:

598 + times = np.arange(0.0, stop=duration, step=1 / max_fps)

599 + times = np.concatenate([times, [duration]], axis=0)

600 + assert len(times) <= max_frames

601 + else:

602 + times = np.linspace(

603 + 0, duration, num=max_frames, endpoint=True, dtype=np.float64

604 + )

605 + return times

606 + else:

607 + raise NotImplementedError(frame_sample_mode)

608 +

609 + @classmethod

610 + def _sample_frames(

611 + cls,

612 + total_num_frames: int,

613 + video_fps: float,

614 + duration: float,

615 + frame_sample_mode: str,

616 + num_frames: int,

617 + max_fps: int,

618 + sampling_fps: int,

619 + ) -> npt.NDArray:

620 + if frame_sample_mode == "uniform_last_frame" and max_fps is not None:

621 + if total_num_frames <= 2:

622 + indices = np.arange(total_num_frames).astype(int)

623 + elif duration > (num_frames - 1) / max_fps: # -1 to include the last frame

624 + # uniform fallback

625 + indices = np.linspace(

626 + 0,

627 + total_num_frames - 1,

628 + num=min(num_frames, total_num_frames),

629 + endpoint=True,

630 + ).astype(int)

631 + else:

632 + float_indices = np.arange(

633 + 0.0,

634 + stop=total_num_frames - 1,

635 + step=float(video_fps / max_fps),

636 + )

637 + if np.round(float_indices[-1]) != total_num_frames - 1:

638 + float_indices = np.concatenate(

639 + [float_indices, [total_num_frames - 1]], axis=0

640 + )

641 + indices = np.round(float_indices).astype(int)

642 + assert indices[-1] < total_num_frames

643 + assert len(float_indices) <= num_frames

644 + elif frame_sample_mode == "uniform_last_frame":

645 + indices = np.linspace(

646 + 0,

647 + total_num_frames - 1,

648 + num=min(num_frames, total_num_frames),

649 + endpoint=True,

650 + ).astype(int)

651 + elif frame_sample_mode == "fps":

652 + candidate_target_fps = cls.get_candidate_target_fps(video_fps, sampling_fps)

653 + selected_target_fps = cls.get_target_fps(

654 + video_fps,

655 + num_frames,

656 + total_num_frames,

657 + frame_sample_mode,

658 + candidate_target_fps,

659 + )

660 + _, indices = cls.get_frame_times_and_chosen_fps(

661 + selected_target_fps,

662 + total_num_frames,

663 + num_frames,

664 + video_fps,

665 + )

666 + else:

667 + raise NotImplementedError(frame_sample_mode)

668 +

669 + return indices

670 +

671 + @classmethod

672 + def load_bytes_opencv(

673 + cls,

674 + data: bytes,

675 + frame_sample_mode: str | None = None,

676 + num_frames: int = -1,

677 + max_fps: int = 2,

678 + sampling_fps: int = 2,

679 + **kwargs,

680 + ) -> tuple[npt.NDArray, dict[str, Any]]:

681 + import cv2

682 +

683 + backend = cls().get_cv2_video_api()

684 + cap = cv2.VideoCapture(BytesIO(data), backend, [])

685 + if not cap.isOpened():

686 + raise ValueError("Could not open video stream")

687 +

688 + total_frames_num = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))

689 + original_fps = cap.get(cv2.CAP_PROP_FPS)

690 + duration = total_frames_num / original_fps if original_fps > 0 else 0

691 +

692 + if frame_sample_mode is None:

693 + # Use transformers transformers.video_utils.VideoMetadata format

694 + frame_idx = list(range(0, total_frames_num))

695 + frame_idx_set = set(frame_idx)

696 + frames, valid_num_frames, valid_frame_indices = cls._read_frames(

697 + cap, frame_idx_set, total_frames_num, max(frame_idx)

698 + )

699 + do_sample_frames = valid_num_frames == total_frames_num

700 + metadata = {

701 + "total_num_frames": total_frames_num,

702 + "fps": original_fps,

703 + "duration": duration,

704 + "video_backend": "opencv",

705 + "do_sample_frames": do_sample_frames,

706 + }

707 + if not do_sample_frames:

708 + metadata["frames_indices"] = valid_frame_indices

709 + return frames, metadata

710 +

711 + frame_idx = cls._sample_frames(

712 + total_frames_num,

713 + original_fps,

714 + duration,

715 + frame_sample_mode,

716 + num_frames,

717 + max_fps,

718 + sampling_fps,

719 + ).tolist()

720 +

721 + frames, valid_num_frames, valid_frame_indices = cls._read_frames(

722 + cap,

723 + set(frame_idx),

724 + len(frame_idx),

725 + total_frames_num - 1,

726 + )

727 +

728 + metadata = {

729 + "total_num_frames": total_frames_num,

730 + "fps": original_fps,

731 + "duration": duration,

732 + "video_backend": "opencv",

733 + "frames_indices": valid_frame_indices,

734 + "do_sample_frames": False,

735 + }

736 +

737 + return frames, metadata

738 +

739 + @classmethod

740 + def load_bytes(

741 + cls,

742 + data: bytes,

743 + num_frames: int = -1,

744 + **kwargs,

745 + ) -> tuple[npt.NDArray, dict[str, Any]]:

746 + frame_sample_mode = cast(str | None, kwargs.pop("frame_sample_mode", None))

747 + max_fps = cast(int, kwargs.pop("max_fps", 2))

748 + sampling_fps = cast(int, kwargs.pop("sampling_fps", 2))

749 + out = cls.load_bytes_opencv(

750 + data,

751 + frame_sample_mode,

752 + num_frames,

753 + max_fps,

754 + sampling_fps,

755 + **kwargs,

756 + )

757 + return out

758 +

759 +

760 class VideoMediaIO(MediaIO[tuple[npt.NDArray, dict[str, Any]]]):

761 def __init__(

762 self,