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>
106 lines
3.8 KiB
Python
106 lines
3.8 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
|
|
"""End-to-end WAN check for the typed pre-D2H media contract.
|
|
|
|
Runs the whole engine twice with the same seed, once per flag state, and compares
|
|
the frames a client receives. WAN denormalizes historical bfloat16 output in the
|
|
input dtype while device preparation computes in float32, so the paths may differ
|
|
by one 255th.
|
|
|
|
Run with a WAN2.2 checkpoint::
|
|
|
|
VLLM_OMNI_DEVICE_POSTPROCESS_MODEL=Wan-AI/Wan2.2-TI2V-5B-Diffusers \
|
|
pytest -s tests/e2e/offline_inference/test_device_postprocess_equivalence.py
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import os
|
|
|
|
import numpy as np
|
|
import pytest
|
|
|
|
from tests.helpers.mark import hardware_test
|
|
from vllm_omni.inputs.data import OmniDiffusionSamplingParams
|
|
|
|
MODEL = os.environ.get("VLLM_OMNI_DEVICE_POSTPROCESS_MODEL")
|
|
|
|
pytestmark = [
|
|
pytest.mark.full_model,
|
|
pytest.mark.diffusion,
|
|
pytest.mark.gpu,
|
|
pytest.mark.skipif(
|
|
not MODEL,
|
|
reason="Set VLLM_OMNI_DEVICE_POSTPROCESS_MODEL to a WAN2.2 checkpoint to run this.",
|
|
),
|
|
]
|
|
|
|
_SEED = 1234
|
|
_STEPS = int(os.environ.get("VLLM_OMNI_DEVICE_POSTPROCESS_STEPS", "2"))
|
|
_SIZE = int(os.environ.get("VLLM_OMNI_DEVICE_POSTPROCESS_SIZE", "256"))
|
|
_FRAMES = int(os.environ.get("VLLM_OMNI_DEVICE_POSTPROCESS_FRAMES", "17"))
|
|
|
|
|
|
def _generate(enable_device_postprocess: bool) -> np.ndarray:
|
|
from vllm_omni.entrypoints.omni import Omni
|
|
|
|
engine = Omni(
|
|
model=MODEL,
|
|
num_gpus=1,
|
|
video_output_transport={"enable_device_postprocess": enable_device_postprocess},
|
|
)
|
|
try:
|
|
# output_type must travel with the request: the pipelines read it from the
|
|
# sampling params, and the reduction only covers the "np" path.
|
|
sampling_params = OmniDiffusionSamplingParams(
|
|
output_type="np",
|
|
seed=_SEED,
|
|
num_inference_steps=_STEPS,
|
|
height=_SIZE,
|
|
width=_SIZE,
|
|
num_frames=_FRAMES,
|
|
)
|
|
outputs = engine.generate({"prompt": "a robot waving"}, sampling_params)
|
|
finally:
|
|
engine.close()
|
|
|
|
video = outputs[0].images[0]
|
|
return video.numpy() if hasattr(video, "numpy") else np.asarray(video)
|
|
|
|
|
|
def _as_uint8_frames(video: np.ndarray) -> np.ndarray:
|
|
"""Reproduce what the API server does before encoding.
|
|
|
|
A float payload is scaled and rounded there, so that is the only fair place to
|
|
compare it against a payload the worker already reduced.
|
|
"""
|
|
if video.dtype == np.uint8:
|
|
return video
|
|
return np.rint(np.clip(video, 0.0, 1.0) * 255.0).astype(np.uint8)
|
|
|
|
|
|
# A dtype difference moves ~20% of pixels by one; dropping the round moves ~50%.
|
|
_MAX_DIFFERING_FRACTION = 0.35
|
|
|
|
|
|
@hardware_test(res={"cuda": ["H100", "B200"]}, num_cards=1)
|
|
def test_device_postprocess_produces_equivalent_frames() -> None:
|
|
float_path = _generate(False)
|
|
device_path = _generate(True)
|
|
|
|
# Guard against the run silently taking the float path twice, which would make
|
|
# the comparison meaningless.
|
|
assert float_path.dtype != np.uint8, "the float path returned uint8; the gate may be stuck open"
|
|
assert device_path.dtype == np.uint8, "the WAN typed-media path did not prepare uint8 frames"
|
|
assert float_path.shape == device_path.shape
|
|
|
|
expected = _as_uint8_frames(float_path)
|
|
deviation = np.abs(device_path.astype(np.int16) - expected.astype(np.int16))
|
|
differing = float((deviation > 0).mean())
|
|
print(f"maxdiff={deviation.max()} differing={differing:.4f} ({int((deviation > 0).sum())}/{deviation.size})")
|
|
|
|
assert deviation.max() <= 1, "the reduction moved a pixel by more than one 255th"
|
|
assert differing < _MAX_DIFFERING_FRACTION, (
|
|
f"{differing:.1%} of pixels differ, more than a denormalize-precision difference explains"
|
|
)
|