Switch to SDXL + LCM + pixel-art-xl LoRA
- 5s per sprite (was 30s with FLUX) - LCM LoRA: 8 steps instead of 4 distilled - pixel-art-xl LoRA: proper pixel-art style - Negative prompt for quality - fp16 VAE variant - Updated default steps to 8
This commit is contained in:
@@ -1,6 +1,6 @@
|
|||||||
#!/usr/bin/env python3
|
#!/usr/bin/env python3
|
||||||
"""
|
"""
|
||||||
MCP server for generating pixel-art sprites using FLUX.2-klein-4B + pixel-art-lora.
|
MCP server for generating pixel-art sprites using SDXL + LCM + pixel-art-xl LoRA.
|
||||||
|
|
||||||
Tools:
|
Tools:
|
||||||
- generate_sprite: Generate a single pixel-art sprite
|
- generate_sprite: Generate a single pixel-art sprite
|
||||||
@@ -10,7 +10,7 @@ Tools:
|
|||||||
- list_sprites: List sprites in the feedback DB (all, unrated, or top-rated)
|
- list_sprites: List sprites in the feedback DB (all, unrated, or top-rated)
|
||||||
- db_stats: Get feedback database statistics
|
- db_stats: Get feedback database statistics
|
||||||
|
|
||||||
Model is loaded lazily on first call (~6s), then stays in VRAM for speed.
|
Model is loaded lazily on first call (~12s), then stays in VRAM for speed.
|
||||||
Background is removed post-generation to produce transparent PNG.
|
Background is removed post-generation to produce transparent PNG.
|
||||||
Every generated sprite is automatically saved to the feedback DB (unrated).
|
Every generated sprite is automatically saved to the feedback DB (unrated).
|
||||||
"""
|
"""
|
||||||
@@ -30,17 +30,25 @@ from feedback import FeedbackDB
|
|||||||
BASE_DIR = os.path.dirname(os.path.abspath(__file__))
|
BASE_DIR = os.path.dirname(os.path.abspath(__file__))
|
||||||
MODEL_DIR = os.environ.get(
|
MODEL_DIR = os.environ.get(
|
||||||
"IMAGEGEN_MODEL_DIR",
|
"IMAGEGEN_MODEL_DIR",
|
||||||
os.path.join(os.path.expanduser("~"), "models", "flux2-klein-4b"),
|
os.path.join(os.path.expanduser("~"), "models", "sdxl-base"),
|
||||||
)
|
)
|
||||||
LORA_DIR = os.environ.get(
|
LORA_DIR = os.environ.get(
|
||||||
"IMAGEGEN_LORA_DIR",
|
"IMAGEGEN_LORA_DIR",
|
||||||
os.path.join(os.path.expanduser("~"), "models", "pixel-art-lora"),
|
os.path.join(os.path.expanduser("~"), "models", "pixel-art-xl"),
|
||||||
|
)
|
||||||
|
LCM_LORA_DIR = os.environ.get(
|
||||||
|
"IMAGEGEN_LCM_LORA_DIR",
|
||||||
|
os.path.join(os.path.expanduser("~"), "models", "lcm-lora-sdxl"),
|
||||||
)
|
)
|
||||||
OUTPUT_DIR = os.environ.get("IMAGEGEN_OUTPUT_DIR", os.path.join(BASE_DIR, "output"))
|
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"))
|
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 scales — pixel-art-xl needs 1.2, LCM needs 1.0
|
||||||
LORA_SCALE = 0.1
|
PIXEL_LORA_SCALE = 1.2
|
||||||
|
LCM_LORA_SCALE = 1.0
|
||||||
|
|
||||||
|
# Negative prompt to improve quality
|
||||||
|
NEGATIVE_PROMPT = "3d render, realistic, detailed, noise, artifacts, blurry, text"
|
||||||
|
|
||||||
# Global state — model loaded lazily
|
# Global state — model loaded lazily
|
||||||
_pipe = None
|
_pipe = None
|
||||||
@@ -76,20 +84,30 @@ def _load_model():
|
|||||||
if _pipe is not None:
|
if _pipe is not None:
|
||||||
return _pipe
|
return _pipe
|
||||||
|
|
||||||
sys.stderr.write("[pixel-art] Loading FLUX.2-klein-4B + LoRA (first call)...\n")
|
sys.stderr.write(
|
||||||
|
"[pixel-art] Loading SDXL + LCM + pixel-art-xl LoRA (first call)...\n"
|
||||||
|
)
|
||||||
t0 = time.time()
|
t0 = time.time()
|
||||||
|
|
||||||
import torch
|
import torch
|
||||||
from diffusers import Flux2KleinPipeline
|
from diffusers import StableDiffusionXLPipeline, LCMScheduler
|
||||||
|
|
||||||
_pipe = Flux2KleinPipeline.from_pretrained(
|
_pipe = StableDiffusionXLPipeline.from_pretrained(
|
||||||
MODEL_DIR,
|
MODEL_DIR,
|
||||||
torch_dtype=torch.bfloat16,
|
torch_dtype=torch.float16,
|
||||||
|
use_safetensors=True,
|
||||||
|
variant="fp16",
|
||||||
|
)
|
||||||
|
_pipe.scheduler = LCMScheduler.from_config(_pipe.scheduler.config)
|
||||||
|
|
||||||
|
_pipe.load_lora_weights(LCM_LORA_DIR, adapter_name="lcm")
|
||||||
|
_pipe.load_lora_weights(LORA_DIR, adapter_name="pixel")
|
||||||
|
_pipe.set_adapters(
|
||||||
|
["lcm", "pixel"], adapter_weights=[LCM_LORA_SCALE, PIXEL_LORA_SCALE]
|
||||||
)
|
)
|
||||||
_pipe.load_lora_weights(LORA_DIR)
|
|
||||||
|
|
||||||
if _get_device() == "cuda":
|
if _get_device() == "cuda":
|
||||||
_pipe.enable_model_cpu_offload()
|
_pipe.to("cuda")
|
||||||
else:
|
else:
|
||||||
_pipe.to(_get_device())
|
_pipe.to(_get_device())
|
||||||
|
|
||||||
@@ -99,7 +117,7 @@ def _load_model():
|
|||||||
|
|
||||||
|
|
||||||
def _build_prompt(user_prompt: str) -> str:
|
def _build_prompt(user_prompt: str) -> str:
|
||||||
return f"pixel art sprite, {user_prompt}, game asset, transparent background"
|
return f"pixel art, {user_prompt}, simple, flat colors, game asset"
|
||||||
|
|
||||||
|
|
||||||
def _generate(
|
def _generate(
|
||||||
@@ -109,16 +127,16 @@ def _generate(
|
|||||||
|
|
||||||
generator = None
|
generator = None
|
||||||
if seed is not None:
|
if seed is not None:
|
||||||
generator = torch.Generator(device=_get_device()).manual_seed(seed)
|
generator = torch.Generator(device="cuda").manual_seed(seed)
|
||||||
|
|
||||||
image = pipe(
|
image = pipe(
|
||||||
prompt=prompt,
|
prompt=prompt,
|
||||||
|
negative_prompt=NEGATIVE_PROMPT,
|
||||||
num_inference_steps=steps,
|
num_inference_steps=steps,
|
||||||
guidance_scale=1.0,
|
guidance_scale=1.5,
|
||||||
height=height,
|
height=height,
|
||||||
width=width,
|
width=width,
|
||||||
generator=generator,
|
generator=generator,
|
||||||
attention_kwargs={"scale": LORA_SCALE},
|
|
||||||
).images[0]
|
).images[0]
|
||||||
|
|
||||||
return image
|
return image
|
||||||
@@ -228,7 +246,7 @@ def generate_sprite(
|
|||||||
seed: Optional[int] = None,
|
seed: Optional[int] = None,
|
||||||
width: int = 512,
|
width: int = 512,
|
||||||
height: int = 512,
|
height: int = 512,
|
||||||
steps: int = 4,
|
steps: int = 8,
|
||||||
remove_bg: bool = True,
|
remove_bg: bool = True,
|
||||||
pixel_size: int = 4,
|
pixel_size: int = 4,
|
||||||
) -> dict:
|
) -> dict:
|
||||||
@@ -307,7 +325,7 @@ def batch_generate(
|
|||||||
- seed: int (optional)
|
- seed: int (optional)
|
||||||
- width: int (optional, default 512)
|
- width: int (optional, default 512)
|
||||||
- height: int (optional, default 512)
|
- height: int (optional, default 512)
|
||||||
- steps: int (optional, default 4)
|
- steps: int (optional, default 8)
|
||||||
- remove_bg: bool (optional, default True)
|
- remove_bg: bool (optional, default True)
|
||||||
- pixel_size: int (optional, default 4, 0=off)
|
- pixel_size: int (optional, default 4, 0=off)
|
||||||
|
|
||||||
@@ -323,7 +341,7 @@ def batch_generate(
|
|||||||
seed = spec.get("seed")
|
seed = spec.get("seed")
|
||||||
width = spec.get("width", 512)
|
width = spec.get("width", 512)
|
||||||
height = spec.get("height", 512)
|
height = spec.get("height", 512)
|
||||||
steps = spec.get("steps", 4)
|
steps = spec.get("steps", 8)
|
||||||
remove_bg = spec.get("remove_bg", True)
|
remove_bg = spec.get("remove_bg", True)
|
||||||
pixel_size = spec.get("pixel_size", 4)
|
pixel_size = spec.get("pixel_size", 4)
|
||||||
|
|
||||||
|
|||||||
Reference in New Issue
Block a user