test_vision_embeds.py 1.86 KB
Newer Older
1
2
3
4
5
6
7
8
9
10
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project

import base64

import numpy as np
import pytest
import requests
import torch

11
from vllm.utils.serial_utils import tensor2base64
12

13
from ...utils import RemoteOpenAIServer
14
15


16
def _terratorch_dummy_messages():
17
18
19
    pixel_values = torch.full((6, 512, 512), 1.0, dtype=torch.float16)
    location_coords = torch.full((1, 2), 1.0, dtype=torch.float16)

20
21
22
23
24
25
26
27
28
29
30
31
32
33
    return [
        {
            "role": "user",
            "content": [
                {
                    "type": "image_embeds",
                    "image_embeds": {
                        "pixel_values": tensor2base64(pixel_values),
                        "location_coords": tensor2base64(location_coords),
                    },
                }
            ],
        }
    ]
34
35


36
37
38
39
@pytest.mark.parametrize(
    "model_name", ["ibm-nasa-geospatial/Prithvi-EO-2.0-300M-TL-Sen1Floods11"]
)
def test_single_request(model_name: str):
40
41
42
43
44
    args = [
        "--runner",
        "pooling",
        # use half precision for speed and memory savings in CI environment
        "--dtype",
45
        "float16",
46
47
48
49
50
51
52
53
54
55
        "--enforce-eager",
        "--trust-remote-code",
        "--max-num-seqs",
        "32",
        "--model-impl",
        "terratorch",
        "--skip-tokenizer-init",
        "--enable-mm-embeds",
    ]

56
57
58
59
60
61
62
63
64
    with RemoteOpenAIServer(model_name, args) as server:
        response = requests.post(
            server.url_for("pooling"),
            json={
                "model": model_name,
                "messages": _terratorch_dummy_messages(),
                "encoding_format": "base64",
            },
        )
65
        response.raise_for_status()
66

67
        output = response.json()["data"][0]["data"]
68

69
70
        np_response = np.frombuffer(base64.b64decode(output), dtype=np.float32)
        assert len(np_response) == 524288