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

# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project

"""
Run vLLM-Omni with diffusers backend, then run diffusers directly. Compare the output similarity.
Also, do a quick performance check on the end-to-end latency.
The latency check is not put under /tests/dfx/perf/ because this is not a complete performance check to be monitored.
"""

import base64
import gc
import io
import time
from pathlib import Path
from types import SimpleNamespace
from typing import Any

import numpy as np
import pytest
import requests
import torch
from diffusers.pipelines.pipeline_utils import DiffusionPipeline
from diffusers.utils import export_to_video  # pyright: ignore[reportPrivateImportUsage]
from PIL import Image

from benchmarks.accuracy.common import pil_to_base64
from tests.e2e.accuracy.helpers import (
    apply_ftfy_mock,
    assert_similarity,
    env_to_apply_ftfy_mock_in_subproc,
    model_output_dir,
    resolve_device_threshold,
)
from tests.e2e.accuracy.helpers import (
    parse_psnr_score as _parse_psnr_score,
)
from tests.e2e.accuracy.helpers import (
    parse_ssim_score as _parse_ssim_score,
)
from tests.e2e.accuracy.helpers import (
    run_ffmpeg_similarity as _run_ffmpeg_similarity,
)
from tests.helpers.clean import cleanup_test_environment
from tests.helpers.mark import hardware_test
from tests.helpers.runtime import OmniServer, OnlineOmniClient
from vllm_omni.diffusion.models.diffusers_adapter.pipeline_diffusers_adapter import (
    CUDA_FLASH_ATTENTION_BACKEND_ATTEMPTS,
)

pytestmark = [pytest.mark.full_model, pytest.mark.diffusion]


def _flash_attempt_backends() -> list[str]:
    """Backends in ``CUDA_FLASH_ATTENTION_BACKEND_ATTEMPTS`` that can actually launch.

    Those names are Hopper FA2/FA3. ``set_attention_backend`` only checks import,
    so a newer GPU still "selects" them and then dies with "no kernel image".
    Do not probe by launching: a missing cubin can leave the CUDA context dead.
    Capability >= 10 is outside that wheel, including future SKUs in that family.
    """
    if not torch.cuda.is_available():
        return []
    major, _ = torch.cuda.get_device_capability()
    if major >= 10:
        return []
    return list(CUDA_FLASH_ATTENTION_BACKEND_ATTEMPTS)


def _set_matched_attention_backend(pipe: DiffusionPipeline) -> None:
    """Pick a Diffusers attention backend the reference run can actually execute."""
    for backend in _flash_attempt_backends():
        try:
            pipe.transformer.set_attention_backend(backend)
            return
        except Exception:
            continue
    pipe.transformer.set_attention_backend("native")


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

# A single-shot latency measurement on this benchmark is dominated by CI variance
# (see issue #5108). Warm up once, then report the median of several timed runs.
NUM_TIMED_RUNS = 2

VIDEO_PROMPT = "The bear in the image dances happily"
VIDEO_WIDTH = 832
VIDEO_HEIGHT = 480
VIDEO_NUM_INFERENCE_STEPS = 10
NUM_FRAMES = 9
FPS = 4
GUIDANCE_SCALE = 4.0
GUIDANCE_SCALE_2 = 1.0
BOUNDARY_RATIO = 0.875
FLOW_SHIFT = 12.0
VIDEO_SSIM_THRESHOLD = 0.95
VIDEO_PSNR_THRESHOLD = 30.0


def _run_vllm_omni_wan22_i2v(
    *,
    model: str,
    output_path: Path,
    conditioning_image: Image.Image,
) -> float:
    server_args = [
        "--num-gpus",
        "1",
        "--stage-init-timeout",
        "400",
        "--init-timeout",
        "900",
        "--diffusion-load-format",
        "diffusers",
        "--boundary-ratio",
        str(BOUNDARY_RATIO),
        "--flow-shift",
        str(FLOW_SHIFT),
    ]
    form_data = {
        "prompt": VIDEO_PROMPT,
        "negative_prompt": NEGATIVE_PROMPT,
        "height": VIDEO_HEIGHT,
        "width": VIDEO_WIDTH,
        "num_inference_steps": VIDEO_NUM_INFERENCE_STEPS,
        "num_frames": NUM_FRAMES,
        "fps": FPS,
        "guidance_scale": GUIDANCE_SCALE,
        "guidance_scale_2": GUIDANCE_SCALE_2,
        "seed": SEED,
    }
    with OmniServer(model, server_args, env_dict=env_to_apply_ftfy_mock_in_subproc(), use_omni=True) as omni_server:
        client = OnlineOmniClient(
            host=omni_server.host,
            port=omni_server.port,
            run_level="full_model",
        )
        request_config = {
            "model": omni_server.model,
            "form_data": form_data,
            "image_reference": f"data:image/png;base64,{pil_to_base64(conditioning_image, 'png')}",
        }
        result = client.send_video_diffusion_request(request_config)[0]
        videos = result.videos
        assert videos is not None
        video_bytes = videos[0]
        output_path.write_bytes(video_bytes)
        latency = result.e2e_latency
        assert latency is not None
        return latency


def _run_diffusers_wan22_i2v(*, model: str, output_path: Path, conditioning_image: Image.Image) -> float:
    from diffusers import WanImageToVideoPipeline  # pyright: ignore[reportPrivateImportUsage]
    from diffusers.schedulers.scheduling_unipc_multistep import UniPCMultistepScheduler

    cleanup_test_environment()
    apply_ftfy_mock()
    pipe: WanImageToVideoPipeline | None = None
    try:
        pipe = WanImageToVideoPipeline.from_pretrained(
            model,
            torch_dtype=torch.bfloat16,
            trust_remote_code=True,
            boundary_ratio=BOUNDARY_RATIO,
        )
        pipe.scheduler = UniPCMultistepScheduler.from_config(pipe.scheduler.config, flow_shift=FLOW_SHIFT)
        pipe.to("cuda")
        _set_matched_attention_backend(pipe)

        _diffusers_dummy_run(pipe)

        generator = torch.Generator(device="cuda").manual_seed(SEED)
        with torch.inference_mode():
            start_time = time.perf_counter()
            result = pipe(
                image=conditioning_image,
                prompt=VIDEO_PROMPT,
                negative_prompt=NEGATIVE_PROMPT,
                height=VIDEO_HEIGHT,
                width=VIDEO_WIDTH,
                num_inference_steps=VIDEO_NUM_INFERENCE_STEPS,
                num_frames=NUM_FRAMES,
                guidance_scale=GUIDANCE_SCALE,
                guidance_scale_2=GUIDANCE_SCALE_2,
                generator=generator,
            )
            end_time = time.perf_counter()
        e2e_latency = end_time - start_time
        frames = result.frames[0]  # pyright: ignore[reportAttributeAccessIssue]
        export_to_video(frames, str(output_path), fps=FPS)  # pyright: ignore[reportArgumentType]
        return e2e_latency
    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()
        cleanup_test_environment()


def _run_vllm_omni_qwen_image(*, model: str, output_path: Path) -> tuple[Image.Image, float]:
    server_args = [
        "--num-gpus",
        "1",
        "--stage-init-timeout",
        "400",
        "--init-timeout",
        "900",
        "--diffusion-load-format",
        "diffusers",
        # "--enable-diffusion-pipeline-profiler",
    ]
    with OmniServer(model, server_args, use_omni=True) as omni_server:
        url = f"http://{omni_server.host}:{omni_server.port}/v1/images/generations"
        request_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,
        }

        def _timed_request() -> tuple[requests.Response, float]:
            start_time = time.perf_counter()
            response = requests.post(url, json=request_json, timeout=600)
            end_time = time.perf_counter()
            response.raise_for_status()
            return response, end_time - start_time

        # Warm up the client request path so first-request lazy init is not charged
        # to the measurement, mirroring _diffusers_dummy_run on the diffusers side.
        _timed_request()

        latencies: list[float] = []
        response: requests.Response | None = None
        for _ in range(NUM_TIMED_RUNS):
            response, latency = _timed_request()
            latencies.append(latency)
        e2e_latency = float(np.median(latencies))

        assert response is not None
        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, e2e_latency


def _run_diffusers_qwen_image(*, model: str, output_path: Path) -> tuple[Image.Image, float]:
    cleanup_test_environment()
    pipe: DiffusionPipeline | None = None
    try:
        pipe = DiffusionPipeline.from_pretrained(
            model,
            torch_dtype=torch.bfloat16,
            trust_remote_code=True,
        ).to("cuda")
        _set_matched_attention_backend(pipe)
        _diffusers_dummy_run(pipe)

        latencies: list[float] = []
        result = None
        with torch.inference_mode():
            for _ in range(NUM_TIMED_RUNS):
                generator = torch.Generator(device="cuda").manual_seed(SEED)
                start_time = time.perf_counter()
                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,
                )
                end_time = time.perf_counter()
                latencies.append(end_time - start_time)
        e2e_latency = float(np.median(latencies))
        assert result is not None
        output_image = result.images[0].convert("RGB")
        output_image.save(output_path)
        return output_image, e2e_latency
    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()
        cleanup_test_environment()


@pytest.mark.benchmark
@hardware_test(res={"cuda": ["H100", "B200"]}, num_cards=1)
@pytest.mark.parametrize("model_id", ["Qwen/Qwen-Image"])
def test_diffusers_backend_t2i_matches_diffusers(model_id: str, accuracy_artifact_root: Path) -> None:
    output_dir = model_output_dir(accuracy_artifact_root, model_id + "-diffusers-backend")

    vllm_output, vllm_latency = _run_vllm_omni_qwen_image(model=model_id, output_path=output_dir / "vllm_omni.png")
    diffusers_output, diffusers_latency = _run_diffusers_qwen_image(
        model=model_id, output_path=output_dir / "diffusers.png"
    )
    vllm_latency = vllm_latency * 1000
    diffusers_latency = diffusers_latency * 1000
    latency_threshold_factor = 0.2
    latency_threshold = diffusers_latency * (1 + latency_threshold_factor)
    print(f"{model_id} latency metrics:")
    print(
        f"  Latency={vllm_latency:.2f}ms, threshold<={latency_threshold:.2f}ms, diffusers latency={diffusers_latency:.2f}ms, lower is better"
    )

    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,
    )
    # Latency is checked against a warmed median of NUM_TIMED_RUNS runs (see the
    # _run_* helpers), so the assertion reflects steady-state rather than a single
    # cold shot -- which is what previously made this check flaky (issue #5108).
    # Correctness is additionally gated by the SSIM/PSNR assertions above.
    assert vllm_latency <= latency_threshold, (
        f"VLLM latency ({vllm_latency:.2f}ms) is greater than "
        f"{latency_threshold_factor * 100}% more than Diffusers latency "
        f"({diffusers_latency:.2f}ms)."
    )


@pytest.mark.benchmark
@hardware_test(res={"cuda": ["H100", "B200"]}, num_cards=1)
@pytest.mark.parametrize(
    "model_id",
    [
        pytest.param(
            "Wan-AI/Wan2.2-I2V-A14B-Diffusers",
        ),
    ],
)
def test_diffusers_backend_i2v_matches_diffusers(
    model_id: str,
    accuracy_artifact_root: Path,
    qwen_bear_image: Image.Image,
) -> None:
    output_dir = model_output_dir(accuracy_artifact_root, model_id + "-diffusers-backend")

    resized_image = qwen_bear_image.resize((VIDEO_WIDTH, VIDEO_HEIGHT), Image.Resampling.LANCZOS)

    vllm_path = output_dir / "vllm_omni.mp4"
    vllm_latency = _run_vllm_omni_wan22_i2v(model=model_id, output_path=vllm_path, conditioning_image=resized_image)
    vllm_latency = vllm_latency * 1000

    diffusers_path = output_dir / "diffusers.mp4"
    diffusers_latency = _run_diffusers_wan22_i2v(
        model=model_id, output_path=diffusers_path, conditioning_image=resized_image
    )
    diffusers_latency = diffusers_latency * 1000
    # H100 keeps the historical 30% slack. B200 measured ~35.8% (6008 vs 4425 ms).
    gpu_key, latency_threshold_factor = resolve_device_threshold(
        {"H100": 0.3, "B200": 0.36},
        label="latency threshold factor",
    )
    latency_threshold = diffusers_latency * (1 + latency_threshold_factor)

    ssim_output = _run_ffmpeg_similarity("ssim", vllm_path, diffusers_path)
    psnr_output = _run_ffmpeg_similarity("psnr", vllm_path, diffusers_path)
    ssim_score = _parse_ssim_score(ssim_output)
    psnr_score = _parse_psnr_score(psnr_output)
    print(f"{model_id} latency metrics:")
    print(
        f"  Latency={vllm_latency:.2f}ms, threshold<={latency_threshold:.2f}ms, "
        f"diffusers latency={diffusers_latency:.2f}ms, slack={latency_threshold_factor:.0%} ({gpu_key}), lower is better"
    )
    print(f"{model_id} similarity metrics:")
    print(f"  SSIM: value={ssim_score:.6f}, threshold>={VIDEO_SSIM_THRESHOLD:.6f}, range=[-1, 1], higher is better")
    print(
        f"  PSNR: value={psnr_score:.6f} dB, threshold>={VIDEO_PSNR_THRESHOLD:.6f} dB, range=[0, +inf), higher is better"
    )

    assert ssim_score >= VIDEO_SSIM_THRESHOLD, (
        f"SSIM below threshold for {model_id}: got {ssim_score:.6f}, expected >= {VIDEO_SSIM_THRESHOLD:.6f}."
    )
    assert psnr_score >= VIDEO_PSNR_THRESHOLD, (
        f"PSNR below threshold for {model_id}: got {psnr_score:.6f}, expected >= {VIDEO_PSNR_THRESHOLD:.6f}."
    )
    assert vllm_latency <= latency_threshold, (
        f"VLLM latency ({vllm_latency:.2f}ms) is greater than {latency_threshold_factor * 100}% more than Diffusers latency ({diffusers_latency:.2f}ms)."
    )


def _diffusers_dummy_run(pipe: DiffusionPipeline) -> None:
    from vllm_omni.diffusion.io_support import supports_multimodal_input

    supports_image_input, supports_audio_input = supports_multimodal_input(
        SimpleNamespace(
            diffusion_load_format="diffusers",
            diffusers_pipeline_cls=pipe.__class__,
            model_class_name="DiffusersAdapterPipeline",
        ),  # pyright: ignore[reportArgumentType]
    )
    height = 512
    width = 512
    kwargs: dict[str, Any] = {
        "prompt": "dummy run",
        "height": height,
        "width": width,
        "num_inference_steps": 1,
    }
    if supports_image_input:
        # Provide a dummy image input if the model supports it
        dummy_image = Image.new("RGB", (width, height))
        kwargs["image"] = dummy_image
    if supports_audio_input:
        audio_sr = 16000
        dummy_audio = np.random.randn(audio_sr * 2).astype(np.float32)
        kwargs["audio"] = dummy_audio
    with torch.inference_mode():
        pipe(**kwargs)  # pyright: ignore[reportCallIssue]