"""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)