| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | """Tests for Step3-VL precomputed image embedding inputs.""" |
| 4 | |
| 5 | import pytest |
| 6 | import torch |
| 7 | |
| 8 | from vllm.model_executor.models.step3_vl import ( |
| 9 | Step3VLForConditionalGeneration, |
| 10 | Step3VLImageEmbeddingInputs, |
| 11 | ) |
| 12 | |
| 13 | |
| 14 | class _FakeStep3VL: |
| 15 | @staticmethod |
| 16 | def _process_image_features(image_features: torch.Tensor) -> torch.Tensor: |
| 17 | return image_features |
| 18 | |
| 19 | |
| 20 | def test_image_embedding_inputs_construction(): |
| 21 | """Step3VLImageEmbeddingInputs should store embeddings in the data field.""" |
| 22 | image_embeds = torch.randn(2, 16, 64) |
| 23 | |
| 24 | inputs = Step3VLImageEmbeddingInputs( |
| 25 | type="image_embeds", |
| 26 | data=image_embeds, |
| 27 | ) |
| 28 | |
| 29 | assert inputs["type"] == "image_embeds" |
| 30 | assert torch.equal(inputs["data"], image_embeds) |
| 31 | assert torch.equal(inputs.data, image_embeds) |
| 32 | |
| 33 | |
| 34 | def test_image_embedding_inputs_validation_rejects_wrong_rank(): |
| 35 | """Validation should reject tensors with wrong rank.""" |
| 36 | with pytest.raises(ValueError, match="rank"): |
| 37 | Step3VLImageEmbeddingInputs( |
| 38 | type="image_embeds", |
| 39 | data=torch.randn(16, 64), |
| 40 | ) |
| 41 | |
| 42 | |
| 43 | def test_process_image_embeds_does_not_require_pixel_input_fields(): |
| 44 | """The image_embeds branch should not reference patch pixel metadata.""" |
| 45 | image_embeds = torch.randn(2, 4, 8) |
| 46 | image_input = Step3VLImageEmbeddingInputs( |
| 47 | type="image_embeds", |
| 48 | data=image_embeds, |
| 49 | ) |
| 50 | |
| 51 | outputs = Step3VLForConditionalGeneration._process_image_input( |
| 52 | _FakeStep3VL(), |
| 53 | image_input, |
| 54 | ) |
| 55 | |
| 56 | assert len(outputs) == 2 |
| 57 | assert torch.equal(outputs[0], image_embeds[0]) |
| 58 | assert torch.equal(outputs[1], image_embeds[1]) |