aadebuger-ComfyUI-LensTurbo/cli/lens_turbo_batch.py

#!/usr/bin/env python3
"""lens_turbo_batch.py — batch generate images from a JSON prompt list.

Input JSON shape (list of objects):
    [
      {"word": "apple", "sentence": "...", "prompt": "..."},
      {"word": "cat",   "sentence": "...", "prompt": "..."},
      ...
    ]

Each entry's `prompt` is sent to ComfyUI's Lens-Turbo (microsoft/Lens-Turbo)
and the resulting PNG is saved as `<out_dir>/<word>.png`. Existing files are
skipped by default, so you can re-run to resume a partially-completed batch.

⚠️ 前置:
  - lens_turbo_cli.py 在同一目录 (本脚本 import 它的 ComfyClient + build_workflow)
  - ComfyUI 已运行,LensTurboLoader/LensTurboSampler 已注册
  - microsoft/Lens-Turbo + openai/gpt-oss-20b + black-forest-labs/FLUX.2-dev/vae/
    已 hf download 到 $HF_HOME (~58 GB)

依赖:
    uv pip install websocket-client

用法:
    # 基础 — 读 JSON,写到 ./out/
    uv run python lens_turbo_batch.py vocab.json --ip 10.8.0.12

    # 自定义输出目录 + 远程 GPU 选择
    uv run python lens_turbo_batch.py vocab.json \\
        --ip 10.8.0.12 --gpu 0 \\
        --out-dir vocab_images

    # 覆盖已有文件
    uv run python lens_turbo_batch.py vocab.json --ip 10.8.0.12 --force

    # 只处理前 3 个 (干跑)
    uv run python lens_turbo_batch.py vocab.json --ip 10.8.0.12 --limit 3

    # 固定 seed (re-runs deterministic)
    uv run python lens_turbo_batch.py vocab.json --ip 10.8.0.12 --seed 42

    # 1440x1440, 4 步 turbo 默认即可
    uv run python lens_turbo_batch.py vocab.json --ip 10.8.0.12 \\
        --base-resolution 1440 --aspect-ratio 1:1

JSON 也支持 {entries: [...]} 包裹形式,或 {data: [...]} / {items: [...]}.
"""
from __future__ import annotations

import argparse
import json
import random
import re
import sys
import time
import urllib.error
import urllib.request
from pathlib import Path

# import from neighbor module (lens_turbo_cli.py 必须在同一目录)
try:
    from lens_turbo_cli import ComfyClient, build_workflow
except ImportError as e:
    sys.exit(f"❌ 不能 import lens_turbo_cli.py。把它放本脚本同一目录。\n  {e}")


def sanitize_filename(word: str) -> str:
    """Turn a word into a safe filename stem.

    Examples:
        "a/an"   -> "a_an"
        "don't"  -> "don_t"
        "  hi "  -> "hi"
        "学习"   -> "学习" (CJK 保留)
    """
    s = word.strip().lower()
    # 保留字母数字、CJK、下划线、连字符
    s = re.sub(r"[^\w\-\u4e00-\u9fff]+", "_", s)
    s = re.sub(r"_+", "_", s).strip("_")
    return s or "unnamed"


def build_parser() -> argparse.ArgumentParser:
    p = argparse.ArgumentParser(
        description="Lens-Turbo 批量文生图 CLI",
        formatter_class=argparse.RawDescriptionHelpFormatter,
        epilog="""\
JSON 格式示例:
    [
      {"word": "apple",  "prompt": "a fresh red apple, studio light, 4k"},
      {"word": "sunset", "prompt": "warm sunset over the ocean, cinematic"}
    ]

注意:
- 优先用每个 entry 的 prompt 字段
- 输出文件名由 word 字段决定 (清洗成安全文件名)
- 已存在的输出会跳过,加 --force 覆盖
- 第一张会触发 ComfyUI 端加载 LensPipeline,~2 分钟冷启动是正常的
  (服务端 cache 命中后续直接复用,后面每张 ~12 秒)
""",
    )
    p.add_argument("json_file", help="JSON 文件,list of {word, prompt} 或 {entries: [...]}")
    p.add_argument("--out-dir", default="out", help="PNG 输出目录 (默认 ./out)")

    # 生成参数 (Turbo defaults)
    p.add_argument("--base-resolution", "-R", type=int, default=1024,
                   help="长边 (默认 1024,Turbo 推荐 1024/1440)")
    p.add_argument("--aspect-ratio", "-A", default="1:1",
                   choices=["1:1", "16:9", "9:16", "4:3", "3:4",
                            "3:2", "2:3", "21:9", "9:21"])
    p.add_argument("--steps", type=int, default=4,
                   help="采样步数 (Lens-Turbo 4 步,Lens 20 步,Lens-Base 50 步)")
    p.add_argument("--cfg", type=float, default=1.0,
                   help="CFG (Turbo 1.0,Lens/Base 5.0)")
    p.add_argument("--seed", type=int, default=None,
                   help="固定 seed (省略 = 每张随机)")
    p.add_argument("--negative-prompt", default="",
                   help="negative prompt (CFG=1 时被忽略)")

    # 模型 / 加载策略
    p.add_argument("--repo-id", default="microsoft/Lens-Turbo")
    p.add_argument("--dtype", default="bf16", choices=["bf16", "fp16", "fp32"])
    p.add_argument("--gpu", type=int, default=0,
                   help="cuda:N (负数 = cpu)")
    p.add_argument("--offload-mode", default="cpu_text_encoder",
                   choices=["cpu_text_encoder", "sequential", "model", "none"],
                   help="VRAM 策略 (默认 cpu_text_encoder,24 GB 卡跑得动)")
    p.add_argument("--no-cpu-construct", action="store_true",
                   help="跳过 cpu_construct 构造守卫")

    # 流程控制
    p.add_argument("--force", action="store_true", help="覆盖已存在文件")
    p.add_argument("--limit", type=int, default=None, help="只处理前 N 条")
    p.add_argument("--retry", type=int, default=2, help="失败重试次数")
    p.add_argument("--prompt-key", default="prompt",
                   help="JSON 里的 prompt 字段名 (默认 'prompt')")
    p.add_argument("--word-key", default="word",
                   help="JSON 里的 word/filename 字段名 (默认 'word')")

    # 服务器
    p.add_argument("--ip", default="127.0.0.1")
    p.add_argument("--port", type=int, default=8188)
    p.add_argument("--server", default=None,
                   help="完整 host:port,覆盖 --ip/--port")
    p.add_argument("--timeout", type=float, default=900,
                   help="单张超时 (默认 900s,首张可能慢)")
    p.add_argument("--ws-recv-timeout", type=float, default=1200,
                   help="WS 闲置超时 (首张冷加载用,默认 1200s)")

    return p


def health_check(server: str) -> bool:
    """Quick check ComfyUI 在不在 + 节点是否注册"""
    try:
        with urllib.request.urlopen(f"http://{server}/system_stats", timeout=5) as r:
            stats = json.loads(r.read())
            sys_info = stats.get("system", {})
            print(f"✓ server {server} OK  "
                  f"comfyui={sys_info.get('comfyui_version', '?')}  "
                  f"python={sys_info.get('python_version', '?').split()[0]}")
    except (urllib.error.URLError, ConnectionRefusedError) as e:
        print(f"✗ server unreachable at {server}: {e}", file=sys.stderr)
        return False
    except Exception as e:
        # /system_stats 可能不存在(老版本),换个端点
        try:
            with urllib.request.urlopen(f"http://{server}/queue", timeout=5) as r:
                r.read()
            print(f"✓ server {server} OK (no system_stats endpoint)")
        except Exception as e2:
            print(f"✗ server unreachable: {e2}", file=sys.stderr)
            return False

    # 确认 Lens 节点注册
    try:
        with urllib.request.urlopen(f"http://{server}/object_info", timeout=10) as r:
            info = json.loads(r.read())
        missing = [n for n in ("LensTurboLoader", "LensTurboSampler") if n not in info]
        if missing:
            print(f"✗ ComfyUI 缺节点: {missing}. 检查 ComfyUI-LensTurbo 是否部署。",
                  file=sys.stderr)
            return False
        print("✓ LensTurboLoader / LensTurboSampler 都注册了")
        return True
    except Exception as e:
        print(f"! /object_info 查询失败 ({e}),继续尝试", file=sys.stderr)
        return True  # 不阻断,让用户的 prompt 自己撞


def extract_first_image(history: dict) -> dict | None:
    """从 /history 返回里挖出第一张图的 {filename, subfolder, type}"""
    outputs = history.get("outputs", {})
    # 我们的 workflow 里 SaveImage 节点 id 是 'save'
    for node_id in ("save",) + tuple(outputs.keys()):
        node_out = outputs.get(node_id, {})
        images = node_out.get("images") or []
        if images:
            return images[0]
    return None


def generate_one(
    client: ComfyClient,
    *,
    prompt: str,
    args,
    seed: int,
) -> bytes | None:
    """跑一个 prompt,返回 PNG bytes (None = 全部重试失败)"""
    device = "cpu" if args.gpu < 0 else f"cuda:{args.gpu}"
    last_err = ""
    for attempt in range(args.retry + 1):
        try:
            wf = build_workflow(
                prompt=prompt,
                negative_prompt=args.negative_prompt,
                base_resolution=args.base_resolution,
                aspect_ratio=args.aspect_ratio,
                steps=args.steps,
                cfg=args.cfg,
                seed=seed,
                device=device,
                dtype=args.dtype,
                offload_mode=args.offload_mode,
                cpu_construct=not args.no_cpu_construct,
                repo_id=args.repo_id,
                force_reload=False,  # 批量永远不要 reload,靠 Loader 的 cache
            )
            history = client.run(wf, timeout_s=args.timeout)
            meta = extract_first_image(history)
            if meta is None:
                last_err = f"history 里没有 save.images: {list(history.get('outputs', {}).keys())}"
            else:
                return client.view(meta["filename"],
                                   meta.get("subfolder", ""),
                                   meta.get("type", "output"))
        except Exception as e:
            last_err = f"{type(e).__name__}: {e}"

        if attempt < args.retry:
            wait = 2 * (attempt + 1)
            print(f"\n    ! {last_err[:160]} — retry in {wait}s")
            time.sleep(wait)

    print(f"\n    ✗ 放弃: {last_err}", file=sys.stderr)
    return None


def load_entries(path: Path, prompt_key: str) -> list[dict]:
    raw = json.loads(path.read_text(encoding="utf-8"))
    if isinstance(raw, dict):
        # 兼容 {entries: [...]} / {data: [...]} / {items: [...]} 形式
        raw = raw.get("entries") or raw.get("data") or raw.get("items") or [raw]
    if not isinstance(raw, list):
        sys.exit("JSON 根必须是 list,或是 {entries: [...]} 包裹的对象")
    return [e for e in raw if isinstance(e, dict) and prompt_key in e]


def main() -> int:
    args = build_parser().parse_args()

    in_path = Path(args.json_file)
    if not in_path.is_file():
        sys.exit(f"❌ JSON 不存在: {in_path}")

    try:
        entries = load_entries(in_path, args.prompt_key)
    except json.JSONDecodeError as e:
        sys.exit(f"❌ JSON 解析失败: {e}")

    if not entries:
        sys.exit(f"❌ {in_path.name} 里没有带 '{args.prompt_key}' 字段的 entry")

    if args.limit:
        entries = entries[: args.limit]

    server = args.server or f"{args.ip}:{args.port}"
    if not health_check(server):
        return 1

    out_dir = Path(args.out_dir)
    out_dir.mkdir(parents=True, exist_ok=True)

    print(f"→ {len(entries)} 个 entry,输出 → {out_dir.resolve()}/")
    print(f"  size={args.base_resolution}@{args.aspect_ratio}  "
          f"steps={args.steps}  cfg={args.cfg}  "
          f"offload={args.offload_mode}  "
          f"seed={'fixed=' + str(args.seed) if args.seed is not None else 'random'}")
    print()

    client = ComfyClient(server=server, ws_recv_timeout=args.ws_recv_timeout)
    stats = {"ok": 0, "skip": 0, "fail": 0}
    t_start = time.time()
    total = len(entries)

    for i, entry in enumerate(entries, 1):
        word = str(entry.get(args.word_key, "")).strip()
        prompt = str(entry.get(args.prompt_key, "")).strip()

        if not word or not prompt:
            print(f"[{i}/{total}] ! 跳过 (缺 {args.word_key} 或 {args.prompt_key}): "
                  f"{json.dumps(entry, ensure_ascii=False)[:80]}")
            stats["fail"] += 1
            continue

        stem = sanitize_filename(word)
        target = out_dir / f"{stem}.png"

        if target.exists() and not args.force:
            print(f"[{i}/{total}] ⊙ {word!r:<20} → {target.name}  (exists, skip)")
            stats["skip"] += 1
            continue

        seed = args.seed if args.seed is not None else random.randint(0, 2**31 - 1)

        print(f"[{i}/{total}] ▸ {word!r:<20} ", end="", flush=True)
        t0 = time.time()
        png = generate_one(client, prompt=prompt, args=args, seed=seed)
        dt = time.time() - t0

        if png is None:
            print(f" ✗ failed ({dt:.1f}s)")
            stats["fail"] += 1
            continue

        target.write_bytes(png)
        print(f"\r[{i}/{total}] ✓ {word!r:<20} → {target.name}  "
              f"({dt:.1f}s, {len(png)//1024} KB, seed={seed})" + " " * 10)
        stats["ok"] += 1

    total_t = time.time() - t_start
    print()
    print(f"done in {total_t:.1f}s — "
          f"{stats['ok']} ok, {stats['skip']} skipped, {stats['fail']} failed")
    if stats["ok"]:
        avg = total_t / stats["ok"]
        print(f"avg {avg:.1f}s/image")
    print(f"images in: {out_dir.resolve()}")
    return 0 if stats["fail"] == 0 else 1


if __name__ == "__main__":
    try:
        raise SystemExit(main())
    except KeyboardInterrupt:
        sys.exit(130)