Some reformatting to make the pre-commit hooks succeed

Signed-off-by: Christian Pinto <christian.pinto@ibm.com>

Author: christian-pinto

examples/offline_inference/prithvi_geospatial_mae.py

@@ -144,7 +144,7 @@ def __init__(self):
             model=os.path.join(os.path.dirname(__file__), "./model"),
             skip_tokenizer_init=True,
             dtype="float16",
-            enforce_eager=True
+            enforce_eager=True,
         )
 
     def run(self, input_data, location_coords):

tests/models/multimodal/pooling/test_prithvi_mae.py

@@ -0,0 +1,44 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+import pytest
+import torch
+
+from ....conftest import VllmRunner
+
+def generate_test_mm_data():
+    mm_data = {
+        "pixel_values": torch.full((6, 512, 512), 1.0, dtype=torch.float16),
+        "location_coords": torch.full((1, 2), 1.0, dtype=torch.float16),
+    }
+    return mm_data
+     
+def _run_test(
+    vllm_runner: type[VllmRunner],
+    model: str,
+) -> None:   
+
+    mm_data = generate_test_mm_data()
+    prompt = {
+        # This model deals with no text input
+        "prompt_token_ids": [1],
+        "multi_modal_data": mm_data
+    }
+    with vllm_runner(model, task="embed",
+                    dtype=torch.float16,
+                    enforce_eager=True,
+                    skip_tokenizer_init=True) as vllm_model:
+        output = vllm_model.encode(prompt)
+            
+MODELS=["christian-pinto/Prithvi-EO-2.0-300M-TL-VLLM"]
+@pytest.mark.parametrize("model", MODELS)
+def test_models_image(
+    hf_runner,
+    vllm_runner,
+    image_assets,
+    model: str,
+) -> None:
+    _run_test(
+        vllm_runner,
+        model,
+    )

vllm/config.py

@@ -642,8 +642,10 @@ def __post_init__(self) -> None:
         self.original_max_model_len = self.max_model_len
         self.max_model_len = self.get_and_verify_max_len(self.max_model_len)
         self.multimodal_config = self._init_multimodal_config()
-        self.is_pooling_model = self.registry.is_pooling_model(self.architectures)
-        self.model_supports_multimodal_raw_input = self._init_model_supports_multimodal_raw_input()
+        self.is_pooling_model = self.registry.is_pooling_model(
+            self.architectures)
+        self.model_supports_multimodal_raw_input = (
+            self._init_model_supports_multimodal_raw_input())
         if not self.skip_tokenizer_init:
             self._verify_tokenizer_mode()
 

vllm/model_executor/models/interfaces.py

@@ -165,18 +165,21 @@ class _SupportsMultiModalWithRawInput(Protocol):
 
 
 @overload
-def supports_multimodal_raw_input(model: object) -> TypeIs[SupportsMultiModalWithRawInput]:
+def supports_multimodal_raw_input(
+        model: object) -> TypeIs[SupportsMultiModalWithRawInput]:
     ...
 
 
 @overload
-def supports_multimodal_raw_input(model: type[object]) -> TypeIs[type[SupportsMultiModalWithRawInput]]:
+def supports_multimodal_raw_input(
+        model: type[object]) -> TypeIs[type[SupportsMultiModalWithRawInput]]:
     ...
 
 
 def supports_multimodal_raw_input(
     model: Union[type[object], object]
-) -> Union[TypeIs[type[SupportsMultiModalWithRawInput]], TypeIs[SupportsMultiModalWithRawInput]]:
+) -> Union[TypeIs[type[SupportsMultiModalWithRawInput]],
+           TypeIs[SupportsMultiModalWithRawInput]]:
     if isinstance(model, type):
         return isinstance(model, _SupportsMultiModalWithRawInput)
 

vllm/model_executor/models/prithvi_geospatial_mae.py

@@ -25,13 +25,14 @@
 
 from vllm.config import VllmConfig
 from vllm.model_executor.model_loader.weight_utils import default_weight_loader
-from vllm.model_executor.models.interfaces import (IsAttentionFree,
-                                                   SupportsMultiModalWithRawInput)
+from vllm.model_executor.models.interfaces import (
+    IsAttentionFree, SupportsMultiModalWithRawInput)
 from vllm.model_executor.models.utils import AutoWeightsLoader
 from vllm.model_executor.pooling_metadata import PoolingMetadata
 from vllm.multimodal import MULTIMODAL_REGISTRY
-from vllm.multimodal.inputs import (MultiModalDataDict, MultiModalFieldConfig, MultiModalFieldElem,
-                                    MultiModalInputs, MultiModalKwargs, MultiModalKwargsItem,
+from vllm.multimodal.inputs import (MultiModalDataDict, MultiModalFieldConfig,
+                                    MultiModalFieldElem, MultiModalInputs,
+                                    MultiModalKwargs, MultiModalKwargsItem,
                                     MultiModalSharedField, PlaceholderRange)
 from vllm.multimodal.parse import MultiModalDataItems
 from vllm.multimodal.processing import (BaseMultiModalProcessor,
@@ -62,7 +63,8 @@ def get_dummy_mm_data(
         # The size of pixel_values might change in the cases where we resize
         # the input but never exceeds the dimensions below.
         return {
-            "pixel_values": torch.full((6, 512, 512), 1.0, dtype=torch.float16),
+            "pixel_values": torch.full((6, 512, 512), 1.0,
+                                       dtype=torch.float16),
             "location_coords": torch.full((1, 2), 1.0, dtype=torch.float16),
         }
 
@@ -75,8 +77,10 @@ def _get_mm_fields_config(
         hf_processor_mm_kwargs: Mapping[str, object],
     ) -> Mapping[str, MultiModalFieldConfig]:
         return dict(
-            pixel_values=MultiModalFieldConfig.shared(batch_size=1, modality="image"),
-            location_coords=MultiModalFieldConfig.shared(batch_size=1, modality="image"),
+            pixel_values=MultiModalFieldConfig.shared(batch_size=1,
+                                                      modality="image"),
+            location_coords=MultiModalFieldConfig.shared(batch_size=1,
+                                                         modality="image"),
         )
 
     def _get_prompt_updates(
@@ -99,15 +103,16 @@ def apply(
 
         for k, v in mm_data.items():
             mm_kwargs[k] = v
-        mm_place_holders = {
-            "image": [PlaceholderRange(offset=0, length=0)]
-        }
+        mm_place_holders = {"image": [PlaceholderRange(offset=0, length=0)]}
 
         multimodal_kwargs_items = [
-            MultiModalKwargsItem.from_elems(
-                [MultiModalFieldElem(modality="image", key=key, data=data, field=MultiModalSharedField(1))
-                 for key, data in mm_kwargs.items()]
-            )
+            MultiModalKwargsItem.from_elems([
+                MultiModalFieldElem(modality="image",
+                                    key=key,
+                                    data=data,
+                                    field=MultiModalSharedField(1))
+                for key, data in mm_kwargs.items()
+            ])
         ]
 
         return MultiModalInputs(
@@ -124,7 +129,8 @@ def apply(
     PrithviGeoSpatialMAEMultiModalProcessor,
     info=PrithviGeoSpatialMAEProcessingInfo,
     dummy_inputs=PrithviGeoSpatialMAEInputBuilder)
-class PrithviGeoSpatialMAE(nn.Module, IsAttentionFree, SupportsMultiModalWithRawInput):
+class PrithviGeoSpatialMAE(nn.Module, IsAttentionFree,
+                           SupportsMultiModalWithRawInput):
     """ Prithvi Masked Autoencoder"""
 
     @classmethod
@@ -189,13 +195,14 @@ def _parse_and_validate_multimodal_data(
             location_coords = None
 
         return pixel_values, location_coords
-    
+
     def get_input_embeddings(self, input_ids: torch.Tensor) -> torch.Tensor:
-        # We do not really use any input tokens and therefore no embeddings to be calculated
-        # However, due to the mandatory token ids in the input prompt we pass one token and the
-        # size of the dummy embedding tensors must reflect that.
+        # We do not really use any input tokens and therefore no embeddings
+        # to be calculated. However, due to the mandatory token ids in
+        # the input prompt we pass one token and the size of the dummy
+        #  embedding tensors must reflect that.
         return torch.empty(input_ids.shape)
-    
+
     def forward(
         self,
         input_ids: Optional[torch.Tensor],
@@ -217,7 +224,10 @@ def pooler(
         hidden_states: torch.Tensor,
         pooling_metadata: PoolingMetadata,
     ) -> Optional[PoolerOutput]:
-        return PoolerOutput([PoolingSequenceGroupOutput(hidden_state) for hidden_state in hidden_states])
+        return PoolerOutput([
+            PoolingSequenceGroupOutput(hidden_state)
+            for hidden_state in hidden_states
+        ])
 
     def load_weights(self, weights: Iterable[tuple[str,
                                                    torch.Tensor]]) -> set[str]:

vllm/model_executor/models/registry.py

@@ -23,8 +23,7 @@
 from .interfaces import (has_inner_state, has_noops, is_attention_free,
                          is_hybrid, supports_cross_encoding,
                          supports_multimodal, supports_multimodal_raw_input,
-                         supports_pp, supports_transcription,
-                         supports_v0_only)
+                         supports_pp, supports_transcription, supports_v0_only)
 from .interfaces_base import is_text_generation_model
 
 logger = init_logger(__name__)
@@ -539,7 +538,7 @@ def is_multimodal_model(
     ) -> bool:
         model_cls, _ = self.inspect_model_cls(architectures)
         return model_cls.supports_multimodal
-    
+
     def supports_multimodal_raw_input(
         self,
         architectures: Union[str, list[str]],

vllm/v1/engine/core.py

@@ -134,17 +134,8 @@ def _initialize_kv_caches(
             self, vllm_config: VllmConfig) -> tuple[int, int, KVCacheConfig]:
         start = time.time()
 
-        #TODO: CP start from here
-        if vllm_config.model_config.is_attention_free:
-            # No need for initializing anything related to KV cache if the model
-            # is attention free.
-            kv_cache_specs = []
-            kv_cache_configs = [
-                    KVCacheConfig(num_blocks=0, kv_cache_tensors={}, kv_cache_groups=[])
-                ]
-        else:
-            # Get all kv cache needed by the model
-            kv_cache_specs = self.model_executor.get_kv_cache_specs()
+        # Get all kv cache needed by the model
+        kv_cache_specs = self.model_executor.get_kv_cache_specs()
 
         # Profiles the peak memory usage of the model to determine how much
         # memory can be allocated for kv cache.
@@ -161,8 +152,8 @@ def _initialize_kv_caches(
         kv_cache_configs = [
             get_kv_cache_config(vllm_config, kv_cache_spec_one_worker,
                                 available_gpu_memory_one_worker)
-            for kv_cache_spec_one_worker, available_gpu_memory_one_worker in
-            zip(kv_cache_specs, available_gpu_memory)
+            for kv_cache_spec_one_worker, available_gpu_memory_one_worker
+            in zip(kv_cache_specs, available_gpu_memory)
         ]
 
         # Since we use a shared centralized controller, we need the

vllm/v1/engine/llm_engine.py

@@ -82,8 +82,7 @@ def __init__(
             self.dp_group = None
         self.should_execute_dummy_batch = False
 
-        
-        if not self.vllm_config.model_config.skip_tokenizer_init:     
+        if not self.vllm_config.model_config.skip_tokenizer_init:
             # Tokenizer (+ ensure liveness if running in another process).
             self.tokenizer = init_tokenizer_from_configs(
                 model_config=vllm_config.model_config,

vllm/v1/engine/output_processor.py

@@ -330,14 +330,13 @@ def add_request(
         tokenizer = None if not self.tokenizer else \
             self.tokenizer.get_lora_tokenizer(request.lora_request)
 
-        req_state = RequestState.from_new_request(
-            tokenizer=tokenizer,
-            request=request,
-            prompt=prompt,
-            parent_req=parent_req,
-            request_index=request_index,
-            queue=queue,
-            log_stats=self.log_stats)
+        req_state = RequestState.from_new_request(tokenizer=tokenizer,
+                                                  request=request,
+                                                  prompt=prompt,
+                                                  parent_req=parent_req,
+                                                  request_index=request_index,
+                                                  queue=queue,
+                                                  log_stats=self.log_stats)
         self.request_states[request_id] = req_state
         self.lora_states.add_request(req_state)
         if parent_req:

vllm/v1/engine/processor.py

@@ -395,7 +395,8 @@ def _validate_model_input(
         if tokenizer:
             max_input_id = max(prompt_ids, default=0)
             if max_input_id > tokenizer.max_token_id:
-                raise ValueError(f"Token id {max_input_id} is out of vocabulary")
+                raise ValueError(
+                    f"Token id {max_input_id} is out of vocabulary")
 
         max_prompt_len = self.model_config.max_model_len
         if len(prompt_ids) > max_prompt_len: