| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | |
| 4 | # Copyright 2025 The vLLM team. |
| 5 | # Copyright 2025 IBM. |
| 6 | # |
| 7 | # Licensed under the Apache License, Version 2.0 (the "License"); |
| 8 | # you may not use this file except in compliance with the License. |
| 9 | # You may obtain a copy of the License at |
| 10 | # |
| 11 | # http://www.apache.org/licenses/LICENSE-2.0 |
| 12 | # |
| 13 | # Unless required by applicable law or agreed to in writing, software |
| 14 | # distributed under the License is distributed on an "AS IS" BASIS, |
| 15 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. |
| 16 | # See the License for the specific language governing permissions and |
| 17 | # limitations under the License. |
| 18 | """Wrapper around `Terratorch` models.""" |
| 19 | |
| 20 | from collections import OrderedDict |
| 21 | from collections.abc import Iterable, Mapping, Sequence |
| 22 | from functools import cached_property |
| 23 | from typing import Any |
| 24 | |
| 25 | import torch |
| 26 | import torch.nn as nn |
| 27 | from terratorch.vllm import ( |
| 28 | DummyDataGenerator, |
| 29 | InferenceRunner, |
| 30 | InputDefinition, |
| 31 | InputTypeEnum, |
| 32 | ) |
| 33 | from transformers import BatchFeature |
| 34 | |
| 35 | from vllm.config import VllmConfig |
| 36 | from vllm.config.multimodal import MultiModalDummyOptions |
| 37 | from vllm.inputs import ModalityData, MultiModalDataDict, MultiModalInput, mm_input |
| 38 | from vllm.logger import init_logger |
| 39 | from vllm.model_executor.layers.pooler import IdentityPooler |
| 40 | from vllm.model_executor.model_loader.weight_utils import default_weight_loader |
| 41 | from vllm.model_executor.models.utils import AutoWeightsLoader |
| 42 | from vllm.multimodal import MULTIMODAL_REGISTRY |
| 43 | from vllm.multimodal.inputs import ( |
| 44 | ImageItem, |
| 45 | MultiModalFieldConfig, |
| 46 | MultiModalKwargsItems, |
| 47 | PlaceholderRange, |
| 48 | ) |
| 49 | from vllm.multimodal.parse import ( |
| 50 | DictEmbeddingItems, |
| 51 | ModalityDataItems, |
| 52 | MultiModalDataItems, |
| 53 | MultiModalDataParser, |
| 54 | ) |
| 55 | from vllm.multimodal.processing import ( |
| 56 | BaseDummyInputsBuilder, |
| 57 | BaseMultiModalProcessor, |
| 58 | BaseProcessingInfo, |
| 59 | ProcessorInputs, |
| 60 | PromptUpdate, |
| 61 | TimingContext, |
| 62 | ) |
| 63 | from vllm.sequence import IntermediateTensors |
| 64 | from vllm.utils.gpu_sync_debug import gpu_sync_allowed |
| 65 | |
| 66 | from .interfaces import IsAttentionFree, MultiModalEmbeddings, SupportsMultiModal |
| 67 | from .interfaces_base import attn_type |
| 68 | |
| 69 | logger = init_logger(__name__) |
| 70 | |
| 71 | |
| 72 | def _terratorch_field_names(input_definition: InputDefinition): |
| 73 | return set(input_definition.data.keys()) |
| 74 | |
| 75 | |
| 76 | def _terratorch_field_factory( |
| 77 | input_definition: InputDefinition, |
| 78 | *, |
| 79 | is_shared: bool = True, # True for unprocessed data, False for processed data |
| 80 | ): |
| 81 | def _terratorch_field_config( |
| 82 | hf_inputs: Mapping[str, torch.Tensor], |
| 83 | ) -> Mapping[str, MultiModalFieldConfig]: |
| 84 | fields = dict[str, MultiModalFieldConfig]() |
| 85 | for name, input in input_definition.data.items(): |
| 86 | modality = "image" |
| 87 | if input.type == InputTypeEnum.tensor: |
| 88 | fields[name] = ( |
| 89 | MultiModalFieldConfig.shared(modality, batch_size=1) |
| 90 | if is_shared |
| 91 | else MultiModalFieldConfig.batched(modality) |
| 92 | ) |
| 93 | |
| 94 | return fields |
| 95 | |
| 96 | return _terratorch_field_config |
| 97 | |
| 98 | |
| 99 | class TerratorchMultiModalDataParser(MultiModalDataParser): |
| 100 | def __init__(self, input_definition: InputDefinition, *args, **kwargs): |
| 101 | super().__init__(*args, **kwargs) |
| 102 | |
| 103 | self.input_definition = input_definition |
| 104 | |
| 105 | def _parse_image_data( |
| 106 | self, |
| 107 | data: dict[str, torch.Tensor] | ModalityData[ImageItem], |
| 108 | ) -> ModalityDataItems[Any, Any] | None: |
| 109 | if isinstance(data, dict): |
| 110 | return DictEmbeddingItems( |
| 111 | data, |
| 112 | modality="image", |
| 113 | required_fields=_terratorch_field_names(self.input_definition), |
| 114 | fields_factory=_terratorch_field_factory(self.input_definition), |
| 115 | ) |
| 116 | |
| 117 | return super()._parse_image_data(data) |
| 118 | |
| 119 | def parse_mm_data(self, mm_data: MultiModalDataDict) -> MultiModalDataItems: |
| 120 | if "image" not in mm_data: |
| 121 | mm_data = {"image": mm_data} |
| 122 | |
| 123 | return super().parse_mm_data(mm_data) |
| 124 | |
| 125 | |
| 126 | class TerratorchProcessingInfo(BaseProcessingInfo): |
| 127 | @cached_property |
| 128 | def input_definition(self) -> InputDefinition: |
| 129 | pretrained_cfg = self.get_hf_config().to_dict()["pretrained_cfg"] |
| 130 | return InputDefinition(**pretrained_cfg["input"]) |
| 131 | |
| 132 | def get_data_parser(self): |
| 133 | return TerratorchMultiModalDataParser( |
| 134 | self.input_definition, |
| 135 | expected_hidden_size=self._get_expected_hidden_size(), |
| 136 | ) |
| 137 | |
| 138 | def get_supported_mm_limits(self) -> Mapping[str, int | None]: |
| 139 | return {"image": None} |
| 140 | |
| 141 | |
| 142 | class TerratorchInputBuilder(BaseDummyInputsBuilder[TerratorchProcessingInfo]): |
| 143 | def __init__(self, info: TerratorchProcessingInfo): |
| 144 | super().__init__(info) |
| 145 | self.dummy_data_generator = DummyDataGenerator( |
| 146 | self.info.get_hf_config().to_dict()["pretrained_cfg"] |
| 147 | ) |
| 148 | |
| 149 | def get_dummy_text(self, mm_counts: Mapping[str, int]) -> str: |
| 150 | return "" |
| 151 | |
| 152 | def get_dummy_mm_data( |
| 153 | self, |
| 154 | seq_len: int, |
| 155 | mm_counts: Mapping[str, int], |
| 156 | mm_options: MultiModalDummyOptions, |
| 157 | ) -> MultiModalDataDict: |
| 158 | # Dummy data is generated based on the 'input' section |
| 159 | # defined in the HF configuration file |
| 160 | |
| 161 | if mm_options: |
| 162 | logger.warning( |
| 163 | "Configurable multimodal profiling " |
| 164 | "options are not supported for Terratorch. " |
| 165 | "They are ignored for now." |
| 166 | ) |
| 167 | |
| 168 | return self.dummy_data_generator.get_dummy_mm_data() |
| 169 | |
| 170 | |
| 171 | class TerratorchMultiModalProcessor(BaseMultiModalProcessor[TerratorchProcessingInfo]): |
| 172 | # Terratorch models, including Prithvi, consume tokenizer-free inputs. |
| 173 | requires_tokenizer = False |
| 174 | |
| 175 | def _get_mm_fields_config( |
| 176 | self, |
| 177 | hf_inputs: BatchFeature, |
| 178 | hf_processor_mm_kwargs: Mapping[str, object], |
| 179 | *, |
| 180 | is_shared: bool = True, |
| 181 | ) -> Mapping[str, MultiModalFieldConfig]: |
| 182 | factory = _terratorch_field_factory( |
| 183 | self.info.input_definition, |
| 184 | is_shared=is_shared, |
| 185 | ) |
| 186 | return factory(hf_inputs) |
| 187 | |
| 188 | def _get_prompt_updates( |
| 189 | self, |
| 190 | mm_items: MultiModalDataItems, |
| 191 | hf_processor_mm_kwargs: Mapping[str, object], |
| 192 | out_mm_kwargs: MultiModalKwargsItems, |
| 193 | ) -> Sequence[PromptUpdate]: |
| 194 | return [] |
| 195 | |
| 196 | def apply( |
| 197 | self, |
| 198 | inputs: ProcessorInputs, |
| 199 | timing_ctx: TimingContext, |
| 200 | ) -> MultiModalInput: |
| 201 | mm_items = inputs.mm_data_items |
| 202 | hf_processor_mm_kwargs = inputs.hf_processor_mm_kwargs |
| 203 | |
| 204 | with timing_ctx.record("apply_hf_processor"): |
| 205 | passthrough_data = self._get_hf_mm_inputs( |
| 206 | mm_items, hf_processor_mm_kwargs |
| 207 | ).passthrough_data |
| 208 | mm_processed_data = BatchFeature( |
| 209 | { |
| 210 | k: torch.as_tensor(v).unsqueeze(0) |
| 211 | for k, v in passthrough_data.items() |
| 212 | }, |
| 213 | tensor_type="pt", |
| 214 | ) |
| 215 | |
| 216 | mm_kwargs = MultiModalKwargsItems.from_hf_inputs( |
| 217 | mm_processed_data, |
| 218 | self._get_mm_fields_config( |
| 219 | mm_processed_data, |
| 220 | hf_processor_mm_kwargs, |
| 221 | is_shared=False, |
| 222 | ), |
| 223 | ) |
| 224 | |
| 225 | with timing_ctx.record("get_mm_hashes"): |
| 226 | mm_hashes = inputs.get_mm_hashes( |
| 227 | self.info.model_id, |
| 228 | self.info.ctx.get_mm_config().mm_hasher_algorithm, |
| 229 | ) |
| 230 | |
| 231 | mm_placeholders = {"image": [PlaceholderRange(offset=0, length=0)]} |
| 232 | |
| 233 | return mm_input( |
| 234 | prompt_token_ids=[1], |
| 235 | mm_kwargs=mm_kwargs, |
| 236 | mm_hashes=mm_hashes, |
| 237 | mm_placeholders=mm_placeholders, |
| 238 | ) |
| 239 | |
| 240 | |
| 241 | @attn_type("attention_free") |
| 242 | @MULTIMODAL_REGISTRY.register_processor( |
| 243 | TerratorchMultiModalProcessor, |
| 244 | info=TerratorchProcessingInfo, |
| 245 | dummy_inputs=TerratorchInputBuilder, |
| 246 | ) |
| 247 | class Terratorch(nn.Module, IsAttentionFree, SupportsMultiModal): |
| 248 | supports_multimodal_raw_input_only = True |
| 249 | is_pooling_model = True |
| 250 | |
| 251 | @classmethod |
| 252 | def get_placeholder_str(cls, modality: str, i: int) -> str | None: |
| 253 | if modality.startswith("image"): |
| 254 | return None |
| 255 | |
| 256 | raise ValueError("Only image modality is supported") |
| 257 | |
| 258 | def __init__(self, vllm_config: VllmConfig, prefix: str = ""): |
| 259 | super().__init__() |
| 260 | |
| 261 | config = vllm_config.model_config.hf_config.to_dict()["pretrained_cfg"] |
| 262 | |
| 263 | self.inference_runner = InferenceRunner(config) |
| 264 | self.model = self.inference_runner.model |
| 265 | |
| 266 | self.pooler = IdentityPooler() |
| 267 | |
| 268 | def embed_input_ids( |
| 269 | self, |
| 270 | input_ids: torch.Tensor, |
| 271 | multimodal_embeddings: MultiModalEmbeddings | None = None, |
| 272 | *, |
| 273 | is_multimodal: torch.Tensor | None = None, |
| 274 | ) -> torch.Tensor: |
| 275 | # We do not really use any input tokens and therefore no embeddings |
| 276 | # to be calculated. However, due to the mandatory token ids in |
| 277 | # the input prompt we pass one token and the size of the dummy |
| 278 | # embedding tensors must reflect that. |
| 279 | return torch.empty((input_ids.shape[0], 0)) |
| 280 | |
| 281 | def forward( |
| 282 | self, |
| 283 | input_ids: torch.Tensor | None, |
| 284 | positions: torch.Tensor, |
| 285 | intermediate_tensors: IntermediateTensors | None = None, |
| 286 | inputs_embeds: torch.Tensor | None = None, |
| 287 | **kwargs: object, |
| 288 | ): |
| 289 | # terratorch's forward has internal GPU syncs. |
| 290 | with gpu_sync_allowed(): |
| 291 | model_output = self.inference_runner.forward(**kwargs) |
| 292 | return model_output.output |
| 293 | |
| 294 | def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: |
| 295 | params_list = [] |
| 296 | model_buffers = dict(self.named_buffers()) |
| 297 | loaded_buffers: list[str] = [] |
| 298 | for key, value in weights: |
| 299 | if isinstance(value, (dict, OrderedDict)): |
| 300 | if key == "state_dict": |
| 301 | weights_to_parse = value |
| 302 | for name, weight in weights_to_parse.items(): |
| 303 | name = f"inference_runner.{name}" |
| 304 | |
| 305 | if "pos_embed" in name: |
| 306 | continue |
| 307 | |
| 308 | if "_timm_module." in name: |
| 309 | name = name.replace("_timm_module.", "") |
| 310 | |
| 311 | # this model requires a couple of buffers to be loaded |
| 312 | # that are not loadable with the AutoWeightsLoader |
| 313 | if name in model_buffers: |
| 314 | if "_timm_module." in name: |
| 315 | name = name.replace("_timm_module.", "") |
| 316 | buffer = model_buffers[name] |
| 317 | weight_loader = getattr( |
| 318 | buffer, "weight_loader", default_weight_loader |
| 319 | ) |
| 320 | weight_loader(buffer, weight) |
| 321 | loaded_buffers.append(name) |
| 322 | else: |
| 323 | params_list.append((name, weight)) |
| 324 | break |
| 325 | |
| 326 | elif isinstance(value, torch.Tensor): |
| 327 | params_list.append((f"inference_runner.model.{key}", value)) |
| 328 | |
| 329 | # Load the remaining model parameters |
| 330 | loader = AutoWeightsLoader(self) |
| 331 | autoloaded_weights = loader.load_weights(params_list) |
| 332 | |
| 333 | return autoloaded_weights.union(set(loaded_buffers)) |