#!/usr/bin/env python3 """ comfy_workflows_lib.py — custom ComfyUI workflow library for drjones' website art. Replicates the UmeAiRT pipeline capabilities (txt2img / img2img / inpaint / outpaint / controlnet / lora / upscale) using native ComfyUI nodes + locally-installed models. Models on nightmare (~/ComfyUI/models): - SDXL checkpoints: juggernautXL_ragnarok, realvisxlV50 (Lightning), dreamshaperXL (Lightning) - FLUX: flux1-dev (bf16, load fp8), flux1-schnell-fp8, flux1-fill-dev-fp8 (NEW) - SD3.5: sd3.5_large_fp8_scaled - text encoders: clip_l.safetensors, clip_g.safetensors, t5xxl_fp8_e4m3fn.safetensors - vae: ae.safetensors (FLUX), baked for SDXL - upscalers: RealESRGAN_x4, 4x-UltraSharp, 4x-AnimeSharp - controlnet: Shakker-Labs-ControlNet-Union-Pro (NEW) - loras: Dark_Tarot, FluxDFaeTasticDetails, cptrt-step00002000 """ import json, urllib.request, urllib.error, urllib.parse, time, os, random COMFY = "http://10.30.20.128:8188" HYPERSWAP = "http://10.30.20.128:9090" # --------------------------------------------------------------------------- # model constants # --------------------------------------------------------------------------- FLUX_DEV = "flux1-dev.safetensors" # bf16, loaded as fp8_e4m3fn FLUX_FILL = "flux1-fill-dev-fp8.safetensors" # NEW — inpaint/outpaint FLUX_SCHNELL = "flux1-schnell-fp8.safetensors" SD35 = "sd3.5_large_fp8_scaled.safetensors" CLIP_L = "clip_l.safetensors" CLIP_G = "clip_g.safetensors" T5XXL_FP8 = "t5xxl_fp8_e4m3fn.safetensors" FLUX_VAE = "ae.safetensors" SDXL_CHECKPOINTS = { "juggernaut": "juggernautXL_ragnarok.safetensors", "realvis": "realvisxlV50_v50LightningBakedvae.safetensors", "dreamshaper": "dreamshaperXL_lightningDPMSDE.safetensors", } UPSCALERS = ["4x-UltraSharp.pth", "RealESRGAN_x4.pth", "4x-AnimeSharp.pth"] LORAS = ["Dark_Tarot.safetensors", "FluxDFaeTasticDetails.safetensors", "cptrt-step00002000.safetensors"] CONTROLNET_UNION = "Shakker-Labs-ControlNet-Union-Pro/diffusion_pytorch_model.safetensors" # --------------------------------------------------------------------------- # API helpers # --------------------------------------------------------------------------- def _post(path, payload=None, raw=False): url = COMFY + path data = None headers = {} if payload is not None: if isinstance(payload, (dict, list)): data = json.dumps(payload).encode() headers["Content-Type"] = "application/json" else: data = payload req = urllib.request.Request(url, data=data, headers=headers) try: r = urllib.request.urlopen(req, timeout=60) b = r.read() return b if raw else json.loads(b) except urllib.error.HTTPError as e: return json.loads(e.read()) if not raw else e.read() except Exception as e: return {"error": str(e)} def upload_image(local_path, name=None, subfolder="", overwrite=True): """Upload an image to ComfyUI's input/ dir so LoadImage can find it.""" if name is None: name = os.path.basename(local_path) boundary = "----WebKitFormBoundary7MA4YWxkTrZu0gW" with open(local_path, "rb") as f: img = f.read() parts = [] parts.append(("--" + boundary).encode()) parts.append(b'Content-Disposition: form-data; name="image"; filename="%s"' % name.encode()) parts.append(b"Content-Type: image/png") parts.append(b"") parts.append(img) parts.append(("--" + boundary).encode()) parts.append(b'Content-Disposition: form-data; name="overwrite"') parts.append(b"") parts.append(b"true" if overwrite else b"false") parts.append(("--" + boundary + "--").encode()) body = b"\r\n".join(parts) req = urllib.request.Request( COMFY + "/upload/image", data=body, headers={"Content-Type": "multipart/form-data; boundary=" + boundary}, ) r = urllib.request.urlopen(req, timeout=60) return json.loads(r.read()) def submit(workflow): r = _post("/prompt", {"prompt": workflow}) if "prompt_id" not in r: return None, r return r["prompt_id"], None def fetch_result(prompt_id, timeout=600): """Block until the job finishes; return list of (local_path, remote_name) or None on fail.""" deadline = time.time() + timeout while time.time() < deadline: h = _post(f"/history/{prompt_id}") if prompt_id in h: outs = h[prompt_id].get("outputs", {}) images = [] for node_id, o in outs.items(): for im in o.get("images", []): images.append((im["filename"], im.get("subfolder", ""), im.get("type", "output"))) if images: local_paths = [] for fn, sub, typ in images: q = urllib.parse.urlencode({"filename": fn, "subfolder": sub, "type": typ}) url = COMFY + "/view?" + q local = os.path.join(os.path.expanduser("~"), "auto-publisher/core/generated_art", fn) os.makedirs(os.path.dirname(local), exist_ok=True) urllib.request.urlretrieve(url, local) local_paths.append(local) return local_paths time.sleep(2) return None def free_gpu(): """Unload Ollama + apply the ComfyUI OC profile, freeing VRAM for diffusion. Called by the web worker before a render. Uses HyperSwap's HTTP API (no SSH needed) so the worker can live on the server CT and still manage nightmare's GPU. free-vram unloads ONE model per call, so loop to clear everything loaded. """ for _ in range(4): try: urllib.request.urlopen(urllib.request.Request( HYPERSWAP + "/api/free-vram", data=b"", method="POST"), timeout=30).read() except Exception: break time.sleep(1) try: body = json.dumps({"profile": "comfy"}).encode() urllib.request.urlopen(urllib.request.Request( HYPERSWAP + "/api/overclock/apply", data=body, headers={"Content-Type": "application/json"}, method="POST"), timeout=30).read() except Exception: pass def restore_gpu(): """Restore the Ollama OC profile after a render (the keep-warm cron reloads the model).""" try: body = json.dumps({"profile": "ollama"}).encode() urllib.request.urlopen(urllib.request.Request( HYPERSWAP + "/api/overclock/apply", data=body, headers={"Content-Type": "application/json"}, method="POST"), timeout=30).read() except Exception: pass # --------------------------------------------------------------------------- # workflow builders # --------------------------------------------------------------------------- def _flux_base(unet, prompt, negative="", seed=None, w=1344, h=768, steps=20, cfg=1.0, sampler="euler", scheduler="simple"): if seed is None: seed = random.randint(0, 2**63) return { "1": {"class_type": "UNETLoader", "inputs": {"unet_name": unet, "weight_dtype": "fp8_e4m3fn"}}, "2": {"class_type": "DualCLIPLoader", "inputs": {"clip_name1": CLIP_L, "clip_name2": T5XXL_FP8, "type": "flux"}}, "3": {"class_type": "VAELoader", "inputs": {"vae_name": FLUX_VAE}}, "4": {"class_type": "CLIPTextEncode", "inputs": {"text": prompt, "clip": ["2", 0]}}, "5": {"class_type": "CLIPTextEncode", "inputs": {"text": negative, "clip": ["2", 0]}}, "6": {"class_type": "EmptySD3LatentImage", "inputs": {"width": w, "height": h, "batch_size": 1}}, "7": {"class_type": "KSampler", "inputs": {"model": ["1", 0], "positive": ["4", 0], "negative": ["5", 0], "latent_image": ["6", 0], "seed": seed, "steps": steps, "cfg": cfg, "sampler_name": sampler, "scheduler": scheduler, "denoise": 1.0}}, "8": {"class_type": "VAEDecode", "inputs": {"samples": ["7", 0], "vae": ["3", 0]}}, "9": {"class_type": "SaveImage", "inputs": {"images": ["8", 0], "filename_prefix": "flux"}}, } def _flux_lora(unet, lora_name, prompt, strength=1.0, seed=None, w=1024, h=1024, steps=20): if seed is None: seed = random.randint(0, 2**63) return { "1": {"class_type": "UNETLoader", "inputs": {"unet_name": unet, "weight_dtype": "fp8_e4m3fn"}}, "2": {"class_type": "DualCLIPLoader", "inputs": {"clip_name1": CLIP_L, "clip_name2": T5XXL_FP8, "type": "flux"}}, "3": {"class_type": "VAELoader", "inputs": {"vae_name": FLUX_VAE}}, "4": {"class_type": "LoraLoader", "inputs": {"lora_name": lora_name, "strength_model": strength, "strength_clip": strength, "model": ["1", 0], "clip": ["2", 0]}}, "5": {"class_type": "CLIPTextEncode", "inputs": {"text": prompt, "clip": ["4", 1]}}, "6": {"class_type": "CLIPTextEncode", "inputs": {"text": "", "clip": ["4", 1]}}, "7": {"class_type": "EmptySD3LatentImage", "inputs": {"width": w, "height": h, "batch_size": 1}}, "8": {"class_type": "KSampler", "inputs": {"model": ["4", 0], "positive": ["5", 0], "negative": ["6", 0], "latent_image": ["7", 0], "seed": seed, "steps": steps, "cfg": 1.0, "sampler_name": "euler", "scheduler": "simple", "denoise": 1.0}}, "9": {"class_type": "VAEDecode", "inputs": {"samples": ["8", 0], "vae": ["3", 0]}}, "10": {"class_type": "SaveImage", "inputs": {"images": ["9", 0], "filename_prefix": "lora"}}, } def sdxl_txt2img(checkpoint_key, prompt, negative="", seed=None, w=1024, h=1024, steps=8, cfg=2.0, upscale_model=None): """SDXL (Lightning checkpoints = 4-8 steps). Optional 4x upscale.""" if seed is None: seed = random.randint(0, 2**63) ckpt = SDXL_CHECKPOINTS[checkpoint_key] wf = { "1": {"class_type": "CheckpointLoaderSimple", "inputs": {"ckpt_name": ckpt}}, "2": {"class_type": "CLIPTextEncode", "inputs": {"text": prompt, "clip": ["1", 1]}}, "3": {"class_type": "CLIPTextEncode", "inputs": {"text": negative, "clip": ["1", 1]}}, "4": {"class_type": "EmptyLatentImage", "inputs": {"width": w, "height": h, "batch_size": 1}}, "5": {"class_type": "KSampler", "inputs": {"model": ["1", 0], "positive": ["2", 0], "negative": ["3", 0], "latent_image": ["4", 0], "seed": seed, "steps": steps, "cfg": cfg, "sampler_name": "dpmpp_2m", "scheduler": "karras", "denoise": 1.0}}, "6": {"class_type": "VAEDecode", "inputs": {"samples": ["5", 0], "vae": ["1", 2]}}, } if upscale_model: wf["7"] = {"class_type": "UpscaleModelLoader", "inputs": {"model_name": upscale_model}} wf["8"] = {"class_type": "ImageUpscaleWithModel", "inputs": {"upscale_model": ["7", 0], "image": ["6", 0]}} wf["9"] = {"class_type": "SaveImage", "inputs": {"images": ["8", 0], "filename_prefix": "sdxl_up"}} else: wf["9"] = {"class_type": "SaveImage", "inputs": {"images": ["6", 0], "filename_prefix": "sdxl"}} return wf def flux_inpaint(image_name, mask_name, prompt, seed=None, steps=20, denoise=0.7): """FLUX Fill inpainting — redraw the masked region of image_name.""" if seed is None: seed = random.randint(0, 2**63) return { "1": {"class_type": "UNETLoader", "inputs": {"unet_name": FLUX_FILL, "weight_dtype": "fp8_e4m3fn"}}, "2": {"class_type": "DualCLIPLoader", "inputs": {"clip_name1": CLIP_L, "clip_name2": T5XXL_FP8, "type": "flux"}}, "3": {"class_type": "VAELoader", "inputs": {"vae_name": FLUX_VAE}}, "4": {"class_type": "LoadImage", "inputs": {"image": image_name}}, "5": {"class_type": "LoadImage", "inputs": {"image": mask_name}}, "5b": {"class_type": "ImageToMask", "inputs": {"image": ["5", 0], "channel": "red"}}, "6": {"class_type": "CLIPTextEncode", "inputs": {"text": prompt, "clip": ["2", 0]}}, "7": {"class_type": "CLIPTextEncode", "inputs": {"text": "", "clip": ["2", 0]}}, "8": {"class_type": "InpaintModelConditioning", "inputs": {"positive": ["6", 0], "negative": ["7", 0], "vae": ["3", 0], "pixels": ["4", 0], "mask": ["5b", 0], "noise_mask": True}}, "9": {"class_type": "KSampler", "inputs": {"model": ["1", 0], "positive": ["8", 0], "negative": ["8", 1], "latent_image": ["8", 2], "seed": seed, "steps": steps, "cfg": 1.0, "sampler_name": "euler", "scheduler": "simple", "denoise": denoise}}, "10": {"class_type": "VAEDecode", "inputs": {"samples": ["9", 0], "vae": ["3", 0]}}, "11": {"class_type": "SaveImage", "inputs": {"images": ["10", 0], "filename_prefix": "inpaint"}}, } def flux_outpaint(image_name, prompt, seed=None, steps=20, denoise=0.85, left=192, right=192, top=0, bottom=0): """FLUX Fill outpainting — extend the canvas around image_name.""" if seed is None: seed = random.randint(0, 2**63) return { "1": {"class_type": "UNETLoader", "inputs": {"unet_name": FLUX_FILL, "weight_dtype": "fp8_e4m3fn"}}, "2": {"class_type": "DualCLIPLoader", "inputs": {"clip_name1": CLIP_L, "clip_name2": T5XXL_FP8, "type": "flux"}}, "3": {"class_type": "VAELoader", "inputs": {"vae_name": FLUX_VAE}}, "4": {"class_type": "LoadImage", "inputs": {"image": image_name}}, "5": {"class_type": "ImagePadForOutpaint", "inputs": {"image": ["4", 0], "left": left, "top": top, "right": right, "bottom": bottom, "feathering": 40}}, "6": {"class_type": "CLIPTextEncode", "inputs": {"text": prompt, "clip": ["2", 0]}}, "7": {"class_type": "CLIPTextEncode", "inputs": {"text": "", "clip": ["2", 0]}}, "8": {"class_type": "InpaintModelConditioning", "inputs": {"positive": ["6", 0], "negative": ["7", 0], "vae": ["3", 0], "pixels": ["5", 0], "mask": ["5", 1], "noise_mask": True}}, "9": {"class_type": "KSampler", "inputs": {"model": ["1", 0], "positive": ["8", 0], "negative": ["8", 1], "latent_image": ["8", 2], "seed": seed, "steps": steps, "cfg": 1.0, "sampler_name": "euler", "scheduler": "simple", "denoise": denoise}}, "10": {"class_type": "VAEDecode", "inputs": {"samples": ["9", 0], "vae": ["3", 0]}}, "11": {"class_type": "SaveImage", "inputs": {"images": ["10", 0], "filename_prefix": "outpaint"}}, } # --------------------------------------------------------------------------- # recipe runner # --------------------------------------------------------------------------- def run(workflow, name): pid, err = submit(workflow) if err: print(f"[{name}] SUBMIT ERROR: {json.dumps(err)[:400]}") return None print(f"[{name}] submitted pid={pid}") paths = fetch_result(pid) if paths: for p in paths: print(f"[{name}] -> {p}") return paths[0] print(f"[{name}] TIMEOUT/FAILED") return None