scripts/build_replicate_faceinpaint.py

"""Replicate masked face-inpaint beautify (gated) — keep hair/bg, edit face.

Uses a Replicate masked face-inpaint model (default
`lucataco/ip_adapter-face-inpaint`) which auto-masks the face, regenerates only
the face with IP-Adapter identity, and keeps the hair / clothing / background.
Runs on the EXISTING Replicate token (no new account).

It reads the model's input SCHEMA first (free GET) and auto-matches field names
(image / prompt / strength), so it adapts to whatever inpaint model you point it
at instead of guessing and hitting a 422. Then GFPGAN sharpen + watermark.
Needs REPLICATE_API_TOKEN (this shell only). Output -> runtime/ (gitignored).
Pilot Ready: NOT CONFIRMED.
"""
from __future__ import annotations

import argparse
import base64
import json
import sys
from pathlib import Path

sys.path.insert(0, str(Path(__file__).resolve().parents[1]))

from io import BytesIO  # noqa: E402

from PIL import Image  # noqa: E402

from app.services.face_pipeline import face_enhance  # noqa: E402
from app.services.gemini_client import (  # noqa: E402
    normalize_image_orientation_bytes,
    upscale_to_original_and_enhance,
)
from app.services.replicate_client import (  # noqa: E402
    ReplicateUnavailable,
    fetch_model_schema,
    replicate_real_enabled,
    run_replicate_model,
)
from app.services.watermark import apply_ai_watermark  # noqa: E402

DEFAULT_MODEL = "lucataco/ip_adapter-face-inpaint"
DEFAULT_PROMPT = (
    "the same Korean person, beautified to a polished K-idol / actor aesthetic: "
    "clear smooth skin, brighter clearer eyes, refined nose, slimmer V-line jaw, "
    "balanced attractive features; natural photorealistic; keep the same identity"
)
_STRENGTH_FIELDS = {"strength", "denoise", "denoising_strength", "prompt_strength", "image_strength"}


def _data_uri(b: bytes) -> str:
    return "data:image/png;base64," + base64.b64encode(b).decode("ascii")


def _downscale(image_bytes: bytes, max_side: int) -> bytes:
    """Shrink so the longest side <= max_side (SD1.5 OOMs on full-res selfies)."""
    img = Image.open(BytesIO(image_bytes)).convert("RGB")
    w, h = img.size
    longest = max(w, h)
    if longest > max_side:
        s = max_side / float(longest)
        img = img.resize((max(1, int(w * s)), max(1, int(h * s))), Image.LANCZOS)
    buf = BytesIO()
    img.save(buf, format="PNG")
    return buf.getvalue()


def _is_image_field(name: str, spec: dict) -> bool:
    if "mask" in name.lower():
        return False
    if spec.get("format") == "uri":
        return True
    n = name.lower()
    return n == "image" or n.endswith("_image") or n in {"source_image", "input_image", "face_image", "subject"}


def _build_input(schema: dict, source_uri: str, prompt: str, strength: float | None) -> dict:
    """Auto-fill only fields that exist in the schema (avoids 422 on unknowns)."""
    inp: dict = {}
    for name, spec in schema.items():
        ln = name.lower()
        if _is_image_field(name, spec):
            inp[name] = source_uri                 # identity + source are the same person
        elif ln == "prompt":
            inp[name] = prompt
        elif ln == "negative_prompt":
            inp[name] = "deformed, distorted, disfigured, plastic, cartoon, different person"
        elif ln in _STRENGTH_FIELDS and strength is not None:
            inp[name] = float(strength)
        elif ln in {"num_outputs", "num_images"}:
            inp[name] = 1
    return inp


def main(argv: list[str] | None = None) -> int:
    ap = argparse.ArgumentParser(description="Replicate masked face-inpaint beautify")
    ap.add_argument("--source", required=True)
    ap.add_argument("--prompt", default=DEFAULT_PROMPT)
    ap.add_argument("--strength", type=float, default=0.6)
    ap.add_argument("--max-size", type=int, default=768,
                    help="downscale longest side before sending (SD1.5 OOMs on full-res)")
    ap.add_argument("--model", default=DEFAULT_MODEL)
    ap.add_argument("--evidence-root", default="runtime/gemini-smoke-evidence")
    ap.add_argument("--tag", default="rfi-01")
    args = ap.parse_args(argv)

    src = Path(args.source)
    if not src.exists():
        print(f"REFUSED: source not found: {src}")
        return 2
    if not replicate_real_enabled():
        print("REFUSED: REPLICATE_API_TOKEN not set (no network call made).")
        return 2

    out_dir = Path(args.evidence_root) / "gemini-smoke" / "replicate-faceinpaint" / args.tag
    out_dir.mkdir(parents=True, exist_ok=True)

    source_bytes = normalize_image_orientation_bytes(src.read_bytes())
    model_in_bytes = _downscale(source_bytes, args.max_size)   # SD1.5-friendly size
    source_uri = _data_uri(model_in_bytes)

    def _fail(msg: str) -> int:
        (out_dir / "error.txt").write_text(msg, encoding="utf-8")
        print(f"RUN_FAILED: {msg}")
        return 1

    try:
        schema = fetch_model_schema(args.model)
    except ReplicateUnavailable as exc:
        return _fail(f"schema fetch: {exc}")
    if not schema:
        return _fail(f"could not read input schema for {args.model}")

    model_input = _build_input(schema, source_uri, args.prompt, args.strength)
    img_fields = [k for k in model_input if model_input[k] is source_uri or model_input[k] == source_uri]
    print(f"replicate-faceinpaint: model={args.model} fields={list(model_input)} ...")

    try:
        raw = run_replicate_model(args.model, model_input)
    except ReplicateUnavailable as exc:
        return _fail(f"{exc} | schema_fields={sorted(schema)}")
    except Exception as exc:  # noqa: BLE001 - capture anything else for diagnosis
        import traceback
        return _fail(f"UNEXPECTED {type(exc).__name__}: {exc}\n{traceback.format_exc()}")
    (out_dir / "faceinpaint-raw.png").write_bytes(raw)

    # Upscale the (downscaled) inpaint result back to the original size + GFPGAN.
    final = upscale_to_original_and_enhance(source_bytes, raw)
    final = apply_ai_watermark(final)
    (out_dir / "faceinpaint-output.png").write_bytes(final)

    metrics = {
        "hosted_model": args.model,
        "strength": args.strength,
        "image_fields_used": img_fields,
        "face_enhanced": face_enhance.face_enhancer_available(),
        "watermark_applied": True,
        "tag": args.tag,
        "human_qa_required": "YES",
        "pilot_ready": "NOT CONFIRMED",
    }
    (out_dir / "faceinpaint-summary.json").write_text(
        json.dumps(metrics, indent=2, ensure_ascii=False), encoding="utf-8"
    )
    print(f"replicate-faceinpaint: DONE tag={args.tag}")
    print(f"  final : {out_dir / 'faceinpaint-output.png'}")
    print("Pilot Ready: NOT CONFIRMED.")
    return 0


if __name__ == "__main__":
    sys.exit(main())