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>
132 lines
4.0 KiB
Python
132 lines
4.0 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
|
|
|
|
from __future__ import annotations
|
|
|
|
import base64
|
|
import io
|
|
import json
|
|
import os
|
|
from pathlib import Path
|
|
|
|
import pytest
|
|
import requests
|
|
import torch
|
|
from PIL import Image
|
|
|
|
from tests.e2e.accuracy.helpers import model_output_dir
|
|
from tests.helpers.mark import hardware_test
|
|
from tests.helpers.runtime import OmniServer
|
|
|
|
pytestmark = [pytest.mark.full_model, pytest.mark.diffusion]
|
|
|
|
MODEL_ENV_VAR = "VLLM_TEST_COSMOS3_MODEL"
|
|
MODEL_ID = "cosmos3"
|
|
PROMPT = "A small warehouse robot moves a blue box across a clean floor."
|
|
NEGATIVE_PROMPT = "blurry, distorted, low quality"
|
|
SEED = 42
|
|
WIDTH = HEIGHT = 256
|
|
NUM_INFERENCE_STEPS = 2
|
|
|
|
|
|
def _model_name() -> str:
|
|
model = os.environ.get(MODEL_ENV_VAR)
|
|
if not model:
|
|
pytest.skip(f"Set {MODEL_ENV_VAR} to run Cosmos3 full-model smoke tests.")
|
|
if not torch.cuda.is_available():
|
|
pytest.skip("Cosmos3 full-model smoke tests require CUDA.")
|
|
assert model
|
|
return model
|
|
|
|
|
|
def _server_args() -> list[str]:
|
|
return [
|
|
"--num-gpus",
|
|
"1",
|
|
"--stage-init-timeout",
|
|
"900",
|
|
"--init-timeout",
|
|
"1200",
|
|
]
|
|
|
|
|
|
def _image_data_url(image: Image.Image) -> str:
|
|
buf = io.BytesIO()
|
|
image.save(buf, format="PNG")
|
|
return f"data:image/png;base64,{base64.b64encode(buf.getvalue()).decode('ascii')}"
|
|
|
|
|
|
@pytest.mark.benchmark
|
|
@hardware_test(res={"cuda": ["H100", "B200"]}, num_cards=1)
|
|
def test_cosmos3_t2i_serving_smoke(accuracy_artifact_root: Path) -> None:
|
|
output_dir = model_output_dir(accuracy_artifact_root, MODEL_ID)
|
|
with OmniServer(_model_name(), _server_args(), use_omni=True) as server:
|
|
response = requests.post(
|
|
f"http://{server.host}:{server.port}/v1/images/generations",
|
|
json={
|
|
"model": server.model,
|
|
"prompt": PROMPT,
|
|
"negative_prompt": NEGATIVE_PROMPT,
|
|
"size": f"{WIDTH}x{HEIGHT}",
|
|
"n": 1,
|
|
"response_format": "b64_json",
|
|
"num_inference_steps": NUM_INFERENCE_STEPS,
|
|
"guidance_scale": 1.0,
|
|
"seed": SEED,
|
|
},
|
|
timeout=1800,
|
|
)
|
|
|
|
response.raise_for_status()
|
|
payload = response.json()
|
|
assert len(payload["data"]) == 1
|
|
image = Image.open(io.BytesIO(base64.b64decode(payload["data"][0]["b64_json"]))).convert("RGB")
|
|
image.save(output_dir / "cosmos3_t2i.png")
|
|
assert image.size == (WIDTH, HEIGHT)
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
("name", "prompt", "num_frames", "image_reference"),
|
|
[
|
|
("t2v", PROMPT, "1", None),
|
|
(
|
|
"i2v",
|
|
"The blue rectangle moves slowly forward.",
|
|
"5",
|
|
Image.new("RGB", (96, 64), color=(40, 80, 160)),
|
|
),
|
|
],
|
|
)
|
|
@pytest.mark.benchmark
|
|
@hardware_test(res={"cuda": ["H100", "B200"]}, num_cards=1)
|
|
def test_cosmos3_video_serving_smoke(
|
|
accuracy_artifact_root: Path,
|
|
name: str,
|
|
prompt: str,
|
|
num_frames: str,
|
|
image_reference: Image.Image | None,
|
|
) -> None:
|
|
output_dir = model_output_dir(accuracy_artifact_root, MODEL_ID)
|
|
data = {
|
|
"model": "",
|
|
"prompt": prompt,
|
|
"negative_prompt": NEGATIVE_PROMPT,
|
|
"size": f"{WIDTH}x{HEIGHT}",
|
|
"num_frames": num_frames,
|
|
"fps": "1",
|
|
"num_inference_steps": str(NUM_INFERENCE_STEPS),
|
|
"guidance_scale": "1.0",
|
|
"seed": str(SEED),
|
|
}
|
|
if image_reference is not None:
|
|
data["image_reference"] = json.dumps({"image_url": _image_data_url(image_reference)})
|
|
|
|
with OmniServer(_model_name(), _server_args(), use_omni=True) as server:
|
|
data["model"] = server.model
|
|
response = requests.post(f"http://{server.host}:{server.port}/v1/videos/sync", data=data, timeout=1800)
|
|
|
|
response.raise_for_status()
|
|
assert response.headers["content-type"].startswith("video/mp4")
|
|
assert response.content
|
|
(output_dir / f"cosmos3_{name}.mp4").write_bytes(response.content)
|