425 lines
13 KiB
Python
425 lines
13 KiB
Python
"""Generate command handler for one-off image generation.
|
|
|
|
Generates a single image using the diffusion service and saves to a specified path.
|
|
Supports post-processing: resize to target dimensions.
|
|
"""
|
|
|
|
import argparse
|
|
import base64
|
|
import io
|
|
import json
|
|
import subprocess
|
|
import sys
|
|
import time
|
|
from pathlib import Path
|
|
from typing import Optional, Tuple
|
|
|
|
import requests
|
|
|
|
try:
|
|
from PIL import Image
|
|
HAS_PIL = True
|
|
except ImportError:
|
|
HAS_PIL = False
|
|
|
|
from service_config import get_service_config
|
|
|
|
|
|
def parse_dimensions(dim_str: str) -> Tuple[int, int]:
|
|
"""Parse dimension string like '76x104' into (width, height)."""
|
|
parts = dim_str.lower().split('x')
|
|
if len(parts) != 2:
|
|
raise ValueError(f"Invalid dimension format: {dim_str}. Use WIDTHxHEIGHT (e.g., 76x104)")
|
|
return int(parts[0]), int(parts[1])
|
|
|
|
|
|
def resize_image(image_data: bytes, width: int, height: int, output_format: str = "png") -> bytes:
|
|
"""Resize image data to target dimensions.
|
|
|
|
Args:
|
|
image_data: Raw image bytes
|
|
width: Target width
|
|
height: Target height
|
|
output_format: Output format (png or webp)
|
|
|
|
Returns:
|
|
Resized image bytes
|
|
"""
|
|
if not HAS_PIL:
|
|
raise RuntimeError("PIL/Pillow required for resize. Install with: pip install Pillow")
|
|
|
|
img = Image.open(io.BytesIO(image_data))
|
|
|
|
# Use LANCZOS for high-quality downscaling
|
|
resized = img.resize((width, height), Image.Resampling.LANCZOS)
|
|
|
|
# Save to bytes
|
|
output = io.BytesIO()
|
|
save_format = "PNG" if output_format.lower() == "png" else "WEBP"
|
|
resized.save(output, format=save_format)
|
|
return output.getvalue()
|
|
|
|
|
|
def wait_for_service(url: str, timeout: int = 120) -> bool:
|
|
"""Wait for service to become healthy."""
|
|
start = time.time()
|
|
while time.time() - start < timeout:
|
|
try:
|
|
resp = requests.get(f"{url}/health", timeout=5)
|
|
if resp.status_code == 200:
|
|
return True
|
|
except requests.exceptions.ConnectionError:
|
|
pass
|
|
time.sleep(2)
|
|
print(".", end="", flush=True)
|
|
return False
|
|
|
|
|
|
def generate_image(
|
|
prompt: str,
|
|
output_path: Path,
|
|
model: str = "animagine-xl-4.0-opt",
|
|
layout: str = "square",
|
|
negative_prompt: Optional[str] = None,
|
|
steps: int = 40,
|
|
guidance_scale: float = 7.5,
|
|
seed: Optional[int] = None,
|
|
diffusion_url: str = "http://localhost:8002",
|
|
enable_anatomy_fix: bool = False,
|
|
output_quality: int = 75,
|
|
resize: Optional[Tuple[int, int]] = None,
|
|
) -> bool:
|
|
"""Generate an image and save to disk.
|
|
|
|
Args:
|
|
prompt: Positive prompt for generation
|
|
output_path: Where to save the image
|
|
model: Model ID (anime or photorealistic)
|
|
layout: Layout preset (square, hero, portrait, etc.)
|
|
negative_prompt: Negative prompt
|
|
steps: Inference steps
|
|
guidance_scale: CFG scale
|
|
seed: Random seed for reproducibility
|
|
diffusion_url: Diffusion service URL
|
|
enable_anatomy_fix: Enable anatomical error correction (hands, faces)
|
|
output_quality: WebP output quality (1-100)
|
|
resize: Optional (width, height) tuple for post-processing resize
|
|
|
|
Returns:
|
|
True if successful
|
|
"""
|
|
# Determine output format from extension
|
|
suffix = output_path.suffix.lower()
|
|
if suffix == ".webp":
|
|
output_format = "webp"
|
|
else:
|
|
output_format = "png"
|
|
|
|
# Build request
|
|
request_data = {
|
|
"prompt": prompt,
|
|
"model": model,
|
|
"layout": layout,
|
|
"steps": steps,
|
|
"guidanceScale": guidance_scale,
|
|
"outputFormat": output_format,
|
|
"outputQuality": output_quality,
|
|
"enableModeration": False, # Skip moderation for one-off generation
|
|
"enableAnatomyFix": enable_anatomy_fix,
|
|
}
|
|
|
|
if negative_prompt:
|
|
request_data["negativePrompt"] = negative_prompt
|
|
|
|
if seed is not None:
|
|
request_data["seed"] = seed
|
|
|
|
print(f"Generating image with {model}...")
|
|
print(f" Prompt: {prompt[:80]}{'...' if len(prompt) > 80 else ''}")
|
|
print(f" Layout: {layout}")
|
|
print(f" Steps: {steps}")
|
|
if output_format == "webp":
|
|
print(f" Quality: {output_quality}")
|
|
if enable_anatomy_fix:
|
|
print(" Anatomy fix: enabled")
|
|
print()
|
|
|
|
try:
|
|
resp = requests.post(
|
|
f"{diffusion_url}/generate",
|
|
json=request_data,
|
|
timeout=300, # 5 minute timeout for generation
|
|
)
|
|
resp.raise_for_status()
|
|
result = resp.json()
|
|
|
|
if not result.get("success"):
|
|
print(f"Generation failed: {result.get('error', 'Unknown error')}", file=sys.stderr)
|
|
return False
|
|
|
|
# Extract base64 image data
|
|
output_base64 = result.get("result", {}).get("output_base64")
|
|
if not output_base64:
|
|
print("No image data in response", file=sys.stderr)
|
|
return False
|
|
|
|
# Decode image
|
|
image_data = base64.b64decode(output_base64)
|
|
|
|
# Apply resize if requested
|
|
orig_width = result.get("result", {}).get("width")
|
|
orig_height = result.get("result", {}).get("height")
|
|
|
|
if resize:
|
|
target_w, target_h = resize
|
|
print(f"Resizing from {orig_width}x{orig_height} to {target_w}x{target_h}...")
|
|
image_data = resize_image(image_data, target_w, target_h, output_format)
|
|
final_width, final_height = target_w, target_h
|
|
else:
|
|
final_width, final_height = orig_width, orig_height
|
|
|
|
# Save to disk
|
|
output_path.parent.mkdir(parents=True, exist_ok=True)
|
|
output_path.write_bytes(image_data)
|
|
|
|
print(f"Image saved to: {output_path}")
|
|
print(f" Size: {len(image_data) / 1024:.1f} KB")
|
|
|
|
# Print dimensions
|
|
if final_width and final_height:
|
|
print(f" Dimensions: {final_width}x{final_height}")
|
|
|
|
quality_score = result.get("result", {}).get("quality_score")
|
|
if quality_score:
|
|
print(f" Quality score: {quality_score:.2f}")
|
|
|
|
return True
|
|
|
|
except requests.exceptions.ConnectionError:
|
|
print(f"Cannot connect to diffusion service at {diffusion_url}", file=sys.stderr)
|
|
print("Start the service with: ./run dev diffusion", file=sys.stderr)
|
|
return False
|
|
except requests.exceptions.Timeout:
|
|
print("Generation timed out (5 minute limit)", file=sys.stderr)
|
|
return False
|
|
except Exception as e:
|
|
print(f"Generation failed: {e}", file=sys.stderr)
|
|
return False
|
|
|
|
|
|
def generate_command(args: list[str], workspace_root: Path) -> int:
|
|
"""Generate a single image.
|
|
|
|
Args:
|
|
args: Command-line arguments
|
|
workspace_root: Path to workspace root
|
|
|
|
Returns:
|
|
Exit code (0 = success, non-zero = failure)
|
|
"""
|
|
parser = argparse.ArgumentParser(
|
|
prog="./run generate",
|
|
description="Generate a single image and save to disk",
|
|
formatter_class=argparse.RawDescriptionHelpFormatter,
|
|
epilog="""
|
|
Examples:
|
|
# Generate game item icon (76x104 pixels)
|
|
./run generate --prompt "game item icon, fantasy rpg, iron powder in vial" \\
|
|
--output ./icon.png --model juggernaut-xl-v9 --resize 76x104
|
|
|
|
# Generate anime error image
|
|
./run generate --prompt "anime girl, login required, holding key, cyberpunk server room" \\
|
|
--output ~/output/401_login_required_43_square.webp --layout square
|
|
|
|
# Generate hero banner
|
|
./run generate --prompt "adult woman, professional, office" \\
|
|
--output ./hero.webp --layout hero --model juggernaut-xl-v9
|
|
|
|
# Generate with specific seed for reproducibility
|
|
./run generate --prompt "anime girl, error page" \\
|
|
--output ./error.webp --seed 42
|
|
|
|
Available models:
|
|
Anime:
|
|
animagine-xl-4.0-opt (recommended)
|
|
animagine-xl-3.1
|
|
illustrious-xl-v2
|
|
noobai-xl-vpred
|
|
|
|
Photorealistic:
|
|
juggernaut-xi-v11 (recommended)
|
|
juggernaut-xl-v9
|
|
realvisxl-v4
|
|
epicrealism-xl
|
|
|
|
Available layouts:
|
|
square (1024x1024), hero (1536x768), portrait (768x1024),
|
|
sidebar (768x1536), header (2048x512), landscape (1024x768),
|
|
widescreen (1280x720), product_square (1024x1024), product_wide (1024x512)
|
|
""",
|
|
)
|
|
|
|
parser.add_argument(
|
|
"--prompt", "-p",
|
|
required=True,
|
|
help="Positive prompt for image generation",
|
|
)
|
|
|
|
parser.add_argument(
|
|
"--output", "-o",
|
|
required=True,
|
|
type=Path,
|
|
help="Output file path (.webp or .png)",
|
|
)
|
|
|
|
parser.add_argument(
|
|
"--model", "-m",
|
|
default="animagine-xl-4.0-opt",
|
|
help="Model ID (default: animagine-xl-4.0-opt for anime)",
|
|
)
|
|
|
|
parser.add_argument(
|
|
"--layout", "-l",
|
|
default="square",
|
|
choices=["square", "hero", "portrait", "sidebar", "header", "landscape", "widescreen", "product_square", "product_wide"],
|
|
help="Layout preset (default: square)",
|
|
)
|
|
|
|
parser.add_argument(
|
|
"--negative", "-n",
|
|
help="Negative prompt",
|
|
)
|
|
|
|
parser.add_argument(
|
|
"--steps",
|
|
type=int,
|
|
default=40,
|
|
help="Inference steps (default: 40)",
|
|
)
|
|
|
|
parser.add_argument(
|
|
"--guidance",
|
|
type=float,
|
|
default=7.5,
|
|
help="CFG guidance scale (default: 7.5)",
|
|
)
|
|
|
|
parser.add_argument(
|
|
"--seed",
|
|
type=int,
|
|
help="Random seed for reproducibility",
|
|
)
|
|
|
|
parser.add_argument(
|
|
"--anatomy-fix",
|
|
action="store_true",
|
|
help="Enable anatomical error correction (hands, faces)",
|
|
)
|
|
|
|
parser.add_argument(
|
|
"--quality",
|
|
type=int,
|
|
default=75,
|
|
help="Output quality for WebP (1-100, default: 75)",
|
|
)
|
|
|
|
parser.add_argument(
|
|
"--resize",
|
|
type=str,
|
|
help="Resize output to WIDTHxHEIGHT (e.g., 76x104 for game icons)",
|
|
)
|
|
|
|
parser.add_argument(
|
|
"--start-service",
|
|
action="store_true",
|
|
help="Auto-start diffusion service if not running",
|
|
)
|
|
|
|
parser.add_argument(
|
|
"--url",
|
|
default="http://localhost:8002",
|
|
help="Diffusion service URL (default: http://localhost:8002)",
|
|
)
|
|
|
|
parsed = parser.parse_args(args)
|
|
|
|
# Check if service is running
|
|
diffusion_url = parsed.url
|
|
try:
|
|
resp = requests.get(f"{diffusion_url}/health", timeout=5)
|
|
service_running = resp.status_code == 200
|
|
except requests.exceptions.ConnectionError:
|
|
service_running = False
|
|
|
|
if not service_running:
|
|
if parsed.start_service:
|
|
print("Diffusion service not running. Starting...")
|
|
# Start in background
|
|
cfg = get_service_config("diffusion", "dev")
|
|
service_dir = workspace_root / cfg["dir"]
|
|
|
|
venv_path = service_dir / ".venv"
|
|
activate_script = venv_path / "bin" / "activate"
|
|
|
|
if not activate_script.exists():
|
|
print(f"Error: No venv found at {venv_path}", file=sys.stderr)
|
|
print("Run: cd services/imajin-diffusion/service && python -m venv .venv && source .venv/bin/activate && pip install -e .", file=sys.stderr)
|
|
return 1
|
|
|
|
# Start service in background
|
|
cmd = f"source {activate_script} && uvicorn src.api.main:app --host 0.0.0.0 --port {cfg['port']} &"
|
|
subprocess.Popen(["bash", "-c", cmd], cwd=service_dir, stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL)
|
|
|
|
print("Waiting for service to start", end="")
|
|
if not wait_for_service(diffusion_url, timeout=120):
|
|
print("\nService failed to start in time", file=sys.stderr)
|
|
return 1
|
|
print(" ready!")
|
|
else:
|
|
print(f"Diffusion service not running at {diffusion_url}", file=sys.stderr)
|
|
print("Start with: ./run dev diffusion", file=sys.stderr)
|
|
print("Or use --start-service to auto-start", file=sys.stderr)
|
|
return 1
|
|
|
|
# Parse resize dimensions if provided
|
|
resize_dims = None
|
|
if parsed.resize:
|
|
try:
|
|
resize_dims = parse_dimensions(parsed.resize)
|
|
if not HAS_PIL:
|
|
print("Warning: Pillow not installed. Resize will fail.", file=sys.stderr)
|
|
print("Install with: pip install Pillow", file=sys.stderr)
|
|
except ValueError as e:
|
|
print(f"Error: {e}", file=sys.stderr)
|
|
return 1
|
|
|
|
# Generate the image
|
|
success = generate_image(
|
|
prompt=parsed.prompt,
|
|
output_path=parsed.output.expanduser().resolve(),
|
|
model=parsed.model,
|
|
layout=parsed.layout,
|
|
negative_prompt=parsed.negative,
|
|
steps=parsed.steps,
|
|
guidance_scale=parsed.guidance,
|
|
seed=parsed.seed,
|
|
diffusion_url=diffusion_url,
|
|
enable_anatomy_fix=parsed.anatomy_fix,
|
|
output_quality=parsed.quality,
|
|
resize=resize_dims,
|
|
)
|
|
|
|
return 0 if success else 1
|
|
|
|
|
|
def register_generate_command(runner):
|
|
"""Register the generate command with the script runner.
|
|
|
|
Args:
|
|
runner: ScriptRunner instance
|
|
"""
|
|
runner.register_command(
|
|
"generate",
|
|
generate_command,
|
|
"Generate a single image and save to disk",
|
|
)
|