imajin/scripts/run/generate_command.py

"""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,
    enable_background_removal: 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)
        enable_background_removal: Remove background for transparent PNG output
        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,
        "enableBackgroundRemoval": enable_background_removal,
    }

    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")
    if enable_background_removal:
        print("  Background removal: enabled (transparent PNG)")
    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(
        "--transparent", "-t",
        action="store_true",
        help="Remove background for transparent PNG output (icons, stickers, product images)",
    )

    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,
        enable_background_removal=parsed.transparent,
        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",
    )