[Bugfix] Fix OpenPangu multimodal embedding merge (#55941)

Signed-off-by: Taneem Ibrahim <taneem.ibrahim@gmail.com>

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)

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"]