wangyuandHuang, Zeyu
6f2ad7b403
Signed-off-by: wangyu <410167048@qq.com> Signed-off-by: Huang, Zeyu <11222265+fhfuih@users.noreply.github.com> Signed-off-by: wangyu <53896905+yenuo26@users.noreply.github.com> Co-authored-by: Huang, Zeyu <11222265+fhfuih@users.noreply.github.com>
34 lines
1.1 KiB
Python
34 lines
1.1 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
|
|
|
import pytest
|
|
|
|
from tests.helpers.mark import hardware_test
|
|
from tests.helpers.runtime import OmniRunnerHandler
|
|
from vllm_omni.inputs.data import OmniDiffusionSamplingParams
|
|
|
|
MODEL = "Tongyi-MAI/Z-Image-Turbo"
|
|
PROMPT = "A high-detail studio photo of an orange tabby cat sitting on a laptop keyboard."
|
|
|
|
|
|
@pytest.mark.advanced_model
|
|
@pytest.mark.diffusion
|
|
@hardware_test(res={"cuda": "L4", "rocm": "MI325", "xpu": "B60"}, num_cards={"cuda": 1, "rocm": 1, "xpu": 2})
|
|
@pytest.mark.parametrize("omni_runner", [(MODEL, None)], indirect=True)
|
|
def test_zimage(omni_runner_handler: OmniRunnerHandler):
|
|
# high resolution may cause OOM on L4
|
|
sampling = OmniDiffusionSamplingParams(
|
|
height=256,
|
|
width=256,
|
|
num_inference_steps=2,
|
|
guidance_scale=0.0,
|
|
seed=42,
|
|
num_outputs_per_prompt=2,
|
|
)
|
|
request_config = {
|
|
"model": MODEL,
|
|
"prompt": PROMPT,
|
|
"sampling_params": sampling,
|
|
}
|
|
omni_runner_handler.send_diffusion_request(request_config)
|