commit e28b466c8a1fa870613db7b9b8a989368a888d31 Author: Emil Date: Sun Jun 28 14:50:58 2026 +0300 Initial commit: MCP server for pixel-art sprite generation diff --git a/.gitignore b/.gitignore new file mode 100644 index 0000000..4278879 --- /dev/null +++ b/.gitignore @@ -0,0 +1,33 @@ +# venv +venv/ +.venv/ + +# Python +__pycache__/ +*.pyc +*.pyo +*.egg-info/ +dist/ +build/ + +# Output (generated sprites) +output/ + +# IDE +.idea/ +.vscode/ +*.swp +*.swo + +# OS +.DS_Store +Thumbs.db + +# Models (not included — download separately) +models/ +*.safetensors +*.bin +*.ckpt + +# opencode local config (keep .opencode/opencode.json) +.opencode/session* diff --git a/.opencode/opencode.json b/.opencode/opencode.json new file mode 100644 index 0000000..62a9a95 --- /dev/null +++ b/.opencode/opencode.json @@ -0,0 +1,10 @@ +{ + "$schema": "https://opencode.ai/config.json", + "mcp": { + "pixel-art": { + "type": "local", + "command": ["./venv/bin/python", "server.py"], + "enabled": true + } + } +} diff --git a/LICENSE b/LICENSE new file mode 100644 index 0000000..b7c47c8 --- /dev/null +++ b/LICENSE @@ -0,0 +1,21 @@ +MIT License + +Copyright (c) 2026 emil + +Permission is hereby granted, free of charge, to any person obtaining a copy +of this software and associated documentation files (the "Software"), to deal +in the Software without restriction, including without limitation the rights +to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +copies of the Software, and to permit persons to whom the Software is +furnished to do so, subject to the following conditions: + +The above copyright notice and this permission notice shall be included in all +copies or substantial portions of the Software. + +THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE +SOFTWARE. diff --git a/README.md b/README.md new file mode 100644 index 0000000..c24cd25 --- /dev/null +++ b/README.md @@ -0,0 +1,112 @@ +# Imagegen + +MCP server for generating pixel-art sprites using FLUX.2-klein-4B + [pixel-art-lora](https://huggingface.co/Limbicnation/pixel-art-lora). + +## Features + +- **Text-to-sprite generation** — describe any character, get a pixel-art PNG +- **Transparent background** — automatic background removal via flood-fill +- **Pixel-art effect** — downscale/upscale with NEAREST interpolation +- **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.) + +## Quick Start + +### 1. Install dependencies + +```bash +python3 -m venv venv +source venv/bin/activate +pip install -r requirements.txt +``` + +### 2. Download models + +```bash +# Create models directory +mkdir -p ~/models + +# Download base model (~23 GB) +huggingface-cli download black-forest-labs/FLUX.2-klein-4b \ + --local-dir ~/models/flux2-klein-4b + +# Download LoRA adapter (~625 MB) +huggingface-cli download Limbicnation/pixel-art-lora \ + --local-dir ~/models/pixel-art-lora +``` + +### 3. Configure paths (optional) + +By default, models are expected at `~/models/`. Override with environment variables: + +```bash +export IMAGEGEN_MODEL_DIR=/path/to/flux2-klein-4b +export IMAGEGEN_LORA_DIR=/path/to/pixel-art-lora +export IMAGEGEN_OUTPUT_DIR=/path/to/output +``` + +### 4. Run as MCP server + +```bash +./venv/bin/python server.py +``` + +Or configure in your MCP client: + +```json +{ + "mcp": { + "pixel-art": { + "type": "local", + "command": ["./venv/bin/python", "server.py"], + "enabled": true + } + } +} +``` + +## Tools + +### `generate_sprite` + +Generate a single pixel-art sprite. + +| Parameter | Type | Default | Description | +|---|---|---|---| +| `prompt` | str | required | Sprite description (e.g. "a brave knight in armor") | +| `output_path` | str | required | PNG save path (relative to output dir or absolute) | +| `seed` | int? | null | Seed for reproducibility | +| `width` | int | 512 | Image width | +| `height` | int | 512 | Image height | +| `steps` | int | 4 | Inference steps (FLUX.2-klein is distilled) | +| `remove_bg` | bool | true | Remove background, make transparent | +| `pixel_size` | int | 4 | Pixel block size (0 = off, 4 = chunky pixel-art) | + +### `batch_generate` + +Generate multiple sprites in one call. Accepts a list of specs with the same parameters. + +## How It Works + +1. **Generation** — FLUX.2-klein-4B (4B params, distilled to 4 steps) with pixel-art LoRA (scale 0.1 for rsLoRA compatibility) +2. **Pixelation** — downscale with LANCZOS, upscale with NEAREST → chunky pixel-art blocks +3. **Background removal** — detect border color, normalize to magenta fill, flood-fill from edges → transparent PNG + +## Requirements + +- **GPU:** NVIDIA with >= 8 GB VRAM (uses CPU offload) +- **Python:** 3.12+ +- **CUDA:** 12.0+ + +## Credits + +- Base model: [FLUX.2-klein-4B](https://huggingface.co/black-forest-labs/FLUX.2-klein-4B) by Black Forest Labs (Apache 2.0) +- LoRA: [pixel-art-lora](https://huggingface.co/Limbicnation/pixel-art-lora) by Limbicnation (Apache 2.0) +- MCP SDK: [modelcontextprotocol/python-sdk](https://github.com/modelcontextprotocol/python-sdk) + +## License + +MIT — see [LICENSE](LICENSE) + +Model licenses are separate (Apache 2.0). Check model cards for details. diff --git a/requirements.txt b/requirements.txt new file mode 100644 index 0000000..41a982f --- /dev/null +++ b/requirements.txt @@ -0,0 +1,6 @@ +mcp>=1.27,<2 +diffusers>=0.37.0 +transformers +accelerate +torch +pillow diff --git a/server.py b/server.py new file mode 100644 index 0000000..2a62289 --- /dev/null +++ b/server.py @@ -0,0 +1,331 @@ +#!/usr/bin/env python3 +""" +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 + +Model is loaded lazily on first call (~6s), then stays in VRAM for speed. +Background is removed post-generation to produce transparent PNG. +""" + +import os +import sys +import time +from typing import Optional + +import numpy as np +from PIL import Image +from mcp.server.fastmcp import FastMCP + +# Paths — models live in a shared location +BASE_DIR = os.path.dirname(os.path.abspath(__file__)) +MODEL_DIR = os.environ.get( + "IMAGEGEN_MODEL_DIR", + os.path.join(os.path.expanduser("~"), "models", "flux2-klein-4b"), +) +LORA_DIR = os.environ.get( + "IMAGEGEN_LORA_DIR", + os.path.join(os.path.expanduser("~"), "models", "pixel-art-lora"), +) +OUTPUT_DIR = os.environ.get("IMAGEGEN_OUTPUT_DIR", os.path.join(BASE_DIR, "output")) + +# rsLoRA requires much lower scale in diffusers — 1.0 produces black images +LORA_SCALE = 0.1 + +# Global state — model loaded lazily +_pipe = None +_device = None + + +def _get_device(): + global _device + if _device is None: + import torch + + if torch.cuda.is_available(): + _device = "cuda" + else: + _device = "cpu" + sys.stderr.write( + "[pixel-art] WARNING: CUDA not available, using CPU (very slow)\n" + ) + return _device + + +def _load_model(): + global _pipe + if _pipe is not None: + return _pipe + + sys.stderr.write("[pixel-art] Loading FLUX.2-klein-4B + LoRA (first call)...\n") + t0 = time.time() + + import torch + from diffusers import Flux2KleinPipeline + + _pipe = Flux2KleinPipeline.from_pretrained( + MODEL_DIR, + torch_dtype=torch.bfloat16, + ) + _pipe.load_lora_weights(LORA_DIR) + + if _get_device() == "cuda": + _pipe.enable_model_cpu_offload() + else: + _pipe.to(_get_device()) + + elapsed = time.time() - t0 + sys.stderr.write(f"[pixel-art] Model loaded in {elapsed:.1f}s\n") + return _pipe + + +def _build_prompt(user_prompt: str) -> str: + return f"pixel art sprite, {user_prompt}, game asset, transparent background" + + +def _generate( + pipe, prompt: str, seed: Optional[int], width: int, height: int, steps: int +): + import torch + + generator = None + if seed is not None: + generator = torch.Generator(device=_get_device()).manual_seed(seed) + + image = pipe( + prompt=prompt, + num_inference_steps=steps, + guidance_scale=1.0, + height=height, + width=width, + generator=generator, + attention_kwargs={"scale": LORA_SCALE}, + ).images[0] + + return image + + +def _remove_background(image: Image.Image, threshold: int = 30) -> Image.Image: + """Remove background using flood-fill from edges. + + Two-pass approach: + 1. Detect border color, replace all near-border pixels with a flat fill color + 2. Flood-fill from edges to remove the flat color cleanly + + This normalizes gradient/noisy backgrounds into one solid color, + making flood-fill removal much cleaner. + """ + from collections import deque + + rgb = image.convert("RGB") + arr = np.array(rgb).astype(int) + h, w = arr.shape[:2] + + # Sample border colors from all 4 edges + border_colors = [] + for x in range(w): + border_colors.append(arr[0, x]) + border_colors.append(arr[h - 1, x]) + for y in range(h): + border_colors.append(arr[y, 0]) + border_colors.append(arr[y, w - 1]) + + border_colors = np.array(border_colors) + bg_color = np.median(border_colors, axis=0).astype(int) + + # Pass 1: normalize background — replace all pixels within threshold + # of border color with a flat fill color (pure magenta, unlikely in sprites) + fill_color = np.array([255, 0, 255], dtype=int) + dist_to_bg = np.abs(arr - bg_color).sum(axis=2) + bg_mask = dist_to_bg < threshold * 3 + arr[bg_mask] = fill_color + + # Pass 2: flood-fill from edges to remove connected fill_color regions + alpha = np.full((h, w), 255, dtype=np.uint8) + visited = np.zeros((h, w), dtype=bool) + queue = deque() + + fill_dist_threshold = 30 # tolerance for near-fill pixels + + # Seed from all border pixels + for x in range(w): + for y in [0, h - 1]: + if not visited[y, x]: + queue.append((y, x)) + visited[y, x] = True + for y in range(h): + for x in [0, w - 1]: + if not visited[y, x]: + queue.append((y, x)) + visited[y, x] = True + + # BFS flood-fill + while queue: + y, x = queue.popleft() + dist = np.abs(arr[y, x] - fill_color).sum() + if dist > fill_dist_threshold: + continue + alpha[y, x] = 0 + + for dy, dx in [(-1, 0), (1, 0), (0, -1), (0, 1)]: + ny, nx = y + dy, x + dx + if 0 <= ny < h and 0 <= nx < w and not visited[ny, nx]: + visited[ny, nx] = True + queue.append((ny, nx)) + + # Clean up: any remaining near-magenta pixels that weren't flood-filled + # (small isolated background pockets) get removed too + remaining_bg = np.abs(arr - fill_color).sum(axis=2) < fill_dist_threshold + alpha[remaining_bg] = 0 + + rgba = np.dstack([arr.astype(np.uint8), alpha]) + return Image.fromarray(rgba, mode="RGBA") + + +def _pixelate(image: Image.Image, pixel_size: int = 8) -> Image.Image: + """Downscale then upscale with NEAREST to create chunky pixel-art effect. + + pixel_size=8 means each "pixel" in the result is an 8x8 block. + """ + w, h = image.size + small = image.resize((w // pixel_size, h // pixel_size), Image.LANCZOS) + return small.resize((w, h), Image.NEAREST) + + +def _ensure_dir(path: str): + dir_path = os.path.dirname(path) + if dir_path: + os.makedirs(dir_path, exist_ok=True) + + +# Create MCP server +mcp = FastMCP("pixel-art") + + +@mcp.tool() +def generate_sprite( + prompt: str, + output_path: str, + seed: Optional[int] = None, + width: int = 512, + height: int = 512, + steps: int = 4, + remove_bg: bool = True, + pixel_size: int = 4, +) -> dict: + """Generate a pixel-art sprite and save it as PNG with transparent background. + + Args: + prompt: Description of the sprite (e.g. "a crystal warrior with geometric armor") + output_path: Where to save the PNG file (relative to output dir or absolute) + seed: Optional seed for reproducibility + width: Image width in pixels (default 512) + 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) + + Returns: + Dict with output_path, seed_used, generation_time, prompt, size. + """ + pipe = _load_model() + full_prompt = _build_prompt(prompt) + + if not os.path.isabs(output_path): + output_path = os.path.join(OUTPUT_DIR, output_path) + + _ensure_dir(output_path) + + t0 = time.time() + image = _generate(pipe, full_prompt, seed, width, height, steps) + + if pixel_size > 0: + image = _pixelate(image, pixel_size) + + if remove_bg: + image = _remove_background(image) + + image.save(output_path) + elapsed = time.time() - t0 + + return { + "output_path": output_path, + "seed_used": seed, + "generation_time": f"{elapsed:.1f}s", + "prompt": full_prompt, + "size": f"{width}x{height}", + "transparent": remove_bg, + "pixel_size": pixel_size, + } + + +@mcp.tool() +def batch_generate( + specs: list[dict], +) -> list[dict]: + """Generate multiple pixel-art sprites in one call. + + Args: + specs: List of dicts, each with: + - prompt: str (required) — sprite description + - output_path: str (required) — PNG save path + - seed: int (optional) + - width: int (optional, default 512) + - height: int (optional, default 512) + - steps: int (optional, default 4) + - remove_bg: bool (optional, default True) + - pixel_size: int (optional, default 4, 0=off) + + Returns: + List of dicts with output_path, seed_used, generation_time, prompt, size, transparent. + """ + pipe = _load_model() + results = [] + + for spec in specs: + prompt = spec["prompt"] + output_path = spec["output_path"] + seed = spec.get("seed") + width = spec.get("width", 512) + height = spec.get("height", 512) + steps = spec.get("steps", 4) + remove_bg = spec.get("remove_bg", True) + pixel_size = spec.get("pixel_size", 4) + + full_prompt = _build_prompt(prompt) + + if not os.path.isabs(output_path): + output_path = os.path.join(OUTPUT_DIR, output_path) + + _ensure_dir(output_path) + + t0 = time.time() + image = _generate(pipe, full_prompt, seed, width, height, steps) + + if pixel_size > 0: + image = _pixelate(image, pixel_size) + + if remove_bg: + image = _remove_background(image) + + image.save(output_path) + elapsed = time.time() - t0 + + results.append( + { + "output_path": output_path, + "seed_used": seed, + "generation_time": f"{elapsed:.1f}s", + "prompt": full_prompt, + "size": f"{width}x{height}", + "transparent": remove_bg, + } + ) + + return results + + +if __name__ == "__main__": + mcp.run(transport="stdio")