Vibe-HD-2D/tools/pixellab.py

#!/usr/bin/env python3
"""PixelLab asset pipeline.

Reads tools/assets.json, generates every asset whose output file is missing
(or all of them with --force / the ones named on the command line), and writes
PNGs into assets/textures/. The API key is read from PIXELLAB_API_KEY in the
environment or from the repo's .env file.

  python3 tools/pixellab.py                 # generate missing assets
  python3 tools/pixellab.py tall_grass      # (re)generate named assets
  python3 tools/pixellab.py --list          # list manifest entries
  python3 tools/pixellab.py --balance       # show account balance

Manifest entry kinds:
  "image"  -> POST /create-image-pixflux (synchronous)
  "tiles"  -> POST /create-tiles-pro, polls GET /tiles-pro/{id}; writes
              <out>_<n>.png for every variation plus <out> for "pick" index
"""
import base64
import json
import os
import sys
import time
import urllib.error
import urllib.request
from concurrent.futures import ThreadPoolExecutor
from pathlib import Path

ROOT = Path(__file__).resolve().parent.parent
API = "https://api.pixellab.ai/v2"
MANIFEST = ROOT / "tools" / "assets.json"
OUT_DIR = ROOT / "assets" / "textures"


def api_key() -> str:
    key = os.environ.get("PIXELLAB_API_KEY")
    if key:
        return key.strip()
    env = ROOT / ".env"
    if env.exists():
        for line in env.read_text().splitlines():
            line = line.strip()
            if line.startswith("PIXELLAB_API_KEY="):
                return line.split("=", 1)[1].strip().strip('"').strip("'")
    sys.exit("PIXELLAB_API_KEY not set (environment or .env)")


def request(method: str, path: str, body=None, timeout=180):
    data = json.dumps(body).encode() if body is not None else None
    req = urllib.request.Request(
        API + path,
        data=data,
        method=method,
        headers={
            "Authorization": f"Bearer {api_key()}",
            "Content-Type": "application/json",
        },
    )
    try:
        with urllib.request.urlopen(req, timeout=timeout) as r:
            return r.status, json.loads(r.read() or b"null")
    except urllib.error.HTTPError as e:
        return e.code, e.read().decode(errors="replace")


def save_b64(b64: str, path: Path):
    path.parent.mkdir(parents=True, exist_ok=True)
    if "," in b64[:64]:  # data URL prefix
        b64 = b64.split(",", 1)[1]
    path.write_bytes(base64.b64decode(b64))


# Manifest keys consumed by this script, never sent to the API.
LOCAL_KEYS = {"kind", "out", "pick", "tile_id"}


def api_body(spec):
    return {k: v for k, v in spec.items() if k not in LOCAL_KEYS}


def gen_image(name, spec):
    body = api_body(spec)
    status, res = request("POST", "/create-image-pixflux", body)
    if status != 200:
        raise RuntimeError(f"{name}: HTTP {status}: {res}")
    out = OUT_DIR / spec["out"]
    save_b64(res["image"]["base64"], out)
    return [out]


def gen_tiles(name, spec):
    # A manifest "tile_id" re-downloads an existing (already paid for) tile set
    # instead of generating a new one. Delete it to regenerate.
    tile_id = spec.get("tile_id")
    if not tile_id:
        body = api_body(spec)
        status, res = request("POST", "/create-tiles-pro", body)
        if status not in (200, 202):
            raise RuntimeError(f"{name}: HTTP {status}: {res}")
        tile_id = res["tile_id"]
        print(f"[..] {name}: tile_id {tile_id} (add to manifest to reuse)", flush=True)
    deadline = time.time() + 600
    while time.time() < deadline:
        status, res = request("GET", f"/tiles-pro/{tile_id}")
        if status == 423:  # still processing
            time.sleep(6)
            continue
        if status != 200:
            raise RuntimeError(f"{name}: poll HTTP {status}: {res}")
        if res.get("storage_urls"):
            break
        if res.get("status") == "failed":
            raise RuntimeError(f"{name}: generation failed: {res}")
        time.sleep(6)
    else:
        raise RuntimeError(f"{name}: timed out")
    stem = Path(spec["out"]).stem
    outs = []
    urls = res["storage_urls"]
    keys = sorted(urls, key=lambda k: int("".join(c for c in k if c.isdigit()) or 0))
    for i, key in enumerate(keys):
        p = OUT_DIR / f"{stem}_{i}.png"
        p.parent.mkdir(parents=True, exist_ok=True)
        # The storage host rejects urllib's default User-Agent with 403.
        dl = urllib.request.Request(urls[key], headers={"User-Agent": "hd2d-asset-pipeline/1.0"})
        with urllib.request.urlopen(dl, timeout=120) as r:
            p.write_bytes(r.read())
        outs.append(p)
    pick = spec.get("pick", 0)
    if outs:
        (OUT_DIR / spec["out"]).write_bytes(outs[min(pick, len(outs) - 1)].read_bytes())
    return outs


def generate(name, spec):
    kind = spec.get("kind", "image")
    fn = {"image": gen_image, "tiles": gen_tiles}[kind]
    t0 = time.time()
    outs = fn(name, spec)
    print(f"[ok] {name} ({time.time() - t0:.1f}s) -> {', '.join(p.name for p in outs)}", flush=True)


def main(argv):
    manifest = json.loads(MANIFEST.read_text())
    if "--balance" in argv:
        print(request("GET", "/balance")[1])
        return
    if "--list" in argv:
        for name, spec in manifest.items():
            have = (OUT_DIR / spec["out"]).exists()
            print(f"{'x' if have else ' '} {name:20s} {spec['out']}")
        return
    force = "--force" in argv
    names = [a for a in argv if not a.startswith("--")]
    unknown = [n for n in names if n not in manifest]
    if unknown:
        sys.exit(f"unknown asset(s): {unknown}")
    todo = {
        n: s
        for n, s in manifest.items()
        if (n in names) or (not names and (force or not (OUT_DIR / s["out"]).exists()))
    }
    if not todo:
        print("all assets present")
        return
    failures = 0
    with ThreadPoolExecutor(max_workers=4) as pool:
        futs = {pool.submit(generate, n, s): n for n, s in todo.items()}
        for f, n in futs.items():
            try:
                f.result()
            except Exception as e:  # report every failure, keep going
                failures += 1
                print(f"[fail] {n}: {e}", file=sys.stderr, flush=True)
    sys.exit(1 if failures else 0)


if __name__ == "__main__":
    main(sys.argv[1:])