ai-agent/scripts/avatar_pipeline.py

#!/usr/bin/env python3
"""Post-process a nanobanana 2x2 sprite grid into an aligned sprite-sheet strip.

Per clip: split the grid into 4 quadrant frames, strip each frame's background
via the brain `bgremove` model (EVOX2), then compose a 4x1 horizontal strip.
Alignment: one uniform scale per clip (largest frame fits the cell), each frame
bottom-anchored on a common baseline so feet stay planted across frames.

Usage:
  avatar_pipeline.py <grid.png> <clip-name> <out-dir> [--cell 512] [--keep-bg]

Needs Pillow + requests (run under a venv, see gen-avatar-sprites.sh) and
BRAIN_API_KEY in the environment (gen-avatar-sprites.sh sources the repo .env).
Writes <out-dir>/sheets/<clip>.png, per-frame cutouts under
<out-dir>/cutouts/<clip>/, and updates <out-dir>/manifest.json.
"""
from __future__ import annotations

import argparse
import base64
import io
import json
import os
import sys
import time
from pathlib import Path

import requests
from PIL import Image

BRAIN_URL = os.environ.get("BRAIN_URL", "http://100.93.171.39:8010").rstrip("/")
BRAIN_KEY = os.environ.get("BRAIN_API_KEY", "")

# Playback defaults per clip; anything absent uses DEFAULT_META.
DEFAULT_META = {"frames": 4, "fps": 6, "loop": True}
CLIP_META = {
    "celebrate": {"fps": 8},
    "walk": {"fps": 8},
    "oops": {"fps": 4},
    "thinking": {"fps": 3},
    "idle": {"fps": 3},
    "lifted": {"fps": 4},
}


def bgremove(png_bytes: bytes) -> Image.Image:
    """Send one frame through the brain bgremove model, return RGBA cutout."""
    if not BRAIN_KEY:
        sys.exit("BRAIN_API_KEY not set")
    h = {"X-API-Key": BRAIN_KEY}
    r = requests.post(
        f"{BRAIN_URL}/api/jobs",
        headers={**h, "Content-Type": "application/json"},
        json={"model": "bgremove",
              "input": {"image_b64": base64.b64encode(png_bytes).decode()}},
        timeout=30,
    )
    r.raise_for_status()
    job_id = r.json()["job_id"]
    deadline = time.time() + 300
    while time.time() < deadline:
        s = requests.get(f"{BRAIN_URL}/api/jobs/{job_id}", headers=h, timeout=30).json()
        state = s.get("state") or s.get("status")
        if state in ("done", "succeeded", "completed"):
            break
        if state in ("failed", "error", "cancelled"):
            sys.exit(f"bgremove job {job_id} failed: {s}")
        time.sleep(2)
    else:
        sys.exit(f"bgremove job {job_id} timed out")
    res = requests.get(f"{BRAIN_URL}/api/jobs/{job_id}/result", headers=h, timeout=60)
    res.raise_for_status()
    return Image.open(io.BytesIO(res.content)).convert("RGBA")


def split_grid(grid: Image.Image) -> list[Image.Image]:
    w, h = grid.size
    cw, ch = w // 2, h // 2
    boxes = [(0, 0), (cw, 0), (0, ch), (cw, ch)]  # TL TR BL BR = frames 1..4
    return [grid.crop((x, y, x + cw, y + ch)) for x, y in boxes]


def alpha_bbox(img: Image.Image) -> tuple[int, int, int, int]:
    bbox = img.getchannel("A").getbbox()
    return bbox or (0, 0, img.width, img.height)


def compose_strip(frames: list[Image.Image], cell: int) -> Image.Image:
    bboxes = [alpha_bbox(f) for f in frames]
    # One scale for the whole clip so inter-frame motion is preserved.
    biggest = max(max(x1 - x0, y1 - y0) for x0, y0, x1, y1 in bboxes)
    scale = (cell * 0.92) / biggest if biggest else 1.0
    baseline = int(cell * 0.97)
    strip = Image.new("RGBA", (cell * len(frames), cell), (0, 0, 0, 0))
    for i, (frame, (x0, y0, x1, y1)) in enumerate(zip(frames, bboxes)):
        cut = frame.crop((x0, y0, x1, y1))
        cut = cut.resize((max(1, int(cut.width * scale)),
                          max(1, int(cut.height * scale))), Image.LANCZOS)
        px = i * cell + (cell - cut.width) // 2
        py = baseline - cut.height
        strip.paste(cut, (px, max(0, py)), cut)
    return strip


def main() -> None:
    ap = argparse.ArgumentParser()
    ap.add_argument("grid")
    ap.add_argument("clip")
    ap.add_argument("out")
    ap.add_argument("--cell", type=int, default=512)
    ap.add_argument("--keep-bg", action="store_true",
                    help="skip bgremove (debugging the crop/align only)")
    args = ap.parse_args()

    out = Path(args.out)
    cutdir = out / "cutouts" / args.clip
    cutdir.mkdir(parents=True, exist_ok=True)
    (out / "sheets").mkdir(parents=True, exist_ok=True)

    grid = Image.open(args.grid).convert("RGBA")
    frames = []
    for i, quad in enumerate(split_grid(grid), 1):
        if args.keep_bg:
            cut = quad
        else:
            buf = io.BytesIO()
            quad.save(buf, "PNG")
            print(f"[{args.clip}] bgremove frame {i}/4 …", flush=True)
            cut = bgremove(buf.getvalue())
        cut.save(cutdir / f"{i}.png")
        frames.append(cut)

    strip = compose_strip(frames, args.cell)
    sheet = out / "sheets" / f"{args.clip}.png"
    strip.save(sheet)

    manifest_path = out / "manifest.json"
    manifest = json.loads(manifest_path.read_text()) if manifest_path.exists() else {}
    manifest[args.clip] = {**DEFAULT_META, **CLIP_META.get(args.clip, {}),
                           "cell": args.cell}
    manifest_path.write_text(json.dumps(manifest, indent=2, sort_keys=True) + "\n")
    print(f"[{args.clip}] sheet -> {sheet}")


if __name__ == "__main__":
    main()