Add feedback system: SQLite DB, rating tools, reference search
- feedback.py: SQLite-backed FeedbackDB (add, rate, search, export) - server.py: auto-save to DB on generate, 4 new MCP tools: - rate_sprite: rate 1-5 stars with feedback - get_reference_sprites: find high-rated similar sprites - list_sprites: list all/unrated/top - db_stats: database statistics - Updated README with feedback tool docs
This commit is contained in:
Emil
@@ -13,6 +13,10 @@ build/
|
||||
# Output (generated sprites)
|
||||
output/
|
||||
|
||||
# Feedback database
|
||||
feedback.db
|
||||
feedback_dataset.jsonl
|
||||
|
||||
# IDE
|
||||
.idea/
|
||||
.vscode/
|
||||
|
||||
@@ -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
|
||||
|
||||
|
||||
+192
@@ -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
|
||||
]
|
||||
@@ -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")
|
||||
|
||||
Reference in New Issue
Block a user