#!/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())