Initial commit: MCP server for pixel-art sprite generation
This commit is contained in:
Emil
@@ -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")
|
||||
Reference in New Issue
Block a user