3b8c88ce48
LORA_DIRS searches ~/Documents/localmodels/Lora (John's master) then the repose copy (SD_LORA_DIRS to override). Load via filtered state_dict keeping only lora_unet_* keys: this diffusers build IndexErrors on these kohya files' text-encoder half (empty rank_dict in load_lora_into_text_encoder); the UNet half carries the look. Verified: all 3 anatomy LoRAs load + generate. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
98 lines
4.2 KiB
Python
98 lines
4.2 KiB
Python
"""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)
|