feat: Unified controller with the sampler.
- Also added a version for the Adv. KSampler
This commit is contained in:
Brian Blankenship
+11
-16
@@ -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"]
|
||||
|
||||
+18
-20
@@ -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);
|
||||
}
|
||||
}
|
||||
|
||||
+96
-55
@@ -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
|
||||
|
||||
Reference in New Issue
Block a user