111 lines
4.5 KiB
Python
111 lines
4.5 KiB
Python
"""Minimal PixelLab v2 client for style-consistent pixel-art generation."""
|
|
|
|
import base64
|
|
import json
|
|
import os
|
|
import time
|
|
from io import BytesIO
|
|
from typing import Any, Optional
|
|
from urllib.error import HTTPError, URLError
|
|
from urllib.request import Request, urlopen
|
|
|
|
from PIL import Image
|
|
|
|
|
|
API_URL = "https://api.pixellab.ai/v2"
|
|
|
|
|
|
class PixelLabError(RuntimeError):
|
|
pass
|
|
|
|
|
|
class PixelLabClient:
|
|
def __init__(self, api_key: Optional[str] = None, api_url: str = API_URL):
|
|
self.api_key = api_key or os.environ.get("PIXELLAB_API_KEY")
|
|
self.api_url = api_url.rstrip("/")
|
|
|
|
def _request(self, method: str, path: str, payload: Optional[dict] = None) -> dict:
|
|
if not self.api_key:
|
|
raise PixelLabError("PIXELLAB_API_KEY is not configured.")
|
|
request = Request(
|
|
f"{self.api_url}{path}",
|
|
data=json.dumps(payload).encode() if payload else None,
|
|
headers={"Authorization": f"Bearer {self.api_key}", "Content-Type": "application/json"},
|
|
method=method,
|
|
)
|
|
try:
|
|
with urlopen(request, timeout=125) as response:
|
|
return json.loads(response.read().decode())
|
|
except HTTPError as exc:
|
|
raise PixelLabError(f"PixelLab API returned HTTP {exc.code}: {exc.read().decode(errors='replace')}") from exc
|
|
except (URLError, TimeoutError) as exc:
|
|
raise PixelLabError(f"Could not reach PixelLab API: {exc}") from exc
|
|
|
|
@staticmethod
|
|
def _reference(value: str) -> dict:
|
|
if value.startswith("data:"):
|
|
raw = base64.b64decode(value.split(",", 1)[1])
|
|
elif value.startswith(("https://", "http://")):
|
|
with urlopen(value, timeout=60) as response:
|
|
raw = response.read()
|
|
else:
|
|
with open(os.path.expanduser(value), "rb") as image_file:
|
|
raw = image_file.read()
|
|
with Image.open(BytesIO(raw)) as image:
|
|
image = image.convert("RGBA")
|
|
image.thumbnail((512, 512), Image.Resampling.LANCZOS)
|
|
width, height = image.size
|
|
output = BytesIO()
|
|
image.save(output, format="PNG")
|
|
encoded = "data:image/png;base64," + base64.b64encode(output.getvalue()).decode()
|
|
return {"image": {"base64": encoded}, "width": width, "height": height}
|
|
|
|
def create_image(self, description: str, width: int, height: int, *, seed: Optional[int] = None,
|
|
reference_images: Optional[list[str]] = None, no_background: bool = True) -> dict:
|
|
refs = [self._reference(ref) for ref in (reference_images or [])]
|
|
if refs:
|
|
payload: dict[str, Any] = {
|
|
"description": description, "image_size": {"width": width, "height": height},
|
|
"style_images": refs[:4], "no_background": no_background,
|
|
}
|
|
if seed is not None:
|
|
payload["seed"] = seed
|
|
return self._request("POST", "/generate-with-style-v2", payload)
|
|
payload = {"description": description, "image_size": {"width": width, "height": height},
|
|
"no_background": no_background}
|
|
if seed is not None:
|
|
payload["seed"] = seed
|
|
return self._request("POST", "/generate-image-v2", payload)
|
|
|
|
def get_status(self, job_id: str) -> dict:
|
|
return self._request("GET", f"/background-jobs/{job_id}")
|
|
|
|
def wait_for_completion(self, job_id: str, timeout_seconds: int = 115) -> dict:
|
|
deadline = time.monotonic() + timeout_seconds
|
|
while True:
|
|
status = self.get_status(job_id)
|
|
state = status.get("status", "").lower()
|
|
if state == "completed":
|
|
return status
|
|
if state in {"failed", "error", "cancelled"}:
|
|
raise PixelLabError(str(status.get("last_response") or state))
|
|
if time.monotonic() >= deadline:
|
|
return status
|
|
time.sleep(2)
|
|
|
|
@staticmethod
|
|
def image_sources(response: dict) -> list[str]:
|
|
sources: list[str] = []
|
|
def visit(value: Any) -> None:
|
|
if isinstance(value, dict):
|
|
for key, item in value.items():
|
|
if key in {"url", "base64", "b64_json"} and isinstance(item, str):
|
|
sources.append(item)
|
|
else:
|
|
visit(item)
|
|
elif isinstance(value, list):
|
|
for item in value:
|
|
visit(item)
|
|
visit(response.get("last_response", response))
|
|
return list(dict.fromkeys(sources))
|