[Quantization][tools] Add diffusion quantization output comparison tool (#3175)

Signed-off-by: david6666666 <530634352@qq.com>
Signed-off-by: David Chen <530634352@qq.com>

docs/user_guide/quantization/overview.md

@@ -138,3 +138,98 @@ build_quant_config({"method": "gguf", "gguf_model": "/path/to/model.gguf"})
 build_quant_config({"transformer": {"method": "fp8"}, "vae": None})
 build_quant_config(None)
 ```
+
+## Output Similarity Comparison Tool
+
+Use `vllm_omni.quantization.tools.compare_diffusion_trajectory_similarity`
+to compare a reference diffusion run with a quantized candidate run using the
+same prompt, seed, resolution, scheduler settings, and inference steps. The
+tool compares final decoded images or video frames, and also reports generation
+latency and worker-reported peak memory when available.
+
+This is useful when validating whether online quantization, an offline
+pre-quantized checkpoint, or a new `ignored_layers` choice keeps generation
+quality close to the BF16 reference.
+
+### Online Quantization Example
+
+```bash
+python -m vllm_omni.quantization.tools.compare_diffusion_trajectory_similarity \
+  --task t2i \
+  --model Qwen/Qwen-Image \
+  --candidate-quantization fp8 \
+  --ignored-layers img_mlp \
+  --prompt "a cup of coffee on the table" \
+  --height 512 --width 512 \
+  --num-inference-steps 20 \
+  --seed 142 \
+  --output-json /tmp/qwen_image_fp8_similarity/result.json \
+  --save-output-dir /tmp/qwen_image_fp8_similarity/images \
+  --enforce-eager
+```
+
+### Offline Checkpoint Example
+
+Use `--candidate-model` when the candidate is already quantized or lives at a
+different model path:
+
+```bash
+python -m vllm_omni.quantization.tools.compare_diffusion_trajectory_similarity \
+  --task t2i \
+  --reference-model Qwen/Qwen-Image \
+  --candidate-model /path/to/qwen-image-fp8-checkpoint \
+  --prompt "a cup of coffee on the table" \
+  --height 512 --width 512 \
+  --num-inference-steps 20 \
+  --seed 142 \
+  --output-json /tmp/qwen_image_fp8_checkpoint_similarity/result.json
+```
+
+If the checkpoint does not include a loadable quantization config, pass one
+explicitly:
+
+```bash
+--candidate-quantization-config-json '{"method":"fp8"}'
+```
+
+### Output Metrics
+
+The output JSON includes `output_metrics`, `reference_generation`, and
+`candidate_generation`.
+
+| Metric | Direction | Meaning |
+|--------|-----------|---------|
+| `cosine_similarity` | Higher is better | Vector direction similarity between output pixels or frames. Useful as a broad sanity check. |
+| `mae` | Lower is better | Mean absolute pixel or frame error. For decoded outputs, values are in uint8 pixel units. |
+| `mse` / `rmse` | Lower is better | Squared error and its square root. These penalize localized large differences more than `mae`. |
+| `max_abs` | Lower is better | Worst single-element absolute error. Treat it as an outlier/debug signal, not as a release gate. |
+| `l2` / `relative_l2` | Lower is better | Absolute and reference-normalized L2 distance. `relative_l2` is easier to compare across resolutions. |
+| `psnr_db` | Higher is better | Pixel-space signal-to-noise ratio in dB for uint8 images or frames. |
+| `avg_generation_time_s` | Lower is better | Average wall-clock generation time across measured runs. |
+| `max_peak_memory_mb` | Lower is better | Maximum worker-reported peak device memory across measured runs, when the worker reports it. |
+
+Recommended starting thresholds for same-seed diffusion comparisons:
+
+| Metric | Smoke threshold | Stricter target | Notes |
+|--------|-----------------|-----------------|-------|
+| `psnr_db` | `>= 20.0` | `>= 25.0` | Good for quick image or frame regression checks. |
+| `mae` | `<= 12.0` | `<= 6.0` | Interpreted in decoded uint8 pixel units. |
+| `cosine_similarity` | `>= 0.98` | `>= 0.995` | Less sensitive to global scale than L2-style metrics. |
+| `relative_l2` | `<= 0.20` | `<= 0.08` | Useful when comparing across prompts or resolutions. |
+
+These thresholds are heuristics. Tune them by model family, task, resolution,
+quantization method, and deployment tolerance. For release gating, pair the
+numeric report with visual inspection of saved reference and candidate outputs.
+
+The tool intentionally reports separate quality, latency, and memory metrics
+instead of a single consolidated similarity score. A single score can hide
+important tradeoffs, for example a candidate with good PSNR but a meaningful
+memory regression, or a candidate with low average error but localized visual
+artifacts. If you need a project-specific pass/fail gate, define it as an
+explicit policy over the individual metrics.
+
+Pixel-level metrics do not measure semantic consistency. For higher-cost
+evaluation, you can complement this report with a vision-language judge that
+describes the reference and candidate outputs and compares those descriptions.
+Keep that semantic check separate from this lightweight tool so users can
+choose whether the additional model cost and latency are appropriate.

tests/diffusion/quantization/test_trajectory_similarity_tool.py

@@ -0,0 +1,126 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+import math
+import sys
+from types import SimpleNamespace
+
+import numpy as np
+import pytest
+import torch
+from PIL import Image
+
+from vllm_omni.quantization.tools.compare_diffusion_trajectory_similarity import (
+    VariantRun,
+    _build_variant_config,
+    _request_peak_memory_mb,
+    _run_summary,
+    compute_tensor_metrics,
+    compute_uint8_image_metrics,
+    metric_guidance,
+    parse_args,
+    summarize_output_image_metrics,
+)
+
+
+def test_compute_tensor_metrics_identical_tensors():
+    metrics = compute_tensor_metrics(torch.ones(2, 3), torch.ones(2, 3))
+
+    assert metrics["cosine_similarity"] == pytest.approx(1.0)
+    assert metrics["mae"] == 0.0
+    assert metrics["mse"] == 0.0
+    assert metrics["rmse"] == 0.0
+    assert metrics["max_abs"] == 0.0
+    assert metrics["l2"] == 0.0
+    assert metrics["relative_l2"] == 0.0
+
+
+def test_compute_uint8_image_metrics_adds_psnr():
+    lhs = np.zeros((2, 2, 3), dtype=np.uint8)
+    rhs = np.zeros((2, 2, 3), dtype=np.uint8)
+
+    metrics = compute_uint8_image_metrics(lhs, rhs)
+
+    assert math.isinf(metrics["psnr_db"])
+
+
+def test_summarize_output_image_metrics_stacks_pil_images():
+    reference = [Image.fromarray(np.zeros((2, 2, 3), dtype=np.uint8))]
+    candidate = [Image.fromarray(np.ones((2, 2, 3), dtype=np.uint8))]
+
+    summary = summarize_output_image_metrics(reference, candidate)
+
+    assert summary["num_images"] == 1
+    assert summary["image0_metrics"]["mae"] == 1.0
+    assert summary["all_images_metrics"]["mse"] == 1.0
+
+
+def test_run_summary_reports_worker_peak_memory():
+    summary = _run_summary(
+        VariantRun(
+            label="candidate",
+            result=object(),
+            generation_times_s=[1.0, 3.0],
+            peak_memory_mb=[100.0, 150.0],
+        )
+    )
+
+    assert summary["peak_memory_mb"] == 150.0
+    assert summary["avg_peak_memory_mb"] == 125.0
+    assert summary["max_peak_memory_mb"] == 150.0
+    assert summary["per_run_peak_memory_mb"] == [100.0, 150.0]
+
+
+def test_request_peak_memory_prefers_inner_diffusion_output():
+    inner = type("InnerOutput", (), {"peak_memory_mb": 321.0})()
+    outer = type("OuterOutput", (), {"request_output": inner, "peak_memory_mb": 123.0})()
+
+    assert _request_peak_memory_mb(outer) == 321.0
+
+
+def test_metric_guidance_describes_thresholds():
+    guidance = metric_guidance()
+
+    assert "cosine_similarity" in guidance["descriptions"]
+    assert guidance["recommended_thresholds"]["output_images_or_frames_uint8"]["psnr_db"]["recommended_min"] == 20.0
+    assert (
+        guidance["recommended_thresholds"]["performance"]["max_peak_memory_ratio_candidate_over_reference"][
+            "recommended_max"
+        ]
+        == 1.00
+    )
+
+
+def test_candidate_model_can_point_to_offline_checkpoint_without_online_quantization():
+    args = SimpleNamespace(
+        model="Qwen/Qwen-Image",
+        candidate_model="Qwen/Qwen-Image-FP8",
+        candidate_quantization=None,
+        candidate_quantization_config_json=None,
+        candidate_ignored_layers=None,
+        ignored_layers=None,
+        candidate_gguf_model=None,
+        candidate_load_format="default",
+    )
+
+    config = _build_variant_config(args, "candidate")
+
+    assert config.model == "Qwen/Qwen-Image-FP8"
+    assert config.quantization is None
+    assert config.quantization_config is None
+
+
+def test_step_execution_defaults_to_false(monkeypatch):
+    monkeypatch.setattr(
+        sys,
+        "argv",
+        [
+            "compare_diffusion_trajectory_similarity.py",
+            "--output-json",
+            "result.json",
+        ],
+    )
+
+    args = parse_args()
+
+    assert args.step_execution is False

vllm_omni/quantization/tools/compare_diffusion_trajectory_similarity.py

@@ -0,0 +1,637 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Compare diffusion quantization quality with final output metrics.
+
+This tool runs a reference diffusion configuration and a candidate
+configuration with the same prompt, seed, resolution, and step count. It reports
+final image or video frame similarity metrics plus generation latency and memory.
+
+The main metrics are:
+
+* output images: the same pixel-space metrics plus PSNR
+* runtime: per-run and average generation latency, and worker-reported peak
+  device memory when available
+
+Online FP8 Qwen-Image example:
+
+    python -m vllm_omni.quantization.tools.compare_diffusion_trajectory_similarity \\
+        --model Qwen/Qwen-Image \\
+        --candidate-quantization fp8 \\
+        --ignored-layers img_mlp \\
+        --prompt "a cup of coffee on the table" \\
+        --height 1024 --width 1024 --num-inference-steps 20 --seed 42 \\
+        --output-json /mnt/data4/cwq/tmp/qwen_image_fp8_similarity.json
+
+Offline checkpoint example:
+
+    python -m vllm_omni.quantization.tools.compare_diffusion_trajectory_similarity \\
+        --reference-model Qwen/Qwen-Image \\
+        --candidate-model /path/to/quantized/qwen-image-checkpoint \\
+        --candidate-quantization-config-json '{"method":"inc","weight_bits":4}' \\
+        --prompt "a cup of coffee on the table" \\
+        --output-json /mnt/data4/cwq/tmp/qwen_image_checkpoint_similarity.json
+
+Offline FP8 model-id example:
+
+    python -m vllm_omni.quantization.tools.compare_diffusion_trajectory_similarity \\
+        --model Qwen/Qwen-Image \\
+        --candidate-model Qwen/Qwen-Image-FP8 \\
+        --prompt "a cup of coffee on the table" \\
+        --height 1024 --width 1024 --num-inference-steps 20 --seed 42 \\
+        --output-json /mnt/data4/cwq/tmp/qwen_image_fp8_checkpoint_similarity.json
+
+If the checkpoint does not carry a loadable quantization config, add
+--candidate-quantization-config-json, for example '{"method":"fp8"}'.
+"""
+
+from __future__ import annotations
+
+import argparse
+import gc
+import json
+import math
+import time
+from dataclasses import dataclass
+from pathlib import Path
+from typing import Any
+
+import numpy as np
+import torch
+import torch.nn.functional as F
+
+
+@dataclass
+class VariantConfig:
+    label: str
+    model: str
+    quantization: str | None
+    quantization_config: dict[str, Any] | None
+    diffusion_load_format: str
+
+
+@dataclass
+class VariantRun:
+    label: str
+    result: Any
+    generation_times_s: list[float]
+    peak_memory_mb: list[float | None]
+
+
+def metric_guidance() -> dict[str, Any]:
+    """Human-readable metric notes and default heuristic thresholds."""
+    return {
+        "scope": (
+            "Heuristic defaults for same-model comparisons with identical prompt, seed, resolution, scheduler, "
+            "and inference steps. Tune thresholds per model family, task, and quantization method; visual review is "
+            "still recommended for release gating."
+        ),
+        "descriptions": {
+            "output_metrics": (
+                "Pixel-space metrics computed on generated RGB images or video frames after decode. These are often "
+                "the most actionable signal for image/video quality regressions."
+            ),
+            "cosine_similarity": (
+                "Higher is better. Measures vector direction similarity and is less sensitive to scale."
+            ),
+            "mae": (
+                "Lower is better. Mean absolute error in the tensor's native scale; uint8 output MAE is in "
+                "pixel values."
+            ),
+            "mse": "Lower is better. Mean squared error; more sensitive to large localized differences than MAE.",
+            "rmse": "Lower is better. Square root of MSE, in the same unit as the compared tensor.",
+            "max_abs": (
+                "Lower is better. Worst absolute element error; useful for debugging outliers, not stable as a gate."
+            ),
+            "l2": "Lower is better. Absolute L2 distance; shape- and scale-dependent.",
+            "relative_l2": "Lower is better. L2 distance normalized by the reference L2 norm.",
+            "psnr_db": "Higher is better. Pixel-space signal-to-noise ratio in dB for uint8 images/frames.",
+            "generation_time_s": "Lower is better. Wall-clock time for the last measured generation run.",
+            "avg_generation_time_s": "Lower is better. Mean wall-clock time across measured runs.",
+            "peak_memory_mb": (
+                "Lower is better. Worker-reported peak device memory for the last measured generation run."
+            ),
+            "max_peak_memory_mb": "Lower is better. Maximum worker-reported peak memory across measured runs.",
+        },
+        "recommended_thresholds": {
+            "output_images_or_frames_uint8": {
+                "psnr_db": {
+                    "recommended_min": 20.0,
+                    "good_min": 25.0,
+                    "strict_min": 30.0,
+                    "note": (
+                        "For stochastic diffusion outputs, 20 dB is a practical smoke threshold; 25+ dB is usually "
+                        "close."
+                    ),
+                },
+                "mae": {
+                    "recommended_max": 12.0,
+                    "good_max": 6.0,
+                    "strict_max": 3.0,
+                    "note": "Measured on uint8 RGB pixels or frames.",
+                },
+                "cosine_similarity": {
+                    "recommended_min": 0.98,
+                    "strict_min": 0.995,
+                    "note": "Useful as a broad image/frame similarity sanity check.",
+                },
+                "relative_l2": {
+                    "recommended_max": 0.20,
+                    "strict_max": 0.08,
+                    "note": "Pixel-space relative L2 is generally more interpretable than latent-space relative L2.",
+                },
+            },
+            "performance": {
+                "avg_generation_time_ratio_candidate_over_reference": {
+                    "recommended_max": 1.10,
+                    "note": (
+                        "A value above 1.10 suggests a latency regression unless expected for the quantization mode."
+                    ),
+                },
+                "max_peak_memory_ratio_candidate_over_reference": {
+                    "recommended_max": 1.00,
+                    "target_max": 0.95,
+                    "note": (
+                        "Quantization should usually not increase peak memory; <=0.95 indicates a meaningful saving."
+                    ),
+                },
+            },
+        },
+    }
+
+
+def _split_csv(value: str | None) -> list[str]:
+    if not value:
+        return []
+    return [item.strip() for item in value.split(",") if item.strip()]
+
+
+def _load_json_arg(value: str | None) -> dict[str, Any] | None:
+    if value is None:
+        return None
+    path = Path(value).expanduser()
+    if path.exists():
+        value = path.read_text(encoding="utf-8")
+    parsed = json.loads(value)
+    if not isinstance(parsed, dict):
+        raise ValueError("quantization config JSON must decode to an object.")
+    return parsed
+
+
+def _cosine_similarity(flat_a: torch.Tensor, flat_b: torch.Tensor) -> float:
+    norm_a = torch.linalg.vector_norm(flat_a).item()
+    norm_b = torch.linalg.vector_norm(flat_b).item()
+    if norm_a == 0.0 and norm_b == 0.0:
+        return 1.0
+    if norm_a == 0.0 or norm_b == 0.0:
+        return 0.0
+    cosine = float(F.cosine_similarity(flat_a, flat_b, dim=0).item())
+    return max(-1.0, min(1.0, cosine))
+
+
+def compute_tensor_metrics(lhs: Any, rhs: Any) -> dict[str, float]:
+    lhs_tensor = torch.as_tensor(lhs).detach().cpu().double()
+    rhs_tensor = torch.as_tensor(rhs).detach().cpu().double()
+    if lhs_tensor.shape != rhs_tensor.shape:
+        raise ValueError(f"Metric shape mismatch: {tuple(lhs_tensor.shape)} vs {tuple(rhs_tensor.shape)}")
+
+    diff = lhs_tensor - rhs_tensor
+    mse = float(diff.square().mean().item())
+    rmse = float(math.sqrt(mse))
+    mae = float(diff.abs().mean().item())
+    max_abs = float(diff.abs().max().item())
+    l2 = float(torch.linalg.vector_norm(diff).item())
+    ref_l2 = float(torch.linalg.vector_norm(lhs_tensor).item())
+    relative_l2 = 0.0 if ref_l2 == 0.0 and l2 == 0.0 else float("inf") if ref_l2 == 0.0 else l2 / ref_l2
+    cosine = _cosine_similarity(lhs_tensor.reshape(-1), rhs_tensor.reshape(-1))
+    return {
+        "cosine_similarity": cosine,
+        "mae": mae,
+        "mse": mse,
+        "rmse": rmse,
+        "max_abs": max_abs,
+        "l2": l2,
+        "relative_l2": relative_l2,
+    }
+
+
+def compute_uint8_image_metrics(lhs: Any, rhs: Any) -> dict[str, float]:
+    metrics = compute_tensor_metrics(lhs, rhs)
+    mse = metrics["mse"]
+    metrics["psnr_db"] = float("inf") if mse == 0.0 else 20 * math.log10(255.0) - 10 * math.log10(mse)
+    return metrics
+
+
+def summarize_output_frame_metrics(reference_frames: list[Any], candidate_frames: list[Any]) -> dict[str, Any]:
+    if len(reference_frames) != len(candidate_frames):
+        raise ValueError(f"Output frame count mismatch: {len(reference_frames)} vs {len(candidate_frames)}")
+    if not reference_frames:
+        raise ValueError("No output frames available for comparison.")
+
+    ref_stack = np.stack([np.asarray(frame.convert("RGB")) for frame in reference_frames], axis=0)
+    cand_stack = np.stack([np.asarray(frame.convert("RGB")) for frame in candidate_frames], axis=0)
+
+    frame0_metrics = compute_uint8_image_metrics(ref_stack[0], cand_stack[0])
+    mid_index = len(reference_frames) // 2
+    mid_metrics = compute_uint8_image_metrics(ref_stack[mid_index], cand_stack[mid_index])
+    all_metrics = compute_uint8_image_metrics(ref_stack, cand_stack)
+
+    return {
+        "num_frames": len(reference_frames),
+        "frame0_metrics": frame0_metrics,
+        "mid_frame_index": mid_index,
+        "mid_frame_metrics": mid_metrics,
+        "all_frames_metrics": all_metrics,
+        # Backward-compatible aliases for existing image callers.
+        "num_images": len(reference_frames),
+        "image0_metrics": frame0_metrics,
+        "mid_image_index": mid_index,
+        "mid_image_metrics": mid_metrics,
+        "all_images_metrics": all_metrics,
+    }
+
+
+def summarize_output_image_metrics(reference_images: list[Any], candidate_images: list[Any]) -> dict[str, Any]:
+    return summarize_output_frame_metrics(reference_images, candidate_images)
+
+
+def _extract_inner_output(outputs: Any) -> Any:
+    if isinstance(outputs, list):
+        if len(outputs) != 1:
+            raise ValueError(f"Expected one request output, got {len(outputs)}.")
+        outputs = outputs[0]
+    inner = getattr(outputs, "request_output", None)
+    if inner is not None:
+        return inner
+    return outputs
+
+
+def _request_peak_memory_mb(result: Any) -> float | None:
+    diffusion_output = getattr(result, "request_output", None)
+    peak_memory_mb = getattr(diffusion_output, "peak_memory_mb", None)
+    if peak_memory_mb is None:
+        peak_memory_mb = getattr(result, "peak_memory_mb", None)
+    if peak_memory_mb is None:
+        return None
+    return float(peak_memory_mb)
+
+
+def _build_quantization_config(
+    *,
+    quantization: str | None,
+    quantization_config_json: str | None,
+    ignored_layers: str | None,
+    gguf_model: str | None,
+) -> dict[str, Any] | None:
+    explicit = _load_json_arg(quantization_config_json)
+    if explicit is not None:
+        return explicit
+    if quantization is None:
+        return None
+
+    config: dict[str, Any] = {"method": quantization}
+    ignored = _split_csv(ignored_layers)
+    if ignored:
+        config["ignored_layers"] = ignored
+    if gguf_model is not None:
+        config["gguf_model"] = gguf_model
+    return config
+
+
+def _build_variant_config(args: argparse.Namespace, variant: str) -> VariantConfig:
+    common_ignored_layers = args.ignored_layers if variant == "candidate" else None
+    ignored_layers = getattr(args, f"{variant}_ignored_layers") or common_ignored_layers
+    quantization = getattr(args, f"{variant}_quantization")
+    quantization_config = _build_quantization_config(
+        quantization=quantization,
+        quantization_config_json=getattr(args, f"{variant}_quantization_config_json"),
+        ignored_layers=ignored_layers,
+        gguf_model=getattr(args, f"{variant}_gguf_model"),
+    )
+    return VariantConfig(
+        label=variant,
+        model=getattr(args, f"{variant}_model") or args.model,
+        quantization=quantization if quantization_config is None else None,
+        quantization_config=quantization_config,
+        diffusion_load_format=getattr(args, f"{variant}_load_format"),
+    )
+
+
+def _build_omni_kwargs(args: argparse.Namespace, config: VariantConfig) -> dict[str, Any]:
+    from vllm_omni.diffusion.data import DiffusionParallelConfig
+
+    parallel_config = DiffusionParallelConfig(
+        ulysses_degree=args.ulysses_degree,
+        ring_degree=args.ring_degree,
+        cfg_parallel_size=args.cfg_parallel_size,
+        tensor_parallel_size=args.tensor_parallel_size,
+        vae_patch_parallel_size=args.vae_patch_parallel_size,
+    )
+    kwargs: dict[str, Any] = {
+        "model": config.model,
+        "mode": "text-to-video" if args.task == "t2v" else "text-to-image",
+        "parallel_config": parallel_config,
+        "enforce_eager": args.enforce_eager,
+        "enable_cpu_offload": args.enable_cpu_offload,
+        "enable_layerwise_offload": args.enable_layerwise_offload,
+        "vae_use_slicing": args.vae_use_slicing,
+        "vae_use_tiling": args.vae_use_tiling,
+        "step_execution": args.step_execution if args.task == "t2i" else False,
+        "diffusion_load_format": config.diffusion_load_format,
+    }
+    if args.model_class_name:
+        kwargs["model_class_name"] = args.model_class_name
+    if args.flow_shift is not None:
+        kwargs["flow_shift"] = args.flow_shift
+    if config.quantization_config is not None:
+        kwargs["quantization_config"] = config.quantization_config
+    elif config.quantization is not None:
+        kwargs["quantization"] = config.quantization
+    return kwargs
+
+
+def _build_sampling_params(args: argparse.Namespace, seed: int):
+    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(seed)
+    kwargs: dict[str, Any] = {
+        "height": args.height,
+        "width": args.width,
+        "generator": generator,
+        "seed": seed,
+        "guidance_scale": args.guidance_scale,
+        "guidance_scale_2": args.guidance_scale_2,
+        "num_inference_steps": args.num_inference_steps,
+        "return_frames": True,
+        "extra_args": {
+            "timesteps_shift": args.timesteps_shift,
+            "cfg_schedule": args.cfg_schedule,
+            "use_norm": args.use_norm,
+            "use_system_prompt": args.use_system_prompt,
+            "system_prompt": args.system_prompt,
+        },
+    }
+    if args.task == "t2v":
+        kwargs["num_frames"] = args.num_frames
+        kwargs["fps"] = args.fps
+        kwargs["frame_rate"] = args.frame_rate
+    else:
+        kwargs["true_cfg_scale"] = args.cfg_scale
+        kwargs["num_outputs_per_prompt"] = args.num_images_per_prompt
+    return OmniDiffusionSamplingParams(**kwargs)
+
+
+def _run_variant(args: argparse.Namespace, config: VariantConfig) -> VariantRun:
+    from vllm_omni.entrypoints.omni import Omni
+    from vllm_omni.platforms import current_omni_platform
+
+    omni = Omni(**_build_omni_kwargs(args, config))
+    measured_results: list[Any] = []
+    generation_times: list[float] = []
+    peak_memory_mb: list[float | None] = []
+    try:
+        for _ in range(args.warmup_runs):
+            sampling_params = _build_sampling_params(args, args.seed)
+            omni.generate(
+                {"prompt": args.prompt, "negative_prompt": args.negative_prompt},
+                sampling_params,
+                use_tqdm=False,
+            )
+
+        for _ in range(args.measure_runs):
+            sampling_params = _build_sampling_params(args, args.seed)
+            start = time.perf_counter()
+            outputs = omni.generate(
+                {"prompt": args.prompt, "negative_prompt": args.negative_prompt},
+                sampling_params,
+                use_tqdm=False,
+            )
+            current_omni_platform.synchronize()
+            generation_times.append(time.perf_counter() - start)
+            result = _extract_inner_output(outputs)
+            peak_memory_mb.append(_request_peak_memory_mb(result))
+            measured_results.append(result)
+    finally:
+        omni.close()
+        del omni
+        gc.collect()
+        current_omni_platform.empty_cache()
+
+    if not measured_results:
+        raise RuntimeError("No measured results were produced.")
+    return VariantRun(
+        label=config.label,
+        result=measured_results[-1],
+        generation_times_s=generation_times,
+        peak_memory_mb=peak_memory_mb,
+    )
+
+
+def _get_output_frames(result: Any) -> list[Any]:
+    images = list(getattr(result, "images", []) or [])
+    if len(images) == 1 and isinstance(images[0], list):
+        return list(images[0])
+    return images
+
+
+def _save_outputs(result: Any, output_dir: Path, label: str, task: str, fps: int) -> list[str]:
+    frames = _get_output_frames(result)
+    saved: list[str] = []
+    if not frames:
+        return saved
+    variant_dir = output_dir / label
+    variant_dir.mkdir(parents=True, exist_ok=True)
+    if task == "t2v":
+        from diffusers.utils import export_to_video
+
+        video_path = variant_dir / "video.mp4"
+        export_to_video(frames, str(video_path), fps=fps)
+        saved.append(str(video_path))
+        for idx in (0, len(frames) // 2, len(frames) - 1):
+            path = variant_dir / f"frame_{idx}.png"
+            frames[idx].save(path)
+            saved.append(str(path))
+        return saved
+    for idx, image in enumerate(frames):
+        path = variant_dir / f"image_{idx}.png"
+        image.save(path)
+        saved.append(str(path))
+    return saved
+
+
+def _run_summary(run: VariantRun) -> dict[str, Any]:
+    valid_peak = [value for value in run.peak_memory_mb if value is not None]
+    return {
+        "generation_time_s": run.generation_times_s[-1],
+        "avg_generation_time_s": sum(run.generation_times_s) / len(run.generation_times_s),
+        "per_run_generation_time_s": run.generation_times_s,
+        "peak_memory_mb": valid_peak[-1] if valid_peak else None,
+        "avg_peak_memory_mb": sum(valid_peak) / len(valid_peak) if valid_peak else None,
+        "max_peak_memory_mb": max(valid_peak) if valid_peak else None,
+        "per_run_peak_memory_mb": run.peak_memory_mb,
+    }
+
+
+def _to_jsonable(result: dict[str, Any]) -> dict[str, Any]:
+    return json.loads(json.dumps(result, allow_nan=True))
+
+
+def run_comparison(args: argparse.Namespace) -> dict[str, Any]:
+    if args.warmup_runs < 0:
+        raise ValueError("--warmup-runs must be >= 0.")
+    if args.measure_runs <= 0:
+        raise ValueError("--measure-runs must be >= 1.")
+
+    reference_config = _build_variant_config(args, "reference")
+    candidate_config = _build_variant_config(args, "candidate")
+
+    reference_run = _run_variant(args, reference_config)
+    candidate_run = _run_variant(args, candidate_config)
+
+    save_output_paths: dict[str, list[str]] = {}
+    if args.save_output_dir:
+        save_dir = Path(args.save_output_dir).expanduser().resolve()
+        save_output_paths["reference"] = _save_outputs(reference_run.result, save_dir, "reference", args.task, args.fps)
+        save_output_paths["candidate"] = _save_outputs(candidate_run.result, save_dir, "candidate", args.task, args.fps)
+
+    result = {
+        "model": args.model,
+        "prompt": args.prompt,
+        "negative_prompt": args.negative_prompt,
+        "seed": args.seed,
+        "warmup_runs": args.warmup_runs,
+        "measure_runs": args.measure_runs,
+        "variant_configs": {
+            "reference": reference_config.__dict__,
+            "candidate": candidate_config.__dict__,
+        },
+        "sampling_kwargs": {
+            "task": args.task,
+            "width": args.width,
+            "height": args.height,
+            "num_frames": args.num_frames if args.task == "t2v" else None,
+            "num_inference_steps": args.num_inference_steps,
+            "guidance_scale": args.guidance_scale,
+            "cfg_scale": args.cfg_scale,
+            "guidance_scale_2": args.guidance_scale_2,
+            "flow_shift": args.flow_shift,
+            "num_images_per_prompt": args.num_images_per_prompt,
+            "step_execution": args.step_execution if args.task == "t2i" else False,
+        },
+        "metric_guidance": metric_guidance(),
+        "reference_generation": _run_summary(reference_run),
+        "candidate_generation": _run_summary(candidate_run),
+        "output_metrics": summarize_output_image_metrics(
+            _get_output_frames(reference_run.result),
+            _get_output_frames(candidate_run.result),
+        ),
+    }
+    if save_output_paths:
+        result["saved_outputs"] = save_output_paths
+    return _to_jsonable(result)
+
+
+def parse_args() -> argparse.Namespace:
+    parser = argparse.ArgumentParser(
+        description=__doc__,
+        formatter_class=argparse.RawDescriptionHelpFormatter,
+    )
+    parser.add_argument("--task", choices=["t2i", "t2v"], default="t2i")
+    parser.add_argument("--model", default="Qwen/Qwen-Image", help="Default model for both variants.")
+    parser.add_argument("--model-class-name", help="Optional diffusion pipeline class override.")
+    parser.add_argument("--reference-model", help="Reference model path or HF ID. Defaults to --model.")
+    parser.add_argument("--candidate-model", help="Candidate model path or HF ID. Defaults to --model.")
+    parser.add_argument("--reference-quantization", help="Reference online quantization method.")
+    parser.add_argument("--candidate-quantization", help="Candidate online quantization method, e.g. fp8.")
+    parser.add_argument(
+        "--reference-quantization-config-json",
+        help="Reference quantization_config JSON string or file.",
+    )
+    parser.add_argument(
+        "--candidate-quantization-config-json",
+        help="Candidate quantization_config JSON string or file.",
+    )
+    parser.add_argument("--reference-ignored-layers", help="Comma-separated reference ignored layer patterns.")
+    parser.add_argument("--candidate-ignored-layers", help="Comma-separated candidate ignored layer patterns.")
+    parser.add_argument("--ignored-layers", help="Alias for --candidate-ignored-layers.")
+    parser.add_argument("--reference-gguf-model", help="Reference GGUF file or HF reference when using gguf.")
+    parser.add_argument("--candidate-gguf-model", help="Candidate GGUF file or HF reference when using gguf.")
+    parser.add_argument("--reference-load-format", default="default", help="Reference diffusion load format.")
+    parser.add_argument("--candidate-load-format", default="default", help="Candidate diffusion load format.")
+    parser.add_argument("--prompt", default="a cup of coffee on the table")
+    parser.add_argument("--negative-prompt")
+    parser.add_argument("--output-json", required=True)
+    parser.add_argument("--save-output-dir")
+    parser.add_argument("--width", type=int, default=1024)
+    parser.add_argument("--height", type=int, default=1024)
+    parser.add_argument("--num-frames", type=int, default=1, help="Number of frames for t2v.")
+    parser.add_argument("--fps", type=int, default=24, help="FPS for saved t2v videos.")
+    parser.add_argument("--frame-rate", type=float, help="Optional generation frame rate for t2v.")
+    parser.add_argument("--num-inference-steps", type=int, default=50)
+    parser.add_argument("--num-images-per-prompt", type=int, default=1)
+    parser.add_argument("--guidance-scale", type=float, default=4.0)
+    parser.add_argument("--cfg-scale", type=float, default=4.0, help="Qwen-Image true CFG scale.")
+    parser.add_argument("--guidance-scale-2", type=float)
+    parser.add_argument("--seed", type=int, default=42)
+    parser.add_argument("--warmup-runs", type=int, default=0)
+    parser.add_argument("--measure-runs", type=int, default=1)
+    parser.add_argument("--ulysses-degree", type=int, default=1)
+    parser.add_argument("--ring-degree", type=int, default=1)
+    parser.add_argument("--cfg-parallel-size", type=int, default=1, choices=[1, 2])
+    parser.add_argument("--tensor-parallel-size", type=int, default=1)
+    parser.add_argument("--vae-patch-parallel-size", type=int, default=1)
+    parser.add_argument("--enforce-eager", action="store_true")
+    parser.add_argument("--enable-cpu-offload", action="store_true")
+    parser.add_argument("--enable-layerwise-offload", action="store_true")
+    parser.add_argument("--vae-use-slicing", action="store_true")
+    parser.add_argument("--vae-use-tiling", action="store_true")
+    parser.add_argument("--flow-shift", type=float)
+    parser.add_argument(
+        "--step-execution",
+        action=argparse.BooleanOptionalAction,
+        default=False,
+        help=(
+            "Use the diffusion step scheduler for text-to-image models. Disabled by default so the tool follows the "
+            "normal offline inference path unless explicitly requested."
+        ),
+    )
+    parser.add_argument("--timesteps-shift", type=float, default=1.0)
+    parser.add_argument("--cfg-schedule", choices=["constant", "linear"], default="constant")
+    parser.add_argument("--use-norm", action="store_true")
+    parser.add_argument(
+        "--use-system-prompt",
+        choices=[
+            "None",
+            "dynamic",
+            "en_vanilla",
+            "en_recaption",
+            "en_think_recaption",
+            "en_unified",
+            "custom",
+        ],
+    )
+    parser.add_argument("--system-prompt")
+    return parser.parse_args()
+
+
+def main() -> None:
+    args = parse_args()
+    output_json = Path(args.output_json).expanduser().resolve()
+    output_json.parent.mkdir(parents=True, exist_ok=True)
+
+    result = run_comparison(args)
+    output_json.write_text(json.dumps(result, indent=2, sort_keys=True), encoding="utf-8")
+
+    image0 = result["output_metrics"]["image0_metrics"]
+    summary = {
+        "output_json": str(output_json),
+        "reference_avg_generation_time_s": result["reference_generation"]["avg_generation_time_s"],
+        "candidate_avg_generation_time_s": result["candidate_generation"]["avg_generation_time_s"],
+        "image0_psnr_db": image0["psnr_db"],
+        "image0_mae": image0["mae"],
+    }
+    print(json.dumps(summary, indent=2))
+
+
+if __name__ == "__main__":
+    main()