diff --git a/tests/models/multimodal/test_openpangu_vl.py b/tests/models/multimodal/test_openpangu_vl.py new file mode 100644 index 0000000000..e8bab16d5d --- /dev/null +++ b/tests/models/multimodal/test_openpangu_vl.py @@ -0,0 +1,55 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +from types import SimpleNamespace + +import pytest +import torch +from torch import nn + +from vllm.model_executor.models.openpangu_vl import ( + OpenPanguVLForConditionalGeneration, +) + + +class _LanguageModel(nn.Module): + def __init__(self, hidden_size: int) -> None: + super().__init__() + self.hidden_size = hidden_size + + def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor: + return input_ids.float().unsqueeze(-1).expand(-1, self.hidden_size).clone() + + +class _OpenPanguVL(OpenPanguVLForConditionalGeneration): + def __init__(self, hidden_size: int) -> None: + nn.Module.__init__(self) + self.config = SimpleNamespace(image_token_id=41, video_token_id=42) + self.language_model = _LanguageModel(hidden_size) + + def get_language_model(self) -> _LanguageModel: + return self.language_model + + +@pytest.mark.parametrize("device_type", ["cpu", "cuda"]) +def test_get_input_embeddings_merges_image_and_video_embeddings( + device_type: str, +) -> None: + if device_type == "cuda" and not torch.cuda.is_available(): + pytest.skip("CUDA is not available") + + device = torch.device(device_type) + hidden_size = 4 + model = _OpenPanguVL(hidden_size).to(device) + input_ids = torch.tensor([7, 41, 41, 8, 42, 9], device=device) + multimodal_embeddings = ( + torch.full((2, hidden_size), -1.0, device=device), + torch.full((1, hidden_size), -2.0, device=device), + ) + + actual = model.get_input_embeddings(input_ids, multimodal_embeddings) + + expected = model.language_model.embed_input_ids(input_ids) + expected[1:3] = -1.0 + expected[4] = -2.0 + torch.testing.assert_close(actual, expected) diff --git a/vllm/model_executor/models/openpangu_vl.py b/vllm/model_executor/models/openpangu_vl.py index 04a6dba41f..10b72f9292 100644 --- a/vllm/model_executor/models/openpangu_vl.py +++ b/vllm/model_executor/models/openpangu_vl.py @@ -1039,17 +1039,20 @@ class OpenPanguVLForConditionalGeneration( def get_input_embeddings( self, input_ids: torch.Tensor, - multimodal_embeddings=None, + multimodal_embeddings: MultiModalEmbeddings | None = None, ) -> torch.Tensor: - inputs_embeds = self.language_model.embed_input_ids(input_ids) - if multimodal_embeddings is not None: - inputs_embeds = self.embed_input_ids( - input_ids, - inputs_embeds, - multimodal_embeddings, - [self.config.image_token_id, self.config.video_token_id], - ) - return inputs_embeds + if multimodal_embeddings is None: + return self.language_model.embed_input_ids(input_ids) + + mm_token_ids = input_ids.new_tensor( + [self.config.image_token_id, self.config.video_token_id] + ) + is_multimodal = torch.isin(input_ids, mm_token_ids) + return self.embed_input_ids( + input_ids, + multimodal_embeddings, + is_multimodal=is_multimodal, + ) def _process_image_input(self, image_input) -> tuple[torch.Tensor, ...]: grid_thw = image_input["image_grid_thw"]