Blame for generate5.py

#!/usr/bin/env python3
"""Reproduce the 5 example images and write outputs/benchmark.json + outputs/prompts.txt.

Same layout as HelloSun/FLUX.2-klein-4B-OpenVINO-INT4 so the two model cards can be
compared side by side. Parameters follow the SD3.5-medium original model card:

    1024x1024, num_inference_steps=40, guidance_scale=4.5, seeds 42-46

Warning: 40 steps on CPU is the expensive part. On a 16 vCPU box expect minutes
per image, not seconds. Use --steps 28 for a quick pass, or --size 512 to
pre-flight the pipeline, then re-run at full settings for the published numbers.

Usage
-----
    python generate5.py
    python generate5.py --steps 28 --size 768 --outdir outputs_preview
"""

from __future__ import annotations

import argparse
import json
import os
import platform
import threading
import time
from pathlib import Path

from PIL import Image

REPO_ID = "HelloSun/SD3.5-medium-OpenVINO-INT4"

# Original model card parameters.
DEFAULT_STEPS = 40
DEFAULT_GUIDANCE = 4.5
DEFAULT_SIZE = 1024

PROMPTS = [
    (
        "01_hanfu",
        42,
        "Young Chinese woman in red Hanfu, intricate embroidery, impeccable makeup, "
        "red floral forehead pattern, elaborate high bun, golden phoenix headdress, "
        "soft-lit outdoor night background, silhouetted tiered pagoda, blurred colorful "
        "distant lights, photorealistic, ultra detailed, 8k",
    ),
    (
        "02_astronaut",
        43,
        "Astronaut in a jungle, cold color palette, muted colors, detailed, "
        "8k, photorealistic, cinematic lighting",
    ),
    (
        "03_taipei",
        44,
        "Cyberpunk street in Taipei at night, heavy rain, neon signs with text 'TAIPEI' "
        "and Chinese characters '台北', reflections on wet asphalt, crowded night market, "
        "cinematic, ultra detailed",
    ),
    (
        "04_shiba",
        45,
        "Cute Shiba Inu wearing a tiny astronaut helmet, sitting in a field of sunflowers "
        "under a starry sky, dreamy illustration, vibrant colors, high quality",
    ),
    (
        "05_ink",
        46,
        "Traditional Chinese ink wash landscape, misty mountains, a small pagoda on a cliff, "
        "cranes flying, minimalist, elegant, high aesthetic quality",
    ),
]


def _vmhwm_mb() -> float:
    """Peak RSS high-water mark in MB from /proc. Exact, and dependency-free.

    VmHWM is monotonic, so a single read at the end gives the true peak for the
    whole process lifetime. Returns 0.0 if unavailable (non-Linux).
    """
    try:
        for line in Path("/proc/self/status").read_text().splitlines():
            if line.startswith("VmHWM:"):
                return int(line.split()[1]) / 1024
    except (OSError, ValueError, IndexError):
        pass
    return 0.0


class PeakRss:
    """Tracks peak resident memory for this process.

    Prefers the kernel's VmHWM counter (exact high-water mark, no sampling
    error, no extra dependency). Falls back to sampling RSS via psutil only if
    VmHWM is unavailable, and to 0.0 if neither works - in which case
    `peak_rss_mb` is reported as 0 rather than being silently wrong.
    """

    def __init__(self, interval: float = 0.2) -> None:
        self.interval = interval
        self.peak_mb = 0.0
        self.source = "none"
        self._stop = threading.Event()
        self._thread: threading.Thread | None = None

    def _sample_psutil(self) -> None:
        import psutil

        proc = psutil.Process(os.getpid())
        while not self._stop.is_set():
            self.peak_mb = max(self.peak_mb, proc.memory_info().rss / 1024**2)
            self._stop.wait(self.interval)

    def _start(self) -> None:
        if _vmhwm_mb() > 0.0:
            self.source = "/proc/self/status VmHWM"
            return
        try:
            import psutil  # noqa: F401
        except ImportError:
            print(
                "WARNING: no VmHWM and no psutil; peak_rss_mb will be reported as 0. "
                "Install psutil to get memory numbers.",
                flush=True,
            )
            return
        self.source = "psutil sampled"
        self._thread = threading.Thread(target=self._sample_psutil, daemon=True)
        self._thread.start()

    def _refresh(self) -> None:
        # VmHWM is monotonic and authoritative, so always fold it in even when
        # psutil sampling is the active source.
        hwm = _vmhwm_mb()
        if hwm > 0.0:
            self.peak_mb = max(self.peak_mb, hwm)
        elif self.source == "none":
            self.source = "unavailable"

    def __enter__(self) -> PeakRss:
        self._start()
        self._refresh()
        return self

    def __exit__(self, *exc) -> None:
        self._stop.set()
        if self._thread:
            self._thread.join(timeout=2 * self.interval)
        self._refresh()


def cpu_name() -> str:
    try:
        for line in Path("/proc/cpuinfo").read_text().splitlines():
            if line.startswith("model name"):
                return line.split(":", 1)[1].strip()
    except OSError:
        pass
    return platform.processor() or "unknown"


def main() -> int:
    parser = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
    parser.add_argument("--model", default=REPO_ID, help="local path or Hub id")
    parser.add_argument("--steps", type=int, default=DEFAULT_STEPS)
    parser.add_argument("--guidance", type=float, default=DEFAULT_GUIDANCE)
    parser.add_argument("--size", type=int, default=DEFAULT_SIZE)
    parser.add_argument("--preview-size", type=int, default=512, help="0 disables the downscaled copies")
    parser.add_argument("--outdir", type=Path, default=Path("outputs"))
    parser.add_argument("--device", default="CPU")
    parser.add_argument("--no-compile", action="store_true")
    args = parser.parse_args()

    import torch
    from optimum.intel import OVDiffusionPipeline

    args.outdir.mkdir(parents=True, exist_ok=True)

    print(f"loading {args.model} ...", flush=True)
    t0 = time.perf_counter()
    pipe = OVDiffusionPipeline.from_pretrained(MODEL_ID, device=args.device, compile=not args.no_compile)
    load_seconds = time.perf_counter() - t0
    print(f"load+compile {load_seconds:.2f}s", flush=True)

    try:
        import openvino

        ov_version = openvino.__version__
    except Exception:
        ov_version = "unknown"

    results = []
    with PeakRss() as rss:
        for name, seed, prompt in PROMPTS:
            print(f"--- {name} (seed {seed}) ---", flush=True)
            started = time.perf_counter()
            image = pipe(
                prompt=prompt,
                num_inference_steps=args.steps,
                guidance_scale=args.guidance,
                width=args.size,
                height=args.size,
                generator=torch.Generator().manual_seed(seed),
            ).images[0]
            elapsed = time.perf_counter() - started

            image.save(args.outdir / f"{name}_{args.size}.png")
            if args.preview_size:
                try:
                    resample = Image.Resampling.LANCZOS  # Pillow >= 9.1
                except AttributeError:
                    resample = Image.LANCZOS  # Pillow < 9.1
                image.resize((args.preview_size, args.preview_size), resample).save(
                    args.outdir / f"{name}_{args.preview_size}.png"
                )

            entry = {
                "name": name,
                "seed": seed,
                "steps": args.steps,
                "guidance": args.guidance,
                "size": f"{args.size}x{args.size}",
                "total_seconds": round(elapsed, 2),
                "seconds_per_step": round(elapsed / max(args.steps, 1), 2),
                "peak_rss_mb": round(rss.peak_mb, 1),
                "prompt": prompt,
            }
            results.append(entry)
            print(json.dumps(entry, ensure_ascii=False), flush=True)

    benchmark = {
        "model_id": REPO_ID,
        "hardware": {
            "cpu": cpu_name(),
            "logical_cores": os.cpu_count(),
            "device": args.device,
        },
        "software": {
            "openvino": ov_version,
            "python": platform.python_version(),
            "platform": platform.platform(),
        },
        "load_compile_seconds": round(load_seconds, 2),
        "params": {
            "num_inference_steps": args.steps,
            "guidance_scale": args.guidance,
            "height": args.size,
            "width": args.size,
        },
        "results": results,
        "peak_rss_mb": round(max(r["peak_rss_mb"] for r in results), 1),
        "peak_rss_source": rss.source,
    }
    (args.outdir / "benchmark.json").write_text(
        json.dumps(benchmark, indent=2, ensure_ascii=False) + "\n", encoding="utf-8"
    )
    (args.outdir / "prompts.txt").write_text(
        "".join(f"{n}\tseed={s}\t{p}\n\n" for n, s, p in PROMPTS), encoding="utf-8"
    )

    print(f"\nwrote {args.outdir / 'benchmark.json'} and {args.outdir / 'prompts.txt'}", flush=True)
    print("paste the table into README.md 實驗數據 when publishing.", flush=True)
    return 0


if __name__ == "__main__":
    raise SystemExit(main())