Batch product shots

"""Product-shot variations across seeds and aspect ratios. Safe to re-run."""

import asyncio
import json
import os
import random
import sys
from decimal import Decimal
from pathlib import Path

import httpx

API = os.environ.get("WINDPAINT_API_URL", "https://api.windpaint.ai").rstrip("/") + "/v1"
KEY = os.environ["WINDPAINT_API_KEY"]
PROJECT = os.environ.get("BATCH_PROJECT", "product-shots")  # used only by this script

PROMPT = (
    "studio product photo of a matte black ceramic coffee mug on a pale oak table, "
    "soft window light from the left, shallow depth of field, clean background"
)
SEEDS = [101, 202, 303]
ASPECT_RATIOS = ["1:1", "4:5", "9:16", "16:9"]
RESOLUTION = "1k"
CONCURRENCY = 4
MAX_ATTEMPTS = 2  # a job that ends `failed` is free, so it gets one more try

OUT = Path("shots")
MANIFEST = OUT / "manifest.json"
TERMINAL = {"completed", "failed", "nsfw", "canceled"}
RUNNABLE = {"pending", "queued"}


class APIError(Exception):
    def __init__(self, response: httpx.Response):
        self.status = response.status_code
        try:
            error = response.json().get("error")
        except ValueError:
            error = response.text
        if isinstance(error, dict):
            self.code, self.message = error.get("code"), error.get("message")
        else:  # 401/403 come back as {"status": false, "error": "Unauthorized"}
            self.code, self.message = None, error
        super().__init__(f"{self.status} {self.code or ''} {self.message}".strip())


async def call(api: httpx.AsyncClient, method: str, path: str, **kwargs) -> dict:
    response = await api.request(method, path, **kwargs)
    if response.is_error:
        raise APIError(response)
    return response.json()


async def get_with_retry(api: httpx.AsyncClient, path: str, **kwargs) -> dict:
    """GETs are safe to retry on network errors and 5xx."""
    delay = 1.0
    for attempt in range(5):
        try:
            return await call(api, "GET", path, **kwargs)
        except (httpx.TransportError, APIError) as exc:
            if isinstance(exc, APIError) and exc.status < 500 or attempt == 4:
                raise
            await asyncio.sleep(delay + random.random())
            delay *= 2


def save(manifest: dict) -> None:
    MANIFEST.write_text(json.dumps(manifest, indent=2))


def load_manifest() -> dict:
    manifest = json.loads(MANIFEST.read_text()) if MANIFEST.exists() else {}
    for seed in SEEDS:
        for aspect in ASPECT_RATIOS:
            key = f"seed{seed}-{aspect.replace(':', 'x')}"
            manifest.setdefault(key, {"seed": seed, "aspect_ratio": aspect, "state": "pending", "attempts": 0})
    for entry in manifest.values():
        if entry["state"] == "rejected":  # never created, never billed: fine to try again
            entry.update(state="pending", error=None)
    return manifest


async def ensure_project(api: httpx.AsyncClient) -> None:
    projects = (await call(api, "GET", "/projects"))["data"]
    if not any(p["slug"] == PROJECT for p in projects):
        await call(api, "POST", "/projects", json={"name": "Product shots", "slug": PROJECT})
        print(f"created project {PROJECT}")


async def reconcile(api: httpx.AsyncClient, manifest: dict) -> None:
    """Resolve submits whose outcome we never learned, by asking the API what exists."""
    unknown = [k for k, e in manifest.items() if e["state"] == "unknown"]
    if not unknown:
        return
    known = {e.get("request_id") for e in manifest.values()}
    recent = (await get_with_retry(api, "/generation/requests", params={"limit": 200}))["data"]
    orphans = [r for r in recent if r["request_id"] not in known]
    if not orphans:
        print(f"{len(unknown)} uncertain submit(s) never reached the API; submitting them again")
        for key in unknown:
            manifest[key].update(state="pending", error=None)
    else:
        # These jobs exist and are billed. Keep them, and leave the uncertain variations alone.
        print(f"found {len(orphans)} job(s) from uncertain submits; downloading them instead of resubmitting")
        for r in orphans:
            manifest[f"orphan-{r['request_id'][:8]}"] = {"request_id": r["request_id"], "state": "queued", "attempts": MAX_ATTEMPTS}
        for key in unknown:
            manifest[key].update(state="skipped", error="may have been created as one of the orphan jobs")
    save(manifest)


async def estimate_total(api: httpx.AsyncClient, todo: list[dict]) -> Decimal:
    total = Decimal(0)
    for aspect in sorted({e["aspect_ratio"] for e in todo}):
        body = {"capability": "image.generate", "prompt": PROMPT, "aspect_ratio": aspect, "resolution": RESOLUTION}
        estimate = (await call(api, "POST", "/generation/estimate", json=body))["data"]
        if estimate["credits"] is None:
            sys.exit(f"{estimate['model']} has no price at {RESOLUTION}")
        count = sum(1 for e in todo if e["aspect_ratio"] == aspect)
        total += Decimal(estimate["credits"]) * count
        print(f"  {aspect:>5}  {estimate['width']}x{estimate['height']}  {count} x {estimate['credits']} credits")
    return total


async def poll(api: httpx.AsyncClient, request_id: str) -> dict:
    delay = 1.5
    while True:
        status = await get_with_retry(api, f"/generation/requests/{request_id}/status")
        if status["status"] in TERMINAL:
            return status
        await asyncio.sleep(delay + random.uniform(0, 0.5))
        delay = min(delay * 1.5, 15)


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


async def submit(api: httpx.AsyncClient, entry: dict, stop: asyncio.Event) -> bool:
    """Submit once. Returns True when we hold a request_id."""
    body = {"prompt": PROMPT, "aspect_ratio": entry["aspect_ratio"], "resolution": RESOLUTION, "seed": entry["seed"]}
    try:
        job = await call(api, "POST", "/generation/capabilities/image.generate", json=body)
    except (httpx.ConnectError, httpx.ConnectTimeout) as exc:
        entry.update(state="rejected", error=repr(exc))  # never reached the API
        return False
    except httpx.TransportError as exc:
        entry.update(state="unknown", error=repr(exc))  # sent, answer lost: the job may exist
        return False
    except APIError as exc:
        if exc.status >= 500:
            entry.update(state="unknown", error=str(exc))
        else:
            entry.update(state="rejected", error=str(exc))  # 4xx: nothing was created or billed
            if exc.code == "billing.insufficient_credits":
                stop.set()
        return False
    entry.update(request_id=job["request_id"], state="queued", attempts=entry["attempts"] + 1, error=None)
    return True


async def run_one(api, sem, key, manifest, stop) -> None:
    entry = manifest[key]
    async with sem:
        while True:
            if not entry.get("request_id"):
                if stop.is_set():
                    return
                ok = await submit(api, entry, stop)
                save(manifest)  # persist the request_id before doing anything else
                if not ok:
                    print(f"{key}: {entry['state']} ({entry['error']})")
                    return
            status = await poll(api, entry["request_id"])
            if status["status"] == "completed":
                path = OUT / f"{key}.png"
                await download(api, status["outputs"][0]["id"], path)
                entry.update(state="completed", file=str(path), credits=status["credits"]["actual"])
                save(manifest)
                print(f"{key}: saved {path}")
                return
            if status["status"] == "failed" and entry["attempts"] < MAX_ATTEMPTS:
                print(f"{key}: failed ({status['error']}), trying once more")
                entry.pop("request_id")
                continue
            entry.update(state=status["status"], error=status["error"])
            save(manifest)
            print(f"{key}: {status['status']} ({status['error']})")
            return


async def main() -> None:
    OUT.mkdir(exist_ok=True)
    manifest = load_manifest()
    async with httpx.AsyncClient(base_url=API, headers={"Authorization": f"Bearer {KEY}"}, timeout=60) as api:
        await ensure_project(api)
        api.headers["X-Windpaint-Project"] = PROJECT
        await reconcile(api, manifest)

        todo = [e for e in manifest.values() if e["state"] == "pending"]
        if todo:
            print(f"{len(todo)} new image(s) at {RESOLUTION}:")
            total = await estimate_total(api, todo)
            available = Decimal((await call(api, "GET", "/billing/balance"))["data"]["available"])
            print(f"total {total} credits, {available} available")
            if total > available:
                sys.exit("not enough credits; top up in Settings -> Billing or shrink the batch")
            if input("Submit? [y/N] ").lower() != "y":
                sys.exit("nothing submitted")

        sem, stop = asyncio.Semaphore(CONCURRENCY), asyncio.Event()
        keys = [k for k, e in manifest.items() if e["state"] in RUNNABLE]
        await asyncio.gather(*(run_one(api, sem, k, manifest, stop) for k in keys))

    states: dict[str, int] = {}
    for entry in manifest.values():
        states[entry["state"]] = states.get(entry["state"], 0) + 1
    spent = sum(Decimal(e["credits"]) for e in manifest.values() if e.get("credits"))
    print(f"done: {states}, {spent} credits charged in total")
    if states.get("unknown") or states.get("rejected"):
        print("re-run the script to resolve unknown and rejected variations")


if __name__ == "__main__":
    asyncio.run(main())