#!/usr/bin/env python3
"""Run the fine-tuned model locally. Weights come from the Hub, not from disk.

    pip install -r requirements-local.txt
    python3 -m src.infer "a steampunk owl inside a glass jar, intricate detail"

    python3 src/infer.py "a cyberpunk portrait" --lora-scale 0 --out base.png
    python3 src/infer.py "a cyberpunk portrait" --lora-scale 1 --out tuned.png
    # same seed, both scales -> a direct before/after against vanilla SD 1.5

Everything is pulled from the Hub on first run and cached under ~/.cache/
huggingface: SD 1.5 (~4 GB) plus the LoRA adapter (~25 MB).

On an M-series Mac this uses the MPS backend. Expect roughly a minute per
image at 25 steps, and keep the batch at 1 - 8 GB of unified memory does not
leave much headroom.
"""

import argparse
import sys
from pathlib import Path

# Keep `python3 src/infer.py ...` working alongside `python3 -m src.infer ...`:
# both need the repo root importable so `src.pipeline` resolves the same way
# it does from app.py.
sys.path.insert(0, str(Path(__file__).resolve().parent.parent))


def main() -> int:
    parser = argparse.ArgumentParser(
        description=__doc__,
        formatter_class=argparse.RawDescriptionHelpFormatter,
    )
    parser.add_argument("prompt")
    parser.add_argument("--negative", default="", help="what to avoid")
    parser.add_argument("--steps", type=int, default=25)
    parser.add_argument("--guidance", type=float, default=7.5)
    parser.add_argument("--lora-scale", type=float, default=1.0,
                        help="0 = vanilla SD 1.5, 1 = as trained")
    parser.add_argument("--seed", type=int, default=-1, help="-1 for random")
    parser.add_argument("--count", type=int, default=1)
    parser.add_argument("--checkpoint", default=None,
                        help="'base' to skip the adapter entirely")
    parser.add_argument("--out", default="out.png")
    args = parser.parse_args()

    # Imported here so --help works without torch installed.
    from src.pipeline import LORA_CHECKPOINT, LORA_REPO, ModelRunner, pick_device

    args.checkpoint = args.checkpoint or LORA_CHECKPOINT
    print(f"device     {pick_device()}")
    print(f"model      {LORA_REPO} @ {args.checkpoint}")
    print("loading (first run downloads ~4 GB) ...", flush=True)

    runner = ModelRunner(checkpoint=args.checkpoint)
    result = runner.generate(
        prompt=args.prompt,
        negative_prompt=args.negative,
        steps=args.steps,
        guidance=args.guidance,
        lora_scale=args.lora_scale,
        seed=None if args.seed < 0 else args.seed,
        count=args.count,
    )

    out = Path(args.out)
    if len(result.images) == 1:
        result.images[0].save(out)
        print(f"seed {result.seed} -> {out}")
    else:
        # Image i used seed+i, so each file names the seed that reproduces
        # it on its own with --count 1.
        for index, image in enumerate(result.images):
            path = out.with_name(f"{out.stem}-{index}{out.suffix}")
            image.save(path)
            print(f"seed {result.seed + index} -> {path}")
    return 0


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