9c7164f77a
Checkpoints are no longer uniform across the fleet (M3 keeps only the LoRA-compatible SD1.5 model; M1 holds the full SDXL archive). The gpu pool routes by operator, not by checkpoint, so a job asking for bigLust could land on a node without it. Query the node's own ComfyUI /object_info and exit with the available list + where the rest lives, instead of a cryptic 400.
177 lines
7.1 KiB
Python
177 lines
7.1 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).
|
|
"""
|
|
import argparse
|
|
import json
|
|
import os
|
|
import random
|
|
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 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")
|
|
subprocess.Popen(
|
|
[str(COMFY / ".venv" / "bin" / "python"), "main.py", "--port", "8188"],
|
|
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 build(seed):
|
|
ckpt = p.get("checkpoint", "Hyper_Realism_1.2_fp16.safetensors")
|
|
lora = (p.get("lora") or "").strip()
|
|
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]
|
|
if lora:
|
|
g["2"] = {"class_type": "LoraLoader", "inputs": {
|
|
"lora_name": lora,
|
|
"strength_model": float(p.get("lora_weight", 0.8)),
|
|
"strength_clip": float(p.get("lora_weight", 0.8)),
|
|
"model": msrc, "clip": csrc}}
|
|
msrc, csrc = ["2", 0], ["2", 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}}
|
|
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": ["3", 0], "negative": ["4", 0],
|
|
"latent_image": ["5", 0]}}
|
|
return g
|
|
|
|
|
|
ensure_server()
|
|
seed = int(p.get("seed", -1))
|
|
if seed < 0:
|
|
seed = random.randint(0, 2**31 - 1)
|
|
CKPT = p.get("checkpoint", "Hyper_Realism_1.2_fp16.safetensors")
|
|
LORA = (p.get("lora") or "").strip()
|
|
print(f"seed={seed} ckpt={CKPT} "
|
|
f"lora={LORA or 'none'}" + (f"@{p.get('lora_weight', 0.8)}" if LORA else ""), flush=True)
|
|
if LORA and ("xl" in CKPT.lower() or "biglust" in CKPT.lower() or "v2-1" in CKPT.lower()):
|
|
print(f"WARNING: {LORA} is SD1.5 but {CKPT} is not — the LoRA will silently do nothing. "
|
|
f"Use Hyper_Realism_1.2_fp16.safetensors (see localmodels/README.md)", flush=True)
|
|
|
|
# Checkpoints are NOT identical across nodes (the M3 keeps only the LoRA-compatible
|
|
# SD1.5 model; the M1 holds the full SDXL archive). Ask this node's ComfyUI what it
|
|
# actually has and fail fast with a useful message rather than a cryptic 400.
|
|
try:
|
|
info = json.load(urllib.request.urlopen(f"{API}/object_info/CheckpointLoaderSimple"))
|
|
have = info["CheckpointLoaderSimple"]["input"]["required"]["ckpt_name"][0]
|
|
if CKPT not in have:
|
|
print(f"ERROR: '{CKPT}' is not on this node.\n"
|
|
f" available here: {', '.join(have) or '(none)'}\n"
|
|
f" The full SDXL set lives on the m1 node; every node has "
|
|
f"Hyper_Realism_1.2_fp16.safetensors (the only LoRA-compatible checkpoint).\n"
|
|
f" Fix: use a checkpoint listed above, or rsync it from m1.", flush=True)
|
|
sys.exit(1)
|
|
if LORA:
|
|
have_l = json.load(urllib.request.urlopen(f"{API}/object_info/LoraLoader"))
|
|
have_l = have_l["LoraLoader"]["input"]["required"]["lora_name"][0]
|
|
if LORA not in have_l:
|
|
print(f"ERROR: LoRA '{LORA}' is not on this node. available: {', '.join(have_l) or '(none)'}", flush=True)
|
|
sys.exit(1)
|
|
except SystemExit:
|
|
raise
|
|
except Exception as e:
|
|
print(f"(could not pre-check model availability: {e}) — continuing", 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": p.get("checkpoint"), "lora": p.get("lora") or None}})
|
|
|
|
(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)
|