wangyuandHongsheng Liu
7619ed6b82
Signed-off-by: [Your Name] <your.email@example.com> Signed-off-by: wangyu <410167048@qq.com> Signed-off-by: wangyu <53896905+yenuo26@users.noreply.github.com> Signed-off-by: Hongsheng Liu <liuhongsheng4@huawei.com> Signed-off-by: [Your Name] <[Your Email]> Co-authored-by: Hongsheng Liu <liuhongsheng4@huawei.com>
37 lines
1.0 KiB
Python
37 lines
1.0 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
|
|
|
"""Tests for Flux1 Schnell."""
|
|
|
|
import pytest
|
|
|
|
from tests.helpers.mark import hardware_test
|
|
from tests.helpers.runtime import OmniRunnerHandler
|
|
from vllm_omni.inputs.data import OmniDiffusionSamplingParams
|
|
|
|
MODEL = "black-forest-labs/FLUX.1-schnell"
|
|
|
|
_OMNI_RUNNER_PARAM = (MODEL, None)
|
|
|
|
pytestmark = [
|
|
pytest.mark.slow,
|
|
pytest.mark.diffusion,
|
|
pytest.mark.parametrize("omni_runner", [_OMNI_RUNNER_PARAM], indirect=True),
|
|
]
|
|
|
|
|
|
@hardware_test(res={"cuda": "H100"}, num_cards=1)
|
|
def test_flux_schnell_text_to_image(omni_runner_handler: OmniRunnerHandler):
|
|
"""Test FLUX.1-schnell text-to-image generation."""
|
|
request_config = {
|
|
"model": MODEL,
|
|
"prompt": "A photo of a cat sitting on a laptop",
|
|
"sampling_params": OmniDiffusionSamplingParams(
|
|
height=512,
|
|
width=512,
|
|
num_inference_steps=2,
|
|
seed=42,
|
|
),
|
|
}
|
|
omni_runner_handler.send_diffusion_request(request_config)
|