"""Generate 5 prompts x 8-step images with OV INT4 Z-Image-Turbo + benchmark. 1024x1024 main + 512x512 control (same prompts/seeds). Turbo: num_inference_steps=9 -> 8 DiT forwards (official README), guidance_scale=0.0. """ import json, time, platform, os, gc, subprocess from pathlib import Path import psutil import torch from optimum.intel import OVDiffusionPipeline MODEL_ID = "HelloSun/Z-Image-Turbo-OpenVINO-INT4" OUT_DIR = Path("outputs") OUT_DIR.mkdir(exist_ok=True) PROMPTS = [ { "id": "01_hanfu_portrait", "prompt": "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", "seed": 42, }, { "id": "02_astronaut_jungle", "prompt": "Astronaut in a jungle, cold color palette, muted colors, detailed, 8k, photorealistic, cinematic lighting", "seed": 43, }, { "id": "03_taipei_cyberpunk", "prompt": "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", "seed": 44, }, { "id": "04_shiba_astronaut", "prompt": "Cute Shiba Inu wearing a tiny astronaut helmet, sitting in a field of sunflowers under a starry sky, dreamy illustration, vibrant colors, high quality", "seed": 45, }, { "id": "05_ink_landscape", "prompt": "Traditional Chinese ink wash landscape, misty mountains, a small pagoda on a cliff, cranes flying, minimalist, elegant, high aesthetic quality", "seed": 46, }, ] HEIGHT = WIDTH = 1024 # Turbo: num_inference_steps=9 -> 8 DiT forwards (official README) NUM_STEPS = 9 GUIDANCE = 0.0 CONTROL_RES = 512 # 512px對照組 (同prompt/seed各一張) def mem_gb(): p = psutil.Process(os.getpid()) return p.memory_info().rss / 1e9 print("=== Loading OV INT4 pipeline ===") t_load0 = time.time() pipe = OVDiffusionPipeline.from_pretrained(MODEL_ID, compile=True) t_load = time.time() - t_load0 print(f"load+compile: {t_load:.1f}s, RSS={mem_gb():.2f} GB") results = [] for item in PROMPTS: pid, prompt, seed = item["id"], item["prompt"], item["seed"] print(f"\n=== {pid} seed={seed} ===") print(prompt) gc.collect() rss_before = mem_gb() # per-step timing via callback step_times = [] t_last = [time.time()] def cb(pipe, step_idx, timestep, callback_kwargs): now = time.time() step_times.append(now - t_last[0]) t_last[0] = now return callback_kwargs t0 = time.time() out = pipe( prompt=prompt, height=HEIGHT, width=WIDTH, num_inference_steps=NUM_STEPS, guidance_scale=GUIDANCE, generator=torch.Generator().manual_seed(seed), callback_on_step_end=cb, callback_on_step_end_tensor_inputs=["latents"], ) total = time.time() - t0 rss_after = mem_gb() img = out.images[0] save_path = OUT_DIR / f"{pid}_seed{seed}_8step_1024.png" img.save(save_path) print(f"saved {save_path} size={img.size} total={total:.1f}s RSS {rss_before:.2f}->{rss_after:.2f} GB") print(f"step times (scheduler steps={len(step_times)}): {[f'{x:.2f}' for x in step_times]}") # DiT forwards = scheduler steps - 1 for this model (first step is init?) # Report mean of step_times[1:] as single DiT step if len>1 entry = { "id": pid, "prompt": prompt, "seed": seed, "height": HEIGHT, "width": WIDTH, "num_inference_steps_param": NUM_STEPS, "scheduler_steps_observed": len(step_times), "step_times_s": [round(x, 3) for x in step_times], "total_s": round(total, 2), "mean_step_s": round(sum(step_times)/len(step_times), 3) if step_times else None, "rss_before_gb": round(rss_before, 3), "rss_after_gb": round(rss_after, 3), "image": str(save_path), } # 512px對照組 (同prompt/seed, 省load/compile重用pipe) gc.collect() c_times = [] c_last = [time.time()] def cb2(pipe, step_idx, timestep, callback_kwargs): now = time.time() c_times.append(now - c_last[0]) c_last[0] = now return callback_kwargs tc0 = time.time() cout = pipe( prompt=prompt, height=CONTROL_RES, width=CONTROL_RES, num_inference_steps=NUM_STEPS, guidance_scale=GUIDANCE, generator=torch.Generator().manual_seed(seed), callback_on_step_end=cb2, callback_on_step_end_tensor_inputs=["latents"], ) c_total = time.time() - tc0 cimg = cout.images[0] c_path = OUT_DIR / f"{pid}_seed{seed}_8step_512.png" cimg.save(c_path) print(f"control 512 saved {c_path} total={c_total:.1f}s mean={sum(c_times)/len(c_times):.2f}s") entry["control_512"] = { "image": str(c_path), "height": CONTROL_RES, "width": CONTROL_RES, "total_s": round(c_total, 2), "mean_step_s": round(sum(c_times)/len(c_times), 3) if c_times else None, "step_times_s": [round(x, 3) for x in c_times], } results.append(entry) with open(OUT_DIR / "benchmark.json", "w") as f: import openvino as ov try: lscpu = subprocess.check_output("lscpu", shell=True, text=True) except Exception as e: lscpu = str(e) sysinfo = { "cpu_model": "Intel(R) Xeon(R) Platinum 8559C", "os_cpu_count": os.cpu_count(), "lscpu": lscpu, "openvino_version": ov.__version__, "settings": f"num_inference_steps={NUM_STEPS}, guidance_scale={GUIDANCE}, 1024 main + 512 control", } json.dump({"model": MODEL_ID, "load_compile_s": round(t_load, 2), "height": HEIGHT, "width": WIDTH, "results": results, "sysinfo": sysinfo}, f, indent=2, ensure_ascii=False) print("\nDone. See outputs/benchmark.json") # prompts.txt with open(OUT_DIR / "prompts.txt", "w") as f: for r in results: f.write(f"[{r['id']}] seed={r['seed']}\n{r['prompt']}\n\n") f.write("\n[control] 512px對照組: 同上5組prompt/seed, height=width=512, num_inference_steps=9, guidance_scale=0.0\n") f.write("文件: outputs/*_8step_512.png (01..05各一張)\n")