diff --git a/__init__.py b/__init__.py index 36108bd..9771c11 100644 --- a/__init__.py +++ b/__init__.py @@ -1,36 +1,31 @@ import server from aiohttp import web -from .pause_nodes import PSampler, Signaler +from .pause_nodes import PSampler, PSamplerAdvanced 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" } - """ +# API Route +@server.PromptServer.instance.routes.post("/comfy/pause_signal") +async def set_pause_command(request): 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}) + print(f"[P-Sampler] 📡 Signal Received: {command}") + return web.json_response({"status": "success"}) -# --- Node Mappings --- +# Node Mappings NODE_CLASS_MAPPINGS = { "PSampler": PSampler, - "Signaler": Signaler + "PSamplerAdvanced": PSamplerAdvanced } NODE_DISPLAY_NAME_MAPPINGS = { - "PSampler": "⏸️ P-Sampler (Pause Enabled)", - "Signaler": "🚦 Signaler (Control)" + "PSampler": "⏸️ Pausable Sampler", + "PSamplerAdvanced": "⏸️ Pausable Sampler (Advanced)" } WEB_DIRECTORY = "./js" -__all__ = ["NODE_CLASS_MAPPINGS", "NODE_DISPLAY_NAME_MAPPINGS", "PAUSE_STATE", "WEB_DIRECTORY"] +__all__ = ["NODE_CLASS_MAPPINGS", "NODE_DISPLAY_NAME_MAPPINGS", "WEB_DIRECTORY"] diff --git a/js/pausing_bridge.js b/js/pausing_bridge.js index aa93cb6..a90c2d0 100644 --- a/js/pausing_bridge.js +++ b/js/pausing_bridge.js @@ -2,44 +2,42 @@ import { app } from "../../scripts/app.js"; import { api } from "../../scripts/api.js"; app.registerExtension({ - name: "Comfy.Pause.Signaler", + name: "Comfy.Pause.Control", async beforeRegisterNodeDef(nodeType, nodeData, app) { - if (nodeData.name === "Signaler") { - console.log("[ComfyUI-Pause] ✅ Found Signaler node definition! Injecting buttons..."); - + // Target both Standard and Advanced nodes + if (nodeData.name === "PSampler" || nodeData.name === "PSamplerAdvanced") { 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"); - }); + // 1. Control Deck + this.addWidget("button", "▶️ RESUME", null, () => sendSignal("PROCEED")); + this.addWidget("button", "⏸️ PAUSE", null, () => sendSignal("PAUSE")); - // 2. Add the Resume Button - this.addWidget("button", "▶️ RESUME", null, () => { - sendBridgeSignal("PROCEED"); - }); + // 2. Precision CFG Fix + const cfgWidget = this.widgets.find(w => w.name === "cfg"); + if (cfgWidget) { + cfgWidget.step = 0.1; + if (!cfgWidget.options) cfgWidget.options = {}; + cfgWidget.options.step = 0.1; + cfgWidget.options.precision = 1; + } - // 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."); + this.setSize([300, nodeData.name === "PSamplerAdvanced" ? 440 : 280]); return r; }; } } }); -async function sendBridgeSignal(command) { - console.log(`[ComfyUI-Pause] 📤 Sending: ${command}`); +async function sendSignal(command) { try { - const response = await api.fetchApi("/comfy/steer", { + await api.fetchApi("/comfy/pause_signal", { method: "POST", headers: { "Content-Type": "application/json" }, body: JSON.stringify({ command: command }), }); } catch (error) { - console.error("[ComfyUI-Pause] ❌ API Error:", error); + console.error("[P-Sampler] API Error:", error); } } diff --git a/pause_nodes.py b/pause_nodes.py index a02887a..873f944 100644 --- a/pause_nodes.py +++ b/pause_nodes.py @@ -1,24 +1,16 @@ -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",) - +# --- NODE 1: STANDARD PAUSABLE SAMPLER --- class PSampler: + """ + The Pausable Sampler (Standard). + Features: Pause & Resume. In-Memory Only. + """ @classmethod def INPUT_TYPES(s): return { @@ -26,72 +18,45 @@ class PSampler: "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}), + "cfg": ("FLOAT", {"default": 8.0, "min": 0.0, "max": 100.0, "step": 0.1, "round": 0.01}), "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" + CATEGORY = "Pausable/Engine" - def sample(self, model, seed, steps, cfg, sampler_name, scheduler, positive, negative, latent_image, denoise=1.0, signal_bus=None): + def sample(self, model, seed, steps, cfg, sampler_name, scheduler, positive, negative, latent_image, denoise): 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"] - + 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 = latent_image.get("noise_mask", None) 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 + human_step = step + 1 + if preview_callback: preview_callback(step, x0, x, total_steps) - # 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": + if PAUSE_STATE.get("command") == "PAUSE": + print(f"[P-Sampler] ⏸️ PAUSED at Step {human_step}. Waiting...") + 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") + print(f"[P-Sampler] ▶️ RESUMED.") - # 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, + disable_noise=False, + start_step=0, last_step=None, force_full_denoise=False, noise_mask=noise_mask, @@ -99,7 +64,83 @@ class PSampler: seed=seed ) except comfy.model_management.InterruptProcessingException: - return (latent_image,) + print("[P-Sampler] 🛑 System Interrupt. Aborting.") + raise + + out = latent_image.copy() + out['samples'] = samples + return (out,) + +# --- NODE 2: ADVANCED PAUSABLE SAMPLER --- +class PSamplerAdvanced: + """ + The Pausable Sampler (Advanced). + Exposes manual control over noise and steps. + """ + @classmethod + def INPUT_TYPES(s): + return { + "required": { + "model": ("MODEL",), + "add_noise": (["enable", "disable"], ), + "noise_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, "step": 0.1, "round": 0.01}), + "sampler_name": (comfy.samplers.KSampler.SAMPLERS, ), + "scheduler": (comfy.samplers.KSampler.SCHEDULERS, ), + "positive": ("CONDITIONING", ), + "negative": ("CONDITIONING", ), + "latent_image": ("LATENT", ), + "start_at_step": ("INT", {"default": 0, "min": 0, "max": 10000}), + "end_at_step": ("INT", {"default": 10000, "min": 0, "max": 10000}), + "return_with_leftover_noise": (["disable", "enable"], ), + } + } + + RETURN_TYPES = ("LATENT",) + FUNCTION = "sample" + CATEGORY = "Pausable/Engine" + + def sample(self, model, add_noise, noise_seed, steps, cfg, sampler_name, scheduler, positive, negative, latent_image, start_at_step, end_at_step, return_with_leftover_noise): + + force_full_denoise = (return_with_leftover_noise == "disable") + disable_noise = (add_noise == "disable") + latent_samples = latent_image["samples"] + + if not disable_noise: + batch_inds = latent_image.get("batch_index") if "batch_index" in latent_image else None + noise = comfy.sample.prepare_noise(latent_samples, noise_seed, batch_inds) + else: + noise = torch.zeros(latent_samples.size(), dtype=latent_samples.dtype, layout=latent_samples.layout, device="cpu") + + noise_mask = latent_image.get("noise_mask", None) + preview_callback = latent_preview.prepare_callback(model, steps, latent_image) + + def intercept_callback(step, x0, x, total_steps): + human_step = step + 1 + if preview_callback: preview_callback(step, x0, x, total_steps) + + if PAUSE_STATE.get("command") == "PAUSE": + print(f"[P-Sampler] ⏸️ PAUSED at Step {human_step}. Waiting...") + while PAUSE_STATE.get("command") == "PAUSE": + time.sleep(0.1) + print(f"[P-Sampler] ▶️ RESUMED.") + + try: + samples = comfy.sample.sample( + model, noise, steps, cfg, sampler_name, scheduler, + positive, negative, latent_samples, + disable_noise=disable_noise, + start_step=start_step, + last_step=end_at_step, + force_full_denoise=force_full_denoise, + noise_mask=noise_mask, + callback=intercept_callback, + seed=noise_seed + ) + except comfy.model_management.InterruptProcessingException: + print("[P-Sampler] 🛑 System Interrupt. Aborting.") + raise out = latent_image.copy() out['samples'] = samples