[CI] qwen image edit L4 accuracy test (#2761)

.buildkite/test-nightly-diffusion.yml

@@ -375,3 +375,43 @@ steps:
                     hostPath:
                       path: /mnt/hf-cache
                       type: DirectoryOrCreate
+
+      - label: ":full_moon: Diffusion · Qwen-Image · Accuracy Test"
+        key: nightly-qwen-image-accuracy
+        timeout_in_minutes: 180
+        if: *nightly_or_pr_label
+        commands:
+          - pytest -s -v tests/e2e/accuracy/test_qwen_image*.py --run-level advanced_model
+        agents:
+          queue: "mithril-h100-pool"
+        plugins:
+          - kubernetes:
+              podSpec:
+                containers:
+                  - image: 936637512419.dkr.ecr.us-west-2.amazonaws.com/vllm-ci-pull-through-cache/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
+                    resources:
+                      limits:
+                        nvidia.com/gpu: 1
+                    volumeMounts:
+                      - name: devshm
+                        mountPath: /dev/shm
+                      - name: hf-cache
+                        mountPath: /root/.cache/huggingface
+                    env:
+                      - name: HF_HOME
+                        value: /root/.cache/huggingface
+                      - name: HF_TOKEN
+                        valueFrom:
+                          secretKeyRef:
+                            name: hf-token-secret
+                            key: token
+                nodeSelector:
+                  node.kubernetes.io/instance-type: gpu-h100-sxm
+                volumes:
+                  - name: devshm
+                    emptyDir:
+                      medium: Memory
+                  - name: hf-cache
+                    hostPath:
+                      path: /mnt/hf-cache
+                      type: DirectoryOrCreate

pyproject.toml

@@ -55,6 +55,7 @@ dev = [
     "pyttsx3>=2.99",
     "opencc>=1.2.0",
     "mistune>=3.2.0", # for example tests
+    "torchmetrics>=1.4.0", # for accuracy similarity metrics
 ]
 
 demo = [

tests/conftest.py

@@ -2397,7 +2397,7 @@ class OpenAIClientHandler:
                                 image_url = item.get("image_url", {}).get("url")
                             else:
                                 image_url_obj = getattr(item, "image_url", None)
-                                image_url = hasattr(image_url_obj, "url", None) if image_url_obj else None
+                                image_url = getattr(image_url_obj, "url", None) if image_url_obj else None
                             if image_url and image_url.startswith("data:image"):
                                 b64_data = image_url.split(",", 1)[1]
                                 img = decode_b64_image(b64_data)
@@ -2703,7 +2703,7 @@ class OpenAIClientHandler:
 
         return responses
 
-    def send_diffusion_request(self, request_config: dict[str, Any], request_num: int = 1) -> list[OmniResponse]:
+    def send_diffusion_request(self, request_config: dict[str, Any], request_num: int = 1) -> list[DiffusionResponse]:
         """
         Send OpenAI requests for diffusion models.
 
@@ -2711,9 +2711,9 @@ class OpenAIClientHandler:
             request_config: Request configuration dictionary containing parameters like model, messages
             request_num: Number of requests to send concurrently, defaults to 1 (single request)
         Returns:
-            List[OmniResponse]: List of response objects
+            List[DiffusionResponse]: List of response objects
         """
-        responses = []
+        responses: list[DiffusionResponse] = []
         stream = request_config.get("stream", False)
         modalities = request_config.get("modalities", omit)  # Most diffusion models don't require modalities param
         extra_body = request_config.get("extra_body", None)
@@ -2876,7 +2876,7 @@ class OpenAIClientHandler:
         return f"{self.base_url.rstrip('/')}/{path.lstrip('/')}"
 
 
-@pytest.fixture
+@pytest.fixture(scope="module")
 def openai_client(request: pytest.FixtureRequest, run_level: str):
     """Create OpenAIClientHandler fixture to facilitate communication with OmniServer
     with encapsulated request sending, concurrent requests, response handling, and validation."""

tests/e2e/accuracy/conftest.py

@@ -5,10 +5,13 @@ import shutil
 import subprocess
 from contextlib import contextmanager
 from dataclasses import dataclass
+from io import BytesIO
 from pathlib import Path
 
 import pytest
+import requests
 import torch
+from PIL import Image
 
 from tests.conftest import OmniServer, OmniServerParams
 
@@ -183,6 +186,28 @@ def accuracy_artifact_root() -> Path:
     return root
 
 
+@pytest.fixture(scope="session")
+def qwen_bear_image(accuracy_artifact_root: Path) -> Image.Image:
+    """Download the Qwen bear image from the URL and save it to the accuracy artifact root."""
+    QWEN_BEAR_IMAGE_URL = "https://vllm-public-assets.s3.us-west-2.amazonaws.com/omni-assets/qwen-bear.png"
+    response = requests.get(QWEN_BEAR_IMAGE_URL, timeout=60)
+    response.raise_for_status()
+    image = Image.open(BytesIO(response.content)).convert("RGB")
+    image.save(accuracy_artifact_root / "qwen_bear.png")
+    return image
+
+
+@pytest.fixture(scope="session")
+def rabbit_image(accuracy_artifact_root: Path) -> Image.Image:
+    """Download the rabbit image from the URL and save it to the accuracy artifact root."""
+    RABBIT_IMAGE_URL = "https://vllm-public-assets.s3.us-west-2.amazonaws.com/omni-assets/rabbit.png"
+    response = requests.get(RABBIT_IMAGE_URL, timeout=60)
+    response.raise_for_status()
+    image = Image.open(BytesIO(response.content)).convert("RGB")
+    image.save(accuracy_artifact_root / "rabbit.png")
+    return image
+
+
 def reset_artifact_dir(path: Path) -> Path:
     if path.exists():
         shutil.rmtree(path)

tests/e2e/accuracy/test_qwen_image_edit.py

@@ -0,0 +1,232 @@
+from __future__ import annotations
+
+import gc
+from pathlib import Path
+
+import pytest
+import requests
+import torch
+from diffusers import QwenImageEditPipeline, QwenImageEditPlusPipeline
+from PIL import Image
+
+from benchmarks.accuracy.common import decode_base64_image, pil_to_png_bytes
+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
+
+SINGLE_MODEL = "Qwen/Qwen-Image-Edit"
+MULTIPLE_MODEL = "Qwen/Qwen-Image-Edit-2509"
+WIDTH = 512
+HEIGHT = 512
+NUM_INFERENCE_STEPS = 20
+TRUE_CFG_SCALE = 4.0
+SEED = 42
+SSIM_THRESHOLD = 0.94
+PSNR_THRESHOLD = 28.0
+
+PROMPT_SINGLE_IMAGE = "The input is a 2D cartoon bear mascot. Restyle it into a painterly oil artwork with warm colors while preserving the main structure."
+PROMPT_MULTIPLE_IMAGE = "Put the cartoon bear mascot and the furry rabbit into one coherent scene with a painterly oil artwork style and consistent lighting."
+NEGATIVE_PROMPT = "low quality, blurry, artifacts, distortion"
+SERVER_ARGS = ["--num-gpus", "1", "--stage-init-timeout", "300", "--init-timeout", "900"]
+
+
+def _run_vllm_omni_image_edit(
+    *,
+    omni_server: OmniServer,
+    prompt: str,
+    input_images: list[Image.Image],
+    output_path: Path,
+) -> Image.Image:
+    response = requests.post(
+        f"http://{omni_server.host}:{omni_server.port}/v1/images/edits",
+        data={
+            "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,
+        },
+        files=[
+            ("image", (f"image_{index}.png", pil_to_png_bytes(image), "image/png"))
+            for index, image in enumerate(input_images)
+        ],
+        timeout=600,
+    )
+    response.raise_for_status()
+    payload = response.json()
+    assert len(payload["data"]) == 1
+    image = decode_base64_image(payload["data"][0]["b64_json"])
+    image.load()
+    image.save(output_path)
+    return image
+
+
+def _run_diffusers_image_edit(
+    *,
+    model: str,
+    pipeline_class: type[QwenImageEditPipeline] | type[QwenImageEditPlusPipeline],
+    prompt: str,
+    input_images: list[Image.Image],
+    output_path: Path,
+) -> Image.Image:
+    _run_pre_test_cleanup(enable_force=True)
+    pipe: QwenImageEditPipeline | QwenImageEditPlusPipeline | None = None
+    device = torch.device("cuda:0")
+    torch.cuda.set_device(device)
+    try:
+        images = input_images[0] if len(input_images) == 1 else input_images
+        pipe = pipeline_class.from_pretrained(
+            model,
+            torch_dtype=torch.bfloat16,
+            trust_remote_code=True,
+        ).to(device)
+        pipe.set_progress_bar_config(disable=False)
+        generator = torch.Generator(device=device).manual_seed(SEED)
+        result = pipe(  # pyright: ignore[reportCallIssue]
+            prompt=prompt,
+            image=images,
+            negative_prompt=NEGATIVE_PROMPT,
+            num_inference_steps=NUM_INFERENCE_STEPS,
+            true_cfg_scale=TRUE_CFG_SCALE,
+            width=WIDTH,
+            height=HEIGHT,
+            generator=generator,
+        )
+        output_image = result.images[0].convert("RGB")  # pyright: ignore[reportAttributeAccessIssue]
+        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)
+
+
+def _vllm_omni_output_single_image(
+    accuracy_artifact_root: Path,
+    qwen_bear_image: Image.Image,
+) -> Image.Image:
+    output_dir = model_output_dir(accuracy_artifact_root, SINGLE_MODEL)
+    output_path = output_dir / "vllm_omni_single.png"
+    with OmniServer(model=SINGLE_MODEL, serve_args=SERVER_ARGS) as server:
+        output = _run_vllm_omni_image_edit(
+            omni_server=server,
+            prompt=PROMPT_SINGLE_IMAGE,
+            input_images=[qwen_bear_image],
+            output_path=output_path,
+        )
+    return output
+
+
+def _diffusers_output_single_image(accuracy_artifact_root: Path, qwen_bear_image: Image.Image) -> Image.Image:
+    output_dir = model_output_dir(accuracy_artifact_root, SINGLE_MODEL)
+    output_path = output_dir / "diffusers_single.png"
+    return _run_diffusers_image_edit(
+        model=SINGLE_MODEL,
+        pipeline_class=QwenImageEditPipeline,
+        prompt=PROMPT_SINGLE_IMAGE,
+        input_images=[qwen_bear_image],
+        output_path=output_path,
+    )
+
+
+def _vllm_omni_output_multiple_image(
+    accuracy_artifact_root: Path,
+    qwen_bear_image: Image.Image,
+    rabbit_image: Image.Image,
+) -> Image.Image:
+    output_dir = model_output_dir(accuracy_artifact_root, MULTIPLE_MODEL)
+    output_path = output_dir / "vllm_omni_multiple.png"
+    with OmniServer(model=MULTIPLE_MODEL, serve_args=SERVER_ARGS) as server:
+        output = _run_vllm_omni_image_edit(
+            omni_server=server,
+            prompt=PROMPT_MULTIPLE_IMAGE,
+            input_images=[qwen_bear_image, rabbit_image],
+            output_path=output_path,
+        )
+    return output
+
+
+def _diffusers_output_multiple_image(
+    accuracy_artifact_root: Path, qwen_bear_image: Image.Image, rabbit_image: Image.Image
+) -> Image.Image:
+    output_dir = model_output_dir(accuracy_artifact_root, MULTIPLE_MODEL)
+    output_path = output_dir / "diffusers_multiple.png"
+    return _run_diffusers_image_edit(
+        model=MULTIPLE_MODEL,
+        pipeline_class=QwenImageEditPlusPipeline,
+        prompt=PROMPT_MULTIPLE_IMAGE,
+        input_images=[qwen_bear_image, rabbit_image],
+        output_path=output_path,
+    )
+
+
+@pytest.mark.advanced_model
+@pytest.mark.benchmark
+@pytest.mark.diffusion
+@hardware_test(res={"cuda": "H100"}, num_cards=1)
+def test_qwen_image_edit_single_matches_diffusers(
+    accuracy_artifact_root: Path,
+    qwen_bear_image: Image.Image,
+) -> None:
+    vllm_image = _vllm_omni_output_single_image(
+        accuracy_artifact_root=accuracy_artifact_root,
+        qwen_bear_image=qwen_bear_image,
+    )
+    diffusers_image = _diffusers_output_single_image(
+        accuracy_artifact_root=accuracy_artifact_root,
+        qwen_bear_image=qwen_bear_image,
+    )
+    assert_similarity(
+        model_name=SINGLE_MODEL,
+        vllm_image=vllm_image,
+        diffusers_image=diffusers_image,
+        width=WIDTH,
+        height=HEIGHT,
+        ssim_threshold=SSIM_THRESHOLD,
+        psnr_threshold=PSNR_THRESHOLD,
+    )
+
+
+@pytest.mark.advanced_model
+@pytest.mark.benchmark
+@pytest.mark.diffusion
+@hardware_test(res={"cuda": "H100"}, num_cards=1)
+@pytest.mark.skip(
+    reason="Skipping as the second image seems to be ignored by the API. Will come back to this later after #2772 is merged."
+)
+def test_qwen_image_edit_multiple_matches_diffusers(
+    accuracy_artifact_root: Path,
+    qwen_bear_image: Image.Image,
+    rabbit_image: Image.Image,
+) -> None:
+    vllm_image = _vllm_omni_output_multiple_image(
+        accuracy_artifact_root=accuracy_artifact_root,
+        qwen_bear_image=qwen_bear_image,
+        rabbit_image=rabbit_image,
+    )
+    diffusers_image = _diffusers_output_multiple_image(
+        accuracy_artifact_root=accuracy_artifact_root,
+        qwen_bear_image=qwen_bear_image,
+        rabbit_image=rabbit_image,
+    )
+    assert_similarity(
+        model_name=MULTIPLE_MODEL,
+        vllm_image=vllm_image,
+        diffusers_image=diffusers_image,
+        width=WIDTH,
+        height=HEIGHT,
+        ssim_threshold=SSIM_THRESHOLD,
+        psnr_threshold=PSNR_THRESHOLD,
+    )

tests/e2e/accuracy/utils.py

@@ -0,0 +1,74 @@
+from __future__ import annotations
+
+from pathlib import Path
+
+import numpy as np
+import pytest
+import torch
+from PIL import Image
+from torchmetrics.image import PeakSignalNoiseRatio, StructuralSimilarityIndexMeasure
+
+
+def model_output_dir(parent_dir: Path, model: str) -> Path:
+    safe_model_name = model.split("/")[-1].replace(".", "_")
+    path = parent_dir / safe_model_name
+    path.mkdir(parents=True, exist_ok=True)
+    return path
+
+
+def assert_similarity(
+    *,
+    model_name: str,
+    vllm_image: Image.Image,
+    diffusers_image: Image.Image,
+    width: int,
+    height: int,
+    ssim_threshold: float,
+    psnr_threshold: float,
+) -> None:
+    requested_size = (width, height)
+    if 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}."
+        )
+
+    assert vllm_image.size == diffusers_image.size, (
+        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)
+    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(
+        f"  PSNR: value={psnr_score:.6f} dB, threshold>={psnr_threshold:.6f} dB, range=[0, +inf), higher_is_better=True"
+    )
+
+    assert ssim_score >= ssim_threshold, (
+        f"SSIM below threshold for {model_name}: got {ssim_score:.6f}, expected >= {ssim_threshold:.6f}."
+    )
+    assert psnr_score >= psnr_threshold, (
+        f"PSNR below threshold for {model_name}: got {psnr_score:.6f}, expected >= {psnr_threshold:.6f}."
+    )
+
+
+def compute_image_ssim_psnr(
+    *,
+    prediction: Image.Image,
+    reference: Image.Image,
+) -> tuple[float, float]:
+    pred_tensor = _pil_to_batched_tensor(prediction)
+    ref_tensor = _pil_to_batched_tensor(reference)
+
+    ssim_metric = StructuralSimilarityIndexMeasure(data_range=1.0)
+    psnr_metric = PeakSignalNoiseRatio(data_range=1.0)
+
+    ssim_value = float(ssim_metric(pred_tensor, ref_tensor).item())
+    psnr_value = float(psnr_metric(pred_tensor, ref_tensor).item())
+    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
+    tensor = torch.from_numpy(array).permute(2, 0, 1).unsqueeze(0)
+    return tensor