76 lines
2.7 KiB
Python
76 lines
2.7 KiB
Python
import pytest
|
|
import torch
|
|
|
|
from tests.helpers.mark import hardware_test
|
|
from vllm_omni.inputs.data import OmniDiffusionSamplingParams
|
|
from vllm_omni.outputs import OmniRequestOutput
|
|
from vllm_omni.platforms import current_omni_platform
|
|
|
|
# Match unprefixed HF id even when MODEL_PREFIX is set (omni_runner resolves full path).
|
|
_QWEN_IMAGE_RANDOM_ID = "riverclouds/qwen_image_random"
|
|
|
|
|
|
def _is_qwen_image_random(model_path: str) -> bool:
|
|
return model_path.rstrip("/").endswith(_QWEN_IMAGE_RANDOM_ID)
|
|
|
|
|
|
models = ["Tongyi-MAI/Z-Image-Turbo", "riverclouds/qwen_image_random"]
|
|
|
|
# Modelscope can't find riverclouds/qwen_image_random
|
|
# TODO: When NPU support is ready, remove this branch.
|
|
if current_omni_platform.is_npu():
|
|
models = ["Tongyi-MAI/Z-Image-Turbo", "Qwen/Qwen-Image"]
|
|
|
|
# omni_runner expects (model, stage_configs_path); single-stage diffusion has no YAML.
|
|
test_params = [(m, None) for m in models]
|
|
|
|
|
|
@pytest.mark.core_model
|
|
@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", test_params, indirect=True)
|
|
def test_diffusion_model(omni_runner, run_level):
|
|
resolved = omni_runner.model_name
|
|
if run_level == "core_model" and not _is_qwen_image_random(resolved):
|
|
pytest.skip()
|
|
|
|
if run_level == "advanced_model" and _is_qwen_image_random(resolved):
|
|
pytest.skip()
|
|
|
|
# high resolution may cause OOM on L4
|
|
height = 256
|
|
width = 256
|
|
sampling = OmniDiffusionSamplingParams(
|
|
height=height,
|
|
width=width,
|
|
num_inference_steps=2,
|
|
guidance_scale=0.0,
|
|
generator=torch.Generator(current_omni_platform.device_type).manual_seed(42),
|
|
num_outputs_per_prompt=2,
|
|
)
|
|
|
|
# OmniRunner.generate() is typed for list[TextPrompt]; diffusion uses Omni.generate(str, ...).
|
|
outputs = omni_runner.omni.generate(
|
|
"a photo of a cat sitting on a laptop keyboard",
|
|
sampling,
|
|
)
|
|
|
|
# Extract images from request_output['images']
|
|
first_output = outputs[0]
|
|
assert first_output.final_output_type == "image"
|
|
if not hasattr(first_output, "request_output") or not first_output.request_output:
|
|
raise ValueError("No request_output found in OmniRequestOutput")
|
|
|
|
req_out = first_output.request_output
|
|
if not isinstance(req_out, OmniRequestOutput) or not hasattr(req_out, "images"):
|
|
raise ValueError("Invalid request_output structure or missing 'images' key")
|
|
|
|
images = req_out.images
|
|
|
|
assert len(images) == 2
|
|
# check image size
|
|
assert images[0].width == width
|
|
assert images[0].height == height
|
|
images[0].save("image_output.png")
|