diff --git a/.env.example b/.env.example new file mode 100644 index 0000000..79a2ecd --- /dev/null +++ b/.env.example @@ -0,0 +1,17 @@ +# Copy this file to .env and set real values there. Never commit .env. + +# Required only for cloud generation through provider="polza". +POLZA_API_KEY= + +# Optional default cloud image model. It can be overridden per MCP tool call. +POLZA_IMAGE_MODEL=openai/gpt-image-1.5 + +# Required only for native pixel-art generation through provider="pixellab". +PIXELLAB_API_KEY= + +# Optional local paths for the local Diffusers provider. +# IMAGEGEN_MODEL_DIR=/path/to/sdxl-base +# IMAGEGEN_LORA_DIR=/path/to/pixel-art-xl +# IMAGEGEN_LCM_LORA_DIR=/path/to/lcm-lora-sdxl +# IMAGEGEN_OUTPUT_DIR=/path/to/output +# IMAGEGEN_DB_PATH=/path/to/feedback.db diff --git a/.gitignore b/.gitignore index a14046c..6609186 100644 --- a/.gitignore +++ b/.gitignore @@ -7,6 +7,7 @@ __pycache__/ *.pyc *.pyo *.egg-info/ +.env dist/ build/ diff --git a/README.md b/README.md index 4f9e29f..f111b71 100644 --- a/README.md +++ b/README.md @@ -13,6 +13,7 @@ MCP server for generating pixel-art sprites with transparent backgrounds. Bring - **Reproducible** — optional seed for consistent results - **Batch generation** — generate multiple sprites in one call - **MCP integration** — works with any MCP-compatible client (opencode, Claude, etc.) +- **Cloud image models** — Polza.ai Media API, including image references and multi-image variations - **Feedback loop** — rate generated sprites, AI uses high-rated ones as reference ## Quick Start @@ -64,6 +65,52 @@ export IMAGEGEN_LORA_DIR=/path/to/your/lora # optional, set empty to disa export IMAGEGEN_OUTPUT_DIR=/path/to/output ``` +### Cloud generation with Polza.ai + +No local GPU or model download is required when using the `polza` provider. Create +an API key in Polza.ai and expose it only to the MCP server process. The server +also loads a local `.env` file automatically; create it from the safe template: + +```bash +cp .env.example .env +# Edit .env and set POLZA_API_KEY=your_key + +# Or provide variables directly when starting the MCP server: +export POLZA_API_KEY=your_key +# Optional default; pass model per tool call to override it. +export POLZA_IMAGE_MODEL=openai/gpt-image-1.5 +``` + +Use `generate_sprite` with `provider="polza"` for one image, or +`generate_images` for up to ten coherent variations in one API request. Both +accept `reference_images`: HTTPS URLs, data URIs, or local files. References are +sent to the provider so an LLM can retain a game's palette, outline treatment, +proportions, and character style across new sprites. + +```json +{ + "prompt": "a forest ranger facing left, idle game sprite", + "output_path": "rangers/idle.png", + "count": 4, + "model": "openai/gpt-image-1.5", + "reference_images": ["/assets/style-guide.png", "https://example.com/hero.png"], + "aspect_ratio": "1:1", + "remove_bg": true, + "pixel_size": 4 +} +``` + +Set `wait=false` for a long-running generation and query its returned ID using +`get_generation_status`. + +### Native pixel art with PixelLab + +Set `PIXELLAB_API_KEY` in `.env` and pass `provider="pixellab"`. PixelLab is +specialised in game-ready pixel art; with `reference_images` it uses up to four +style references and generates style-consistent sprites. `generate_images` will +create one background job per requested variant, so `count` is reliable even +when a model does not offer a multi-image parameter. + ### 4. Run as MCP server ```bash @@ -123,6 +170,20 @@ Returns: `output_path`, `db_id`, `generation_time`, and other metadata. Generate multiple sprites in one call. Each is saved to the feedback DB. +Each spec may also set `provider: "polza"`, `model`, `reference_images`, +`aspect_ratio`, `quality`, `count`, and `wait`. + +#### `generate_images` + +Generate 1–10 variants from one prompt through Polza.ai. It accepts the same +style-reference fields as cloud `generate_sprite` and saves every finished +variant to the feedback DB. + +#### `get_generation_status` + +Check a non-blocking Polza generation by its `generation_id` and retrieve its +status, output sources, usage, warnings, or error. + ### Feedback #### `rate_sprite` @@ -167,6 +228,9 @@ Get database statistics: total sprites, rated, unrated, average rating. | `IMAGEGEN_MODEL_DIR` | `~/models/flux2-klein-4b` | Path to base model | | `IMAGEGEN_LORA_DIR` | `~/models/pixel-art-lora` | Path to LoRA adapter | | `IMAGEGEN_OUTPUT_DIR` | `./output` | Default output directory | +| `POLZA_API_KEY` | — | Polza.ai API key; required for cloud generation | +| `POLZA_IMAGE_MODEL` | `openai/gpt-image-1.5` | Default Polza image model | +| `PIXELLAB_API_KEY` | — | PixelLab API key; required for `pixellab` provider | ### Swapping models diff --git a/assets/plasma-impact.png b/assets/plasma-impact.png new file mode 100644 index 0000000..2719833 Binary files /dev/null and b/assets/plasma-impact.png differ diff --git a/assets/space-nebula.png b/assets/space-nebula.png new file mode 100644 index 0000000..a17f111 Binary files /dev/null and b/assets/space-nebula.png differ diff --git a/demo.html b/demo.html new file mode 100644 index 0000000..bb0275c --- /dev/null +++ b/demo.html @@ -0,0 +1,70 @@ + + + + + + Imagen — Fleet Demo + + + +
+
+
+

Imagen · fleet catalogue

Asterion Patrol

+
SECTOR 04 · 12:00 FORMATION
+
+
+ Asterion vanguard + Asterion escort + Asterion corvette + Asterion scout + Asterion support ship + Asterion tender + Asterion micro scout + Asterion micro interceptor + Asterion micro patrol ship + Asterion micro support ship +
Asterion-01
pixel-perfect vanguard
open glossary →
+
+ +
+ + diff --git a/feedback.py b/feedback.py index 26c6a89..542185c 100644 --- a/feedback.py +++ b/feedback.py @@ -33,6 +33,17 @@ class DBStats: avg_rating: float = 0.0 +@dataclass +class StyleReference: + id: str + image_path: str + name: str + role: str + notes: Optional[str] + priority: int + created_at: str + + class FeedbackDB: """SQLite-backed feedback database for generated sprites.""" @@ -55,6 +66,18 @@ class FeedbackDB: """) conn.execute("CREATE INDEX IF NOT EXISTS idx_prompt ON feedback(prompt)") conn.execute("CREATE INDEX IF NOT EXISTS idx_rating ON feedback(rating)") + conn.execute(""" + CREATE TABLE IF NOT EXISTS style_references ( + id TEXT PRIMARY KEY, + image_path TEXT NOT NULL UNIQUE, + name TEXT NOT NULL, + role TEXT NOT NULL, + notes TEXT, + priority INTEGER NOT NULL DEFAULT 0, + created_at TEXT NOT NULL + ) + """) + conn.execute("CREATE INDEX IF NOT EXISTS idx_style_priority ON style_references(priority DESC)") conn.commit() return cls(conn) @@ -154,6 +177,47 @@ class FeedbackDB: self.conn.execute("DELETE FROM feedback WHERE id = ?", (entry_id,)) self.conn.commit() + def add_style_reference( + self, + image_path: str, + name: str, + role: str, + notes: Optional[str] = None, + priority: int = 0, + ) -> str: + """Register an existing image as a reusable visual style reference.""" + existing = self.conn.execute( + "SELECT id FROM style_references WHERE image_path = ?", (image_path,) + ).fetchone() + if existing: + self.conn.execute( + "UPDATE style_references SET name=?, role=?, notes=?, priority=? WHERE id=?", + (name, role, notes, priority, existing[0]), + ) + self.conn.commit() + return existing[0] + + reference_id = str(uuid.uuid4()) + self.conn.execute( + "INSERT INTO style_references (id, image_path, name, role, notes, priority, created_at) " + "VALUES (?, ?, ?, ?, ?, ?, ?)", + (reference_id, image_path, name, role, notes, priority, str(int(time.time()))), + ) + self.conn.commit() + return reference_id + + def get_style_references(self, limit: int = 8) -> list[StyleReference]: + cursor = self.conn.execute( + "SELECT id, image_path, name, role, notes, priority, created_at " + "FROM style_references ORDER BY priority DESC, created_at DESC LIMIT ?", + (limit,), + ) + return [StyleReference(*row) for row in cursor] + + def delete_style_reference(self, reference_id: str) -> None: + self.conn.execute("DELETE FROM style_references WHERE id = ?", (reference_id,)) + self.conn.commit() + def export_jsonl(self, path: str, min_rating: int = 4) -> int: entries = self.top_rated(10000, min_rating) lines = [] diff --git a/fleet-glossary.html b/fleet-glossary.html new file mode 100644 index 0000000..2ff50ce --- /dev/null +++ b/fleet-glossary.html @@ -0,0 +1,12 @@ + +Asterion Fleet Glossary + +
← return to scene

Asterion Fleet

A working glossary for the newly generated 64×64 fleet. All entries share a pixel-perfect grid, blue-gray armour, cyan propulsion and a strict north-facing top-down silhouette.

+
Needle scoutAST-01 · Recon

Needle Scout

Fast survey craft. Its narrow nose and compact engine pair make it legible at the smallest tactical scale.

+
Mirror interceptorAST-05 · Intercept

Mirror Interceptor

Symmetric pursuit hull with swept wings, designed to anchor the visual language of the patrol wing.

+
Missile corvetteAST-09 · Strike

Missile Corvette

Compact midline striker. The broad central volume makes it a readable combat silhouette without extra effects.

+
Support frigateAST-13 · Support

Support Frigate

Wide, stable frame intended for repair and relay duties; its wings visually distinguish it from the attack craft.

+
Patrol shipAST-17 · Patrol

Diamond Patrol Ship

Heavy symmetric hull for perimeter duty. The diamond centre reads clearly in formation and in isolation.

+
Cargo tenderAST-21 · Logistics

Cargo Tender

Utility vessel for fuel and supplies. Its softer, wider profile adds a practical non-combat role to the fleet.

+

Source batch: PixelLab · 64×64 · transparent PNG · prompt and generation provenance embedded in each file.

diff --git a/game.html b/game.html new file mode 100644 index 0000000..faf7013 --- /dev/null +++ b/game.html @@ -0,0 +1,32 @@ + + + + + Asterion Patrol — Demo Game + + +
+
ASTERION PATROL
SCORE 00000
HULL ■■■
+

WASD/←↑→↓ move · SPACE or click fire · P pause

+

ASTERION PATROL

Defend the nebula corridor.

+
+ diff --git a/pixellab.py b/pixellab.py new file mode 100644 index 0000000..f1b92a3 --- /dev/null +++ b/pixellab.py @@ -0,0 +1,110 @@ +"""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)) diff --git a/polza.py b/polza.py new file mode 100644 index 0000000..fc7ee3b --- /dev/null +++ b/polza.py @@ -0,0 +1,155 @@ +"""Client for Polza.ai's media API. + +The client deliberately uses the standard library so enabling cloud generation +does not add a runtime dependency to the local Diffusers installation. +""" + +import base64 +import json +import mimetypes +import os +import time +from pathlib import Path +from typing import Any, Optional +from urllib.error import HTTPError, URLError +from urllib.request import Request, urlopen + + +API_URL = "https://polza.ai/api/v1" +MAX_REFERENCE_BYTES = 50 * 1024 * 1024 + + +class PolzaError(RuntimeError): + """An API or transport error returned by Polza.ai.""" + + +class PolzaClient: + def __init__(self, api_key: Optional[str] = None, api_url: str = API_URL): + self.api_key = api_key or os.environ.get("POLZA_API_KEY") + self.api_url = api_url.rstrip("/") + + def _headers(self) -> dict[str, str]: + if not self.api_key: + raise PolzaError( + "POLZA_API_KEY is not configured. Set it in the MCP server environment." + ) + return { + "Authorization": f"Bearer {self.api_key}", + "Content-Type": "application/json", + } + + def _request(self, method: str, path: str, payload: Optional[dict] = None) -> dict: + body = json.dumps(payload).encode("utf-8") if payload is not None else None + request = Request( + f"{self.api_url}{path}", data=body, headers=self._headers(), method=method + ) + try: + with urlopen(request, timeout=125) as response: + result = json.loads(response.read().decode("utf-8")) + except HTTPError as exc: + detail = exc.read().decode("utf-8", errors="replace") + try: + detail = json.loads(detail).get("error", {}).get("message", detail) + except json.JSONDecodeError: + pass + raise PolzaError(f"Polza API returned HTTP {exc.code}: {detail}") from exc + except (URLError, TimeoutError) as exc: + raise PolzaError(f"Could not reach Polza API: {exc}") from exc + + if result.get("error"): + error = result["error"] + raise PolzaError(error.get("message", str(error))) + return result + + @staticmethod + def reference_payload(reference: str) -> dict[str, str]: + """Convert an HTTPS URL, data URI, or local image path for Media API.""" + if reference.startswith(("https://", "http://")): + return {"type": "url", "data": reference} + if reference.startswith("data:"): + return {"type": "base64", "data": reference} + + path = Path(reference).expanduser() + if not path.is_file(): + raise PolzaError(f"Reference image does not exist: {reference}") + if path.stat().st_size > MAX_REFERENCE_BYTES: + raise PolzaError(f"Reference image exceeds the 50 MB API limit: {reference}") + mime_type = mimetypes.guess_type(path.name)[0] or "application/octet-stream" + encoded = base64.b64encode(path.read_bytes()).decode("ascii") + return {"type": "base64", "data": f"data:{mime_type};base64,{encoded}"} + + def create_image( + self, + *, + model: str, + prompt: str, + reference_images: Optional[list[str]] = None, + count: int = 1, + aspect_ratio: Optional[str] = None, + seed: Optional[int] = None, + quality: Optional[str] = None, + output_format: str = "png", + background: Optional[str] = None, + provider: Optional[dict] = None, + wait: bool = True, + ) -> dict: + if not 1 <= count <= 10: + raise PolzaError("count must be between 1 and 10") + input_data: dict[str, Any] = { + "prompt": prompt, + "max_images": count, + "output_format": output_format, + } + if reference_images: + input_data["images"] = [self.reference_payload(ref) for ref in reference_images] + for name, value in { + "aspect_ratio": aspect_ratio, + "seed": seed, + "quality": quality, + "background": background, + }.items(): + if value is not None: + input_data[name] = value + + payload: dict[str, Any] = {"model": model, "input": input_data, "async": not wait} + if provider: + payload["provider"] = provider + return self._request("POST", "/media", payload) + + def get_status(self, generation_id: str) -> dict: + return self._request("GET", f"/media/{generation_id}") + + def wait_for_completion( + self, generation_id: str, timeout_seconds: int = 115, poll_seconds: float = 2 + ) -> dict: + deadline = time.monotonic() + timeout_seconds + while True: + status = self.get_status(generation_id) + state = status.get("status", "").lower() + if state in {"completed", "succeeded", "success"}: + return status + if state in {"failed", "error", "cancelled", "canceled"}: + error = status.get("error") or {} + raise PolzaError(error.get("message", f"Generation {generation_id} {state}")) + if time.monotonic() >= deadline: + return status + time.sleep(poll_seconds) + + @staticmethod + def image_sources(response: dict) -> list[str]: + """Extract CDN URLs or base64 images from documented response variants.""" + found: list[str] = [] + + def visit(value: Any) -> None: + if isinstance(value, dict): + for key, item in value.items(): + if key in {"url", "b64_json", "base64"} and isinstance(item, str): + found.append(item) + elif key in {"data", "output", "images", "result", "results"}: + visit(item) + elif isinstance(value, list): + for item in value: + visit(item) + + visit(response) + return list(dict.fromkeys(found)) diff --git a/server.py b/server.py index 399c06a..0d6039f 100644 --- a/server.py +++ b/server.py @@ -18,16 +18,41 @@ Every generated sprite is automatically saved to the feedback DB (unrated). import os import sys import time +import base64 +import json +from io import BytesIO from typing import Optional +from urllib.request import urlopen import numpy as np -from PIL import Image +from PIL import Image, PngImagePlugin from mcp.server.fastmcp import FastMCP from feedback import FeedbackDB +from polza import PolzaClient, PolzaError +from pixellab import PixelLabClient + + +def _load_dotenv(path: str) -> None: + """Load a minimal KEY=value .env file without overriding real environment.""" + if not os.path.isfile(path): + return + with open(path, encoding="utf-8") as env_file: + for raw_line in env_file: + line = raw_line.strip() + if not line or line.startswith("#") or "=" not in line: + continue + key, value = line.split("=", 1) + key, value = key.strip(), value.strip() + if value[:1] == value[-1:] and value[:1] in {"'", '"'}: + value = value[1:-1] + if key: + os.environ.setdefault(key, value) + # Paths — models live in a shared location BASE_DIR = os.path.dirname(os.path.abspath(__file__)) +_load_dotenv(os.path.join(BASE_DIR, ".env")) MODEL_DIR = os.environ.get( "IMAGEGEN_MODEL_DIR", os.path.join(os.path.expanduser("~"), "models", "sdxl-base"), @@ -42,6 +67,7 @@ LCM_LORA_DIR = os.environ.get( ) OUTPUT_DIR = os.environ.get("IMAGEGEN_OUTPUT_DIR", os.path.join(BASE_DIR, "output")) DB_PATH = os.environ.get("IMAGEGEN_DB_PATH", os.path.join(BASE_DIR, "feedback.db")) +POLZA_MODEL = os.environ.get("POLZA_IMAGE_MODEL", "openai/gpt-image-1.5") # LoRA scales — pixel-art-xl needs 1.2, LCM needs 1.0 PIXEL_LORA_SCALE = 1.2 @@ -222,8 +248,214 @@ def _ensure_dir(path: str): os.makedirs(dir_path, exist_ok=True) +def _save_sprite(image: Image.Image, path: str, prompt: str, metadata: dict) -> None: + """Save a PNG with portable, machine-readable generation provenance.""" + _ensure_dir(path) + png_info = PngImagePlugin.PngInfo() + png_info.add_text("prompt", prompt) + png_info.add_text("imagen", json.dumps(metadata, ensure_ascii=False)) + image.save(path, format="PNG", pnginfo=png_info) + + +def _output_paths(output_path: str, count: int) -> list[str]: + """Keep a requested filename for one result; suffix variants for many.""" + if not os.path.isabs(output_path): + output_path = os.path.join(OUTPUT_DIR, output_path) + if count == 1: + return [output_path] + stem, extension = os.path.splitext(output_path) + extension = extension or ".png" + return [f"{stem}_{index:02d}{extension}" for index in range(1, count + 1)] + + +def _image_from_source(source: str) -> Image.Image: + if source.startswith("data:"): + source = source.split(",", 1)[-1] + if source.startswith(("http://", "https://")): + with urlopen(source, timeout=60) as response: + data = response.read() + else: + data = base64.b64decode(source) + return Image.open(BytesIO(data)).convert("RGBA") + + +def _polza_generate( + prompt: str, + output_path: str, + *, + count: int = 1, + model: Optional[str] = None, + reference_images: Optional[list[str]] = None, + aspect_ratio: Optional[str] = None, + seed: Optional[int] = None, + quality: Optional[str] = None, + remove_bg: bool = True, + pixel_size: int = 4, + wait: bool = True, +) -> list[dict]: + """Generate and persist image variants through Polza Media API.""" + client = PolzaClient() + full_prompt = _build_prompt(prompt) + started = time.time() + selected_model = model or POLZA_MODEL + response = client.create_image( + model=selected_model, + prompt=full_prompt, + reference_images=reference_images, + count=count, + aspect_ratio=aspect_ratio, + seed=seed, + quality=quality, + # The transparent-background option is specific to GPT Image. Other + # models still get transparent PNGs through local post-processing. + background="transparent" if remove_bg and "gpt-image" in selected_model else None, + wait=wait, + ) + task_id = response.get("id") + state = response.get("status", "").lower() + if wait and task_id and state in {"pending", "queued", "processing", "running"}: + response = client.wait_for_completion(task_id) + + sources = client.image_sources(response) + if not sources: + return [{ + "generation_id": task_id, + "status": response.get("status", "pending"), + "provider": "polza", + "model": response.get("model", model or POLZA_MODEL), + "message": "Generation is still running; call get_generation_status later.", + }] + + paths = _output_paths(output_path, len(sources)) + results = [] + for index, (source, path) in enumerate(zip(sources, paths), start=1): + _ensure_dir(path) + image = _image_from_source(source) + if pixel_size > 0: + image = _pixelate(image, pixel_size) + if remove_bg: + image = _remove_background(image) + _save_sprite(image, path, prompt, {"provider": "polza", "full_prompt": full_prompt, + "generation_id": task_id, "model": response.get("model", model or POLZA_MODEL)}) + params = { + "provider": "polza", + "generation_id": task_id, + "model": response.get("model", model or POLZA_MODEL), + "seed": seed, + "reference_images": reference_images or [], + "aspect_ratio": aspect_ratio, + "quality": quality, + "remove_bg": remove_bg, + "pixel_size": pixel_size, + "full_prompt": full_prompt, + "usage": response.get("usage"), + "warnings": response.get("warnings", []), + } + entry_id = _get_db().add(prompt=prompt, params=params, image_path=path) + results.append({ + "output_path": path, + "source_url": source if source.startswith(("http://", "https://")) else None, + "variant": index, + "db_id": entry_id, + "generation_id": task_id, + "generation_time": f"{time.time() - started:.1f}s", + "prompt": full_prompt, + "provider": "polza", + "model": response.get("model", model or POLZA_MODEL), + "usage": response.get("usage"), + "warnings": response.get("warnings", []), + "rated": False, + }) + return results + + +def _pixellab_generate( + prompt: str, output_path: str, *, count: int = 1, reference_images: Optional[list[str]] = None, + seed: Optional[int] = None, width: int = 128, height: int = 128, remove_bg: bool = True, + pixel_size: int = 0, wait: bool = True, +) -> list[dict]: + """Generate native pixel-art sprites through PixelLab v2.""" + client = PixelLabClient() + full_prompt = _build_prompt(prompt) + responses = [client.create_image(full_prompt, width, height, seed=seed, reference_images=reference_images, + no_background=remove_bg) for _ in range(count)] + completed = [] + for response in responses: + job_id = response.get("background_job_id") + if wait and job_id: + response = client.wait_for_completion(job_id) + completed.append((job_id, response)) + source_records = [(job_id, response, source) for job_id, response in completed + for source in client.image_sources(response)] + if not source_records: + return [{"generation_id": job_id, "status": response.get("status", "processing"), + "provider": "pixellab", "usage": response.get("usage")} + for job_id, response in completed] + results = [] + for index, ((job_id, response, source), path) in enumerate( + zip(source_records, _output_paths(output_path, len(source_records))), start=1 + ): + _ensure_dir(path) + image = _image_from_source(source) + if pixel_size > 0: + image = _pixelate(image, pixel_size) + if remove_bg: + image = _remove_background(image) + _save_sprite(image, path, prompt, {"provider": "pixellab", "full_prompt": full_prompt, + "generation_id": job_id, "width": width, "height": height}) + db_id = _get_db().add(prompt=prompt, image_path=path, params={ + "provider": "pixellab", "generation_id": job_id, "reference_images": reference_images or [], + "seed": seed, "width": width, "height": height, "remove_bg": remove_bg, + "pixel_size": pixel_size, "full_prompt": full_prompt, "usage": response.get("usage"), + }) + results.append({"output_path": path, "db_id": db_id, "generation_id": job_id, "variant": index, + "provider": "pixellab", "usage": response.get("usage"), "rated": False}) + return results + + +# Instructions are part of the MCP initialization response, so every connected +# agent receives the asset-generation rules before it chooses a tool. +MCP_INSTRUCTIONS = """ +Generate game assets deliberately and keep the dataset reusable. + +Style workflow: +1. Call get_project_style_guide first. Its registered references live in the + server database; pass 1–4 selected paths as reference_images to + provider='pixellab' for a coherent style family. get_reference_sprites is + useful for additional rated examples. +2. Generate one semantic asset role per prompt (for example only a missile + corvette). Never put a list of roles in one prompt when individual metadata + must identify each image: all outputs of a multi-image batch inherit the + same prompt and cannot reliably be labelled afterwards. +3. State camera, orientation, silhouette, palette, background, and exclusions + explicitly. For top-down fleets, say: 'orthographic top-down; nose at 12 + o'clock; engines at 6 o'clock'. Add exact vertical mirror symmetry only + when it is artistically required. + +Pixel-art rules: +- For pixel-perfect assets use PixelLab at a native square size (normally + 64x64 or 128x128), transparent background, pixel_size=0, and request no + anti-aliasing, blur, sub-pixel shading, or semi-transparent edges. +- In HTML/game clients display these files at an integer multiple with + image-rendering: pixelated; never use arbitrary CSS scaling. +- Use provider='pixellab' for native game sprites and style references. + Use provider='polza' for broader concept art or backgrounds. Local is an + optional Diffusers/GPU provider. + +Operations and provenance: +- Use a semantic output_path and preserve every returned db_id. Finished PNGs + embed 'prompt' and 'imagen' JSON metadata (provider, full prompt and + generation settings); the feedback database stores the same provenance. +- generate_images is for variants of one precise asset. batch_generate is for + several independent, precisely described assets. For async work use wait=false + and get_generation_status with the matching provider. +- Never expose, request, or store API keys in prompts, output paths, feedback, + or image metadata. +""".strip() + + # Create MCP server -mcp = FastMCP("pixel-art") +mcp = FastMCP("pixel-art", instructions=MCP_INSTRUCTIONS) @mcp.tool() @@ -236,8 +468,18 @@ def generate_sprite( steps: int = 8, remove_bg: bool = True, pixel_size: int = 4, + provider: str = "local", + model: Optional[str] = None, + reference_images: Optional[list[str]] = None, + aspect_ratio: Optional[str] = None, + quality: Optional[str] = None, ) -> dict: - """Generate a pixel-art sprite and save it as PNG with transparent background. + """Generate one precisely described pixel-art sprite and save it as PNG. + + For a reusable game dataset, describe only one semantic role per call and + state view/orientation/palette explicitly. Use provider='pixellab' plus + 1–4 reference_images for style-consistent native pixel art. Each PNG + embeds the prompt and generation provenance in its metadata. Args: prompt: Description of the sprite (e.g. "a crystal warrior with geometric armor") @@ -247,11 +489,36 @@ def generate_sprite( height: Image height in pixels (default 512) steps: Inference steps (default 4, FLUX.2-klein is distilled) remove_bg: Remove background and make transparent (default True) - pixel_size: Size of each pixel block for pixel-art effect (default 4, 0=off) + pixel_size: Size of each pixel block for pixel-art effect (default 4, 0=off) + provider: "local" (SDXL) or "polza" (cloud image models) + model: Polza model ID; defaults to POLZA_IMAGE_MODEL + reference_images: Style references: HTTPS URLs, data URIs, or local image paths + aspect_ratio: Cloud aspect ratio such as "1:1" or "16:9" + quality: Cloud model quality setting Returns: Dict with output_path, seed_used, generation_time, prompt, size. """ + if provider == "polza": + return _polza_generate( + prompt, + output_path, + model=model, + reference_images=reference_images, + aspect_ratio=aspect_ratio, + seed=seed, + quality=quality, + remove_bg=remove_bg, + pixel_size=pixel_size, + )[0] + if provider == "pixellab": + return _pixellab_generate(prompt, output_path, reference_images=reference_images, seed=seed, + width=width, height=height, remove_bg=remove_bg, pixel_size=pixel_size)[0] + if provider != "local": + raise ValueError("provider must be 'local', 'polza', or 'pixellab'") + if reference_images: + raise ValueError("reference_images require provider='polza'") + pipe = _load_model() full_prompt = _build_prompt(prompt) @@ -269,7 +536,8 @@ def generate_sprite( if remove_bg: image = _remove_background(image) - image.save(output_path) + _save_sprite(image, output_path, prompt, {"provider": "local", "full_prompt": full_prompt, + "seed": seed, "width": width, "height": height, "steps": steps}) elapsed = time.time() - t0 db = _get_db() @@ -312,7 +580,12 @@ def generate_sprite( def batch_generate( specs: list[dict], ) -> list[dict]: - """Generate multiple pixel-art sprites in one call. + """Generate multiple independent, precisely described sprites in one call. + + Each spec must represent one asset role. Do not ask one spec for mixed + categories if you need per-image semantic metadata: every result inherits + that spec's single prompt. For a game style family, give every PixelLab spec + the same selected reference_images and native 64x64 or 128x128 dimensions. Args: specs: List of dicts, each with: @@ -328,7 +601,7 @@ def batch_generate( Returns: List of dicts with output_path, seed_used, generation_time, prompt, size, transparent. """ - pipe = _load_model() + pipe = None results = [] for spec in specs: @@ -340,6 +613,51 @@ def batch_generate( steps = spec.get("steps", 8) remove_bg = spec.get("remove_bg", True) pixel_size = spec.get("pixel_size", 4) + provider = spec.get("provider", "local") + model = spec.get("model") + reference_images = spec.get("reference_images") + aspect_ratio = spec.get("aspect_ratio") + quality = spec.get("quality") + + if provider == "polza": + results.extend( + _polza_generate( + prompt, + output_path, + count=spec.get("count", 1), + model=model, + reference_images=reference_images, + aspect_ratio=aspect_ratio, + seed=seed, + quality=quality, + remove_bg=remove_bg, + pixel_size=pixel_size, + wait=spec.get("wait", True), + ) + ) + continue + if provider == "pixellab": + results.extend( + _pixellab_generate( + prompt, + output_path, + count=spec.get("count", 1), + reference_images=reference_images, + seed=seed, + width=width, + height=height, + remove_bg=remove_bg, + pixel_size=pixel_size, + wait=spec.get("wait", True), + ) + ) + continue + if provider != "local": + raise ValueError("provider must be 'local', 'polza', or 'pixellab'") + if reference_images: + raise ValueError("reference_images require provider='polza'") + if pipe is None: + pipe = _load_model() full_prompt = _build_prompt(prompt) @@ -357,7 +675,8 @@ def batch_generate( if remove_bg: image = _remove_background(image) - image.save(output_path) + _save_sprite(image, output_path, prompt, {"provider": "local", "full_prompt": full_prompt, + "seed": seed, "width": width, "height": height, "steps": steps}) elapsed = time.time() - t0 db = _get_db() @@ -399,6 +718,80 @@ def batch_generate( return results +@mcp.tool() +def generate_images( + prompt: str, + output_path: str, + count: int = 1, + model: Optional[str] = None, + reference_images: Optional[list[str]] = None, + aspect_ratio: Optional[str] = None, + seed: Optional[int] = None, + quality: Optional[str] = None, + remove_bg: bool = True, + pixel_size: int = 4, + wait: bool = True, + provider: str = "polza", +) -> list[dict]: + """Generate cloud-image variations of one precise asset request. + + Use ``reference_images`` for a game's style guide, existing characters, or + tiles. Each value may be an HTTPS URL, data URI, or a path readable by this + MCP server. Results use output_path_01.png, output_path_02.png, etc. + When ``wait`` is false, the returned generation_id can be passed to + get_generation_status later. Use separate calls (or batch_generate specs) + for different roles such as scout, corvette, and freighter, so embedded + prompt metadata remains meaningful for every image. For pixel-perfect + assets choose provider='pixellab'; it uses a native 128x128 output here. + """ + if provider == "pixellab": + # PixelLab produces native pixel art; its practical size range is 16–512. + # aspect_ratio is not applicable because its API accepts explicit dimensions. + return _pixellab_generate( + prompt, output_path, count=count, reference_images=reference_images, seed=seed, + remove_bg=remove_bg, pixel_size=pixel_size, wait=wait, + ) + if provider != "polza": + raise ValueError("provider must be 'polza' or 'pixellab'") + return _polza_generate( + prompt, + output_path, + count=count, + model=model, + reference_images=reference_images, + aspect_ratio=aspect_ratio, + seed=seed, + quality=quality, + remove_bg=remove_bg, + pixel_size=pixel_size, + wait=wait, + ) + + +@mcp.tool() +def get_generation_status(generation_id: str, provider: str = "polza") -> dict: + """Get a Polza or PixelLab cloud-generation status and finished sources.""" + if provider == "pixellab": + status = PixelLabClient().get_status(generation_id) + return { + "generation_id": status.get("id", generation_id), "status": status.get("status"), + "provider": "pixellab", "image_sources": PixelLabClient.image_sources(status), + "usage": status.get("usage"), "error": status.get("error"), + } + if provider != "polza": + raise ValueError("provider must be 'polza' or 'pixellab'") + status = PolzaClient().get_status(generation_id) + return { + "generation_id": status.get("id", generation_id), + "status": status.get("status"), + "model": status.get("model"), + "image_sources": PolzaClient.image_sources(status), + "usage": status.get("usage"), + "warnings": status.get("warnings", []), + "error": status.get("error"), + } + + @mcp.tool() def rate_sprite( db_id: str, @@ -468,6 +861,65 @@ def get_reference_sprites( ] +@mcp.tool() +def get_project_style_guide(limit: int = 4) -> dict: + """Return the server's current reusable style guide and reference images. + + Call this before generating a coordinated asset family. The references are + registered in the MCP server's SQLite DB, rather than hard-coded in a + client. Pass the returned recommended_reference_images to PixelLab. + """ + references = _get_db().get_style_references(limit) + return { + "rules": { + "pixel_art": "Use native 64x64 or 128x128 PNG and integer client scaling.", + "top_down": "State orthographic camera, nose at 12 o'clock, engines at 6 o'clock.", + "metadata": "Use one semantic asset role per generation prompt.", + }, + "recommended_reference_images": [reference.image_path for reference in references], + "references": [ + { + "reference_id": reference.id, + "image_path": reference.image_path, + "name": reference.name, + "role": reference.role, + "notes": reference.notes, + "priority": reference.priority, + } + for reference in references + ], + } + + +@mcp.tool() +def register_style_reference( + image_path: str, + name: str, + role: str, + notes: Optional[str] = None, + priority: int = 0, +) -> dict: + """Register an existing generated image as a reusable project style reference. + + Use a stable, well-reviewed PNG. Higher priority references are returned + first by get_project_style_guide. Register 1–4 complementary examples, not + many near-duplicates. + """ + if not os.path.isfile(image_path): + raise ValueError(f"Style reference image does not exist: {image_path}") + reference_id = _get_db().add_style_reference( + image_path, name, role, notes, priority + ) + return {"reference_id": reference_id, "status": "saved", "image_path": image_path} + + +@mcp.tool() +def remove_style_reference(reference_id: str) -> dict: + """Remove a style-reference registration; the image file is not deleted.""" + _get_db().delete_style_reference(reference_id) + return {"reference_id": reference_id, "status": "removed"} + + @mcp.tool() def list_sprites( filter: str = "all", diff --git a/test_feedback.py b/test_feedback.py index 8d07809..d3daa0b 100644 --- a/test_feedback.py +++ b/test_feedback.py @@ -10,7 +10,7 @@ import tempfile import pytest sys.path.insert(0, os.path.dirname(os.path.abspath(__file__))) -from feedback import FeedbackDB, FeedbackEntry, DBStats, _tokenize +from feedback import FeedbackDB, FeedbackEntry, DBStats, StyleReference, _tokenize @pytest.fixture @@ -283,6 +283,30 @@ class TestDelete: assert entries[0].prompt == "archer" +class TestStyleReferences: + def test_add_and_list_style_reference(self, db): + reference_id = db.add_style_reference( + "/tmp/style.png", "Asterion vanguard", "fleet style", "cyan engines", 10 + ) + references = db.get_style_references() + assert references[0].id == reference_id + assert isinstance(references[0], StyleReference) + assert references[0].name == "Asterion vanguard" + + def test_style_reference_upserts_by_image_path(self, db): + first = db.add_style_reference("/tmp/style.png", "old", "style") + second = db.add_style_reference("/tmp/style.png", "new", "style", priority=5) + assert first == second + references = db.get_style_references() + assert len(references) == 1 + assert references[0].name == "new" + + def test_delete_style_reference(self, db): + reference_id = db.add_style_reference("/tmp/style.png", "ship", "style") + db.delete_style_reference(reference_id) + assert db.get_style_references() == [] + + class TestExportJsonl: def test_export_jsonl(self, db, tmp_path): _add_sample(db, "knight", rating=5, feedback="great") diff --git a/test_pixellab.py b/test_pixellab.py new file mode 100644 index 0000000..84bb208 --- /dev/null +++ b/test_pixellab.py @@ -0,0 +1,24 @@ +from pixellab import PixelLabClient + + +def test_image_sources_reads_completed_job_images(): + response = { + "last_response": { + "images": [{"base64": "data:image/png;base64,AAAA"}], + "url": "https://example.com/preview.png", + } + } + assert PixelLabClient.image_sources(response) == [ + "data:image/png;base64,AAAA", + "https://example.com/preview.png", + ] + + +def test_create_image_uses_style_endpoint(monkeypatch): + client = PixelLabClient(api_key="test") + monkeypatch.setattr(client, "_reference", lambda _: {"image": {"base64": "x"}, "width": 64, "height": 64}) + captured = {} + monkeypatch.setattr(client, "_request", lambda method, path, payload: captured.update(method=method, path=path, payload=payload) or {}) + client.create_image("ship", 64, 64, reference_images=["style.png"]) + assert captured["path"] == "/generate-with-style-v2" + assert captured["payload"]["style_images"][0]["width"] == 64 diff --git a/test_polza.py b/test_polza.py new file mode 100644 index 0000000..44f59a8 --- /dev/null +++ b/test_polza.py @@ -0,0 +1,68 @@ +import base64 +from pathlib import Path + +import pytest + +from polza import PolzaClient, PolzaError + + +def test_reference_payload_accepts_url_and_data_uri(): + assert PolzaClient.reference_payload("https://example.com/style.png") == { + "type": "url", + "data": "https://example.com/style.png", + } + data_uri = "data:image/png;base64,AAAA" + assert PolzaClient.reference_payload(data_uri) == {"type": "base64", "data": data_uri} + + +def test_reference_payload_encodes_local_image(tmp_path): + path = tmp_path / "style.png" + path.write_bytes(b"image-bytes") + + payload = PolzaClient.reference_payload(str(path)) + + assert payload["type"] == "base64" + assert payload["data"].startswith("data:image/png;base64,") + assert base64.b64decode(payload["data"].split(",", 1)[1]) == b"image-bytes" + + +def test_reference_payload_rejects_unknown_path(): + with pytest.raises(PolzaError, match="does not exist"): + PolzaClient.reference_payload("missing-style.png") + + +def test_create_image_sends_references_and_variants(monkeypatch): + client = PolzaClient(api_key="test-key") + captured = {} + + def fake_request(method, path, payload=None): + captured.update(method=method, path=path, payload=payload) + return {"id": "gen_1", "status": "pending"} + + monkeypatch.setattr(client, "_request", fake_request) + client.create_image( + model="seedream-3", + prompt="game sprite", + reference_images=["https://example.com/style.png"], + count=3, + aspect_ratio="1:1", + seed=42, + wait=False, + ) + + assert captured["method"] == "POST" + assert captured["path"] == "/media" + assert captured["payload"]["async"] is True + assert captured["payload"]["input"]["max_images"] == 3 + assert captured["payload"]["input"]["images"][0]["type"] == "url" + + +def test_image_sources_handles_urls_and_base64(): + response = { + "data": [{"url": "https://cdn.example/one.png"}], + "result": {"images": [{"b64_json": "aGVsbG8="}]}, + } + assert PolzaClient.image_sources(response) == [ + "https://cdn.example/one.png", + "aGVsbG8=", + ] diff --git a/test_server.py b/test_server.py index eb81bf2..f022bd4 100644 --- a/test_server.py +++ b/test_server.py @@ -152,6 +152,17 @@ class TestEnsureDir: server._ensure_dir("test.png") +class TestPngMetadata: + def test_save_sprite_embeds_prompt_and_provenance(self, tmp_path): + path = tmp_path / "sprite.png" + image = Image.new("RGBA", (4, 4), (1, 2, 3, 255)) + server._save_sprite(image, str(path), "blue scout", {"provider": "test"}) + + with Image.open(path) as saved: + assert saved.text["prompt"] == "blue scout" + assert '"provider": "test"' in saved.text["imagen"] + + class TestEnvPaths: def test_default_model_dir(self): assert "models" in server.MODEL_DIR