@tduncan / Nvidia-Optimization-Stack-VLLM
Code Commits Issues Pulls Cloud Agent CI/CD
main
main
py 113 lines 3.72 KB
Raw
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4 import os
5
6 import pytest
7
8 from vllm.assets.video import VideoAsset
9 from vllm.config import ModelConfig
10 from vllm.multimodal import MULTIMODAL_REGISTRY
11
12 from ....conftest import ImageTestAssets
13 from ...utils import build_model_context
14
15 MODEL_ID = "nvidia/Cosmos3-Edge"
16 IMAGE_PLACEHOLDER = "<|vision_start|><|image_pad|><|vision_end|>"
17 VIDEO_PLACEHOLDER = "<|vision_start|><|video_pad|><|vision_end|>"
18 LOCAL_MODEL_PATH = os.getenv("COSMOS3_EDGE_MODEL_PATH")
19
20
21 @pytest.fixture(scope="module")
22 def processor():
23 if LOCAL_MODEL_PATH:
24 model_config = ModelConfig(
25 LOCAL_MODEL_PATH,
26 tokenizer=LOCAL_MODEL_PATH,
27 max_model_len=4096,
28 limit_mm_per_prompt={"image": 2, "video": 1},
29 )
30 else:
31 ctx = build_model_context(
32 MODEL_ID,
33 limit_mm_per_prompt={"image": 2, "video": 1},
34 )
35 model_config = ctx.model_config
36
37 return MULTIMODAL_REGISTRY.create_processor(model_config)
38
39
40 def _assert_image_outputs(processor, processed, num_images: int) -> None:
41 mm_data = processed["mm_kwargs"].get_data()
42 grid_thw = mm_data["image_grid_thw"]
43 pixel_values = mm_data["pixel_values"]
44
45 assert grid_thw.shape == (num_images, 3)
46 assert pixel_values.shape[0] == int(grid_thw.prod(dim=-1).sum())
47
48 merge_size = processor.info.get_hf_config().vision_config.spatial_merge_size
49 expected_tokens = (grid_thw.prod(dim=-1) // merge_size**2).tolist()
50 image_placeholders = processed["mm_placeholders"]["image"]
51
52 assert len(image_placeholders) == num_images
53 assert [placeholder.length for placeholder in image_placeholders] == (
54 expected_tokens
55 )
56
57 image_token_id = processor.info.get_hf_processor().image_token_id
58 assert processed["prompt_token_ids"].count(image_token_id) == sum(expected_tokens)
59
60
61 def _assert_video_outputs(processor, processed) -> None:
62 mm_data = processed["mm_kwargs"].get_data()
63 grid_thw = mm_data["video_grid_thw"]
64 pixel_values = mm_data["pixel_values_videos"]
65
66 assert grid_thw.shape == (1, 3)
67 assert pixel_values.shape[0] == int(grid_thw.prod())
68 assert len(processed["mm_placeholders"]["video"]) == 1
69
70 merge_size = processor.info.get_hf_config().vision_config.spatial_merge_size
71 expected_tokens = int(grid_thw.prod()) // merge_size**2
72 video_token_id = processor.info.get_hf_config().video_token_id
73 prompt_token_ids = processed["prompt_token_ids"]
74 assert prompt_token_ids.count(video_token_id) == expected_tokens
75
76 hf_processor = processor.info.get_hf_processor()
77 expected_frame_wrappers = int(grid_thw[:, 0].sum())
78 assert (
79 prompt_token_ids.count(hf_processor.vision_start_token_id)
80 == expected_frame_wrappers
81 )
82 assert (
83 prompt_token_ids.count(hf_processor.vision_end_token_id)
84 == expected_frame_wrappers
85 )
86
87
88 @pytest.mark.parametrize("num_images", [1, 2])
89 def test_process_images(
90 processor,
91 image_assets: ImageTestAssets,
92 num_images: int,
93 ) -> None:
94 images = [asset.pil_image for asset in image_assets[:num_images]]
95 processed = processor(
96 IMAGE_PLACEHOLDER * num_images,
97 mm_items=processor.info.parse_mm_data({"image": images}),
98 hf_processor_mm_kwargs={},
99 )
100
101 _assert_image_outputs(processor, processed, num_images)
102
103
104 def test_process_video(processor) -> None:
105 video_asset = VideoAsset(name="baby_reading", num_frames=8)
106 video = (video_asset.np_ndarrays, video_asset.metadata)
107 processed = processor(
108 VIDEO_PLACEHOLDER,
109 mm_items=processor.info.parse_mm_data({"video": [video]}),
110 hf_processor_mm_kwargs={},
111 )
112
113 _assert_video_outputs(processor, processed)