[CI] add qwen image and layered accuracy test (#2772)

Signed-off-by: david6666666 <530634352@qq.com>

tests/e2e/accuracy/test_qwen_image.py

@@ -0,0 +1,124 @@
+from __future__ import annotations
+
+import base64
+import gc
+import io
+import os
+from pathlib import Path
+
+import pytest
+import requests
+import torch
+from diffusers.pipelines.pipeline_utils import DiffusionPipeline
+from PIL import Image
+
+from tests.conftest import (
+    OmniServer,
+    _run_post_test_cleanup,
+    _run_pre_test_cleanup,
+)
+from tests.e2e.accuracy.utils import assert_similarity, model_output_dir
+from tests.utils import hardware_test
+
+MODEL_ID = "Qwen/Qwen-Image"
+MODEL_ENV_VAR = "QWEN_IMAGE_MODEL"
+PROMPT = "A photo of a cat sitting on a laptop keyboard, digital art style."
+NEGATIVE_PROMPT = "blurry, low quality"
+WIDTH = 512
+HEIGHT = 512
+NUM_INFERENCE_STEPS = 20
+TRUE_CFG_SCALE = 4.0
+SEED = 42
+SSIM_THRESHOLD = 0.97
+PSNR_THRESHOLD = 30.0
+
+
+def _model_name() -> str:
+    return os.environ.get(MODEL_ENV_VAR, MODEL_ID)
+
+
+def _local_files_only(model: str) -> bool:
+    return Path(model).exists()
+
+
+def _run_vllm_omni_qwen_image(*, model: str, output_path: Path) -> Image.Image:
+    server_args = ["--num-gpus", "1", "--stage-init-timeout", "300", "--init-timeout", "900"]
+    with OmniServer(model, server_args, use_omni=True) as omni_server:
+        response = requests.post(
+            f"http://{omni_server.host}:{omni_server.port}/v1/images/generations",
+            json={
+                "model": omni_server.model,
+                "prompt": PROMPT,
+                "size": f"{WIDTH}x{HEIGHT}",
+                "n": 1,
+                "response_format": "b64_json",
+                "negative_prompt": NEGATIVE_PROMPT,
+                "num_inference_steps": NUM_INFERENCE_STEPS,
+                "true_cfg_scale": TRUE_CFG_SCALE,
+                "seed": SEED,
+            },
+            timeout=600,
+        )
+        response.raise_for_status()
+        payload = response.json()
+        assert len(payload["data"]) == 1
+        image_bytes = base64.b64decode(payload["data"][0]["b64_json"])
+        image = Image.open(io.BytesIO(image_bytes)).convert("RGB")
+        image.load()
+        image.save(output_path)
+        return image
+
+
+def _run_diffusers_qwen_image(*, model: str, output_path: Path) -> Image.Image:
+    _run_pre_test_cleanup(enable_force=True)
+    pipe: DiffusionPipeline | None = None
+    try:
+        pipe = DiffusionPipeline.from_pretrained(
+            model,
+            torch_dtype=torch.bfloat16,
+            trust_remote_code=True,
+            local_files_only=_local_files_only(model),
+        ).to("cuda")
+        generator = torch.Generator(device="cuda").manual_seed(SEED)
+        result = pipe(  # pyright: ignore[reportCallIssue]
+            prompt=PROMPT,
+            negative_prompt=NEGATIVE_PROMPT,
+            width=WIDTH,
+            height=HEIGHT,
+            num_inference_steps=NUM_INFERENCE_STEPS,
+            true_cfg_scale=TRUE_CFG_SCALE,
+            generator=generator,
+        )
+        output_image = result.images[0].convert("RGB")
+        output_image.save(output_path)
+        return output_image
+    finally:
+        if pipe is not None and hasattr(pipe, "maybe_free_model_hooks"):
+            pipe.maybe_free_model_hooks()
+        del pipe
+        gc.collect()
+        if torch.cuda.is_available():
+            torch.cuda.empty_cache()
+        _run_post_test_cleanup(enable_force=True)
+
+
+@pytest.mark.advanced_model
+@pytest.mark.benchmark
+@pytest.mark.diffusion
+@hardware_test(res={"cuda": "H100"}, num_cards=1)
+def test_qwen_image_matches_diffusers(accuracy_artifact_root: Path) -> None:
+    model = _model_name()
+    output_dir = model_output_dir(accuracy_artifact_root, MODEL_ID)
+
+    vllm_output = _run_vllm_omni_qwen_image(model=model, output_path=output_dir / "vllm_omni.png")
+    diffusers_output = _run_diffusers_qwen_image(model=model, output_path=output_dir / "diffusers.png")
+
+    assert_similarity(
+        model_name=MODEL_ID,
+        vllm_image=vllm_output,
+        diffusers_image=diffusers_output,
+        width=WIDTH,
+        height=HEIGHT,
+        ssim_threshold=SSIM_THRESHOLD,
+        psnr_threshold=PSNR_THRESHOLD,
+    )

tests/e2e/accuracy/test_qwen_image_layered.py

@@ -0,0 +1,151 @@
+from __future__ import annotations
+
+import base64
+import gc
+import io
+import os
+from pathlib import Path
+
+import pytest
+import requests
+import torch
+from diffusers.pipelines.pipeline_utils import DiffusionPipeline
+from PIL import Image
+
+from tests.conftest import (
+    OmniServer,
+    _run_post_test_cleanup,
+    _run_pre_test_cleanup,
+)
+from tests.e2e.accuracy.utils import assert_image_sequence_similarity, model_output_dir
+from tests.utils import hardware_test
+
+MODEL_ID = "Qwen/Qwen-Image-Layered"
+MODEL_ENV_VAR = "QWEN_IMAGE_LAYERED_MODEL"
+PROMPT = "decompose into layers"
+NEGATIVE_PROMPT = " "
+NUM_INFERENCE_STEPS = 20
+TRUE_CFG_SCALE = 4.0
+SEED = 777
+LAYERS = 3
+RESOLUTION = 640
+SSIM_THRESHOLD = 0.97
+PSNR_THRESHOLD = 30.0
+
+
+def _model_name() -> str:
+    return os.environ.get(MODEL_ENV_VAR, MODEL_ID)
+
+
+def _local_files_only(model: str) -> bool:
+    return Path(model).exists()
+
+
+def _normalize_layered_images(images: object) -> list[Image.Image]:
+    if not isinstance(images, list) or not images:
+        raise AssertionError(f"Unexpected layered output container: {type(images).__name__}")
+
+    first_item = images[0]
+    if isinstance(first_item, Image.Image):
+        return [image.convert("RGBA") for image in images if isinstance(image, Image.Image)]
+    if isinstance(first_item, (list, tuple)):
+        return [image.convert("RGBA") for image in first_item if isinstance(image, Image.Image)]
+    raise AssertionError(f"Unexpected layered image element type: {type(first_item).__name__}")
+
+
+def _run_vllm_omni_qwen_image_layered(*, model: str, input_image: Image.Image, output_dir: Path) -> list[Image.Image]:
+    input_image.save(output_dir / "input.png")
+    server_args = ["--num-gpus", "1", "--stage-init-timeout", "300", "--init-timeout", "900"]
+    with OmniServer(model, server_args, use_omni=True) as omni_server:
+        buffer = io.BytesIO()
+        input_image.save(buffer, format="PNG")
+        buffer.seek(0)
+        response = requests.post(
+            f"http://{omni_server.host}:{omni_server.port}/v1/images/edits",
+            data={
+                "model": omni_server.model,
+                "prompt": PROMPT,
+                "size": "auto",
+                "n": 1,
+                "response_format": "b64_json",
+                "negative_prompt": NEGATIVE_PROMPT,
+                "num_inference_steps": NUM_INFERENCE_STEPS,
+                "true_cfg_scale": TRUE_CFG_SCALE,
+                "seed": SEED,
+                "layers": LAYERS,
+                "resolution": RESOLUTION,
+            },
+            files=[("image", ("input.png", buffer, "image/png"))],
+            timeout=600,
+        )
+        response.raise_for_status()
+        payload = response.json()
+        assert len(payload["data"]) == LAYERS
+        output_images = []
+        for item in payload["data"]:
+            image_bytes = base64.b64decode(item["b64_json"])
+            image = Image.open(io.BytesIO(image_bytes)).convert("RGBA")
+            image.load()
+            output_images.append(image)
+        for index, image in enumerate(output_images, start=1):
+            image.save(output_dir / f"vllm_omni_layer_{index}.png")
+        return output_images
+
+
+def _run_diffusers_qwen_image_layered(*, model: str, input_image: Image.Image, output_dir: Path) -> list[Image.Image]:
+    _run_pre_test_cleanup(enable_force=True)
+    pipe: DiffusionPipeline | None = None
+    try:
+        pipe = DiffusionPipeline.from_pretrained(
+            model,
+            torch_dtype=torch.bfloat16,
+            trust_remote_code=True,
+            local_files_only=_local_files_only(model),
+        ).to("cuda")
+        generator = torch.Generator(device="cuda").manual_seed(SEED)
+        result = pipe(  # pyright: ignore[reportCallIssue]
+            image=input_image,
+            prompt=PROMPT,
+            negative_prompt=NEGATIVE_PROMPT,
+            num_inference_steps=NUM_INFERENCE_STEPS,
+            true_cfg_scale=TRUE_CFG_SCALE,
+            generator=generator,
+            num_images_per_prompt=1,
+            layers=LAYERS,
+            resolution=RESOLUTION,
+        )
+        output_images = _normalize_layered_images(result.images)
+        assert len(output_images) == LAYERS, f"Expected {LAYERS} diffusers layers, got {len(output_images)}"
+        for index, image in enumerate(output_images, start=1):
+            image.save(output_dir / f"diffusers_layer_{index}.png")
+        return output_images
+    finally:
+        if pipe is not None and hasattr(pipe, "maybe_free_model_hooks"):
+            pipe.maybe_free_model_hooks()
+        del pipe
+        gc.collect()
+        if torch.cuda.is_available():
+            torch.cuda.empty_cache()
+        _run_post_test_cleanup(enable_force=True)
+
+
+@pytest.mark.advanced_model
+@pytest.mark.benchmark
+@pytest.mark.diffusion
+@hardware_test(res={"cuda": "H100"}, num_cards=1)
+def test_qwen_image_layered_matches_diffusers(accuracy_artifact_root: Path, qwen_bear_image: Image.Image) -> None:
+    model = _model_name()
+    output_dir = model_output_dir(accuracy_artifact_root, MODEL_ID)
+    input_image = qwen_bear_image.convert("RGBA")
+
+    vllm_outputs = _run_vllm_omni_qwen_image_layered(model=model, input_image=input_image, output_dir=output_dir)
+    diffusers_outputs = _run_diffusers_qwen_image_layered(model=model, input_image=input_image, output_dir=output_dir)
+
+    assert_image_sequence_similarity(
+        model_name=MODEL_ID,
+        vllm_images=vllm_outputs,
+        diffusers_images=diffusers_outputs,
+        ssim_threshold=SSIM_THRESHOLD,
+        psnr_threshold=PSNR_THRESHOLD,
+        compare_mode="RGBA",
+    )

tests/e2e/accuracy/utils.py

@@ -21,13 +21,14 @@ def assert_similarity(
     model_name: str,
     vllm_image: Image.Image,
     diffusers_image: Image.Image,
-    width: int,
-    height: int,
     ssim_threshold: float,
     psnr_threshold: float,
+    width: int | None = None,
+    height: int | None = None,
+    compare_mode: str = "RGB",
 ) -> None:
-    requested_size = (width, height)
-    if diffusers_image.size != requested_size:
+    requested_size = (width, height) if width is not None and height is not None else None
+    if requested_size is not None and diffusers_image.size != requested_size:
         pytest.skip(
             "Skipping as diffusers baseline output is corrupt and not comparable: "
             f"dimensions do not match requested size; requested={requested_size}, got={diffusers_image.size}."
@@ -37,7 +38,11 @@ def assert_similarity(
         f"Online and diffusers output sizes mismatch: online={vllm_image.size}, diffusers={diffusers_image.size}"
     )
 
-    ssim_score, psnr_score = compute_image_ssim_psnr(prediction=vllm_image, reference=diffusers_image)
+    ssim_score, psnr_score = compute_image_ssim_psnr(
+        prediction=vllm_image,
+        reference=diffusers_image,
+        compare_mode=compare_mode,
+    )
     print(f"{model_name} similarity metrics:")
     print(f"  SSIM: value={ssim_score:.6f}, threshold>={ssim_threshold:.6f}, range=[-1, 1], higher_is_better=True")
     print(
@@ -52,13 +57,37 @@ def assert_similarity(
     )
 
 
+def assert_image_sequence_similarity(
+    *,
+    model_name: str,
+    vllm_images: list[Image.Image],
+    diffusers_images: list[Image.Image],
+    ssim_threshold: float,
+    psnr_threshold: float,
+    compare_mode: str = "RGB",
+) -> None:
+    assert len(vllm_images) == len(diffusers_images), (
+        f"Output image count mismatch for {model_name}: online={len(vllm_images)}, diffusers={len(diffusers_images)}"
+    )
+    for index, (vllm_image, diffusers_image) in enumerate(zip(vllm_images, diffusers_images, strict=True), start=1):
+        assert_similarity(
+            model_name=f"{model_name}[layer={index}]",
+            vllm_image=vllm_image,
+            diffusers_image=diffusers_image,
+            ssim_threshold=ssim_threshold,
+            psnr_threshold=psnr_threshold,
+            compare_mode=compare_mode,
+        )
+
+
 def compute_image_ssim_psnr(
     *,
     prediction: Image.Image,
     reference: Image.Image,
+    compare_mode: str = "RGB",
 ) -> tuple[float, float]:
-    pred_tensor = _pil_to_batched_tensor(prediction)
-    ref_tensor = _pil_to_batched_tensor(reference)
+    pred_tensor = _pil_to_batched_tensor(prediction, compare_mode=compare_mode)
+    ref_tensor = _pil_to_batched_tensor(reference, compare_mode=compare_mode)
 
     ssim_metric = StructuralSimilarityIndexMeasure(data_range=1.0)
     psnr_metric = PeakSignalNoiseRatio(data_range=1.0)
@@ -68,7 +97,7 @@ def compute_image_ssim_psnr(
     return ssim_value, psnr_value
 
 
-def _pil_to_batched_tensor(image: Image.Image) -> torch.Tensor:
-    array = np.asarray(image.convert("RGB"), dtype=np.float32) / 255.0
+def _pil_to_batched_tensor(image: Image.Image, *, compare_mode: str) -> torch.Tensor:
+    array = np.asarray(image.convert(compare_mode), dtype=np.float32) / 255.0
     tensor = torch.from_numpy(array).permute(2, 0, 1).unsqueeze(0)
     return tensor