Add generate5_flux.py

generate5_flux.py

@@ -0,0 +1,257 @@
+"""Generate the 5 fixed-seed 1024x1024 examples plus the 512px control set,
+recording per-step latency and process RSS via a diffusers callback.
+For FLUX.1-schnell (1-4 steps, CFG 0.0, max_sequence_length=256).
+
+Outputs -> /home/user/app/outputs_flux/ and a machine readable log in
+/home/user/app/outputs_flux/benchmark.json
+"""
+
+import argparse
+import json
+import os
+import platform
+import statistics
+import time
+from pathlib import Path
+
+import psutil
+import torch
+from optimum.intel import OVFluxPipeline
+
+MODEL_PATH = "/home/user/app/flux-schnell-ov-int4"
+OUTPUT_DIR = Path("/home/user/app/outputs_flux")
+
+NEGATIVE_PROMPT = ""
+
+# name, seed, prompt
+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 rss_mb() -> float:
+    return psutil.Process(os.getpid()).memory_info().rss / 1024**2
+
+
+class StepProfiler:
+    """diffusers callback: records wall time and RSS after every denoising step."""
+
+    def __init__(self, total_steps: int):
+        self.total_steps = total_steps
+        self.times: list[float] = []
+        self.rss: list[float] = []
+        self._last = time.perf_counter()
+        self.t_start = self._last
+
+    def __call__(self, pipe, step_index, timestep, callback_kwargs):
+        now = time.perf_counter()
+        self.times.append(now - self._last)
+        self._last = now
+        self.rss.append(rss_mb())
+        print(
+            f"    step {step_index + 1}/{self.total_steps}: "
+            f"{self.times[-1]:.3f}s  rss={self.rss[-1]:.0f}MB",
+            flush=True,
+        )
+        return callback_kwargs
+
+
+def system_info() -> dict:
+    info = {
+        "platform": platform.platform(),
+        "python": platform.python_version(),
+        "logical_cpus_os_cpu_count": os.cpu_count(),
+        "psutil_physical_cores": psutil.cpu_count(logical=False),
+        "psutil_logical_cores": psutil.cpu_count(logical=True),
+        "total_ram_gb": round(psutil.virtual_memory().total / 1024**3, 1),
+    }
+    try:
+        import openvino
+
+        info["openvino_version"] = openvino.__version__
+    except Exception:
+        pass
+    import importlib.metadata as md
+
+    for pkg in [
+        "optimum",
+        "optimum-intel",
+        "diffusers",
+        "transformers",
+        "tokenizers",
+        "huggingface-hub",
+        "nncf",
+        "torch",
+        "pillow",
+        "psutil",
+    ]:
+        try:
+            info[f"pkg_{pkg}"] = md.version(pkg)
+        except Exception:
+            pass
+    try:
+        out = os.popen("lscpu").read()
+        for line in out.splitlines():
+            if line.startswith("Model name"):
+                info["cpu_model"] = line.split(":", 1)[1].strip()
+    except Exception:
+        pass
+    return info
+
+
+def run_one(pipe, name, seed, prompt, size, steps, guidance, max_seq_len, out_dir):
+    print(f"[{name}] {size[0]}x{size[1]} seed={seed} steps={steps} cfg={guidance}", flush=True)
+    rss_before = rss_mb()
+    profiler = StepProfiler(steps)
+    generator = torch.Generator(device="cpu").manual_seed(seed)
+    t0 = time.perf_counter()
+    result = pipe(
+        prompt=prompt,
+        negative_prompt=NEGATIVE_PROMPT,
+        width=size[0],
+        height=size[1],
+        num_inference_steps=steps,
+        guidance_scale=guidance,
+        max_sequence_length=max_seq_len,
+        generator=generator,
+        callback_on_step_end=profiler,
+    )
+    total = time.perf_counter() - t0
+    rss_peak = max(profiler.rss) if profiler.rss else rss_mb()
+    out_path = out_dir / f"{name}.png"
+    result.images[0].save(out_path)
+    rec = {
+        "name": name,
+        "prompt": prompt,
+        "negative_prompt": NEGATIVE_PROMPT,
+        "seed": seed,
+        "width": size[0],
+        "height": size[1],
+        "num_inference_steps": steps,
+        "guidance_scale": guidance,
+        "max_sequence_length": max_seq_len,
+        "steps": len(profiler.times),
+        "total_time_s": round(total, 3),
+        "step_time_mean_s": round(statistics.mean(profiler.times), 3) if profiler.times else None,
+        "step_time_median_s": round(statistics.median(profiler.times), 3) if profiler.times else None,
+        "step_time_min_s": round(min(profiler.times), 3) if profiler.times else None,
+        "step_time_max_s": round(max(profiler.times), 3) if profiler.times else None,
+        "rss_before_mb": round(rss_before, 1),
+        "rss_peak_mb": round(rss_peak, 1),
+        "rss_after_mb": round(rss_mb(), 1),
+        "per_step_time_s": [round(t, 4) for t in profiler.times],
+        "per_step_rss_mb": [round(r, 1) for r in profiler.rss],
+        "image": str(out_path),
+    }
+    print(
+        f"[{name}] done {total:.1f}s  mean {rec['step_time_mean_s']}s/step  peak RSS {rss_peak:.0f}MB",
+        flush=True,
+    )
+    return rec
+
+
+def main() -> None:
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--model_path", type=str, default=MODEL_PATH)
+    parser.add_argument("--steps", type=int, default=4, help="FLUX.1-schnell: 1-4 steps")
+    parser.add_argument("--guidance_scale", type=float, default=0.0, help="FLUX uses CFG 0.0")
+    parser.add_argument("--max_sequence_length", type=int, default=256)
+    parser.add_argument("--small_steps", type=int, default=4)
+    parser.add_argument("--skip_small", action="store_true")
+    args = parser.parse_args()
+
+    out_dir = OUTPUT_DIR
+    out_dir.mkdir(parents=True, exist_ok=True)
+
+    print("loading + compiling pipeline ...", flush=True)
+    t0 = time.perf_counter()
+    pipe = OVFluxPipeline.from_pretrained(args.model_path, compile=True, device="CPU")
+    load_s = time.perf_counter() - t0
+    rss_after_load = rss_mb()
+    print(f"load+compile {load_s:.1f}s  rss={rss_after_load:.0f}MB", flush=True)
+
+    records = []
+    for name, seed, prompt in PROMPTS:
+        records.append(
+            run_one(pipe, name, seed, prompt, (1024, 1024), args.steps, args.guidance_scale, args.max_sequence_length, out_dir)
+        )
+
+    if not args.skip_small:
+        for name, seed, prompt in PROMPTS:
+            records.append(
+                run_one(
+                    pipe,
+                    f"{name}_512",
+                    seed,
+                    prompt,
+                    (512, 512),
+                    args.small_steps,
+                    args.guidance_scale,
+                    args.max_sequence_length,
+                    out_dir,
+                )
+            )
+
+    fp16_dir = Path("/home/user/app/flux-schnell-ov-fp16")
+    benchmark = {
+        "system": system_info(),
+        "load_compile_time_s": round(load_s, 3),
+        "rss_after_load_mb": round(rss_after_load, 1),
+        "pipeline_dir_size_mb": None,
+        "fp16_dir_size_mb": None,
+        "images": records,
+    }
+    int4_dir = Path(args.model_path)
+    if int4_dir.exists():
+        benchmark["pipeline_dir_size_mb"] = round(
+            sum(f.stat().st_size for f in int4_dir.rglob("*") if f.is_file()) / 1024**2, 1
+        )
+    if fp16_dir.exists():
+        benchmark["fp16_dir_size_mb"] = round(
+            sum(f.stat().st_size for f in fp16_dir.rglob("*") if f.is_file()) / 1024**2, 1
+        )
+
+    with open(out_dir / "benchmark.json", "w") as f:
+        json.dump(benchmark, f, indent=2, ensure_ascii=False)
+    with open(out_dir / "prompts.txt", "w") as f:
+        for name, seed, prompt in PROMPTS:
+            f.write(f"{name} | seed={seed} | 1024x1024\n")
+            f.write(f"  prompt: {prompt}\n")
+            f.write(f"  negative: {NEGATIVE_PROMPT}\n\n")
+    print("benchmark.json + prompts.txt written")
+
+
+if __name__ == "__main__":
+    main()