diff --git a/Bridge.py b/Bridge.py index da76199..f853733 100644 --- a/Bridge.py +++ b/Bridge.py @@ -1,5 +1,3 @@ -cat = "XchanBik/Bridge" - import comfy.model_management as model_management import torch import numpy as np @@ -8,8 +6,13 @@ import sys import folder_paths import copy import json +import threading +import time import traceback from server import PromptServer +from nodes import ExecutionBlocker + +cat = "XchanBik/Bridge" class AnyType(str): def __ne__(self, __value: object) -> bool: @@ -17,19 +20,73 @@ class AnyType(str): any = AnyType("*") -# Global bridge storage +# Enhanced global bridge storage with execution control class BridgeStorage: def __init__(self): - self.storage = {} + self.storage = {} + self.completion_events = {} + self.lock = threading.Lock() + def store(self, name, data): - self.storage[name] = data + with self.lock: + self.storage[name] = data + # Signal completion if there are waiting loads + if name in self.completion_events: + self.completion_events[name].set() + def get(self, name): - return self.storage.get(name) + # Create an event if it doesn't exist + with self.lock: + if name not in self.completion_events: + self.completion_events[name] = threading.Event() + + # If data already exists, return it immediately + if name in self.storage: + return self.storage.get(name) + + # Otherwise return a special object that will block execution + return BridgeExecutionBlocker(name, self.completion_events[name]) + def has(self, name): return name in self.storage + + def reset(self): + with self.lock: + # Reset events but keep data for potential reuse + for event in self.completion_events.values(): + event.clear() + # Create a global bridge storage instance bridge_storage = BridgeStorage() +# Special execution blocker that works with ComfyUI's execution system +class BridgeExecutionBlocker(ExecutionBlocker): + def __init__(self, bridge_name, event): + self.bridge_name = bridge_name + self.completion_event = event + print(f"[SimpleBridge] Load '{bridge_name}' waiting for store...") + + def wait(self): + timeout = 300 # 5 minute timeout + while timeout > 0: + if self.completion_event.wait(1): # Wait up to 1 second + print(f"[SimpleBridge] Load '{self.bridge_name}' received data") + return True + timeout -= 1 + + print(f"[SimpleBridge] Error: Timeout waiting for '{self.bridge_name}' after 5 minutes") + return False + + def get_execution_blockers(self): + return [self] + +# Reset bridge events when workflow starts +@PromptServer.instance.routes.post("/prompt") +async def bridge_reset_on_workflow_start(request): + bridge_storage.reset() + # Let the original handler process the request + return await PromptServer.instance.orig_prompt_route(request) + # Store Bridge Node class SimpleBridgeStoreNode: @classmethod @@ -53,7 +110,7 @@ class SimpleBridgeStoreNode: print(f"[SimpleBridge] Stored value under '{bridge_name}': {type(value)}") else: print(f"[SimpleBridge] Tried to store None under '{bridge_name}'") - return (value,) + return () # Load Bridge Node class SimpleBridgeLoadNode: @@ -71,10 +128,12 @@ class SimpleBridgeLoadNode: CATEGORY = cat def load_value(self, bridge_name): - # Get the value from the bridge storage + # Get the value from the bridge storage - this may return a blocker value = bridge_storage.get(bridge_name) - if value is None: - print(f"[SimpleBridge] '{bridge_name}' has no stored value") + # If we get a normal value, just return it + if not isinstance(value, BridgeExecutionBlocker): + return (value,) + # If we get a blocker, ComfyUI will handle it properly return (value,)