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.
"""

import argparse
import base64
import json
import subprocess
import sys
import time
from pathlib import Path
from typing import Optional

import requests

from service_config import get_service_config


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,
) -> 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)

    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 and save
        image_data = base64.b64decode(output_base64)
        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 additional info
        width = result.get("result", {}).get("width")
        height = result.get("result", {}).get("height")
        if width and height:
            print(f"  Dimensions: {width}x{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 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(
        "--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

    # 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,
    )

    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",
    )