4c05a9daa8
wan_video 5B q8 verified clean on M3 Ultra (148.7s) and M2 Max (503.7s), identical 9.2GB max RSS on both; 81GB 'footprint' on m2max identified as MLX opportunistic cache, not demand. Frame-identical output at same seed across machines. HARDWARE.md gains wan_video and ardy_motion rows. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
278 lines
11 KiB
Python
278 lines
11 KiB
Python
"""SD/SDXL + LoRA generation via a resident ComfyUI (Metal).
|
|
|
|
Pure stdlib — no venv needed (the runner falls back to the node's python3), because
|
|
all the heavy lifting happens inside ComfyUI's own venv. Keeps ComfyUI resident so
|
|
checkpoints stay cached in RAM between jobs (a cold load costs seconds; a warm one
|
|
doesn't).
|
|
|
|
`lora` takes one name or a LIST, each optionally `name=weight`, and they are chained —
|
|
style + texture together, which a single loader could never do. `controlnet` +
|
|
`control_image` add structural conditioning: the headline use is recolouring one garment
|
|
into N colourways that share a byte-identical silhouette, so a single cut-out alpha is
|
|
reusable across every variant.
|
|
"""
|
|
import argparse
|
|
import json
|
|
import os
|
|
import random
|
|
import shutil
|
|
import subprocess
|
|
import sys
|
|
import time
|
|
import urllib.parse
|
|
import urllib.request
|
|
from pathlib import Path
|
|
|
|
ROOT = Path(__file__).resolve().parents[3]
|
|
COMFY = ROOT / "vendor" / "comfyui"
|
|
API = "http://127.0.0.1:8188"
|
|
|
|
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)
|
|
outdir = Path(a.outdir)
|
|
|
|
prompt = (p.get("prompt") or "").strip()
|
|
if not prompt:
|
|
print("ERROR: prompt is required")
|
|
sys.exit(1)
|
|
if not (COMFY / "main.py").exists():
|
|
print(f"ERROR: ComfyUI not installed at {COMFY}. Run scripts/install_comfyui.sh")
|
|
sys.exit(1)
|
|
|
|
|
|
def parse_loras():
|
|
"""-> [(name, weight)]. Accepts 'a', 'a=0.6', ['a','b=0.4'], and a list of weights.
|
|
|
|
A LoRA on the wrong base is a SILENT no-op — no error, the image just comes back
|
|
unchanged and you blame the LoRA. So everything here fails loudly instead.
|
|
"""
|
|
raw = p.get("lora") or p.get("loras") or []
|
|
if isinstance(raw, str):
|
|
raw = [s for s in (x.strip() for x in raw.split(",")) if s]
|
|
elif not isinstance(raw, list):
|
|
raw = []
|
|
w_raw = p.get("lora_weight", 0.8)
|
|
weights = w_raw if isinstance(w_raw, list) else [w_raw] * len(raw)
|
|
out = []
|
|
for i, item in enumerate(raw):
|
|
item = str(item).strip()
|
|
if not item:
|
|
continue
|
|
if "=" in item:
|
|
name, _, w = item.rpartition("=")
|
|
try:
|
|
out.append((name.strip(), float(w)))
|
|
continue
|
|
except ValueError:
|
|
pass # '=' was part of the filename, not a weight
|
|
try:
|
|
w = float(weights[i]) if i < len(weights) else 0.8
|
|
except (TypeError, ValueError):
|
|
w = 0.8
|
|
out.append((item, w))
|
|
return out
|
|
|
|
|
|
LORAS = parse_loras()
|
|
CKPT = p.get("checkpoint", "Hyper_Realism_1.2_fp16.safetensors")
|
|
CTRL = (p.get("controlnet") or "").strip()
|
|
CTRL_STRENGTH = float(p.get("controlnet_strength", 1.0))
|
|
ctrl_file = None # basename inside ComfyUI's input dir
|
|
|
|
|
|
def alive():
|
|
try:
|
|
urllib.request.urlopen(f"{API}/", timeout=3)
|
|
return True
|
|
except Exception:
|
|
return False
|
|
|
|
|
|
def ensure_server():
|
|
if alive():
|
|
print("comfyui: already resident (models stay cached)", flush=True)
|
|
return
|
|
print("comfyui: starting resident server ...", flush=True)
|
|
log = open("/tmp/comfyui.log", "ab")
|
|
# --listen 0.0.0.0: tailnet-only box; lets John drive the same resident
|
|
# instance interactively at http://100.89.131.57:8188 (FLUX park lives in
|
|
# models/ now) while farm jobs keep using it via localhost.
|
|
subprocess.Popen(
|
|
[str(COMFY / ".venv" / "bin" / "python"), "main.py", "--port", "8188",
|
|
"--listen", "0.0.0.0"],
|
|
cwd=str(COMFY), stdout=log, stderr=log, start_new_session=True)
|
|
for _ in range(90):
|
|
if alive():
|
|
print("comfyui: up", flush=True)
|
|
return
|
|
time.sleep(2)
|
|
print("ERROR: ComfyUI did not come up — see /tmp/comfyui.log")
|
|
sys.exit(1)
|
|
|
|
|
|
def node_options(cls, field):
|
|
"""What this node actually has installed — used to fail fast with a useful list."""
|
|
try:
|
|
info = json.load(urllib.request.urlopen(f"{API}/object_info/{cls}"))
|
|
return info[cls]["input"]["required"][field][0]
|
|
except Exception:
|
|
return None
|
|
|
|
|
|
def build(seed):
|
|
g = {
|
|
"1": {"class_type": "CheckpointLoaderSimple", "inputs": {"ckpt_name": CKPT}},
|
|
"5": {"class_type": "EmptyLatentImage", "inputs": {
|
|
"width": int(p.get("width", 512)), "height": int(p.get("height", 512)),
|
|
"batch_size": int(p.get("batch", 1))}},
|
|
"7": {"class_type": "VAEDecode", "inputs": {"samples": ["6", 0], "vae": ["1", 2]}},
|
|
"8": {"class_type": "SaveImage", "inputs": {"filename_prefix": "mb_sd", "images": ["7", 0]}},
|
|
}
|
|
msrc, csrc = ["1", 0], ["1", 1]
|
|
# Chain one LoraLoader per entry so stacking works (style + texture together).
|
|
nid = 100
|
|
for name, w in LORAS:
|
|
g[str(nid)] = {"class_type": "LoraLoader", "inputs": {
|
|
"lora_name": name, "strength_model": w, "strength_clip": w,
|
|
"model": msrc, "clip": csrc}}
|
|
msrc, csrc = [str(nid), 0], [str(nid), 1]
|
|
nid += 1
|
|
# clip_skip 2 == CLIPSetLastLayer -2 (what the local SD1.5 LoRAs were trained at)
|
|
if int(p.get("clip_skip", 2)) == 2:
|
|
g["9"] = {"class_type": "CLIPSetLastLayer", "inputs": {"clip": csrc, "stop_at_clip_layer": -2}}
|
|
csrc = ["9", 0]
|
|
g["3"] = {"class_type": "CLIPTextEncode", "inputs": {"text": prompt, "clip": csrc}}
|
|
g["4"] = {"class_type": "CLIPTextEncode", "inputs": {"text": p.get("negative", ""), "clip": csrc}}
|
|
pos, neg = ["3", 0], ["4", 0]
|
|
if CTRL and ctrl_file:
|
|
g["20"] = {"class_type": "LoadImage", "inputs": {"image": ctrl_file}}
|
|
g["21"] = {"class_type": "ControlNetLoader", "inputs": {"control_net_name": CTRL}}
|
|
g["22"] = {"class_type": "ControlNetApplyAdvanced", "inputs": {
|
|
"positive": pos, "negative": neg, "control_net": ["21", 0], "image": ["20", 0],
|
|
"strength": CTRL_STRENGTH,
|
|
"start_percent": float(p.get("controlnet_start", 0.0)),
|
|
"end_percent": float(p.get("controlnet_end", 1.0))}}
|
|
pos, neg = ["22", 0], ["22", 1]
|
|
g["6"] = {"class_type": "KSampler", "inputs": {
|
|
"seed": seed, "steps": int(p.get("steps", 25)), "cfg": float(p.get("cfg", 7.0)),
|
|
"sampler_name": p.get("sampler", "dpmpp_2m"), "scheduler": p.get("scheduler", "karras"),
|
|
"denoise": 1.0, "model": msrc, "positive": pos, "negative": neg,
|
|
"latent_image": ["5", 0]}}
|
|
return g
|
|
|
|
|
|
ensure_server()
|
|
seed = int(p.get("seed", -1))
|
|
if seed < 0:
|
|
seed = random.randint(0, 2**31 - 1)
|
|
lora_desc = ", ".join(f"{n}@{w}" for n, w in LORAS) or "none"
|
|
print(f"seed={seed} ckpt={CKPT} lora={lora_desc}"
|
|
+ (f" controlnet={CTRL}@{CTRL_STRENGTH}" if CTRL else ""), flush=True)
|
|
|
|
# A ControlNet needs its conditioning image inside ComfyUI's own input dir — LoadImage
|
|
# resolves by basename, not path.
|
|
if CTRL:
|
|
if not a.input:
|
|
print("ERROR: controlnet requested but no input image was supplied "
|
|
"(pass the control image as the job's asset)", flush=True)
|
|
sys.exit(1)
|
|
src = Path(a.input[0])
|
|
if not src.is_file():
|
|
print(f"ERROR: control image not found: {src}", flush=True)
|
|
sys.exit(1)
|
|
indir = COMFY / "input"
|
|
indir.mkdir(parents=True, exist_ok=True)
|
|
ctrl_file = f"mb_ctrl_{seed}{src.suffix or '.png'}"
|
|
shutil.copyfile(src, indir / ctrl_file)
|
|
print(f"controlnet: staged {src.name} -> input/{ctrl_file}", flush=True)
|
|
|
|
# Checkpoints, LoRAs and ControlNets are NOT identical across nodes. Ask this node's
|
|
# ComfyUI what it actually has and fail fast with a useful message rather than a
|
|
# cryptic 400 — or worse, a silent no-op.
|
|
have = node_options("CheckpointLoaderSimple", "ckpt_name")
|
|
if have is not None and CKPT not in have:
|
|
print(f"ERROR: '{CKPT}' is not on this node.\n"
|
|
f" available here: {', '.join(have) or '(none)'}\n"
|
|
f" Fix: use a checkpoint listed above, or rsync it to "
|
|
f"~/Documents/localmodels/Stable-diffusion/ on this node.", flush=True)
|
|
sys.exit(1)
|
|
if LORAS:
|
|
have_l = node_options("LoraLoader", "lora_name")
|
|
if have_l is not None:
|
|
missing = [n for n, _ in LORAS if n not in have_l]
|
|
if missing:
|
|
print(f"ERROR: LoRA(s) not on this node: {', '.join(missing)}\n"
|
|
f" available here: {', '.join(have_l) or '(none)'}", flush=True)
|
|
sys.exit(1)
|
|
if CTRL:
|
|
have_c = node_options("ControlNetLoader", "control_net_name")
|
|
if have_c is not None and CTRL not in have_c:
|
|
print(f"ERROR: ControlNet '{CTRL}' is not on this node.\n"
|
|
f" available here: {', '.join(have_c) or '(none)'}", flush=True)
|
|
sys.exit(1)
|
|
|
|
# Architecture mismatch is the classic silent failure: an SD1.5 LoRA on an SDXL base
|
|
# loads without complaint and does nothing at all.
|
|
xl_ckpt = any(t in CKPT.lower() for t in ("xl", "juggernaut"))
|
|
for n, _ in LORAS:
|
|
xl_lora = "xl" in n.lower()
|
|
if xl_ckpt != xl_lora:
|
|
print(f"WARNING: '{n}' looks {'SDXL' if xl_lora else 'SD1.5'} but the checkpoint "
|
|
f"'{CKPT}' looks {'SDXL' if xl_ckpt else 'SD1.5'} — a wrong-base LoRA is a "
|
|
f"SILENT no-op (no error, no effect). Filenames lie; check the metadata.",
|
|
flush=True)
|
|
|
|
body = json.dumps({"prompt": build(seed)}).encode()
|
|
req = urllib.request.Request(f"{API}/prompt", data=body, headers={"Content-Type": "application/json"})
|
|
try:
|
|
pid = json.load(urllib.request.urlopen(req))["prompt_id"]
|
|
except Exception as e:
|
|
print(f"ERROR: ComfyUI rejected the workflow: {e}")
|
|
sys.exit(1)
|
|
|
|
t0 = time.time()
|
|
images = []
|
|
while time.time() - t0 < 900:
|
|
h = json.load(urllib.request.urlopen(f"{API}/history/{pid}"))
|
|
if pid in h:
|
|
st = h[pid].get("status", {})
|
|
if st.get("status_str") == "error":
|
|
print("ERROR: generation failed — check /tmp/comfyui.log")
|
|
print(json.dumps(st)[:400])
|
|
sys.exit(1)
|
|
if h[pid].get("outputs"):
|
|
for node in h[pid]["outputs"].values():
|
|
images += node.get("images", [])
|
|
break
|
|
time.sleep(2)
|
|
|
|
if not images:
|
|
print("ERROR: no image produced (timeout)")
|
|
sys.exit(1)
|
|
|
|
outputs = []
|
|
for i, im in enumerate(images):
|
|
q = urllib.parse.urlencode({"filename": im["filename"], "subfolder": im.get("subfolder", ""),
|
|
"type": im.get("type", "output")})
|
|
data = urllib.request.urlopen(f"{API}/view?{q}").read()
|
|
name = f"sd_{seed}_{i}.png" if len(images) > 1 else f"sd_{seed}.png"
|
|
dest = outdir / name
|
|
dest.write_bytes(data)
|
|
outputs.append({"path": str(dest), "name": name,
|
|
"meta": {"tool": "comfyui_sd", "seed": seed, "checkpoint": CKPT,
|
|
"lora": lora_desc if LORAS else None,
|
|
"controlnet": CTRL or None}})
|
|
|
|
if ctrl_file:
|
|
try:
|
|
(COMFY / "input" / ctrl_file).unlink()
|
|
except OSError:
|
|
pass
|
|
|
|
(outdir / "result.json").write_text(json.dumps({"outputs": outputs}))
|
|
print(f"done: {len(outputs)} image(s) in {time.time() - t0:.1f}s -> {outputs[0]['path']}", flush=True)
|