| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | """vLLM composition layer for Dots3Note image and audio encoders.""" |
| 4 | |
| 5 | from collections.abc import Iterable |
| 6 | from contextlib import contextmanager |
| 7 | |
| 8 | import torch |
| 9 | from torch import nn |
| 10 | |
| 11 | from vllm.config import VllmConfig |
| 12 | from vllm.model_executor.models.interfaces import ( |
| 13 | MultiModalEmbeddings, |
| 14 | SupportsMultiModal, |
| 15 | SupportsPP, |
| 16 | ) |
| 17 | from vllm.model_executor.models.module_mapping import MultiModelKeys |
| 18 | from vllm.model_executor.models.utils import ( |
| 19 | AutoWeightsLoader, |
| 20 | IntermediateTensors, |
| 21 | WeightsMapper, |
| 22 | maybe_prefix, |
| 23 | ) |
| 24 | from vllm.multimodal import MULTIMODAL_REGISTRY |
| 25 | from vllm.transformers_utils.repo_utils import get_hf_file_to_dict |
| 26 | |
| 27 | from ..common.processor import ( |
| 28 | AUDIO_END, |
| 29 | AUDIO_PAD, |
| 30 | AUDIO_START, |
| 31 | IMAGE_END, |
| 32 | IMAGE_PAD, |
| 33 | IMAGE_START, |
| 34 | VIDEO_PLACEHOLDER, |
| 35 | Dots3NoteDummyInputsBuilder, |
| 36 | Dots3NoteMultiModalProcessor, |
| 37 | Dots3NoteProcessingInfo, |
| 38 | load_note_config_section, |
| 39 | ) |
| 40 | from .audio import Dots3NoteAudioConfig, Dots3NoteAudioModel |
| 41 | from .model import Dots3NoteLanguageModelForCausalLM |
| 42 | from .vision import DotsMoEVitConfig, DotsMoEVitModel |
| 43 | |
| 44 | |
| 45 | @contextmanager |
| 46 | def _skip_linear_init(): |
| 47 | """Skip initialization for VE Linear weights loaded immediately afterward.""" |
| 48 | original_reset_parameters = nn.Linear.reset_parameters |
| 49 | |
| 50 | def _noop_reset_parameters(_linear: nn.Linear) -> None: |
| 51 | return None |
| 52 | |
| 53 | nn.Linear.reset_parameters = _noop_reset_parameters |
| 54 | try: |
| 55 | yield |
| 56 | finally: |
| 57 | nn.Linear.reset_parameters = original_reset_parameters |
| 58 | |
| 59 | |
| 60 | @MULTIMODAL_REGISTRY.register_processor( |
| 61 | Dots3NoteMultiModalProcessor, |
| 62 | info=Dots3NoteProcessingInfo, |
| 63 | dummy_inputs=Dots3NoteDummyInputsBuilder, |
| 64 | ) |
| 65 | class Dots3NoteForCausalLM(nn.Module, SupportsMultiModal, SupportsPP): |
| 66 | """Dots3Note model with optional image and audio towers.""" |
| 67 | |
| 68 | supports_encoder_tp_data = True |
| 69 | |
| 70 | hf_to_vllm_mapper = WeightsMapper( |
| 71 | orig_to_new_prefix={ |
| 72 | "model.": "language_model.model.", |
| 73 | "lm_head.": "language_model.lm_head.", |
| 74 | "mtp.": "language_model.mtp.", |
| 75 | "vision_encoder.": "visual.", |
| 76 | "audio_encoder.": "audio_tower.", |
| 77 | } |
| 78 | ) |
| 79 | |
| 80 | @classmethod |
| 81 | def get_placeholder_str(cls, modality: str, i: int) -> str | None: |
| 82 | if modality.startswith("image"): |
| 83 | return f"{IMAGE_START}{IMAGE_PAD}{IMAGE_END}" |
| 84 | if modality.startswith("audio"): |
| 85 | return f"{AUDIO_START}{AUDIO_PAD}{AUDIO_END}" |
| 86 | if modality.startswith("video"): |
| 87 | return VIDEO_PLACEHOLDER |
| 88 | raise ValueError(f"Unsupported modality: {modality}") |
| 89 | |
| 90 | def __init__(self, *, vllm_config: VllmConfig, prefix: str = "") -> None: |
| 91 | super().__init__() |
| 92 | model_config = vllm_config.model_config |
| 93 | self.config = model_config.hf_config |
| 94 | self.quant_config = vllm_config.quant_config |
| 95 | self.multimodal_config = model_config.multimodal_config |
| 96 | assert self.multimodal_config is not None |
| 97 | |
| 98 | added_tokens = get_hf_file_to_dict( |
| 99 | "added_tokens.json", |
| 100 | model_config.model, |
| 101 | model_config.revision, |
| 102 | ) |
| 103 | if added_tokens is None or IMAGE_PAD not in added_tokens: |
| 104 | raise ValueError("NOTE tokenizer is missing the image padding token") |
| 105 | self.config.image_token_index = int(added_tokens[IMAGE_PAD]) |
| 106 | |
| 107 | vision_config_dict = load_note_config_section( |
| 108 | model_config.model, |
| 109 | model_config.revision, |
| 110 | "vision_config", |
| 111 | ) |
| 112 | audio_config_dict = load_note_config_section( |
| 113 | model_config.model, |
| 114 | model_config.revision, |
| 115 | "audio_config", |
| 116 | ) |
| 117 | video_enabled = ( |
| 118 | vision_config_dict is not None |
| 119 | and self.multimodal_config.get_limit_per_prompt("video") > 0 |
| 120 | ) |
| 121 | image_enabled = vision_config_dict is not None and ( |
| 122 | self.multimodal_config.get_limit_per_prompt("image") > 0 or video_enabled |
| 123 | ) |
| 124 | audio_enabled = audio_config_dict is not None and ( |
| 125 | self.multimodal_config.get_limit_per_prompt("audio") > 0 or video_enabled |
| 126 | ) |
| 127 | |
| 128 | with self._mark_tower_model(vllm_config, {"image", "audio", "video"}): |
| 129 | self.visual: DotsMoEVitModel | None |
| 130 | if image_enabled: |
| 131 | assert vision_config_dict is not None |
| 132 | with _skip_linear_init(): |
| 133 | self.visual = DotsMoEVitModel( |
| 134 | DotsMoEVitConfig(**vision_config_dict) |
| 135 | ) |
| 136 | else: |
| 137 | self.visual = None |
| 138 | self.audio_tower: Dots3NoteAudioModel | None |
| 139 | if audio_enabled: |
| 140 | assert audio_config_dict is not None |
| 141 | self.audio_tower = Dots3NoteAudioModel( |
| 142 | Dots3NoteAudioConfig(**audio_config_dict) |
| 143 | ) |
| 144 | else: |
| 145 | self.audio_tower = None |
| 146 | # The native encoder service casts each complete tower before |
| 147 | # loading its checkpoint. This also converts explicitly-created |
| 148 | # floating buffers (for example RoPE tables and router state), |
| 149 | # which merely constructing under vLLM's default dtype does not. |
| 150 | if self.visual is not None: |
| 151 | self.visual.to(dtype=model_config.dtype) |
| 152 | if self.audio_tower is not None: |
| 153 | self.audio_tower.to(dtype=model_config.dtype) |
| 154 | with self._mark_language_model(vllm_config): |
| 155 | self.language_model = Dots3NoteLanguageModelForCausalLM( |
| 156 | vllm_config=vllm_config, |
| 157 | prefix=maybe_prefix(prefix, "language_model"), |
| 158 | ) |
| 159 | self.make_empty_intermediate_tensors = ( |
| 160 | self.language_model.make_empty_intermediate_tensors |
| 161 | ) |
| 162 | |
| 163 | orig_to_new_prefix = dict[str, None]() |
| 164 | if self.visual is None: |
| 165 | orig_to_new_prefix["visual."] = None |
| 166 | if self.audio_tower is None: |
| 167 | orig_to_new_prefix["audio_tower."] = None |
| 168 | if orig_to_new_prefix: |
| 169 | self.hf_to_vllm_mapper = self.hf_to_vllm_mapper | WeightsMapper( |
| 170 | orig_to_new_prefix=orig_to_new_prefix |
| 171 | ) |
| 172 | |
| 173 | def _process_image_input( |
| 174 | self, |
| 175 | pixel_values: torch.Tensor, |
| 176 | image_grid_thw: torch.Tensor, |
| 177 | ) -> tuple[torch.Tensor, ...]: |
| 178 | if self.visual is None: |
| 179 | return () |
| 180 | merge_size = self.visual.spatial_merge_size |
| 181 | sizes = (image_grid_thw.prod(-1) // merge_size**2).tolist() |
| 182 | meta = self.visual.prepare_meta(image_grid_thw) |
| 183 | image_embeds = self.visual(pixel_values, meta) |
| 184 | return image_embeds.split(sizes) |
| 185 | |
| 186 | def _process_audio_input( |
| 187 | self, |
| 188 | audio_values: torch.Tensor, |
| 189 | audio_lengths: torch.Tensor, |
| 190 | ) -> tuple[torch.Tensor, ...]: |
| 191 | if self.audio_tower is None: |
| 192 | return () |
| 193 | if audio_values.dtype != torch.int32: |
| 194 | raise TypeError( |
| 195 | "NOTE audio values must carry float32 waveform bits as int32, " |
| 196 | f"got {audio_values.dtype}" |
| 197 | ) |
| 198 | waveforms = audio_values.contiguous().view(torch.float32) |
| 199 | audio_embeds, item_lengths = self.audio_tower(waveforms, audio_lengths) |
| 200 | return audio_embeds.split(item_lengths) |
| 201 | |
| 202 | def _process_video_input( |
| 203 | self, |
| 204 | pixel_values: torch.Tensor, |
| 205 | image_grid_thw: torch.Tensor, |
| 206 | audio_values: torch.Tensor, |
| 207 | audio_lengths: torch.Tensor, |
| 208 | modalities: torch.Tensor, |
| 209 | frame_counts: torch.Tensor, |
| 210 | audio_counts: torch.Tensor, |
| 211 | emission_counts: torch.Tensor, |
| 212 | ) -> tuple[torch.Tensor, ...]: |
| 213 | image_embeds = self._process_image_input(pixel_values, image_grid_thw) |
| 214 | audio_embeds = ( |
| 215 | self._process_audio_input(audio_values, audio_lengths) |
| 216 | if audio_lengths.numel() |
| 217 | else () |
| 218 | ) |
| 219 | orders = modalities.tolist() |
| 220 | frame_counts_list = frame_counts.tolist() |
| 221 | audio_counts_list = audio_counts.tolist() |
| 222 | emission_counts_list = emission_counts.tolist() |
| 223 | outputs: list[torch.Tensor] = [] |
| 224 | image_idx = audio_idx = order_idx = 0 |
| 225 | for num_frames, num_audios, num_emissions in zip( |
| 226 | frame_counts_list, |
| 227 | audio_counts_list, |
| 228 | emission_counts_list, |
| 229 | ): |
| 230 | video_order = orders[order_idx : order_idx + num_emissions] |
| 231 | video_parts: list[torch.Tensor] = [] |
| 232 | video_image_start = image_idx |
| 233 | video_audio_start = audio_idx |
| 234 | for modality in video_order: |
| 235 | if modality == 0: |
| 236 | video_parts.append(image_embeds[image_idx]) |
| 237 | image_idx += 1 |
| 238 | elif modality == 1: |
| 239 | if audio_idx >= len(audio_embeds): |
| 240 | raise ValueError("NOTE video audio tower output is missing") |
| 241 | video_parts.append(audio_embeds[audio_idx]) |
| 242 | audio_idx += 1 |
| 243 | else: |
| 244 | raise ValueError(f"Unknown NOTE video modality id: {modality}") |
| 245 | if image_idx - video_image_start != num_frames: |
| 246 | raise ValueError("NOTE video frame order/count mismatch") |
| 247 | if audio_idx - video_audio_start != num_audios: |
| 248 | raise ValueError("NOTE video audio order/count mismatch") |
| 249 | outputs.append(torch.cat(video_parts)) |
| 250 | order_idx += num_emissions |
| 251 | if image_idx != len(image_embeds) or audio_idx != len(audio_embeds): |
| 252 | raise ValueError("NOTE video encoder outputs were not fully consumed") |
| 253 | return tuple(outputs) |
| 254 | |
| 255 | def embed_multimodal(self, **kwargs: object) -> MultiModalEmbeddings: |
| 256 | multimodal_embeddings: list[torch.Tensor] = [] |
| 257 | handled: set[str] = set() |
| 258 | for input_key in kwargs: |
| 259 | if input_key == "pixel_values" and "image" not in handled: |
| 260 | pixel_values = kwargs.get("pixel_values") |
| 261 | image_grid_thw = kwargs.get("image_grid_thw") |
| 262 | if isinstance(pixel_values, torch.Tensor) and isinstance( |
| 263 | image_grid_thw, torch.Tensor |
| 264 | ): |
| 265 | multimodal_embeddings.extend( |
| 266 | self._process_image_input(pixel_values, image_grid_thw) |
| 267 | ) |
| 268 | handled.add("image") |
| 269 | elif input_key == "audio_values" and "audio" not in handled: |
| 270 | audio_values = kwargs.get("audio_values") |
| 271 | audio_lengths = kwargs.get("audio_lengths") |
| 272 | if isinstance(audio_values, torch.Tensor) and isinstance( |
| 273 | audio_lengths, torch.Tensor |
| 274 | ): |
| 275 | multimodal_embeddings.extend( |
| 276 | self._process_audio_input(audio_values, audio_lengths) |
| 277 | ) |
| 278 | handled.add("audio") |
| 279 | elif input_key == "video_pixel_values" and "video" not in handled: |
| 280 | video_inputs = ( |
| 281 | kwargs.get("video_pixel_values"), |
| 282 | kwargs.get("video_image_grid_thw"), |
| 283 | kwargs.get("video_audio_values"), |
| 284 | kwargs.get("video_audio_lengths"), |
| 285 | kwargs.get("video_modalities"), |
| 286 | kwargs.get("video_frame_counts"), |
| 287 | kwargs.get("video_audio_counts"), |
| 288 | kwargs.get("video_emission_counts"), |
| 289 | ) |
| 290 | if all(isinstance(value, torch.Tensor) for value in video_inputs): |
| 291 | multimodal_embeddings.extend( |
| 292 | self._process_video_input(*video_inputs) # type: ignore[arg-type] |
| 293 | ) |
| 294 | handled.add("video") |
| 295 | return tuple(multimodal_embeddings) |
| 296 | |
| 297 | def forward( |
| 298 | self, |
| 299 | input_ids: torch.Tensor | None, |
| 300 | positions: torch.Tensor, |
| 301 | intermediate_tensors: IntermediateTensors | None = None, |
| 302 | inputs_embeds: torch.Tensor | None = None, |
| 303 | **kwargs: object, |
| 304 | ) -> torch.Tensor | IntermediateTensors: |
| 305 | return self.language_model( |
| 306 | input_ids, |
| 307 | positions, |
| 308 | intermediate_tensors, |
| 309 | inputs_embeds, |
| 310 | ) |
| 311 | |
| 312 | def compute_logits(self, hidden_states: torch.Tensor) -> torch.Tensor | None: |
| 313 | return self.language_model.compute_logits(hidden_states) |
| 314 | |
| 315 | def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: |
| 316 | return AutoWeightsLoader(self).load_weights( |
| 317 | weights, |
| 318 | mapper=self.hf_to_vllm_mapper, |
| 319 | ) |
| 320 | |
| 321 | def process_weights_after_loading(self) -> None: |
| 322 | if self.visual is not None: |
| 323 | self.visual.process_weights_after_loading() |
| 324 | |
| 325 | def get_mm_mapping(self) -> MultiModelKeys: |
| 326 | return MultiModelKeys.from_string_field( |
| 327 | language_model="language_model", |
| 328 | tower_model=["visual", "audio_tower"], |
| 329 | ) |