diff --git a/.gitignore b/.gitignore index 4278879..7688996 100644 --- a/.gitignore +++ b/.gitignore @@ -13,6 +13,10 @@ build/ # Output (generated sprites) output/ +# Feedback database +feedback.db +feedback_dataset.jsonl + # IDE .idea/ .vscode/ diff --git a/README.md b/README.md index fe84561..e02146f 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.) +- **Feedback loop** — rate generated sprites, AI uses high-rated ones as reference ## Quick Start @@ -85,9 +86,11 @@ Or configure in your MCP client: ## Tools -### `generate_sprite` +### Generation -Generate a single pixel-art sprite. +#### `generate_sprite` + +Generate a single pixel-art sprite. Automatically saved to feedback DB (unrated). | Parameter | Type | Default | Description | |---|---|---|---| @@ -100,9 +103,46 @@ Generate a single pixel-art sprite. | `remove_bg` | bool | true | Remove background, make transparent | | `pixel_size` | int | 4 | Pixel block size (0 = off, 4 = chunky pixel-art) | -### `batch_generate` +Returns: `output_path`, `db_id`, `generation_time`, and other metadata. -Generate multiple sprites in one call. Accepts a list of specs with the same parameters. +#### `batch_generate` + +Generate multiple sprites in one call. Each is saved to the feedback DB. + +### Feedback + +#### `rate_sprite` + +Rate a generated sprite 1-5 stars with optional feedback. + +| Parameter | Type | Default | Description | +|---|---|---|---| +| `db_id` | str | required | ID returned by generate_sprite / batch_generate | +| `rating` | int | required | 1-5 stars | +| `feedback` | str? | null | Optional text feedback | + +#### `get_reference_sprites` + +Get highly-rated reference sprites for a prompt. The AI uses these as examples when generating similar sprites. + +| Parameter | Type | Default | Description | +|---|---|---|---| +| `prompt` | str | required | Search query (e.g. "knight") | +| `limit` | int | 5 | Max results | +| `min_rating` | int | 4 | Minimum rating threshold | + +#### `list_sprites` + +List sprites in the feedback database. + +| Parameter | Type | Default | Description | +|---|---|---|---| +| `filter` | str | "all" | "all", "unrated", or "top" | +| `limit` | int | 20 | Max results | + +#### `db_stats` + +Get database statistics: total sprites, rated, unrated, average rating. ## Configuration diff --git a/feedback.py b/feedback.py new file mode 100644 index 0000000..a62c57e --- /dev/null +++ b/feedback.py @@ -0,0 +1,192 @@ +""" +Feedback database — SQLite-backed storage for sprite ratings. + +Stores generated sprites with their prompts, params, PNG paths, and user ratings (1-5 stars). +Provides similarity search by prompt keywords for few-shot reference examples. +""" + +import json +import os +import sqlite3 +import time +import uuid +from dataclasses import dataclass, field +from typing import Optional + + +@dataclass +class FeedbackEntry: + id: str + prompt: str + params: dict + rating: int + feedback: Optional[str] + image_path: Optional[str] + created_at: str + + +@dataclass +class DBStats: + total: int = 0 + rated: int = 0 + unrated: int = 0 + avg_rating: float = 0.0 + + +class FeedbackDB: + """SQLite-backed feedback database for generated sprites.""" + + def __init__(self, conn: sqlite3.Connection): + self.conn = conn + + @classmethod + def open(cls, path: str) -> "FeedbackDB": + conn = sqlite3.connect(path) + conn.execute(""" + CREATE TABLE IF NOT EXISTS feedback ( + id TEXT PRIMARY KEY, + prompt TEXT NOT NULL, + params_json TEXT NOT NULL, + rating INTEGER DEFAULT 0, + feedback TEXT, + image_path TEXT, + created_at TEXT NOT NULL + ) + """) + 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.commit() + return cls(conn) + + def add( + self, + prompt: str, + params: dict, + image_path: Optional[str] = None, + ) -> str: + entry_id = str(uuid.uuid4()) + params_json = json.dumps(params) + now = str(int(time.time())) + + self.conn.execute( + "INSERT INTO feedback (id, prompt, params_json, rating, feedback, image_path, created_at) " + "VALUES (?, ?, ?, 0, NULL, ?, ?)", + (entry_id, prompt, params_json, image_path, now), + ) + self.conn.commit() + return entry_id + + def update_rating( + self, + entry_id: str, + rating: int, + feedback: Optional[str] = None, + ) -> None: + rating = max(0, min(5, rating)) + self.conn.execute( + "UPDATE feedback SET rating = ?, feedback = ? WHERE id = ?", + (rating, feedback, entry_id), + ) + self.conn.commit() + + def get_all(self) -> list[FeedbackEntry]: + cursor = self.conn.execute( + "SELECT id, prompt, params_json, rating, feedback, image_path, created_at " + "FROM feedback ORDER BY created_at DESC" + ) + return [self._row_to_entry(row) for row in cursor] + + def get_unrated(self) -> list[FeedbackEntry]: + cursor = self.conn.execute( + "SELECT id, prompt, params_json, rating, feedback, image_path, created_at " + "FROM feedback WHERE rating = 0 ORDER BY created_at DESC" + ) + return [self._row_to_entry(row) for row in cursor] + + def top_rated(self, limit: int = 10, min_rating: int = 1) -> list[FeedbackEntry]: + cursor = self.conn.execute( + "SELECT id, prompt, params_json, rating, feedback, image_path, created_at " + "FROM feedback WHERE rating >= ? ORDER BY rating DESC, created_at DESC LIMIT ?", + (min_rating, limit), + ) + return [self._row_to_entry(row) for row in cursor] + + def search_similar(self, query: str, limit: int = 5) -> list[FeedbackEntry]: + query_keywords = _tokenize(query) + + if not query_keywords: + return self.top_rated(limit, 1) + + all_entries = self.get_all() + + scored = [] + for entry in all_entries: + entry_keywords = _tokenize(entry.prompt) + match_count = sum( + 1 for qk in query_keywords if any(ek == qk for ek in entry_keywords) + ) + if match_count > 0: + scored.append((entry, match_count, entry.rating)) + + scored.sort(key=lambda x: (-x[1], -x[2])) + return [e for e, _, _ in scored[:limit]] + + def stats(self) -> DBStats: + total = self.conn.execute("SELECT COUNT(*) FROM feedback").fetchone()[0] + rated = self.conn.execute( + "SELECT COUNT(*) FROM feedback WHERE rating > 0" + ).fetchone()[0] + avg = ( + self.conn.execute( + "SELECT AVG(rating) FROM feedback WHERE rating > 0" + ).fetchone()[0] + or 0.0 + ) + + return DBStats( + total=total, + rated=rated, + unrated=total - rated, + avg_rating=round(avg, 1), + ) + + def delete(self, entry_id: str) -> None: + self.conn.execute("DELETE FROM feedback WHERE id = ?", (entry_id,)) + self.conn.commit() + + def export_jsonl(self, path: str, min_rating: int = 4) -> int: + entries = self.top_rated(10000, min_rating) + lines = [] + + for entry in entries: + line = json.dumps( + { + "instruction": f"Generate a pixel-art sprite for: {entry.prompt}", + "response": entry.params, + "rating": entry.rating, + } + ) + lines.append(line) + + with open(path, "w") as f: + f.write("\n".join(lines)) + + return len(entries) + + def _row_to_entry(self, row: sqlite3.Row) -> FeedbackEntry: + return FeedbackEntry( + id=row[0], + prompt=row[1], + params=json.loads(row[2]), + rating=max(0, min(5, row[3])), + feedback=row[4], + image_path=row[5], + created_at=row[6], + ) + + +def _tokenize(s: str) -> list[str]: + """Tokenize a prompt into lowercase keywords.""" + return [ + w.lower() for w in s.replace("_", " ").replace("-", " ").split() if len(w) > 1 + ] diff --git a/server.py b/server.py index 2a62289..6a7e814 100644 --- a/server.py +++ b/server.py @@ -3,11 +3,16 @@ MCP server for generating pixel-art sprites using FLUX.2-klein-4B + pixel-art-lora. Tools: - - generate_sprite: Generate a single pixel-art sprite - - batch_generate: Generate multiple sprites in one call + - generate_sprite: Generate a single pixel-art sprite + - batch_generate: Generate multiple sprites in one call + - rate_sprite: Rate a generated sprite (1-5 stars) with optional feedback + - get_reference_sprites: Get highly-rated reference sprites for a prompt + - list_sprites: List sprites in the feedback DB (all, unrated, or top-rated) + - db_stats: Get feedback database statistics Model is loaded lazily on first call (~6s), then stays in VRAM for speed. Background is removed post-generation to produce transparent PNG. +Every generated sprite is automatically saved to the feedback DB (unrated). """ import os @@ -19,6 +24,8 @@ import numpy as np from PIL import Image from mcp.server.fastmcp import FastMCP +from feedback import FeedbackDB + # Paths — models live in a shared location BASE_DIR = os.path.dirname(os.path.abspath(__file__)) MODEL_DIR = os.environ.get( @@ -30,6 +37,7 @@ LORA_DIR = os.environ.get( os.path.join(os.path.expanduser("~"), "models", "pixel-art-lora"), ) 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")) # rsLoRA requires much lower scale in diffusers — 1.0 produces black images LORA_SCALE = 0.1 @@ -37,6 +45,15 @@ LORA_SCALE = 0.1 # Global state — model loaded lazily _pipe = None _device = None +_db: Optional[FeedbackDB] = None + + +def _get_db() -> FeedbackDB: + global _db + if _db is None: + _db = FeedbackDB.open(DB_PATH) + sys.stderr.write(f"[pixel-art] Feedback DB: {DB_PATH}\n") + return _db def _get_device(): @@ -250,6 +267,20 @@ def generate_sprite( image.save(output_path) elapsed = time.time() - t0 + db = _get_db() + entry_id = db.add( + prompt=prompt, + params={ + "seed": seed, + "width": width, + "height": height, + "steps": steps, + "remove_bg": remove_bg, + "pixel_size": pixel_size, + }, + image_path=output_path, + ) + return { "output_path": output_path, "seed_used": seed, @@ -258,6 +289,8 @@ def generate_sprite( "size": f"{width}x{height}", "transparent": remove_bg, "pixel_size": pixel_size, + "db_id": entry_id, + "rated": False, } @@ -313,6 +346,20 @@ def batch_generate( image.save(output_path) elapsed = time.time() - t0 + db = _get_db() + entry_id = db.add( + prompt=prompt, + params={ + "seed": seed, + "width": width, + "height": height, + "steps": steps, + "remove_bg": remove_bg, + "pixel_size": pixel_size, + }, + image_path=output_path, + ) + results.append( { "output_path": output_path, @@ -321,11 +368,137 @@ def batch_generate( "prompt": full_prompt, "size": f"{width}x{height}", "transparent": remove_bg, + "db_id": entry_id, + "rated": False, } ) return results +@mcp.tool() +def rate_sprite( + db_id: str, + rating: int, + feedback: Optional[str] = None, +) -> dict: + """Rate a generated sprite (1-5 stars) with optional feedback text. + + Use this after reviewing a sprite to teach the system what looks good. + The AI uses high-rated sprites as reference when generating similar ones. + + Args: + db_id: The ID returned by generate_sprite or batch_generate + rating: 1-5 stars (5 = excellent, 1 = terrible) + feedback: Optional text feedback (e.g. "great colors, bad proportions") + + Returns: + Dict with db_id, rating, feedback, and status. + """ + db = _get_db() + db.update_rating(db_id, rating, feedback) + + return { + "db_id": db_id, + "rating": rating, + "feedback": feedback, + "status": "saved", + } + + +@mcp.tool() +def get_reference_sprites( + prompt: str, + limit: int = 5, + min_rating: int = 4, +) -> list[dict]: + """Get highly-rated reference sprites from the feedback DB for a given prompt. + + Use these as examples when generating similar sprites to improve quality. + Returns sprites with similar prompt keywords that have been rated >= min_rating. + + Args: + prompt: The prompt to search for (e.g. "knight", "crystal warrior") + limit: Max number of results (default 5) + min_rating: Minimum rating (1-5, default 4) + + Returns: + List of dicts with db_id, prompt, rating, feedback, image_path, params. + """ + db = _get_db() + entries = db.search_similar(prompt, limit * 2) + entries = [e for e in entries if e.rating >= min_rating][:limit] + + if not entries: + return [] + + return [ + { + "db_id": e.id, + "prompt": e.prompt, + "rating": e.rating, + "feedback": e.feedback, + "image_path": e.image_path, + "params": e.params, + } + for e in entries + ] + + +@mcp.tool() +def list_sprites( + filter: str = "all", + limit: int = 20, +) -> list[dict]: + """List sprites in the feedback database. + + Args: + filter: "all" = all sprites, "unrated" = only unrated, "top" = highest rated + limit: Max number of results (default 20) + + Returns: + List of dicts with db_id, prompt, rating, image_path, created_at. + """ + db = _get_db() + + if filter == "unrated": + entries = db.get_unrated() + elif filter == "top": + entries = db.top_rated(limit, 1) + else: + entries = db.get_all() + + entries = entries[:limit] + + return [ + { + "db_id": e.id, + "prompt": e.prompt, + "rating": e.rating, + "image_path": e.image_path, + "created_at": e.created_at, + } + for e in entries + ] + + +@mcp.tool() +def db_stats() -> dict: + """Get feedback database statistics. + + Returns: + Dict with total, rated, unrated, avg_rating. + """ + db = _get_db() + stats = db.stats() + + return { + "total": stats.total, + "rated": stats.rated, + "unrated": stats.unrated, + "avg_rating": stats.avg_rating, + } + + if __name__ == "__main__": mcp.run(transport="stdio")