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>
53 lines
1.4 KiB
Python
53 lines
1.4 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
|
|
|
"""
|
|
End-to-end test for ErnieImage text-to-image generation.
|
|
|
|
Equivalent to running:
|
|
vllm serve baidu/ERNIE-Image --omni
|
|
"""
|
|
|
|
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 = "baidu/ERNIE-Image"
|
|
|
|
# (model, stage_configs_path, extra_omni_kwargs) for ``omni_runner`` indirect parametrize
|
|
_OMNI_RUNNER_PARAM = (
|
|
MODEL,
|
|
None,
|
|
{
|
|
"parallel_config": DiffusionParallelConfig(
|
|
tensor_parallel_size=2,
|
|
),
|
|
"enable_cpu_offload": True,
|
|
},
|
|
)
|
|
|
|
pytestmark = [
|
|
pytest.mark.full_model,
|
|
pytest.mark.diffusion,
|
|
pytest.mark.parametrize("omni_runner", [_OMNI_RUNNER_PARAM], indirect=True),
|
|
]
|
|
|
|
|
|
@hardware_test(res={"cuda": "L4"}, num_cards=2)
|
|
def test_ernie_image_text_to_image(omni_runner_handler: OmniRunnerHandler) -> None:
|
|
request_config = {
|
|
"model": omni_runner_handler.runner.model_name,
|
|
"prompt": "A photo of a cat sitting on a laptop",
|
|
"sampling_params": OmniDiffusionSamplingParams(
|
|
height=512,
|
|
width=512,
|
|
num_inference_steps=2,
|
|
guidance_scale=4.0,
|
|
seed=42,
|
|
),
|
|
}
|
|
omni_runner_handler.send_diffusion_request(request_config)
|