Pixellab based system

pixellab_api.py

@@ -0,0 +1,139 @@
+"""
+Minimal PixelLab v2 client used by pixellab_library.py.
+
+- Token from PIXELLAB_API_KEY (environment or .env next to this file).
+- Retries rate limits (429/529) with backoff, polls async jobs (edit endpoints return a job id),
+  and runs one job at a time across all processes (the account's concurrency limit is 1).
+- Appends one line per paid call to pixellab_src/ledger.jsonl so spend can be tracked across
+  runs and parallel workers: `python pixellab_library.py cost`.
+"""
+import base64
+import fcntl
+import io
+import json
+import os
+import random
+import sys
+import time
+from pathlib import Path
+
+import requests
+from PIL import Image
+
+API = "https://api.pixellab.ai/v2"
+LEDGER = Path(__file__).with_name("pixellab_src") / "ledger.jsonl"
+LOCK = LEDGER.with_name(".api.lock")
+RETRIES = 10
+JOB_POLL_SECONDS = 4
+JOB_TIMEOUT_SECONDS = 900
+
+
+class ApiError(Exception):
+    pass
+
+
+def load_key():
+    key = os.environ.get("PIXELLAB_API_KEY")
+    env = Path(__file__).with_name(".env")
+    if not key and env.exists():
+        for line in env.read_text().splitlines():
+            name, sep, value = line.strip().partition("=")
+            if sep and name.strip() == "PIXELLAB_API_KEY":
+                key = value.strip().strip("'\"")
+    if not key:
+        sys.exit("PIXELLAB_API_KEY not set (environment or .env)")
+    return key
+
+
+def b64(img):
+    buf = io.BytesIO()
+    img.save(buf, format="PNG")
+    return {"type": "base64", "base64": base64.b64encode(buf.getvalue()).decode(), "format": "png"}
+
+
+def decode(field):
+    data = field["base64"]
+    if data.startswith("data:"):
+        data = data.split(",", 1)[1]
+    return Image.open(io.BytesIO(base64.b64decode(data))).convert("RGBA")
+
+
+def find_image(obj):
+    """First image in a response: a Base64Image dict or a data: URL, at any depth."""
+    if isinstance(obj, dict):
+        if isinstance(obj.get("base64"), str):
+            return decode(obj)
+        obj = list(obj.values())
+    if isinstance(obj, list):
+        for v in obj:
+            img = find_image(v)
+            if img is not None:
+                return img
+    if isinstance(obj, str) and obj.startswith("data:image"):
+        return decode({"base64": obj})
+    return None
+
+
+def redact(payload):
+    return {k: ("<image>" if isinstance(v, dict) and v.get("type") == "base64" else v) for k, v in payload.items()}
+
+
+class Client:
+    def __init__(self):
+        self.session = requests.Session()
+        self.session.headers["Authorization"] = f"Bearer {load_key()}"
+
+    def call(self, method, path, payload=None):
+        for attempt in range(RETRIES):
+            try:
+                r = self.session.request(method, f"{API}/{path}", json=payload, timeout=300)
+            except requests.RequestException as e:
+                if attempt == RETRIES - 1:
+                    raise ApiError(str(e))
+                time.sleep(5 + attempt * 5)
+                continue
+            if r.status_code in (429, 529, 502, 503) and attempt < RETRIES - 1:
+                time.sleep(min(60, 6 * (attempt + 1)) + random.uniform(0, 4))
+                continue
+            if not r.ok:
+                raise ApiError(f"{path}: {r.status_code} {r.text[:400]}")
+            return r.json()
+
+    def balance(self):
+        return self.call("GET", "balance")
+
+    def run(self, endpoint, payload, label):
+        """POST, wait for the job if there is one, return (image, usage). Logs the call to the ledger.
+
+        The account allows one job at a time, so jobs from every process (parallel workers, other
+        terminals) queue on a file lock instead of failing with 429."""
+        LOCK.parent.mkdir(parents=True, exist_ok=True)
+        with LOCK.open("a") as lock:
+            fcntl.flock(lock, fcntl.LOCK_EX)
+            try:
+                return self._run(endpoint, payload, label)
+            finally:
+                fcntl.flock(lock, fcntl.LOCK_UN)
+
+    def _run(self, endpoint, payload, label):
+        resp = self.call("POST", endpoint, payload)
+        usage = resp.get("usage")
+        job = resp.get("background_job_id")
+        deadline = time.time() + JOB_TIMEOUT_SECONDS
+        while job:
+            if time.time() > deadline:
+                raise ApiError(f"job {job} still running after {JOB_TIMEOUT_SECONDS}s")
+            time.sleep(JOB_POLL_SECONDS)
+            status = self.call("GET", f"background-jobs/{job}")
+            if status["status"] == "failed":
+                raise ApiError(f"job {job} failed: {json.dumps(status.get('last_response'))[:400]}")
+            if status["status"] == "completed":
+                resp, usage, job = status.get("last_response") or {}, status.get("usage") or usage, None
+        img = find_image({k: v for k, v in resp.items() if k != "usage"})
+        LEDGER.parent.mkdir(parents=True, exist_ok=True)
+        with LEDGER.open("a") as f:
+            f.write(json.dumps({"time": round(time.time()), "label": label, "endpoint": endpoint,
+                                "usage": usage, "ok": img is not None}) + "\n")
+        if img is None:
+            raise ApiError(f"no image in response: {json.dumps(redact(resp))[:400]}")
+        return img, usage