vllm-omni/tests/e2e/accuracy/test_qwen_image_edit.py

from __future__ import annotations

import gc
import os
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.e2e.accuracy.helpers import assert_images_pixel_close, assert_similarity, model_output_dir
from tests.helpers.env import run_post_test_cleanup, run_pre_test_cleanup
from tests.helpers.mark import hardware_test
from tests.helpers.media import get_asset_path
from tests.helpers.runtime import OmniServer

pytestmark = [pytest.mark.diffusion]


SINGLE_MODEL = "Qwen/Qwen-Image-Edit"
MULTIPLE_MODEL = "Qwen/Qwen-Image-Edit-2511"
WIDTH = 512
HEIGHT = 512
NUM_INFERENCE_STEPS = 20
TRUE_CFG_SCALE = 4.0
SEED = 42
SSIM_THRESHOLD = 0.94
PSNR_THRESHOLD = 28.0

EDIT_2511_MODEL = "Qwen/Qwen-Image-Edit-2511"
EDIT_2511_MODEL_ENV_VAR = "QWEN_IMAGE_EDIT_2511_MODEL"
EDIT_2511_WIDTH = 1190
EDIT_2511_HEIGHT = 2032
EDIT_2511_NUM_INFERENCE_STEPS = 40
EDIT_2511_TRUE_CFG_SCALE = 4.0
EDIT_2511_GUIDANCE_SCALE = 1.0
EDIT_2511_SEED = 42
EDIT_2511_MEAN_ABS_DIFF_THRESHOLD = 3e-2
EDIT_2511_P99_ABS_DIFF_THRESHOLD = 4e-1
EDIT_2511_NEGATIVE_PROMPT = " "
EDIT_2511_PROMPT = (
    "将第二张图中人脸/猫狗脸转换为3D卡通形象，质量要求：毛发纹理自然电影级色彩校准，注意保留原图的外貌、"
    "肤色、发型特征，仅替换第一张财神爷/猫狗财神的用绿框框住的头部区域，耳朵样式和原图的3D卡通形象一致， "
    "（帽子紧贴头顶，未贴合部分可调整帽子摆放角度达到完全贴合）， 头部与身体过渡必须自然，无拼接痕迹，"
    "无额外元素出现在头部四周，纯白背景，不要出现绿色圆圈"
)
EDIT_2511_IMAGE_PATHS = [
    get_asset_path("qwen_image_edit/qwen_image_edit_2511_test1.png"),
    get_asset_path("qwen_image_edit/qwen_image_edit_2511_test2.png"),
]

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 _edit_2511_model_name() -> str:
    return os.environ.get(EDIT_2511_MODEL_ENV_VAR, EDIT_2511_MODEL)


def _local_files_only(model: str) -> bool:
    return Path(model).exists()


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()
    pipe: QwenImageEditPipeline | QwenImageEditPlusPipeline | None = None
    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("cuda")
        pipe.set_progress_bar_config(disable=False)
        generator = torch.Generator(device="cuda").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.accelerator.empty_cache()
        run_post_test_cleanup()


def _load_edit_2511_input_images() -> list[Image.Image]:
    return [Image.open(path).convert("RGB") for path in EDIT_2511_IMAGE_PATHS]


def _run_vllm_omni_image_edit_2511(
    *,
    model: str,
    input_images: list[Image.Image],
    output_path: Path,
) -> Image.Image:
    with OmniServer(model=model, serve_args=SERVER_ARGS) as server:
        response = requests.post(
            f"http://{server.host}:{server.port}/v1/images/edits",
            data={
                "model": server.model,
                "prompt": EDIT_2511_PROMPT,
                "size": f"{EDIT_2511_WIDTH}x{EDIT_2511_HEIGHT}",
                "n": 1,
                "response_format": "b64_json",
                "negative_prompt": EDIT_2511_NEGATIVE_PROMPT,
                "num_inference_steps": EDIT_2511_NUM_INFERENCE_STEPS,
                "true_cfg_scale": EDIT_2511_TRUE_CFG_SCALE,
                "guidance_scale": EDIT_2511_GUIDANCE_SCALE,
                "seed": EDIT_2511_SEED,
            },
            files=[
                ("image", (f"image_{index}.png", pil_to_png_bytes(image), "image/png"))
                for index, image in enumerate(input_images)
            ],
            timeout=1200,
        )
        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_2511(
    *,
    model: str,
    input_images: list[Image.Image],
    output_path: Path,
) -> Image.Image:
    run_pre_test_cleanup()
    pipe: QwenImageEditPlusPipeline | None = None
    try:
        pipe = QwenImageEditPlusPipeline.from_pretrained(
            model,
            torch_dtype=torch.bfloat16,
            trust_remote_code=True,
            local_files_only=_local_files_only(model),
        ).to("cuda")
        pipe.set_progress_bar_config(disable=False)
        generator = torch.Generator(device="cuda").manual_seed(EDIT_2511_SEED)
        result = pipe(  # pyright: ignore[reportCallIssue]
            image=input_images,
            prompt=EDIT_2511_PROMPT,
            generator=generator,
            true_cfg_scale=EDIT_2511_TRUE_CFG_SCALE,
            negative_prompt=EDIT_2511_NEGATIVE_PROMPT,
            num_inference_steps=EDIT_2511_NUM_INFERENCE_STEPS,
            guidance_scale=EDIT_2511_GUIDANCE_SCALE,
            num_images_per_prompt=1,
            width=EDIT_2511_WIDTH,
            height=EDIT_2511_HEIGHT,
        )
        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.accelerator.empty_cache()
        run_post_test_cleanup()


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,
    )


@pytest.mark.slow
@pytest.mark.benchmark
@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.slow
@pytest.mark.benchmark
@hardware_test(res={"cuda": "H100"}, num_cards=1)
def test_qwen_image_edit_2511_matches_diffusers_pixelwise(accuracy_artifact_root: Path) -> None:
    model = _edit_2511_model_name()
    output_dir = model_output_dir(accuracy_artifact_root, EDIT_2511_MODEL)
    input_images = _load_edit_2511_input_images()
    input_paths: list[Path] = []
    for index, image in enumerate(input_images, start=1):
        input_path = output_dir / f"edit_2511_input_{index}.png"
        image.save(input_path)
        input_paths.append(input_path)
    vllm_output_path = output_dir / "vllm_omni_edit_2511.png"
    diffusers_output_path = output_dir / "diffusers_edit_2511.png"

    vllm_image = _run_vllm_omni_image_edit_2511(
        model=model,
        input_images=input_images,
        output_path=vllm_output_path,
    )
    diffusers_image = _run_diffusers_image_edit_2511(
        model=model,
        input_images=input_images,
        output_path=diffusers_output_path,
    )

    print(f"{EDIT_2511_MODEL} generated images:")
    for index, input_path in enumerate(input_paths, start=1):
        print(f"  input_{index}: {input_path}")
    print(f"  vllm_omni: {vllm_output_path}")
    print(f"  diffusers: {diffusers_output_path}")

    assert_images_pixel_close(
        model_name=EDIT_2511_MODEL,
        vllm_image=vllm_image,
        diffusers_image=diffusers_image,
        mean_threshold=EDIT_2511_MEAN_ABS_DIFF_THRESHOLD,
        p99_threshold=EDIT_2511_P99_ABS_DIFF_THRESHOLD,
    )