benchmark_pixart.py

"""Benchmark PixArt-Sigma hero inference with PyTorch on a ROCm host."""

from __future__ import annotations

import argparse
import json
from pathlib import Path
from statistics import mean
from time import perf_counter
from typing import Any

from pixart import (
    MODEL_CARD_PARAMETER_COUNT,
    MODEL_ID,
    MODEL_LICENSE,
    _load_pipeline,
    _parameter_counts,
    _percentile,
    _rocm_smi_snapshot,
    _run_pipeline,
    _torch,
    runtime_info,
)


def benchmark_pixart(
    *,
    prompt: str,
    warmups: int,
    runs: int,
    seed: int,
    inference_steps: int,
    guidance_scale: float,
    compile_transformer: bool,
    height: int = 1024,
    width: int = 1024,
) -> dict[str, Any]:
    """Measure synchronized PixArt hero generations with ``torch.utils.benchmark``.

    The timer includes a complete 1024px denoising pass and waits for the ROCm
    device before it stops. Model loading and warmup are reported separately
    from steady-state image latency.
    """

    if warmups < 0 or runs < 1:
        raise ValueError("warmups must be non-negative and runs must be positive.")
    if height <= 0 or width <= 0 or height % 8 or width % 8:
        raise ValueError("PixArt image width and height must be positive multiples of 8.")

    runtime = runtime_info(require_rocm=True)
    torch = _torch()
    from torch.utils.benchmark import Timer

    load_started = perf_counter()
    pipeline = _load_pipeline(MODEL_ID, compile_transformer)
    pipeline_load_ms = (perf_counter() - load_started) * 1000.0

    def run_once() -> None:
        # ROCm work is asynchronous; synchronize so the timer measures generation.
        _run_pipeline(
            pipeline,
            prompt=prompt,
            height=height,
            width=width,
            inference_steps=inference_steps,
            guidance_scale=guidance_scale,
            seed=seed,
        )
        torch.cuda.synchronize()

    for _ in range(warmups):
        run_once()
    torch.cuda.synchronize()
    torch.cuda.reset_peak_memory_stats()

    # One synchronized image generation per sample keeps p50/p95 directly
    # representative of the application's hero-image workload.
    timer = Timer(
        stmt="run_once()",
        globals={"run_once": run_once},
        label="PixArt-Sigma hero inference",
        sub_label=f"{width}x{height}, {inference_steps} steps",
        description="One synchronized image generation",
    )
    measurements = [timer.timeit(number=1) for _ in range(runs)]
    latencies_ms = [measurement.mean * 1000.0 for measurement in measurements]
    mean_ms = mean(latencies_ms)

    return {
        "benchmark": "pixart_sigma_rocm_inference",
        "model": {
            "id": MODEL_ID,
            "model_card_parameter_count": MODEL_CARD_PARAMETER_COUNT,
            "loaded_parameter_counts": _parameter_counts(pipeline),
            "license": MODEL_LICENSE,
        },
        "runtime": runtime,
        "configuration": {
            "width": width,
            "height": height,
            "inference_steps": inference_steps,
            "guidance_scale": guidance_scale,
            "warmups": warmups,
            "runs": runs,
            "seed": seed,
            "compiled_transformer": compile_transformer,
        },
        "pytorch_benchmark": {
            "timer": "torch.utils.benchmark.Timer",
            "statement": "run_once()",
            "calls_per_sample": 1,
            "synchronizes_rocm_after_each_sample": True,
        },
        "latency_ms": {
            "pipeline_load_ms": round(pipeline_load_ms, 2),
            "samples": [round(value, 2) for value in latencies_ms],
            "mean": round(mean_ms, 2),
            "p50": round(_percentile(latencies_ms, 0.50), 2),
            "p95": round(_percentile(latencies_ms, 0.95), 2),
            "throughput_images_per_s": round(1000.0 / mean_ms, 4),
        },
        "memory_mib": {
            "peak_allocated": round(torch.cuda.max_memory_allocated() / (1024**2), 2),
            "peak_reserved": round(torch.cuda.max_memory_reserved() / (1024**2), 2),
        },
        "rocm_smi": _rocm_smi_snapshot(),
        "accuracy": {
            "applicable": False,
            "value": None,
            "reason": "Text-to-image generation has no ground-truth classification accuracy. Use a predefined prompt-adherence or human-quality protocol before reporting a quality score.",
        },
    }


def main() -> None:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--prompt", default="An editorial hero image of a bowl of ramen, warm evening light, no text or logos.")
    parser.add_argument("--warmups", type=int, default=1)
    parser.add_argument("--runs", type=int, default=3)
    parser.add_argument("--steps", type=int, default=20)
    parser.add_argument("--guidance-scale", type=float, default=4.5)
    parser.add_argument("--seed", type=int, default=42)
    parser.add_argument("--height", type=int, default=1024)
    parser.add_argument("--width", type=int, default=1024)
    parser.add_argument("--compile-transformer", action="store_true")
    parser.add_argument("--output", type=Path, default=Path("artifacts/benchmarks/pixart-rocm.json"))
    args = parser.parse_args()

    report = benchmark_pixart(
        prompt=args.prompt,
        warmups=args.warmups,
        runs=args.runs,
        seed=args.seed,
        inference_steps=args.steps,
        guidance_scale=args.guidance_scale,
        compile_transformer=args.compile_transformer,
        height=args.height,
        width=args.width,
    )
    args.output.parent.mkdir(parents=True, exist_ok=True)
    args.output.write_text(json.dumps(report, indent=2), encoding="utf-8")
    print(json.dumps(report, indent=2))
    print(f"\nWrote {args.output}")


if __name__ == "__main__":
    main()