Emil_Shanaty/Imagen
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Initial commit: MCP server for pixel-art sprite generation

Browse Source
This commit is contained in:
Emil
2026-06-28 14:50:58 +03:00
commit e28b466c8a
6 changed files with 513 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files
.gitignore
+33
View File
@@ -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*
.opencode/opencode.json
+10
View File
@@ -0,0 +1,10 @@
{
"$schema": "https://opencode.ai/config.json",
"mcp": {
"pixel-art": {
"type": "local",
"command": ["./venv/bin/python", "server.py"],
"enabled": true
}
}
}
LICENSE
+21
View File
@@ -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.
README.md
+112
View File
@@ -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.
requirements.txt
+6
View File
@@ -0,0 +1,6 @@
mcp>=1.27,<2
diffusers>=0.37.0
transformers
accelerate
torch
pillow
server.py
+331
View File
@@ -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")