59 lines
1.7 KiB
Python
59 lines
1.7 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
|
"""Tests for Step3-VL precomputed image embedding inputs."""
|
|
|
|
import pytest
|
|
import torch
|
|
|
|
from vllm.model_executor.models.step3_vl import (
|
|
Step3VLForConditionalGeneration,
|
|
Step3VLImageEmbeddingInputs,
|
|
)
|
|
|
|
|
|
class _FakeStep3VL:
|
|
@staticmethod
|
|
def _process_image_features(image_features: torch.Tensor) -> torch.Tensor:
|
|
return image_features
|
|
|
|
|
|
def test_image_embedding_inputs_construction():
|
|
"""Step3VLImageEmbeddingInputs should store embeddings in the data field."""
|
|
image_embeds = torch.randn(2, 16, 64)
|
|
|
|
inputs = Step3VLImageEmbeddingInputs(
|
|
type="image_embeds",
|
|
data=image_embeds,
|
|
)
|
|
|
|
assert inputs["type"] == "image_embeds"
|
|
assert torch.equal(inputs["data"], image_embeds)
|
|
assert torch.equal(inputs.data, image_embeds)
|
|
|
|
|
|
def test_image_embedding_inputs_validation_rejects_wrong_rank():
|
|
"""Validation should reject tensors with wrong rank."""
|
|
with pytest.raises(ValueError, match="rank"):
|
|
Step3VLImageEmbeddingInputs(
|
|
type="image_embeds",
|
|
data=torch.randn(16, 64),
|
|
)
|
|
|
|
|
|
def test_process_image_embeds_does_not_require_pixel_input_fields():
|
|
"""The image_embeds branch should not reference patch pixel metadata."""
|
|
image_embeds = torch.randn(2, 4, 8)
|
|
image_input = Step3VLImageEmbeddingInputs(
|
|
type="image_embeds",
|
|
data=image_embeds,
|
|
)
|
|
|
|
outputs = Step3VLForConditionalGeneration._process_image_input(
|
|
_FakeStep3VL(),
|
|
image_input,
|
|
)
|
|
|
|
assert len(outputs) == 2
|
|
assert torch.equal(outputs[0], image_embeds[0])
|
|
assert torch.equal(outputs[1], image_embeds[1])
|