vllm-omni/tests/diffusion/quantization/test_quantization_quality.py

# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""
Quantization quality gate — validates that quantized models stay within
acceptable perceptual distance (LPIPS) of the BF16 baseline.

Developers adding a new quantization method should:
1. Add their method + model to QUALITY_CONFIGS below
2. Set a max_lpips threshold (use 0.15 for image, 0.20 for video as defaults)
3. Run: pytest tests/diffusion/quantization/test_quantization_quality.py -v -m ""
4. Paste the output table into their PR description

The test generates outputs with both BF16 and the quantized method using the
same seed, computes LPIPS, and fails if it exceeds the threshold.

Requirements:
    pip install lpips

Example — run only FP8 tests:
    pytest tests/diffusion/quantization/test_quantization_quality.py -v -m "" -k "fp8"

Example — run a specific model:
    pytest tests/diffusion/quantization/test_quantization_quality.py -v -m "" -k "z_image"
"""

from __future__ import annotations

import gc
from dataclasses import dataclass

import numpy as np
import pytest
import torch

from tests.utils import hardware_marks

# ---------------------------------------------------------------------------
# Configuration — add new quantization methods / models here
# ---------------------------------------------------------------------------


@dataclass
class QualityTestConfig:
    """Defines a single quantization quality test case."""

    id: str  # pytest ID, e.g. "fp8_z_image"
    model: str  # HF model name
    quantization: str  # quantization method, e.g. "fp8"
    task: str  # "t2i" or "t2v"
    prompt: str  # generation prompt
    max_lpips: float  # fail threshold — higher = more lenient
    height: int = 1024
    width: int = 1024
    num_inference_steps: int = 20  # keep low for CI speed
    num_frames: int = 5  # only for t2v
    seed: int = 42
    gpu: str = "H100"  # minimum GPU requirement


# Add new quantization methods / models here.
# Developers: copy a config, change quantization + max_lpips, run the test.
QUALITY_CONFIGS = [
    QualityTestConfig(
        id="fp8_z_image",
        model="Tongyi-MAI/Z-Image-Turbo",
        quantization="fp8",
        task="t2i",
        prompt="a cup of coffee on a wooden table, morning light",
        max_lpips=0.10,
        num_inference_steps=20,
    ),
    QualityTestConfig(
        id="fp8_flux",
        model="black-forest-labs/FLUX.1-dev",
        quantization="fp8",
        task="t2i",
        prompt="a cup of coffee on a wooden table, morning light",
        max_lpips=0.20,
        num_inference_steps=10,
    ),
    QualityTestConfig(
        id="fp8_qwen_image",
        model="Qwen/Qwen-Image",
        quantization="fp8",
        task="t2i",
        prompt="a cup of coffee on a wooden table, morning light",
        max_lpips=0.35,
        seed=142,
        num_inference_steps=20,
    ),
]


# ---------------------------------------------------------------------------
# Helpers
# ---------------------------------------------------------------------------


def _generate_image(omni, config: QualityTestConfig):
    """Generate a single image, return (PIL.Image, peak_mem_gib)."""
    from vllm_omni.inputs.data import OmniDiffusionSamplingParams
    from vllm_omni.platforms import current_omni_platform

    generator = torch.Generator(
        device=current_omni_platform.device_type,
    ).manual_seed(config.seed)
    torch.cuda.reset_peak_memory_stats()

    outputs = omni.generate(
        {"prompt": config.prompt},
        OmniDiffusionSamplingParams(
            height=config.height,
            width=config.width,
            generator=generator,
            num_inference_steps=config.num_inference_steps,
        ),
    )

    peak_mem = torch.cuda.max_memory_allocated() / (1024**3)
    first = outputs[0]
    if hasattr(first, "images") and first.images:
        return first.images[0], peak_mem
    inner = first.request_output
    if inner is not None and hasattr(inner, "images") and inner.images:
        return inner.images[0], peak_mem
    raise ValueError("Could not extract image from output.")


def _generate_video(omni, config: QualityTestConfig):
    """Generate a video, return (np.ndarray [F,H,W,C], peak_mem_gib)."""
    from vllm_omni.inputs.data import OmniDiffusionSamplingParams
    from vllm_omni.outputs import OmniRequestOutput
    from vllm_omni.platforms import current_omni_platform

    generator = torch.Generator(
        device=current_omni_platform.device_type,
    ).manual_seed(config.seed)
    torch.cuda.reset_peak_memory_stats()

    outputs = omni.generate(
        {"prompt": config.prompt, "negative_prompt": ""},
        OmniDiffusionSamplingParams(
            height=config.height,
            width=config.width,
            generator=generator,
            num_inference_steps=config.num_inference_steps,
            num_frames=config.num_frames,
        ),
    )

    peak_mem = torch.cuda.max_memory_allocated() / (1024**3)
    first = outputs[0]
    if hasattr(first, "request_output") and isinstance(first.request_output, list):
        inner = first.request_output[0]
        if isinstance(inner, OmniRequestOutput) and hasattr(inner, "images"):
            frames = inner.images[0] if inner.images else None
        else:
            frames = inner
    elif hasattr(first, "images") and first.images:
        frames = first.images
    else:
        raise ValueError("Could not extract video frames from output.")

    if isinstance(frames, torch.Tensor):
        video = frames.detach().cpu()
        if video.dim() == 5:
            video = video[0]
        if video.dim() == 4 and video.shape[0] in (3, 4):
            video = video.permute(1, 2, 3, 0)
        if video.is_floating_point():
            video = video.clamp(-1, 1) * 0.5 + 0.5
        return video.float().numpy(), peak_mem

    return np.asarray(frames), peak_mem


def _compute_lpips(baseline, quantized, task: str) -> float:
    """Compute LPIPS between baseline and quantized outputs."""
    from benchmarks.diffusion.quantization_quality import (
        compute_lpips_images,
        compute_lpips_video,
    )

    if task == "t2i":
        return compute_lpips_images([baseline], [quantized])[0]
    return compute_lpips_video(baseline, quantized)


def _unload(omni):
    del omni
    gc.collect()
    if torch.cuda.is_available():
        torch.cuda.empty_cache()
        torch.cuda.synchronize()


# ---------------------------------------------------------------------------
# Tests
# ---------------------------------------------------------------------------

_marks = hardware_marks(res={"cuda": "H100"})


@pytest.mark.advanced_model
@pytest.mark.diffusion
@pytest.mark.parametrize(
    "config",
    [pytest.param(c, id=c.id, marks=_marks) for c in QUALITY_CONFIGS],
)
def test_quantization_quality(config: QualityTestConfig):
    """Validate that quantized output stays within LPIPS threshold of BF16."""
    from vllm_omni.entrypoints.omni import Omni

    generate_fn = _generate_video if config.task == "t2v" else _generate_image

    # --- BF16 baseline ---
    omni_bl = Omni(model=config.model)
    baseline_out, bl_mem = generate_fn(omni_bl, config)
    _unload(omni_bl)

    # --- Quantized ---
    omni_qt = Omni(model=config.model, quantization_config=config.quantization)
    quant_out, qt_mem = generate_fn(omni_qt, config)
    _unload(omni_qt)

    # --- LPIPS ---
    lpips_score = _compute_lpips(baseline_out, quant_out, config.task)

    # --- Report ---
    mem_reduction = (bl_mem - qt_mem) / bl_mem * 100 if bl_mem > 0 else 0
    print(f"\n{'=' * 60}")
    print(f"Quantization Quality: {config.id}")
    print(f"{'=' * 60}")
    print(f"  Model:         {config.model}")
    print(f"  Method:        {config.quantization}")
    print(f"  LPIPS:         {lpips_score:.4f}  (threshold: {config.max_lpips})")
    print(f"  BF16 memory:   {bl_mem:.2f} GiB")
    print(f"  Quant memory:  {qt_mem:.2f} GiB  ({mem_reduction:.0f}% reduction)")
    print(f"  Result:        {'PASS' if lpips_score <= config.max_lpips else 'FAIL'}")
    print(f"{'=' * 60}\n")

    assert lpips_score <= config.max_lpips, (
        f"LPIPS {lpips_score:.4f} exceeds threshold {config.max_lpips} for {config.quantization} on {config.model}"
    )