#!/usr/bin/env python3 """ render_role_pool.py — pre-render a role-aware face pool by hitting serve_imagegen.py (localhost:3600/generate) with prompts pulled from the bun server's /headshots/_scenes endpoint (single source of truth for SCENES + SCENES_VERSION). Layout: data/headshots_role_pool/ {band}/ {gender}_{race}/ face_00.webp face_01.webp ... manifest.jsonl Each entry in manifest.jsonl: {"band": "warehouse", "gender": "man", "race": "caucasian", "file": "warehouse/man_caucasian/face_03.webp", "seed": 184729338, "scenes_version": "v1"} Idempotent: a file at the target path is skipped. Re-run with --force to regenerate. SCENES_VERSION is captured per render so the server's pool route can refuse stale renders if the version drifts. """ from __future__ import annotations import argparse import base64 import json import os import sys import time import urllib.request import urllib.error DEFAULT_BANDS = ["warehouse", "production", "trades", "driver", "lead"] DEFAULT_GENDERS = ["man", "woman"] DEFAULT_RACES = ["caucasian", "east_asian", "south_asian", "middle_eastern", "black", "hispanic"] def race_text(r: str) -> str: return { "caucasian": "", "east_asian": "East Asian", "south_asian": "South Asian", "middle_eastern": "Middle Eastern", "black": "Black", "hispanic": "Hispanic", }.get(r, "") def fetch_scenes(mcp_url: str) -> tuple[str, dict]: """Pull canonical SCENES from the bun server. Single source of truth.""" req = urllib.request.Request(f"{mcp_url}/headshots/_scenes") with urllib.request.urlopen(req, timeout=10) as resp: data = json.loads(resp.read()) return data["version"], data["scenes"] def render(comfy_url: str, prompt: str, seed: int, steps: int, timeout: int, dim: int) -> bytes | None: payload = json.dumps({ "prompt": prompt, "width": dim, "height": dim, "steps": steps, "seed": seed, }).encode() req = urllib.request.Request( f"{comfy_url}/generate", data=payload, headers={"Content-Type": "application/json"}, ) try: with urllib.request.urlopen(req, timeout=timeout) as resp: data = json.loads(resp.read()) except urllib.error.HTTPError as e: print(f" HTTP {e.code} from comfy: {e.read()[:200]}", file=sys.stderr) return None except Exception as e: print(f" comfy error: {type(e).__name__}: {e}", file=sys.stderr) return None img_b64 = data.get("image") if not img_b64: print(f" comfy response missing 'image' field: {list(data.keys())}", file=sys.stderr) return None return base64.b64decode(img_b64) def main(): p = argparse.ArgumentParser() p.add_argument("--out", default=os.path.join(os.path.dirname(__file__), "..", "..", "data", "headshots_role_pool")) p.add_argument("--per-bucket", type=int, default=10, help="how many faces per (band × gender × race)") p.add_argument("--mcp", default="http://localhost:3700") p.add_argument("--comfy", default="http://localhost:3600") p.add_argument("--steps", type=int, default=8) p.add_argument("--bands", nargs="*", default=DEFAULT_BANDS) p.add_argument("--genders", nargs="*", default=DEFAULT_GENDERS) p.add_argument("--races", nargs="*", default=DEFAULT_RACES) p.add_argument("--force", action="store_true", help="regenerate existing files") p.add_argument("--age", type=int, default=32) p.add_argument("--timeout", type=int, default=120, help="per-render timeout (1024² takes ~5s on A4000)") p.add_argument("--dim", type=int, default=1024, help="square render dimension (v2 default 1024, v1 was 512)") args = p.parse_args() out_root = os.path.realpath(args.out) os.makedirs(out_root, exist_ok=True) print(f"Fetching canonical SCENES from {args.mcp}/headshots/_scenes…") try: version, scenes = fetch_scenes(args.mcp) except Exception as e: print(f"FATAL: could not fetch scenes ({e}). Is the mcp-server up?", file=sys.stderr) sys.exit(1) print(f" SCENES_VERSION={version}, {len(scenes)} bands available: {list(scenes.keys())}") # v2+ files live at {out}/{version}/{band}/{g}_{r}/face_NN.webp. # v1 lived at {out}/{band}/... — keep that layout intact for # rollback; the server route reads both and prefers current. out = out_root if version == "v1" else os.path.join(out_root, version) os.makedirs(out, exist_ok=True) print(f" writing to: {out}") print(f" render dim: {args.dim}×{args.dim}") # Reject any --bands not in the server's SCENES unknown = [b for b in args.bands if b not in scenes] if unknown: print(f"FATAL: unknown bands {unknown}. Server has: {list(scenes.keys())}", file=sys.stderr) sys.exit(1) manifest_rows = [] todo = [ (band, g, r, n) for band in args.bands for g in args.genders for r in args.races for n in range(args.per_bucket) ] print(f"\nPlanning: {len(todo)} renders ({len(args.bands)} bands × {len(args.genders)} genders × {len(args.races)} races × {args.per_bucket} faces).") print(f"Estimated GPU time at 1.5s/render = {len(todo) * 1.5 / 60:.1f} min.\n") t0 = time.time() rendered = 0 skipped = 0 failed = 0 for i, (band, g, r, n) in enumerate(todo): bucket_dir = os.path.join(out, band, f"{g}_{r}") os.makedirs(bucket_dir, exist_ok=True) fname = f"face_{n:02d}.webp" full = os.path.join(bucket_dir, fname) rel = os.path.relpath(full, out) if os.path.exists(full) and os.path.getsize(full) > 1024 and not args.force: skipped += 1 manifest_rows.append({ "band": band, "gender": g, "race": r, "file": rel, "seed": None, "scenes_version": version, "cached": True, }) continue scene_def = scenes[band] scene_clause = scene_def["scene"] race_clause = race_text(r) gender_clause = g # "man" / "woman" # Match the bun server's prompt builder exactly. If you tweak # one, tweak the other (or factor a /prompt-builder endpoint). # The {role} slot is intentionally a band-typical title here # — the pre-rendered face is shared across roles in the same # band, so we use the band's archetypal role. Specific roles # still hit the on-demand /headshots/generate/:key path with # their actual title. archetype_role = { "warehouse": "warehouse worker", "production": "production worker", "trades": "skilled tradesperson", "driver": "delivery driver", "lead": "shift supervisor", }.get(band, "warehouse worker") prompt = ( f"professional headshot portrait of a {args.age}-year-old " f"{race_clause} {gender_clause} {archetype_role}, {scene_clause}, " f"neutral confident expression, sharp focus, photorealistic" ) # Deterministic seed per slot — same (band, g, r, n) always # gets the same face. Mixing scenes_version means a SCENES # tweak shifts every face slightly; that's the right behavior # (it's how cache invalidation propagates to the pool too). seed_str = f"{band}|{g}|{r}|{n}|{version}" seed_h = 5381 for ch in seed_str: seed_h = ((seed_h << 5) + seed_h + ord(ch)) & 0x7fffffff seed = seed_h bytes_ = render(args.comfy, prompt, seed, args.steps, args.timeout, args.dim) if bytes_ is None: failed += 1 continue with open(full, "wb") as f: f.write(bytes_) rendered += 1 manifest_rows.append({ "band": band, "gender": g, "race": r, "file": rel, "seed": seed, "scenes_version": version, "cached": False, }) if (i + 1) % 10 == 0 or (i + 1) == len(todo): elapsed = time.time() - t0 done = i + 1 rate = done / elapsed if elapsed > 0 else 0 eta = (len(todo) - done) / rate if rate > 0 else 0 print(f" [{done}/{len(todo)}] rendered={rendered} skipped={skipped} failed={failed} " f"rate={rate:.2f}/s eta={eta:.0f}s") # Atomic manifest write manifest_path = os.path.join(out, "manifest.jsonl") tmp = manifest_path + ".tmp" with open(tmp, "w") as f: for row in manifest_rows: f.write(json.dumps(row) + "\n") os.replace(tmp, manifest_path) print(f"\nDone. {rendered} new, {skipped} cached, {failed} failed in {time.time()-t0:.1f}s") print(f"Manifest: {manifest_path} ({len(manifest_rows)} entries)") print(f"\nNext: poke {args.mcp}/headshots/__reload to pick up the new pool.") if __name__ == "__main__": main()