/tests/e2e/accuracy/test_diffusers_backend_similarity.py

# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM 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,
)
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.env import run_post_test_cleanup, run_pre_test_cleanup
from tests.helpers.mark import hardware_test
from tests.helpers.runtime import OmniServer, OpenAIClientHandler

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


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

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 = OpenAIClientHandler(
            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]
        video_bytes = result.videos[0]  # pyright: ignore[reportOptionalSubscript] # Guaranteed not None
        output_path.write_bytes(video_bytes)
        return result.e2e_latency  # pyright: ignore[reportReturnType] # Guaranteed not None


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

    run_pre_test_cleanup()
    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")

        _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()
        run_post_test_cleanup()


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:
        start_time = time.perf_counter()
        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,
        )
        end_time = time.perf_counter()
        e2e_latency = end_time - start_time
        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, e2e_latency


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

        generator = torch.Generator(device="cuda").manual_seed(SEED)
        with torch.inference_mode():
            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()
        e2e_latency = end_time - start_time
        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()
        run_post_test_cleanup()


@pytest.mark.benchmark
@hardware_test(res={"cuda": "H100"}, 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,
    )
    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)."
    )


@pytest.mark.benchmark
@hardware_test(res={"cuda": "H100"}, 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
    latency_threshold_factor = 0.3
    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, diffusers latency={diffusers_latency:.2f}ms, 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]