"""sd_local — SD1.5 Hyper_Realism + anatomy LoRAs (+ optional pose ControlNet / IP-Adapter ref).

Runs UNDER the repose venv (manifest python is that venv, absolute → primary-only).
Generalizes scripts/repose.py from MESHGOD: pose template optional, plain txt2img otherwise.
"""
import argparse
import json
import os
import sys

os.environ.setdefault("PYTORCH_ENABLE_MPS_FALLBACK", "1")

import torch  # noqa: E402
from diffusers import DPMSolverMultistepScheduler  # noqa: E402
from PIL import Image  # noqa: E402

ap = argparse.ArgumentParser()
ap.add_argument("--input", action="append", default=[])
ap.add_argument("--outdir", required=True)
ap.add_argument("--params", default="{}")
a = ap.parse_args()
p = json.loads(a.params)

HERE = os.path.expanduser("~/Documents/repose")
BASE = os.path.join(HERE, "models", "Hyper_Realism_1.2_fp16.safetensors")
# LoRA search path: John's master dir first, the repose copy second (SD_LORA_DIRS to override)
LORA_DIRS = [os.path.expanduser(d) for d in os.environ.get(
    "SD_LORA_DIRS", "~/Documents/localmodels/Lora:~/Documents/repose/models/lora").split(":")]


def find_lora(stem):
    for d in LORA_DIRS:
        cand = os.path.join(d, stem + ".safetensors")
        if os.path.exists(cand):
            return cand
    return None
if not os.path.exists(BASE):
    print(f"ERROR: {BASE} missing — sd_local runs on the primary only")
    sys.exit(1)

dev = "mps" if torch.backends.mps.is_available() else "cpu"
dtype = torch.float16 if dev == "mps" else torch.float32
pose = p.get("pose", "none")
kw = {}

if pose != "none":
    from diffusers import ControlNetModel, StableDiffusionControlNetPipeline
    tpl = os.path.join(HERE, f"tpose_{pose}.png")
    cn = ControlNetModel.from_pretrained("lllyasviel/control_v11p_sd15_openpose", torch_dtype=dtype)
    pipe = StableDiffusionControlNetPipeline.from_single_file(
        BASE, controlnet=cn, torch_dtype=dtype, safety_checker=None, requires_safety_checker=False)
    kw["image"] = Image.open(tpl).convert("RGB")
else:
    from diffusers import StableDiffusionPipeline
    pipe = StableDiffusionPipeline.from_single_file(
        BASE, torch_dtype=dtype, safety_checker=None, requires_safety_checker=False)

pipe.scheduler = DPMSolverMultistepScheduler.from_config(pipe.scheduler.config, use_karras_sigmas=True)
pipe = pipe.to(dev)
# no attention slicing — it crashes IP-Adapter attn-processor injection (see repose.py)

if p.get("loras"):
    from safetensors.torch import load_file
    names, weights = [], []
    for spec in p["loras"]:
        stem, _, w = str(spec).partition("=")
        path = find_lora(stem.strip())
        if not path:
            print(f"[sd] LORA MISS '{stem}' — searched {LORA_DIRS}", flush=True)
            continue
        # UNet-only: this diffusers build trips an IndexError on these kohya files' text-encoder
        # half (empty rank_dict in load_lora_into_text_encoder). The UNet half carries the look.
        sd_l = load_file(path)
        unet_only = {k: v for k, v in sd_l.items() if not k.startswith("lora_te")}
        pipe.load_lora_weights(unet_only, adapter_name=stem.replace(" ", "_"))
        names.append(stem.replace(" ", "_"))
        weights.append(float(w or 0.6))
    if names:
        pipe.set_adapters(names, weights)
    print(f"[sd] loras: {dict(zip(names, weights))}", flush=True)

if a.input:                                     # reference image → IP-Adapter
    pipe.load_ip_adapter("h94/IP-Adapter", subfolder="models", weight_name="ip-adapter_sd15.bin")
    pipe.set_ip_adapter_scale(float(p.get("ip", 0.7)))
    kw["ip_adapter_image"] = Image.open(a.input[0]).convert("RGB")
    print(f"[sd] ip-adapter ref @ {p.get('ip', 0.7)}", flush=True)

gen = torch.Generator("cpu").manual_seed(int(p.get("seed", 7)))
img = pipe(prompt=p.get("prompt", ""), negative_prompt=p.get("negative", ""),
           num_inference_steps=int(p.get("steps", 28)), guidance_scale=float(p.get("cfg", 7.0)),
           generator=gen, width=int(p.get("width", 512)), height=int(p.get("height", 768)),
           **kw).images[0]
out = os.path.join(a.outdir, "sd_local.png")
img.save(out)
json.dump({"outputs": [{"path": "sd_local.png", "meta": {"tool": "sd_local"}}]},
          open(os.path.join(a.outdir, "result.json"), "w"))
print("done: sd_local.png", flush=True)