imajin/scripts/run/repaint_command.py

"""Repaint command — background replacement via SDXL inpainting.

Usage:
    ./run repaint --source photo.jpg --prompt "luxury hotel suite, city view"
    ./run repaint --source photo.jpg --prompt "hotel suite" --count 4 --out ./results/
    ./run repaint --source photo.jpg --prompt "..." --seed 42 --count 3
"""

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

import requests


def _encode_image(path: Path) -> str:
    """Read image file and return raw base64 string."""
    return base64.b64encode(path.read_bytes()).decode()


def _submit_job(
    url: str,
    source_b64: str,
    background_prompt: str,
    negative_prompt: Optional[str],
    steps: int,
    guidance_scale: float,
    seed: int,
    rating: str,
) -> str:
    payload = {
        "sourceImage": source_b64,
        "backgroundPrompt": background_prompt,
        "negativePrompt": negative_prompt,
        "steps": steps,
        "guidanceScale": guidance_scale,
        "seed": seed,
        "maturityRating": rating,
    }
    resp = requests.post(f"{url}/generate/repaint-background/async", json=payload, timeout=30)
    resp.raise_for_status()
    data = resp.json()
    if not data.get("success") or not data.get("jobId"):
        raise RuntimeError(f"Submit failed: {data}")
    return data["jobId"]


def _poll_jobs(url: str, job_ids: list[str], interval: float = 3.0) -> dict[str, dict]:
    """Poll until all jobs are terminal. Returns {job_id: result_data}."""
    pending = set(job_ids)
    results: dict[str, dict] = {}

    while pending:
        time.sleep(interval)
        for job_id in list(pending):
            resp = requests.get(f"{url}/jobs/{job_id}", timeout=10)
            resp.raise_for_status()
            data = resp.json()
            status = data.get("status")
            if status == "completed":
                result_resp = requests.get(f"{url}/jobs/{job_id}/result", timeout=30)
                result_resp.raise_for_status()
                result_data = result_resp.json()
                results[job_id] = result_data
                pending.discard(job_id)
                print(f"  ✓ {job_id[:8]}  done")
            elif status == "failed":
                results[job_id] = {"error": data.get("error", "failed")}
                pending.discard(job_id)
                print(f"  ✗ {job_id[:8]}  failed: {data.get('error', '?')}", file=sys.stderr)

    return results


def repaint_command(args: list[str], workspace_root: Path) -> int:
    parser = argparse.ArgumentParser(
        prog="./run repaint",
        description="Replace image background using SDXL inpainting (BiRefNet segmentation)",
        formatter_class=argparse.RawDescriptionHelpFormatter,
        epilog="""
Examples:
  # Single repaint
  ./run repaint --source photo.jpg --prompt "luxury hotel suite, city skyline"

  # Batch of 4 variants from the same source
  ./run repaint --source photo.jpg --prompt "hotel suite" --count 4 --out ./results/

  # Deterministic batch (seed+0, seed+1, ...)
  ./run repaint --source photo.jpg --prompt "hotel suite" --seed 42 --count 3

  # NSFW rating (removes SFW clothing-lock from prompt)
  ./run repaint --source photo.jpg --prompt "hotel suite" --rating nsfw
        """,
    )
    parser.add_argument("--source", "-s", required=True, type=Path, help="Source photo path")
    parser.add_argument("--prompt", "-p", required=True, help="Background scene description")
    parser.add_argument("--negative", "-n", default=None, help="Negative prompt")
    parser.add_argument("--count", "-c", type=int, default=1, help="Number of variants (default: 1)")
    parser.add_argument("--seed", type=int, default=None, help="Starting seed (increments per variant)")
    parser.add_argument("--steps", type=int, default=35, help="Inference steps (default: 35)")
    parser.add_argument("--guidance", type=float, default=7.5, help="CFG guidance scale (default: 7.5)")
    parser.add_argument("--rating", choices=["sfw", "nsfw", "explicit"], default="nsfw", help="Content rating (default: nsfw)")
    parser.add_argument("--out", "-o", type=Path, default=None, help="Output directory (or file if count=1)")
    parser.add_argument("--url", default="http://localhost:8002", help="Diffusion service URL")

    parsed = parser.parse_args(args)

    source_path = parsed.source.expanduser().resolve()
    if not source_path.exists():
        print(f"Source photo not found: {source_path}", file=sys.stderr)
        return 1

    # Check service health
    try:
        requests.get(f"{parsed.url}/health", timeout=5).raise_for_status()
    except Exception:
        print(f"Diffusion service not reachable at {parsed.url}", file=sys.stderr)
        print("Start with: ./run dev diffusion", file=sys.stderr)
        return 1

    # Determine output paths
    out_path = (parsed.out or Path(".")).expanduser().resolve()
    if parsed.count == 1 and out_path.suffix in (".png", ".jpg", ".webp"):
        out_dir = out_path.parent
        out_template = None
        single_out = out_path
    else:
        out_dir = out_path if out_path.suffix == "" else out_path.parent
        out_template = source_path.stem + "_repaint_{n}.png"
        single_out = None
    out_dir.mkdir(parents=True, exist_ok=True)

    # Generate seeds
    import random
    base_seed = parsed.seed if parsed.seed is not None else random.randint(0, 2**31 - 1)
    seeds = [base_seed + i for i in range(parsed.count)]

    print(f"Repainting {source_path.name} × {parsed.count}")
    print(f"  Prompt: {parsed.prompt[:80]}{'...' if len(parsed.prompt) > 80 else ''}")
    print(f"  Seeds:  {seeds[:5]}{'...' if len(seeds) > 5 else ''}")
    print()

    source_b64 = _encode_image(source_path)

    # Submit all jobs
    job_ids: list[str] = []
    for seed in seeds:
        try:
            job_id = _submit_job(
                parsed.url, source_b64, parsed.prompt, parsed.negative,
                parsed.steps, parsed.guidance, seed, parsed.rating,
            )
            job_ids.append(job_id)
            print(f"  → {job_id[:8]}  seed={seed}")
        except Exception as e:
            print(f"  Submit failed (seed={seed}): {e}", file=sys.stderr)

    if not job_ids:
        print("All submissions failed.", file=sys.stderr)
        return 1

    print(f"\nPolling {len(job_ids)} job(s)...")
    results = _poll_jobs(parsed.url, job_ids)

    # Save results
    saved = 0
    for idx, (job_id, result) in enumerate(results.items()):
        if "error" in result and "output_base64" not in str(result):
            continue

        # Navigate into nested result structure
        r = result.get("result", result)
        b64 = r.get("output_base64", "")
        if not b64:
            continue

        if single_out and idx == 0:
            out_file = single_out
        else:
            out_file = out_dir / (out_template or f"{source_path.stem}_repaint_{idx+1:02d}.png").format(n=idx + 1)

        out_file.write_bytes(base64.b64decode(b64))
        w, h = r.get("width", "?"), r.get("height", "?")
        seed_used = r.get("seed", seeds[idx] if idx < len(seeds) else "?")
        print(f"  Saved {out_file.name}  ({w}×{h}, seed={seed_used})")
        saved += 1

    print(f"\n{saved}/{len(job_ids)} images saved to {out_dir}")
    return 0 if saved > 0 else 1


def register_repaint_command(runner) -> None:
    runner.register_command("repaint", repaint_command, "Replace image background via SDXL inpainting")