x-stride
fa673ce0d4
SDXL checkpoint merging with DARE/TIES deltas, bezier UNet profiles, stack visualizer, example workflow PNG, and API smoke tests.
247 lines
8.2 KiB
Python
Executable File
247 lines
8.2 KiB
Python
Executable File
#!/usr/bin/env python3
|
|
"""API test: DARE TIES full merge on trad SDXL checkpoints + save result."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import glob
|
|
import json
|
|
import os
|
|
import sys
|
|
import time
|
|
import urllib.error
|
|
import urllib.request
|
|
|
|
BASE = "http://127.0.0.1:8188"
|
|
OUTPUT_DIR = os.environ.get("COMFYUI_OUTPUT_DIR", "/opt/ComfyUI/output")
|
|
SAVE_PREFIX = "dare_ties_api_test/trad_merge"
|
|
POLL_SEC = 2.0
|
|
TIMEOUT_SEC = 3600 # merges can run a long time on lowvram
|
|
|
|
# Prefer plain SDXL photoreal checkpoints (not Illustrious/anime paths).
|
|
REF_ORDER = [
|
|
"juggernautXL_version6Rundiffusion.safetensors",
|
|
"epicrealismXL_v7FinalDestination.safetensors",
|
|
"RealVisXL_V5.0_fp16.safetensors",
|
|
"cyberrealisticXL_v20.safetensors",
|
|
"dreamshaperXL_v21TurboDPMSDE.safetensors",
|
|
]
|
|
CHILD_ORDER = [
|
|
"RealVisXL_V5.0_fp16.safetensors",
|
|
"epicrealismXL_v7FinalDestination.safetensors",
|
|
"cyberrealistic_v50.safetensors",
|
|
"cyberrealisticXL_v20.safetensors",
|
|
"SDXL1.0/Realistic Freedom - SFW and NSFW - Ophelia - realisticFreedomSFW_ophelia.safetensors",
|
|
]
|
|
|
|
|
|
def req(method: str, path: str, body: dict | None = None) -> tuple[int, dict | list | str]:
|
|
data = None if body is None else json.dumps(body).encode()
|
|
r = urllib.request.Request(
|
|
BASE + path,
|
|
data=data,
|
|
headers={"Content-Type": "application/json"} if body is not None else {},
|
|
method=method,
|
|
)
|
|
try:
|
|
with urllib.request.urlopen(r, timeout=60) as resp:
|
|
raw = resp.read()
|
|
if not raw:
|
|
return resp.status, ""
|
|
return resp.status, json.loads(raw)
|
|
except urllib.error.HTTPError as e:
|
|
raw = e.read().decode()
|
|
try:
|
|
return e.code, json.loads(raw)
|
|
except json.JSONDecodeError:
|
|
return e.code, raw
|
|
|
|
|
|
def list_checkpoints() -> list[str]:
|
|
code, data = req("GET", "/models/checkpoints")
|
|
if code != 200 or not isinstance(data, list):
|
|
raise RuntimeError(f"checkpoint list failed: {code} {data!r}")
|
|
return data
|
|
|
|
|
|
def pick_trad(all_ckpts: list[str]) -> tuple[str, str, str]:
|
|
avail = set(all_ckpts)
|
|
|
|
def first(candidates: list[str]) -> str | None:
|
|
for name in candidates:
|
|
if name in avail:
|
|
return name
|
|
return None
|
|
|
|
ref = first(REF_ORDER)
|
|
if ref is None:
|
|
ref = next(
|
|
(
|
|
n
|
|
for n in all_ckpts
|
|
if "illustrious" not in n.lower()
|
|
and "anime" not in n.lower()
|
|
and "pony" not in n.lower()
|
|
and n.endswith(".safetensors")
|
|
),
|
|
all_ckpts[0],
|
|
)
|
|
|
|
children: list[str] = []
|
|
for name in CHILD_ORDER:
|
|
if name in avail and name != ref and name not in children:
|
|
children.append(name)
|
|
if len(children) == 2:
|
|
break
|
|
|
|
if len(children) < 2:
|
|
for name in all_ckpts:
|
|
if name == ref or name in children:
|
|
continue
|
|
if any(x in name.lower() for x in ("illustrious", "anime", "pony", "noob")):
|
|
continue
|
|
children.append(name)
|
|
if len(children) == 2:
|
|
break
|
|
|
|
if len(children) < 2:
|
|
raise RuntimeError(f"need 2 child trad checkpoints besides ref; have ref={ref!r}")
|
|
|
|
return ref, children[0], children[1]
|
|
|
|
|
|
def wait_prompt(prompt_id: str) -> dict:
|
|
started = time.time()
|
|
deadline = started + TIMEOUT_SEC
|
|
while time.time() < deadline:
|
|
code, hist = req("GET", f"/history/{prompt_id}")
|
|
if code == 200 and isinstance(hist, dict) and prompt_id in hist:
|
|
return hist[prompt_id]
|
|
code, q = req("GET", "/queue")
|
|
if code == 200 and isinstance(q, dict):
|
|
running = q.get("queue_running") or []
|
|
pending = q.get("queue_pending") or []
|
|
active = any(item[1] == prompt_id for item in running + pending if len(item) > 1)
|
|
if not active and prompt_id not in (hist if isinstance(hist, dict) else {}):
|
|
time.sleep(POLL_SEC)
|
|
code, hist = req("GET", f"/history/{prompt_id}")
|
|
if code == 200 and isinstance(hist, dict) and prompt_id in hist:
|
|
return hist[prompt_id]
|
|
elapsed = int(time.time() - started)
|
|
if elapsed and elapsed % 30 == 0:
|
|
print(f" ... still running ({elapsed}s)")
|
|
time.sleep(POLL_SEC)
|
|
raise TimeoutError(f"prompt {prompt_id} did not finish within {TIMEOUT_SEC}s")
|
|
|
|
|
|
def build_merge_prompt(reference: str, child_a: str, child_b: str) -> dict:
|
|
return {
|
|
"1": {
|
|
"class_type": "CheckpointLoaderSimple",
|
|
"inputs": {"ckpt_name": reference},
|
|
},
|
|
"2": {
|
|
"class_type": "DARETIESModelStack",
|
|
"inputs": {
|
|
"ckpt_name_1": child_a,
|
|
"model_weight_1": 1.0,
|
|
"clip_weight_1": 0.0,
|
|
"enabled_1": True,
|
|
"ckpt_name_2": child_b,
|
|
"model_weight_2": 1.0,
|
|
"clip_weight_2": 0.0,
|
|
"enabled_2": True,
|
|
"ckpt_name_3": "None",
|
|
"model_weight_3": 1.0,
|
|
"clip_weight_3": 0.0,
|
|
"enabled_3": False,
|
|
},
|
|
},
|
|
"3": {
|
|
"class_type": "DARETIESMergeModels",
|
|
"inputs": {
|
|
"reference_model": ["1", 0],
|
|
"reference_clip": ["1", 1],
|
|
"reference_vae": ["1", 2],
|
|
"model_stack": ["2", 0],
|
|
"merge_mode": "dare_delta",
|
|
"drop_rate": 0.10,
|
|
"delta_strength": 1.0,
|
|
"seed": 4242,
|
|
"normalize_model_weights": True,
|
|
"merge_clip": "off",
|
|
"normalize_clip_weights": True,
|
|
"clip_delta_strength": 0.5,
|
|
"ties_density": 0.50,
|
|
"merge_math_gpu": False,
|
|
},
|
|
},
|
|
"4": {
|
|
"class_type": "CheckpointSave",
|
|
"inputs": {
|
|
"model": ["3", 0],
|
|
"clip": ["3", 1],
|
|
"vae": ["3", 2],
|
|
"filename_prefix": SAVE_PREFIX,
|
|
},
|
|
},
|
|
}
|
|
|
|
|
|
def find_saved_checkpoint(since: float) -> str | None:
|
|
pattern = os.path.join(OUTPUT_DIR, SAVE_PREFIX + "*.safetensors")
|
|
candidates = [p for p in glob.glob(pattern) if os.path.getmtime(p) >= since - 1.0]
|
|
if not candidates:
|
|
return None
|
|
return max(candidates, key=os.path.getmtime)
|
|
|
|
|
|
def main() -> int:
|
|
print(f"DARE TIES merge API test @ {BASE}")
|
|
code, _ = req("GET", "/object_info/DARETIESMergeModels")
|
|
if code != 200:
|
|
print(f"FAIL: ComfyUI unreachable or merge node missing ({code})", file=sys.stderr)
|
|
return 1
|
|
|
|
all_ckpts = list_checkpoints()
|
|
reference, child_a, child_b = pick_trad(all_ckpts)
|
|
print(f" reference {reference}")
|
|
print(f" child A {child_a}")
|
|
print(f" child B {child_b}")
|
|
print(" mode dare_delta drop=0.10 clip=off merge_math_gpu=false")
|
|
|
|
prompt = build_merge_prompt(reference, child_a, child_b)
|
|
t0 = time.time()
|
|
code, body = req("POST", "/prompt", {"prompt": prompt})
|
|
if code != 200 or not isinstance(body, dict) or "prompt_id" not in body:
|
|
print(f"FAIL: queue rejected: {code} {body!r}", file=sys.stderr)
|
|
return 1
|
|
|
|
prompt_id = body["prompt_id"]
|
|
print(f" queued prompt_id={prompt_id}")
|
|
|
|
entry = wait_prompt(prompt_id)
|
|
elapsed = time.time() - t0
|
|
status = entry.get("status", {})
|
|
if status.get("status_str") != "success":
|
|
print(f"FAIL: {status}", file=sys.stderr)
|
|
for msg in status.get("messages") or []:
|
|
print(" ", msg, file=sys.stderr)
|
|
return 1
|
|
|
|
outputs = entry.get("outputs", {})
|
|
saved_path = find_saved_checkpoint(t0)
|
|
if saved_path is None:
|
|
print(f"FAIL: no checkpoint under {OUTPUT_DIR}/{SAVE_PREFIX}*.safetensors", file=sys.stderr)
|
|
print(f" history outputs: {outputs!r}", file=sys.stderr)
|
|
return 1
|
|
|
|
size_gb = os.path.getsize(saved_path) / (1024**3)
|
|
print(f" saved {saved_path} ({size_gb:.2f} GiB)")
|
|
print(f" elapsed {elapsed:.1f}s")
|
|
print("PASS")
|
|
return 0
|
|
|
|
|
|
if __name__ == "__main__":
|
|
raise SystemExit(main())
|