roux/scripts/upscale-images.py

#!/usr/bin/env python3
"""
Upscale issue plate images with Upscayl's CLI and the Remacri model.

  python3 scripts/upscale-images.py --issue 02

Copies every content/posts/<issue>/<slug>/<slug>.png into a flat staging dir,
runs `upscayl-bin` once over the whole dir (the model loads a single time), then
writes the results back in place. Optionally clamps the long edge and re-strips
metadata with ImageMagick so the committed PNGs stay a sane size.

Idempotent: a slug already recorded in data/prompts/issue-<issue>.upscale.json
for the same settings is skipped unless --force.

Needs only stdlib. `upscayl-bin` and (for --max-edge) `magick` must be on PATH or
at their default locations.
"""
import argparse
import json
import shutil
import subprocess
import sys
import tempfile
import time
from pathlib import Path

SITE = Path(__file__).resolve().parent.parent

UPSCAYL_CANDIDATES = [
    Path("/usr/share/upscayl/bin/upscayl-bin"),
    Path("/opt/Upscayl/resources/bin/upscayl-bin"),
]
MODELS_CANDIDATES = [
    Path("/usr/share/upscayl/models"),
    Path("/opt/Upscayl/resources/models"),
    Path.home() / ".config" / "upscayl" / "models",
]


def which(name: str) -> str | None:
    return shutil.which(name)


def find_upscayl() -> Path:
    for p in UPSCAYL_CANDIDATES:
        if p.is_file():
            return p
    found = which("upscayl-bin")
    if found:
        return Path(found)
    sys.exit("upscayl-bin not found (looked in /usr/share/upscayl/bin and PATH)")


def find_models(explicit: Path | None) -> Path:
    if explicit:
        if explicit.is_dir():
            return explicit
        sys.exit(f"--models-dir not a directory: {explicit}")
    for p in MODELS_CANDIDATES:
        if p.is_dir():
            return p
    sys.exit("upscayl models dir not found")


def load_manifest(path: Path) -> dict:
    if path.is_file():
        try:
            return {r["slug"]: r for r in json.loads(path.read_text(encoding="utf-8"))}
        except (json.JSONDecodeError, KeyError, TypeError):
            pass
    return {}


def main() -> int:
    ap = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
    ap.add_argument("--issue", required=True)
    ap.add_argument("--out", type=Path, help="output root; default is in-place")
    ap.add_argument("--model", default="remacri-4x", help="upscayl model name")
    ap.add_argument("--models-dir", type=Path, help="override models folder")
    ap.add_argument("--scale", type=int, default=4, choices=[2, 3, 4],
                    help="output scale (default 4)")
    ap.add_argument("--max-edge", type=int, default=4992,
                    help="clamp the long edge with ImageMagick after upscaling "
                         "(0 = keep native scale; default 4992 ≈ issue-01 sources)")
    ap.add_argument("--quality", type=int, default=95,
                    help="ImageMagick PNG quality when --max-edge applies")
    ap.add_argument("--compress", type=int, default=0,
                    help="upscayl -c compression (0-100)")
    ap.add_argument("--gpu-id", help="upscayl -g gpu id (default auto)")
    ap.add_argument("--only", help="comma-separated slugs")
    ap.add_argument("--limit", type=int)
    ap.add_argument("--force", action="store_true")
    ap.add_argument("--dry-run", action="store_true")
    args = ap.parse_args()

    upscayl = find_upscayl()
    models_dir = find_models(args.models_dir)
    if not (models_dir / f"{args.model}.bin").is_file():
        sys.exit(f"model {args.model!r} not in {models_dir}")
    magick = which("magick") or which("convert")
    if args.max_edge and not magick:
        sys.exit("--max-edge needs ImageMagick (magick/convert) on PATH; pass --max-edge 0 to skip")

    content = SITE / "content" / "posts" / args.issue
    if not content.is_dir():
        sys.exit(f"no such issue dir: {content}")

    bundles = sorted(p for p in content.iterdir()
                     if p.is_dir() and (p / f"{p.name}.png").is_file())
    if args.only:
        want = {s.strip() for s in args.only.split(",")}
        bundles = [b for b in bundles if b.name in want]
    if args.limit:
        bundles = bundles[: args.limit]
    if not bundles:
        sys.exit("no matching plate images")

    settings = f"{args.model}|s{args.scale}|e{args.max_edge}|c{args.compress}"
    manifest_path = SITE / "data" / "prompts" / f"issue-{args.issue}.upscale.json"
    done = load_manifest(manifest_path)

    todo = [b for b in bundles
            if args.force or done.get(b.name, {}).get("settings") != settings]
    skipped = len(bundles) - len(todo)
    print(f"{len(bundles)} plates · {len(todo)} to upscale · {skipped} already done ({settings})")
    if not todo:
        return 0

    with tempfile.TemporaryDirectory(prefix="roux-upscale-") as td:
        tin, tout = Path(td) / "in", Path(td) / "out"
        tin.mkdir()
        tout.mkdir()
        for b in todo:
            shutil.copy2(b / f"{b.name}.png", tin / f"{b.name}.png")

        cmd = [str(upscayl), "-i", str(tin), "-o", str(tout),
               "-m", str(models_dir), "-n", args.model,
               "-s", str(args.scale), "-f", "png"]
        if args.compress:
            cmd += ["-c", str(args.compress)]
        if args.gpu_id:
            cmd += ["-g", args.gpu_id]

        if args.dry_run:
            print(" ".join(cmd))
            if args.max_edge:
                print(f"# then: {magick} <img> -resize {args.max_edge}x{args.max_edge}\\> "
                      f"-strip -quality {args.quality} <img>")
            return 0

        print(f"→ upscayl {args.model} x{args.scale} ({len(todo)} images)")
        rc = subprocess.run(cmd).returncode
        if rc != 0:
            sys.exit(f"upscayl-bin exited {rc}")

        out_root = args.out.expanduser() if args.out else content
        ok, failed = [], []
        for b in todo:
            produced = tout / f"{b.name}.png"
            if not produced.is_file():
                failed.append(b.name)
                continue
            if args.max_edge:
                subprocess.run([magick, str(produced),
                                "-resize", f"{args.max_edge}x{args.max_edge}>",
                                "-strip", "-quality", str(args.quality),
                                str(produced)], check=True)
            dest = out_root / b.name / f"{b.name}.png"
            dest.parent.mkdir(parents=True, exist_ok=True)
            shutil.move(str(produced), str(dest))
            ok.append(b.name)
            done[b.name] = {
                "slug": b.name,
                "settings": settings,
                "model": args.model,
                "scale": args.scale,
                "max_edge": args.max_edge,
                "done_at": time.strftime("%Y-%m-%dT%H:%M:%SZ", time.gmtime()),
            }

    # Only the in-place pipeline is tracked; --out is for experimentation.
    if not args.out:
        manifest_path.parent.mkdir(parents=True, exist_ok=True)
        manifest_path.write_text(
            json.dumps([done[k] for k in sorted(done)], indent=2, ensure_ascii=False) + "\n",
            encoding="utf-8")

    print(f"\n{len(ok)} upscaled / {len(failed)} failed / {skipped} skipped")
    if failed:
        print("failed:", ", ".join(sorted(failed)))
        return 1
    return 0


if __name__ == "__main__":
    raise SystemExit(main())