196 lines
10 KiB
Python
196 lines
10 KiB
Python
import uuid
|
|
from pathlib import Path
|
|
from aiohttp import web
|
|
import hashlib # For robust hashing of data
|
|
import torch # For hashing torch tensors
|
|
import json
|
|
import folder_paths
|
|
from server import PromptServer
|
|
from comfy_execution.graph import ExecutionBlocker
|
|
TEMP_SUBFOLDER = "minimal-pause-signals"
|
|
|
|
# Structure: { node_id: { 'data': any, 'data_hash': str } }
|
|
PAUSED_DATA = {}
|
|
|
|
class AlwaysEqualProxy(str):
|
|
def __eq__(self, _):
|
|
return True
|
|
|
|
def __ne__(self, _):
|
|
return False
|
|
|
|
any_type = AlwaysEqualProxy("*")
|
|
|
|
# --- Helper functions ---
|
|
def get_signal_file_path(node_id: str) -> Path:
|
|
"""Gets the path for this node's signal file."""
|
|
temp_dir = Path(folder_paths.get_temp_directory())
|
|
signal_dir = temp_dir / TEMP_SUBFOLDER
|
|
signal_dir.mkdir(parents=True, exist_ok=True)
|
|
return signal_dir / f"continue_{node_id}.signal"
|
|
|
|
def compute_data_hash(data) -> str:
|
|
"""
|
|
Computes a robust hash for the input data.
|
|
Prioritizes specific handling for common ComfyUI types like torch.Tensor.
|
|
Falls back to a string representation hash for other types.
|
|
Returns 'NOT_HASHABLE' if hashing fails significantly.
|
|
"""
|
|
try:
|
|
if isinstance(data, torch.Tensor):
|
|
# For torch.Tensor, convert to numpy array bytes for robust hashing
|
|
# This is generally the most reliable way to hash tensor content.
|
|
# Handle empty tensors to avoid tobytes() errors if data.numel() == 0
|
|
if data.numel() == 0:
|
|
return hashlib.sha256(f"EMPTY_TENSOR_SHAPE:{data.shape}_DTYPE:{data.dtype}".encode('utf-8')).hexdigest()
|
|
else:
|
|
return hashlib.sha256(data.cpu().numpy().tobytes()).hexdigest()
|
|
elif isinstance(data, (str, int, float, bool)):
|
|
# Basic types can be directly converted to string and encoded
|
|
return hashlib.sha256(str(data).encode('utf-8')).hexdigest()
|
|
elif isinstance(data, (list, tuple)):
|
|
# For lists/tuples, convert to a consistent JSON string representation
|
|
# This handles nested structures and ensures consistent order for keys if dicts are nested
|
|
# Note: This requires all elements to be JSON serializable.
|
|
return hashlib.sha256(json.dumps(data, sort_keys=True, default=str).encode('utf-8')).hexdigest()
|
|
elif isinstance(data, dict):
|
|
# For dictionaries, convert to a consistent JSON string representation
|
|
return hashlib.sha256(json.dumps(data, sort_keys=True, default=str).encode('utf-8')).hexdigest()
|
|
else:
|
|
# Fallback for other arbitrary objects: try to use repr() or str()
|
|
# Be aware: repr() can include memory addresses, making hashes inconsistent across runs.
|
|
# str() might not be unique enough for complex objects.
|
|
return hashlib.sha256(repr(data).encode('utf-8')).hexdigest()
|
|
|
|
except Exception as e:
|
|
# Catch any unexpected errors during hashing and log them
|
|
print(f"[WARN] Failed to compute hash for data type {type(data)}: {e}. Returning 'NOT_HASHABLE'.")
|
|
return 'NOT_HASHABLE'
|
|
|
|
# --- HTTP Route ---
|
|
if hasattr(PromptServer, 'instance') and PromptServer.instance is not None:
|
|
routes = PromptServer.instance.routes
|
|
@routes.post('/minimal_pause/signal_continue')
|
|
async def signal_continue_handler(request: web.Request) -> web.Response:
|
|
"""Creates the signal file to indicate continuation."""
|
|
try:
|
|
if request.content_type != 'application/json':
|
|
return web.json_response({"status": "error", "message": "Invalid content type"}, status=400)
|
|
data = await request.json()
|
|
node_id = data.get('node_id')
|
|
|
|
if not node_id:
|
|
return web.json_response({"status": "error", "message": "node_id missing"}, status=400)
|
|
|
|
signal_file = get_signal_file_path(node_id)
|
|
# --- START DEBUG HTTP ---
|
|
# --- END DEBUG HTTP ---
|
|
try:
|
|
signal_file.touch(exist_ok=True) # Create the signal file
|
|
# --- START DEBUG HTTP ---
|
|
|
|
# --- END DEBUG HTTP ---
|
|
|
|
if node_id not in PAUSED_DATA:
|
|
print(f"[INFO HTTP Handler] Node {node_id}: Not found in PAUSED_DATA (might have restarted or never paused in memory).")
|
|
else:
|
|
print(f"[INFO HTTP Handler] Node {node_id}: Found in PAUSED_DATA when signal received.")
|
|
|
|
return web.json_response({"status": "ok", "message": "Continue signal received and file created."})
|
|
except IOError as e:
|
|
print(f"[ERROR HTTP Handler] Node {node_id}: Error creating signal file: {e}")
|
|
return web.json_response({"status": "error", "message": f"Failed to create signal file: {e}"}, status=500)
|
|
|
|
except Exception as e:
|
|
print(f"[ERROR HTTP Handler] General error in signal_continue_handler: {e}")
|
|
return web.json_response({"status": "error", "message": "Internal server error"}, status=500)
|
|
|
|
|
|
class MinimalPauseNode:
|
|
"""A node that pauses workflow execution based on a filesystem signal, triggered by an HTTP signal,
|
|
or passes through if upstream data is unchanged."""
|
|
|
|
@classmethod
|
|
def IS_CHANGED(cls, *args, **kwargs):
|
|
# A unique ID for each execution to ensure UI update, but not for changing internal logic
|
|
return str(uuid.uuid4())
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": { "any": (any_type, {}), },
|
|
"hidden": { "prompt": "PROMPT", "extra_pnginfo": "EXTRA_PNGINFO", "unique_id": "UNIQUE_ID", },
|
|
}
|
|
|
|
RETURN_TYPES = (any_type,)
|
|
RETURN_NAMES = ("any",)
|
|
FUNCTION = "pause_or_continue"
|
|
CATEGORY = "🎤MW/MW-Audio-Tools"
|
|
OUTPUT_NODE = True
|
|
|
|
def pause_or_continue(self, any, unique_id: str, prompt=None, extra_pnginfo=None, **kwargs):
|
|
node_id = unique_id
|
|
signal_file = get_signal_file_path(node_id)
|
|
|
|
current_data_hash = compute_data_hash(any)
|
|
stored_data_hash = PAUSED_DATA.get(node_id, {}).get('data_hash')
|
|
should_continue_from_file = signal_file.exists()
|
|
if should_continue_from_file:
|
|
# Case 1: User explicitly clicked "Continue" via HTTP signal.
|
|
# This is the strongest signal; always resume and pass data.
|
|
resumed_data = None
|
|
if node_id in PAUSED_DATA:
|
|
resumed_data = PAUSED_DATA[node_id]['data']
|
|
print(f"[INFO Node {node_id}] Resuming (explicit continue) with original stored data from PAUSED_DATA.")
|
|
else:
|
|
# Fallback: If PAUSED_DATA cleared, use the current input `any`.
|
|
resumed_data = any
|
|
print(f"[WARN Node {node_id}] Resuming (explicit continue) - in-memory data lost. Using current input `any`.")
|
|
|
|
# Clean up the signal file and in-memory data
|
|
try:
|
|
if signal_file.exists():
|
|
signal_file.unlink()
|
|
print(f"[INFO Node {node_id}] Signal file unlinked.")
|
|
else:
|
|
print(f"[WARN Node {node_id}] Signal file not found during unlink attempt.")
|
|
except OSError as e:
|
|
print(f"[ERROR Node {node_id}] Error deleting signal file on resume: {e}")
|
|
|
|
if node_id in PAUSED_DATA:
|
|
del PAUSED_DATA[node_id]
|
|
print(f"[INFO Node {node_id}] PAUSED_DATA entry deleted.")
|
|
return {"result": (resumed_data,)}
|
|
|
|
else: # Case 2: No explicit continue signal. Decide based on data change.
|
|
|
|
if node_id not in PAUSED_DATA:
|
|
# Case 2a: First time hitting this node in this prompt execution, AND no signal file.
|
|
# This means it should pause. Store current data and its hash.
|
|
PAUSED_DATA[node_id] = { 'data': any, 'data_hash': current_data_hash }
|
|
print(f"[INFO Node {node_id}] Pausing for the first time. Storing data and hash: {current_data_hash}. PAUSED_DATA[{node_id}] now present.")
|
|
if hasattr(PromptServer, 'instance') and PromptServer.instance is not None:
|
|
PromptServer.instance.send_sync("minimal_pause_enable_button", { "node_id": node_id })
|
|
|
|
return {"result": (ExecutionBlocker(None),)}
|
|
else:
|
|
# Case 2b: Node was previously paused (node_id is in PAUSED_DATA), and no signal file.
|
|
# Check if the input data has changed since last pause.
|
|
if current_data_hash != stored_data_hash:
|
|
# Input data has changed. We need to re-pause with the new data.
|
|
PAUSED_DATA[node_id] = { 'data': any, 'data_hash': current_data_hash } # Update stored data and hash
|
|
print(f"[INFO Node {node_id}] Input data CHANGED ({current_data_hash} != {stored_data_hash}). Re-pausing with new data.")
|
|
if hasattr(PromptServer, 'instance') and PromptServer.instance is not None:
|
|
PromptServer.instance.send_sync("minimal_pause_enable_button", { "node_id": node_id })
|
|
|
|
return {"result": (ExecutionBlocker(None),)}
|
|
else:
|
|
# Input data has NOT changed. This node should "pass through" its previously stored output.
|
|
# This is crucial for allowing downstream nodes to execute when upstream is stable.
|
|
resumed_data = PAUSED_DATA[node_id]['data']
|
|
print(f"[INFO Node {node_id}] Input data UNCHANGED. Passing through previous output to allow downstream execution.")
|
|
if hasattr(PromptServer, 'instance') and PromptServer.instance is not None:
|
|
# Re-send enable signal in case UI missed it or state changed (still paused, but passing through)
|
|
PromptServer.instance.send_sync("minimal_pause_enable_button", { "node_id": node_id })
|
|
|
|
return {"result": (resumed_data,)} |