generate5.py

"""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")