104 lines
3.4 KiB
Python
104 lines
3.4 KiB
Python
from __future__ import annotations
|
|
|
|
from io import BytesIO
|
|
from pathlib import Path
|
|
|
|
import pytest
|
|
import requests
|
|
from PIL import Image
|
|
|
|
|
|
def pytest_addoption(parser):
|
|
group = parser.getgroup("accuracy-e2e")
|
|
group.addoption(
|
|
"--wan22-i2v-image-source",
|
|
action="store",
|
|
default=None,
|
|
help="Image source for Wan2.2 I2V accuracy tests. Can be local path or remote URL",
|
|
)
|
|
group.addoption(
|
|
"--wan22-i2v-online-timeout-seconds",
|
|
action="store",
|
|
type=int,
|
|
default=1200,
|
|
help="Online serving timeout in seconds for Wan2.2 I2V accuracy tests.",
|
|
)
|
|
group.addoption(
|
|
"--hunyuanvideo15-i2v-image-source",
|
|
action="store",
|
|
default=None,
|
|
help="Image source for HunyuanVideo-1.5 I2V accuracy tests. Can be local path or remote URL",
|
|
)
|
|
group.addoption(
|
|
"--hunyuanvideo15-online-timeout-seconds",
|
|
action="store",
|
|
type=int,
|
|
default=3600,
|
|
help="Online serving timeout in seconds for HunyuanVideo-1.5 accuracy tests.",
|
|
)
|
|
|
|
|
|
@pytest.fixture(scope="session")
|
|
def wan22_i2v_image_source(request: pytest.FixtureRequest) -> str | None:
|
|
value = request.config.getoption("wan22_i2v_image_source")
|
|
return str(value) if value else None
|
|
|
|
|
|
@pytest.fixture(scope="session")
|
|
def wan22_i2v_online_timeout_seconds(request: pytest.FixtureRequest) -> int:
|
|
return int(request.config.getoption("wan22_i2v_online_timeout_seconds"))
|
|
|
|
|
|
@pytest.fixture(scope="session")
|
|
def hunyuanvideo15_i2v_image_source(request: pytest.FixtureRequest) -> str | None:
|
|
value = request.config.getoption("hunyuanvideo15_i2v_image_source")
|
|
return str(value) if value else None
|
|
|
|
|
|
@pytest.fixture(scope="session")
|
|
def hunyuanvideo15_online_timeout_seconds(request: pytest.FixtureRequest) -> int:
|
|
return int(request.config.getoption("hunyuanvideo15_online_timeout_seconds"))
|
|
|
|
|
|
@pytest.fixture(scope="session")
|
|
def accuracy_artifact_root() -> Path:
|
|
root = Path(__file__).resolve().parent / "artifacts"
|
|
root.mkdir(parents=True, exist_ok=True)
|
|
return root
|
|
|
|
|
|
@pytest.fixture(scope="session")
|
|
def qwen_bear_image(accuracy_artifact_root: Path):
|
|
"""Download the Qwen bear image from the URL and save it to the accuracy artifact root."""
|
|
QWEN_BEAR_IMAGE_URL = "https://vllm-public-assets.s3.us-west-2.amazonaws.com/omni-assets/qwen-bear.png"
|
|
image_path = accuracy_artifact_root / "qwen_bear.png"
|
|
if image_path.exists():
|
|
image = Image.open(image_path).convert("RGB")
|
|
yield image
|
|
image.close()
|
|
return
|
|
response = requests.get(QWEN_BEAR_IMAGE_URL, timeout=60)
|
|
response.raise_for_status()
|
|
image = Image.open(BytesIO(response.content)).convert("RGB")
|
|
image.save(image_path)
|
|
yield image
|
|
image.close()
|
|
|
|
|
|
@pytest.fixture(scope="session")
|
|
def rabbit_image(accuracy_artifact_root: Path):
|
|
"""Download the rabbit image from the URL and save it to the accuracy artifact root."""
|
|
RABBIT_IMAGE_URL = "https://vllm-public-assets.s3.us-west-2.amazonaws.com/omni-assets/rabbit.png"
|
|
image_path = accuracy_artifact_root / "rabbit.png"
|
|
if image_path.exists():
|
|
image = Image.open(image_path).convert("RGB")
|
|
yield image
|
|
image.close()
|
|
return
|
|
response = requests.get(RABBIT_IMAGE_URL, timeout=60)
|
|
response.raise_for_status()
|
|
image = Image.open(BytesIO(response.content)).convert("RGB")
|
|
image.save(image_path)
|
|
yield image
|
|
image.close()
|