64903b8e30
Signed-off-by: wangyu <410167048@qq.com> Signed-off-by: [Your Name] [Your Email] Signed-off-by: [Your Name] <your.email@example.com> Co-authored-by: Cursor <cursoragent@cursor.com> Co-authored-by: amy-why-3459 <wuhaiyan17@huawei.com>
84 lines
2.8 KiB
Python
84 lines
2.8 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
|
|
"""End-to-end smoke test for LingBot-World v2 generation."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import os
|
|
from pathlib import Path
|
|
|
|
import numpy as np
|
|
import pytest
|
|
|
|
from tests.helpers.mark import hardware_test
|
|
from tests.helpers.runtime import OmniRunnerHandler
|
|
from vllm_omni.diffusion.data import DiffusionParallelConfig
|
|
from vllm_omni.inputs.data import OmniDiffusionSamplingParams
|
|
|
|
MODEL = os.environ.get(
|
|
"VLLM_OMNI_LINGBOT_WORLD_V2_CHECKPOINT_PATH",
|
|
"robbyant/lingbot-world-v2-14b-causal-fast-diffusers",
|
|
)
|
|
_IMAGE_PATH = os.environ.get("VLLM_OMNI_LINGBOT_WORLD_V2_IMAGE_PATH")
|
|
_ACTION_DIR_VALUE = os.environ.get("VLLM_OMNI_LINGBOT_WORLD_V2_ACTION_DIR")
|
|
_ACTION_DIR = Path(_ACTION_DIR_VALUE).expanduser().resolve() if _ACTION_DIR_VALUE else None
|
|
_HAS_INPUT_ASSETS = _IMAGE_PATH is not None and _ACTION_DIR is not None
|
|
|
|
_OMNI_RUNNER_PARAM = (
|
|
MODEL,
|
|
None,
|
|
{
|
|
"flow_shift": 5.0,
|
|
"parallel_config": DiffusionParallelConfig(tensor_parallel_size=1),
|
|
"model_config": {
|
|
"lingbot_action_root": str(_ACTION_DIR.parent) if _ACTION_DIR is not None else "",
|
|
},
|
|
},
|
|
)
|
|
|
|
pytestmark = [
|
|
pytest.mark.full_model,
|
|
pytest.mark.diffusion,
|
|
pytest.mark.skipif(
|
|
not _HAS_INPUT_ASSETS,
|
|
reason=("VLLM_OMNI_LINGBOT_WORLD_V2_IMAGE_PATH and VLLM_OMNI_LINGBOT_WORLD_V2_ACTION_DIR are required"),
|
|
),
|
|
pytest.mark.parametrize("omni_runner", [_OMNI_RUNNER_PARAM], indirect=True),
|
|
]
|
|
|
|
|
|
@hardware_test(res={"cuda": ["H100", "B200"]}, num_cards=1)
|
|
def test_lingbot_world_v2_tp1_one_block(omni_runner_handler: OmniRunnerHandler) -> None:
|
|
"""Load the real checkpoint and generate one causal block on one GPU."""
|
|
|
|
assert _IMAGE_PATH is not None
|
|
assert _ACTION_DIR is not None
|
|
response = omni_runner_handler.send_diffusion_request(
|
|
{
|
|
"model": omni_runner_handler.runner.model_name,
|
|
"prompt": "The camera moves slowly forward through the scene.",
|
|
"multi_modal_data": {"image": str(Path(_IMAGE_PATH).expanduser().resolve())},
|
|
"sampling_params": OmniDiffusionSamplingParams(
|
|
height=480,
|
|
width=832,
|
|
num_frames=9,
|
|
max_sequence_length=512,
|
|
seed=42,
|
|
fps=16,
|
|
extra_args={
|
|
"action_path": _ACTION_DIR.name,
|
|
"flow_shift": 5.0,
|
|
},
|
|
),
|
|
}
|
|
)
|
|
|
|
assert response.success
|
|
assert response.images is not None
|
|
assert len(response.images) == 1
|
|
video = np.asarray(response.images[0])
|
|
assert video.shape == (1, 9, 480, 832, 3)
|
|
assert np.isfinite(video).all()
|
|
assert float(np.ptp(video)) > 1e-6
|
|
assert float(np.mean(np.abs(np.diff(video, axis=1)))) > 1e-6
|