（Phase 1）Add ModelOpt FP8 auto-detect support for diffusion checkpoints #2709 (#2913)

Signed-off-by: roG0d <rodgarcas98@gmail.com>
Signed-off-by: roG0d <baonudesifeizhai@gmail.com>
Signed-off-by: baonudesifeizhai <85092850+baonudesifeizhai@users.noreply.github.com>
Co-authored-by: roG0d <rodgarcas98@gmail.com>

tests/diffusion/quantization/test_quantization_quality.py

@@ -1,8 +1,7 @@
 # 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.
+Quantization quality gate for diffusion models.
 
 Developers adding a new quantization method should:
 1. Add their method + model to QUALITY_CONFIGS below
@@ -11,7 +10,7 @@ Developers adding a new quantization method should:
 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.
+same seed, computes similarity metrics, and fails if LPIPS exceeds the threshold.
 
 Requirements:
     pip install lpips
@@ -21,12 +20,22 @@ Example — run only FP8 tests:
 
 Example — run a specific model:
     pytest tests/diffusion/quantization/test_quantization_quality.py -v -m "" -k "z_image"
+
+Example — validate a local BF16 baseline against a local pre-quantized checkpoint:
+    export VLLM_OMNI_QUALITY_CONFIGS=/tmp/modelopt_quality_cases.json
+    pytest tests/diffusion/quantization/test_quantization_quality.py -v -m "" -k "qwen_image_2512"
+
+Optional artifact dump:
+    export VLLM_OMNI_QUALITY_OUTPUT_DIR=/tmp/modelopt_quality_outputs
 """
 
 from __future__ import annotations
 
 import gc
+import json
+import os
 from dataclasses import dataclass
+from pathlib import Path
 
 import numpy as np
 import pytest
@@ -44,17 +53,49 @@ 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
+    model: str | None = None  # HF model name
+    quantization: str | None = None  # quantization method, e.g. "fp8"
+    baseline_model: str | None = None  # explicit BF16/local baseline path
+    quantized_model: str | None = None  # explicit quantized/local model path
     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
+    negative_prompt: str = ""
+    guidance_scale: float | None = None
+
+    def baseline_ref(self) -> str:
+        return self.baseline_model or self.model or ""
+
+    def quantized_ref(self) -> str:
+        if self.quantized_model is not None:
+            return self.quantized_model
+        return self.model or ""
+
+    def quantization_ref(self) -> str | None:
+        if self.quantized_model is not None:
+            return None
+        return self.quantization
+
+    def validate(self) -> None:
+        uses_explicit_models = self.baseline_model is not None or self.quantized_model is not None
+        uses_model_plus_method = self.model is not None or self.quantization is not None
+
+        if uses_explicit_models and uses_model_plus_method:
+            raise ValueError(f"{self.id}: explicit baseline/quantized paths cannot be mixed with model/quantization")
+
+        if uses_explicit_models:
+            if self.baseline_model is None or self.quantized_model is None:
+                raise ValueError(f"{self.id}: baseline_model and quantized_model must be provided together")
+            return
+
+        if self.model is None or self.quantization is None:
+            raise ValueError(f"{self.id}: expected either model+quantization or baseline_model+quantized_model")
 
 
 # Add new quantization methods / models here.
@@ -91,6 +132,53 @@ QUALITY_CONFIGS = [
 ]
 
 
+def _load_extra_quality_configs() -> list[QualityTestConfig]:
+    config_path = os.environ.get("VLLM_OMNI_QUALITY_CONFIGS")
+    if not config_path:
+        return []
+
+    payload = json.loads(Path(config_path).read_text(encoding="utf-8"))
+    if not isinstance(payload, list):
+        raise ValueError("VLLM_OMNI_QUALITY_CONFIGS must point to a JSON list")
+
+    configs: list[QualityTestConfig] = []
+    for item in payload:
+        if not isinstance(item, dict):
+            raise ValueError("Each quality config entry must be a JSON object")
+        config = QualityTestConfig(**item)
+        config.validate()
+        configs.append(config)
+    return configs
+
+
+def _all_quality_configs() -> list[QualityTestConfig]:
+    configs = [*QUALITY_CONFIGS, *_load_extra_quality_configs()]
+    for config in configs:
+        config.validate()
+    return configs
+
+
+def _output_path(output_dir: Path, config: QualityTestConfig, label: str, suffix: str) -> Path:
+    safe_id = config.id.replace("/", "_")
+    return output_dir / f"{safe_id}_{label}{suffix}"
+
+
+def _maybe_save_output(output_dir: Path | None, config: QualityTestConfig, label: str, output) -> None:
+    if output_dir is None:
+        return
+
+    output_dir.mkdir(parents=True, exist_ok=True)
+    if config.task == "t2i":
+        output.save(_output_path(output_dir, config, label, ".png"))
+        return
+
+    if isinstance(output, np.ndarray):
+        np.save(_output_path(output_dir, config, label, ".npy"), output)
+        return
+
+    raise TypeError(f"Unsupported output type for saving: {type(output)!r}")
+
+
 # ---------------------------------------------------------------------------
 # Helpers
 # ---------------------------------------------------------------------------
@@ -107,12 +195,13 @@ def _generate_image(omni, config: QualityTestConfig):
     torch.accelerator.reset_peak_memory_stats()
 
     outputs = omni.generate(
-        {"prompt": config.prompt},
+        {"prompt": config.prompt, "negative_prompt": config.negative_prompt},
         OmniDiffusionSamplingParams(
             height=config.height,
             width=config.width,
             generator=generator,
             num_inference_steps=config.num_inference_steps,
+            guidance_scale=config.guidance_scale,
         ),
     )
 
@@ -138,13 +227,14 @@ def _generate_video(omni, config: QualityTestConfig):
     torch.accelerator.reset_peak_memory_stats()
 
     outputs = omni.generate(
-        {"prompt": config.prompt, "negative_prompt": ""},
+        {"prompt": config.prompt, "negative_prompt": config.negative_prompt},
         OmniDiffusionSamplingParams(
             height=config.height,
             width=config.width,
             generator=generator,
             num_inference_steps=config.num_inference_steps,
             num_frames=config.num_frames,
+            guidance_scale=config.guidance_scale,
         ),
     )
 
@@ -186,6 +276,34 @@ def _compute_lpips(baseline, quantized, task: str) -> float:
     return compute_lpips_video(baseline, quantized)
 
 
+def _to_float_array(output, task: str) -> np.ndarray:
+    if task == "t2i":
+        array = np.asarray(output.convert("RGB"), dtype=np.float32) / 255.0
+    else:
+        array = np.asarray(output, dtype=np.float32)
+        if array.max() > 1.0 or array.min() < 0.0:
+            array = np.clip(array, 0.0, 255.0) / 255.0
+        else:
+            array = np.clip(array, 0.0, 1.0)
+    return array
+
+
+def _compute_psnr_and_mae(baseline, quantized, task: str) -> tuple[float, float]:
+    baseline_array = _to_float_array(baseline, task)
+    quantized_array = _to_float_array(quantized, task)
+    if baseline_array.shape != quantized_array.shape:
+        raise ValueError(
+            "Output shapes do not match for metric computation: "
+            f"baseline={baseline_array.shape}, quantized={quantized_array.shape}"
+        )
+
+    diff = baseline_array - quantized_array
+    mae = float(np.mean(np.abs(diff)))
+    mse = float(np.mean(np.square(diff)))
+    psnr = float("inf") if mse == 0.0 else float(20.0 * np.log10(1.0 / np.sqrt(mse)))
+    return psnr, mae
+
+
 def _unload(omni):
     del omni
     gc.collect()
@@ -199,6 +317,7 @@ def _unload(omni):
 # ---------------------------------------------------------------------------
 
 _marks = hardware_marks(res={"cuda": "H100"})
+_OUTPUT_DIR = Path(os.environ["VLLM_OMNI_QUALITY_OUTPUT_DIR"]) if "VLLM_OMNI_QUALITY_OUTPUT_DIR" in os.environ else None
 
 
 def _quality_param(c: QualityTestConfig):
@@ -214,7 +333,7 @@ def _quality_param(c: QualityTestConfig):
 @pytest.mark.diffusion
 @pytest.mark.parametrize(
     "config",
-    [_quality_param(c) for c in QUALITY_CONFIGS],
+    [_quality_param(c) for c in _all_quality_configs()],
 )
 def test_quantization_quality(config: QualityTestConfig):
     """Validate that quantized output stays within LPIPS threshold of BF16."""
@@ -223,31 +342,44 @@ def test_quantization_quality(config: QualityTestConfig):
     generate_fn = _generate_video if config.task == "t2v" else _generate_image
 
     # --- BF16 baseline ---
-    omni_bl = Omni(model=config.model)
+    omni_bl = Omni(model=config.baseline_ref())
     baseline_out, bl_mem = generate_fn(omni_bl, config)
     _unload(omni_bl)
+    _maybe_save_output(_OUTPUT_DIR, config, "baseline", baseline_out)
 
     # --- Quantized ---
-    omni_qt = Omni(model=config.model, quantization_config=config.quantization)
+    quantization = config.quantization_ref()
+    if quantization is None:
+        omni_qt = Omni(model=config.quantized_ref())
+    else:
+        omni_qt = Omni(model=config.quantized_ref(), quantization_config=quantization)
     quant_out, qt_mem = generate_fn(omni_qt, config)
     _unload(omni_qt)
+    _maybe_save_output(_OUTPUT_DIR, config, "quantized", quant_out)
 
-    # --- LPIPS ---
+    # --- Similarity metrics ---
     lpips_score = _compute_lpips(baseline_out, quant_out, config.task)
+    psnr_score, mae_score = _compute_psnr_and_mae(baseline_out, quant_out, config.task)
+    assert lpips_score <= config.max_lpips, (
+        f"LPIPS {lpips_score:.4f} exceeds threshold {config.max_lpips} "
+        f"for {config.quantization_ref() or 'pre-quantized checkpoint'} on {config.quantized_ref()}"
+    )
 
     # --- 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"  Baseline:      {config.baseline_ref()}")
+    print(f"  Quantized:     {config.quantized_ref()}")
+    print(f"  Method:        {config.quantization_ref() or 'pre-quantized checkpoint'}")
     print(f"  LPIPS:         {lpips_score:.4f}  (threshold: {config.max_lpips})")
+    print(f"  PSNR:          {psnr_score:.4f} dB  (higher is better)")
+    print(f"  MAE:           {mae_score:.6f}  (lower is better)")
     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}"
-    )
+    assert np.isfinite(psnr_score) or np.isinf(psnr_score), f"PSNR is invalid for {config.id}: {psnr_score}"
+    assert np.isfinite(mae_score), f"MAE is not finite for {config.id}: {mae_score}"