diff --git a/docs/user_guide/quantization/overview.md b/docs/user_guide/quantization/overview.md index c4605c8c2..39669bf5d 100644 --- a/docs/user_guide/quantization/overview.md +++ b/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. diff --git a/tests/diffusion/quantization/test_trajectory_similarity_tool.py b/tests/diffusion/quantization/test_trajectory_similarity_tool.py new file mode 100644 index 000000000..45d51d356 --- /dev/null +++ b/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 diff --git a/vllm_omni/quantization/tools/__init__.py b/vllm_omni/quantization/tools/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/vllm_omni/quantization/tools/compare_diffusion_trajectory_similarity.py b/vllm_omni/quantization/tools/compare_diffusion_trajectory_similarity.py new file mode 100644 index 000000000..3d4f68d23 --- /dev/null +++ b/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()