Game asset set

"""A matching set of game item icons, packed into a sprite atlas. Safe to re-run.

    python game_assets.py generate   # submit, wait, download candidates
    python game_assets.py pack       # cut out picks and build the atlas
"""

import json
import math
import os
import sys
import time
import uuid
import zlib
from pathlib import Path

import httpx
from PIL import Image

API = os.environ.get("WINDPAINT_API_URL", "https://api.windpaint.ai").rstrip("/") + "/v1"
TERMINAL = {"completed", "failed", "nsfw", "canceled"}

STYLE = (
    "hand-painted fantasy RPG inventory icon of {subject}, a single object centered with empty space around it, "
    "three-quarter view, bold dark outline, soft cel shading, warm light from the top left, "
    "isolated on a solid flat pure magenta background, no ground shadow, no text"
)
ASSETS = {
    "sword": "a short iron sword with a leather-wrapped grip",
    "shield": "a round wooden shield with an iron rim",
    "potion": "a small glass bottle of glowing red potion with a cork",
    "scroll": "a rolled parchment scroll tied with a blue ribbon",
    "boots": "a pair of worn brown leather boots",
    "key": "an ornate brass key",
}
CANDIDATES = 4  # num_outputs, 1 to 4, billed per image delivered
CELL = 128  # atlas cell size in px
PADDING = 6
KEY_COLOR = (255, 0, 255)
TOLERANCE = 120  # colour distance below which a pixel counts as background

OUT = Path("assets")
MANIFEST = OUT / "manifest.json"
PICKS = OUT / "picks.json"

api = httpx.Client(base_url=API, headers={"Authorization": f"Bearer {os.environ['WINDPAINT_API_KEY']}"}, timeout=60)


def check(response: httpx.Response) -> dict:
    if response.is_success:
        return response.json()
    try:
        error = response.json()["error"]
    except (ValueError, KeyError):
        error = response.text
    if isinstance(error, dict):
        sys.exit(f"{response.status_code} {error['code']}: {error['message']}")
    sys.exit(f"{response.status_code}: {error}")


def load(path: Path) -> dict:
    return json.loads(path.read_text()) if path.exists() else {}


def save(path: Path, data: dict) -> None:
    path.write_text(json.dumps(data, indent=2))


def seed_for(name: str) -> int:
    return zlib.crc32(name.encode())  # stable across runs and machines


def submit(body: dict, key: str) -> str:
    """Submit, retrying with the same Idempotency-Key so a retry never creates a second job."""
    for attempt in range(5):
        try:
            r = api.post("/generation/capabilities/image.generate", json=body, headers={"Idempotency-Key": key})
        except httpx.TransportError:
            time.sleep(2 ** attempt)  # the job may exist; the same key finds it
            continue
        in_progress = r.status_code == 409 and r.json()["error"]["code"] == "idempotency.in_progress"
        if r.status_code == 429 or r.status_code >= 500 or in_progress:
            time.sleep(float(r.headers.get("Retry-After") or 2 ** attempt))
            continue
        job = check(r)
        print(f"  submitted {job['request_id']} ({job['credits_estimate']} credits held)")
        return job["request_id"]
    sys.exit("submit kept failing; re-run within 24 hours to retry with the same key")


def wait(request_id: str) -> dict:
    delay = 1.5
    while True:
        status = check(api.get(f"/generation/requests/{request_id}/status"))
        if status["status"] in TERMINAL:
            return status
        time.sleep(delay)
        delay = min(delay * 1.5, 15)


def download(asset_id: str, path: Path) -> None:
    url = check(api.get(f"/generation/assets/{asset_id}/url"))["data"]["url"]
    response = httpx.get(url, timeout=300)  # signed URL: no Authorization header
    response.raise_for_status()
    path.write_bytes(response.content)


def generate() -> None:
    (OUT / "candidates").mkdir(parents=True, exist_ok=True)
    manifest = load(MANIFEST)
    for name, subject in ASSETS.items():
        body = {
            "prompt": STYLE.format(subject=subject),
            "aspect_ratio": "1:1",
            "resolution": "1k",
            "num_outputs": CANDIDATES,
            "seed": seed_for(name),
        }
        entry = manifest.get(name)
        if not entry or entry["body"] != body:  # new asset, or its prompt changed
            entry = manifest[name] = {"body": body, "key": str(uuid.uuid4())}
            save(MANIFEST, manifest)  # persist the key before the first attempt
        if not entry.get("request_id"):
            print(name)
            entry["request_id"] = submit(body, entry["key"])
            save(MANIFEST, manifest)

    for name in ASSETS:
        entry = manifest[name]
        if entry.get("files"):
            continue
        status = wait(entry["request_id"])
        if status["status"] != "completed":
            print(f"{name}: {status['status']} ({status['error']}); re-run to try again")
            del manifest[name]  # a new attempt needs a new key; the old one replays this job
            save(MANIFEST, manifest)
            continue
        files = []
        for i, output in enumerate(status["outputs"]):
            path = OUT / "candidates" / f"{name}-{i}.png"
            download(output["id"], path)
            files.append(str(path))
        entry.update(files=files, credits=status["credits"]["actual"])
        save(MANIFEST, manifest)
        print(f"{name}: {len(files)} candidates, {entry['credits']} credits")

    picks = load(PICKS)
    for name in ASSETS:
        picks.setdefault(name, 0)
    save(PICKS, picks)
    print(f"look through {OUT / 'candidates'}, set your choices in {PICKS}, then run: python {sys.argv[0]} pack")


def cut_out(image: Image.Image) -> tuple[Image.Image, float]:
    """Key out the background colour with a soft edge. Returns the sprite and the share of pixels removed."""
    kr, kg, kb = KEY_COLOR
    pixels = []
    removed = 0
    for r, g, b, _ in image.convert("RGBA").getdata():
        distance = abs(r - kr) + abs(g - kg) + abs(b - kb)
        alpha = max(0, min(255, (distance - TOLERANCE) * 255 // TOLERANCE))
        removed += alpha == 0
        pixels.append((r, g, b, alpha))
    sprite = Image.new("RGBA", image.size)
    sprite.putdata(pixels)
    bbox = sprite.getchannel("A").getbbox()
    if bbox:
        sprite = sprite.crop(bbox)
    sprite.thumbnail((CELL - 2 * PADDING, CELL - 2 * PADDING), Image.Resampling.LANCZOS)
    return sprite, removed / (image.width * image.height)


def pack() -> None:
    manifest, picks = load(MANIFEST), load(PICKS)
    names = [n for n in ASSETS if manifest.get(n, {}).get("files")]
    if not names:
        sys.exit(f"nothing to pack; run: python {sys.argv[0]} generate")
    (OUT / "sprites").mkdir(exist_ok=True)
    cols = math.ceil(math.sqrt(len(names)))
    rows = math.ceil(len(names) / cols)
    atlas = Image.new("RGBA", (cols * CELL, rows * CELL))
    frames = {}
    for i, name in enumerate(names):
        files = manifest[name]["files"]
        pick = min(int(picks.get(name, 0)), len(files) - 1)
        sprite, removed = cut_out(Image.open(files[pick]))
        if removed < 0.25:
            print(f"{name}: only {removed:.0%} of the background keyed out; try another candidate")
        cell = Image.new("RGBA", (CELL, CELL))
        cell.paste(sprite, ((CELL - sprite.width) // 2, (CELL - sprite.height) // 2))
        cell.save(OUT / "sprites" / f"{name}.png")
        x, y = (i % cols) * CELL, (i // cols) * CELL
        atlas.paste(cell, (x, y))
        frames[f"{name}.png"] = {
            "frame": {"x": x, "y": y, "w": CELL, "h": CELL},
            "rotated": False,
            "trimmed": False,
            "spriteSourceSize": {"x": 0, "y": 0, "w": CELL, "h": CELL},
            "sourceSize": {"w": CELL, "h": CELL},
        }
    atlas.save(OUT / "atlas.png")
    save(OUT / "atlas.json", {
        "frames": frames,
        "meta": {"image": "atlas.png", "format": "RGBA8888", "size": {"w": atlas.width, "h": atlas.height}, "scale": "1"},
    })
    print(f"packed {len(names)} sprites into {OUT / 'atlas.png'} ({atlas.width}x{atlas.height})")


if __name__ == "__main__":
    commands = {"generate": generate, "pack": pack}
    if len(sys.argv) != 2 or sys.argv[1] not in commands:
        sys.exit(f"usage: python {sys.argv[0]} generate|pack")
    commands[sys.argv[1]]()

{
  "sword": 2,
  "shield": 0,
  "potion": 1,
  "scroll": 0,
  "boots": 3,
  "key": 0
}