diff --git a/server/operators/sd_local/run.py b/server/operators/sd_local/run.py index 1fc74ec..941b9c5 100644 --- a/server/operators/sd_local/run.py +++ b/server/operators/sd_local/run.py @@ -23,7 +23,17 @@ p = json.loads(a.params) HERE = os.path.expanduser("~/Documents/repose") BASE = os.path.join(HERE, "models", "Hyper_Realism_1.2_fp16.safetensors") -LORA_DIR = os.path.join(HERE, "models", "lora") +# 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) @@ -50,14 +60,23 @@ 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("=") - pipe.load_lora_weights(LORA_DIR, weight_name=f"{stem}.safetensors", - adapter_name=stem.replace(" ", "_")) + 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)) - pipe.set_adapters(names, weights) + if names: + pipe.set_adapters(names, weights) print(f"[sd] loras: {dict(zip(names, weights))}", flush=True) if a.input: # reference image → IP-Adapter