commit 463fadf6c80daee5185409204d443e1882ea9654 Author: Brian Blankenship Date: Tue Jan 20 13:19:01 2026 -0600 initial commit diff --git a/.gitignore b/.gitignore new file mode 100644 index 0000000..bee8a64 --- /dev/null +++ b/.gitignore @@ -0,0 +1 @@ +__pycache__ diff --git a/__init__.py b/__init__.py new file mode 100644 index 0000000..36108bd --- /dev/null +++ b/__init__.py @@ -0,0 +1,36 @@ +import server +from aiohttp import web +from .pause_nodes import PSampler, Signaler +from .shared import PAUSE_STATE, state_lock + +# --- API Route Registration --- +@server.PromptServer.instance.routes.post("/comfy/steer") +async def set_steering_command(request): + """ + Receives control signals from the frontend (The Cockpit). + Payload example: { "command": "PAUSE" } + """ + json_data = await request.json() + command = json_data.get("command", "PROCEED") + + # Update the global state + with state_lock: + PAUSE_STATE["command"] = command + + print(f"[ComfyUI-Pause] πŸ“‘ Signal Received: {command}") + return web.json_response({"status": "success", "received": command}) + +# --- Node Mappings --- +NODE_CLASS_MAPPINGS = { + "PSampler": PSampler, + "Signaler": Signaler +} + +NODE_DISPLAY_NAME_MAPPINGS = { + "PSampler": "⏸️ P-Sampler (Pause Enabled)", + "Signaler": "🚦 Signaler (Control)" +} + +WEB_DIRECTORY = "./js" + +__all__ = ["NODE_CLASS_MAPPINGS", "NODE_DISPLAY_NAME_MAPPINGS", "PAUSE_STATE", "WEB_DIRECTORY"] diff --git a/gitignore b/gitignore new file mode 100644 index 0000000..bee8a64 --- /dev/null +++ b/gitignore @@ -0,0 +1 @@ +__pycache__ diff --git a/js/pausing_bridge.js b/js/pausing_bridge.js new file mode 100644 index 0000000..aa93cb6 --- /dev/null +++ b/js/pausing_bridge.js @@ -0,0 +1,45 @@ +import { app } from "../../scripts/app.js"; +import { api } from "../../scripts/api.js"; + +app.registerExtension({ + name: "Comfy.Pause.Signaler", + async beforeRegisterNodeDef(nodeType, nodeData, app) { + if (nodeData.name === "Signaler") { + console.log("[ComfyUI-Pause] βœ… Found Signaler node definition! Injecting buttons..."); + + const onNodeCreated = nodeType.prototype.onNodeCreated; + nodeType.prototype.onNodeCreated = function () { + const r = onNodeCreated ? onNodeCreated.apply(this, arguments) : undefined; + + // 1. Add the Pause Button + this.addWidget("button", "πŸ›‘ PAUSE", null, () => { + sendBridgeSignal("PAUSE"); + }); + + // 2. Add the Resume Button + this.addWidget("button", "▢️ RESUME", null, () => { + sendBridgeSignal("PROCEED"); + }); + + // 3. Force the node to be large enough to show buttons + this.setSize([200, 100]); + + console.log("[ComfyUI-Pause] πŸ› οΈ Buttons added to new Signaler node."); + return r; + }; + } + } +}); + +async function sendBridgeSignal(command) { + console.log(`[ComfyUI-Pause] πŸ“€ Sending: ${command}`); + try { + const response = await api.fetchApi("/comfy/steer", { + method: "POST", + headers: { "Content-Type": "application/json" }, + body: JSON.stringify({ command: command }), + }); + } catch (error) { + console.error("[ComfyUI-Pause] ❌ API Error:", error); + } +} diff --git a/pause_nodes.py b/pause_nodes.py new file mode 100644 index 0000000..a02887a --- /dev/null +++ b/pause_nodes.py @@ -0,0 +1,106 @@ +import torch +import time +import nodes +import comfy.samplers +import comfy.sample +import comfy.utils +import latent_preview +from .shared import PAUSE_STATE + +class Signaler: + @classmethod + def INPUT_TYPES(s): + return {"required": {}} + RETURN_TYPES = ("PAUSE_BUS",) + RETURN_NAMES = ("signal_bus",) + FUNCTION = "do_nothing" + CATEGORY = "Pause/Control" + def do_nothing(self): + return ("SIGNAL_BUS_ACTIVE",) + +class PSampler: + @classmethod + def INPUT_TYPES(s): + return { + "required": { + "model": ("MODEL",), + "seed": ("INT", {"default": 0, "min": 0, "max": 0xffffffffffffffff}), + "steps": ("INT", {"default": 20, "min": 1, "max": 10000}), + "cfg": ("FLOAT", {"default": 8.0, "min": 0.0, "max": 100.0}), + "sampler_name": (comfy.samplers.KSampler.SAMPLERS, ), + "scheduler": (comfy.samplers.KSampler.SCHEDULERS, ), + "positive": ("CONDITIONING", ), + "negative": ("CONDITIONING", ), + "latent_image": ("LATENT", ), + "denoise": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 1.0, "step": 0.01}), + }, + "optional": { + "signal_bus": ("PAUSE_BUS", ), + } + } + + RETURN_TYPES = ("LATENT",) + FUNCTION = "sample" + CATEGORY = "Pause/Sampling" + + def sample(self, model, seed, steps, cfg, sampler_name, scheduler, positive, negative, latent_image, denoise=1.0, signal_bus=None): + + latent_samples = latent_image["samples"] + disable_noise = False + + if disable_noise: + noise = torch.zeros(latent_samples.size(), dtype=latent_samples.dtype, layout=latent_samples.layout, device="cpu") + else: + batch_inds = latent_image.get("batch_index") if "batch_index" in latent_image else None + noise = comfy.sample.prepare_noise(latent_samples, seed, batch_inds) + + noise_mask = None + if "noise_mask" in latent_image: + noise_mask = latent_image["noise_mask"] + + preview_callback = latent_preview.prepare_callback(model, steps, latent_image) + + # --- REFINED CALLBACK WITH BETTER LOGGING --- + def intercept_callback(step, x0, x, total_steps): + # 'step' is 0-indexed (0, 1, 2...). + # So step=0 is actually the 1st step. + human_step = step + 1 + + # 1. Update the UI Preview first (Show the result of the step we just finished) + if preview_callback: + preview_callback(step, x0, x, total_steps) + + # 2. Check Bridge + command = PAUSE_STATE.get("command") + if command == "PAUSE": + print(f"\n[PSampler] πŸ›‘ PAUSE TRIGGERED.") + print(f"[PSampler] βœ… Finished Step {human_step}/{total_steps}.") + print(f"[PSampler] ⏸️ Holding execution... (Next up: Step {human_step + 1})") + + # Wait Loop + while PAUSE_STATE.get("command") == "PAUSE": + time.sleep(0.1) + + print(f"[PSampler] ▢️ RESUME SIGNAL RECEIVED.") + print(f"[PSampler] πŸš€ Proceeding to calculate Step {human_step + 1}...\n") + + # Execute + try: + samples = comfy.sample.sample( + model, noise, steps, cfg, sampler_name, scheduler, + positive, negative, latent_samples, + denoise=denoise, + disable_noise=disable_noise, + start_step=None, + last_step=None, + force_full_denoise=False, + noise_mask=noise_mask, + callback=intercept_callback, + seed=seed + ) + except comfy.model_management.InterruptProcessingException: + return (latent_image,) + + out = latent_image.copy() + out['samples'] = samples + return (out,) diff --git a/shared.py b/shared.py new file mode 100644 index 0000000..6250050 --- /dev/null +++ b/shared.py @@ -0,0 +1,9 @@ +import threading + +# The Lock ensures that the API and the GPU don't fight over this variable +state_lock = threading.Lock() + +PAUSE_STATE = { + "command": "PROCEED", + "payload": None +}