Spaces:
Runtime error
Runtime error
File size: 2,445 Bytes
f2dbf59 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 |
import hashlib
import os
class BatchManagerNode:
def __init__(self, frames_per_batch=-1):
print("BatchNode init")
self.frames_per_batch = frames_per_batch
self.inputs = {}
self.outputs = {}
self.unique_id = None
self.has_closed_inputs = False
self.total_frames = float('inf')
def reset(self):
print("BatchNode reset")
self.close_inputs()
for key in self.outputs:
if getattr(self.outputs[key][-1], "gi_suspended", False):
try:
self.outputs[key][-1].send(None)
except StopIteration:
pass
self.__init__(self.frames_per_batch)
def has_open_inputs(self):
return len(self.inputs) > 0
def close_inputs(self):
for key in self.inputs:
if getattr(self.inputs[key][-1], "gi_suspended", False):
try:
self.inputs[key][-1].send(1)
except StopIteration:
pass
self.inputs = {}
@classmethod
def INPUT_TYPES(s):
return {
"required": {"frames_per_batch": ("INT", {"default": 16, "min": 1, "max": 128, "step": 1})},
"hidden": {"prompt": "PROMPT", "unique_id": "UNIQUE_ID"},
}
RETURN_TYPES = ("BatchManager",)
RETURN_NAMES = ("meta_batch",)
CATEGORY = "tbox/Video"
FUNCTION = "update_batch"
def update_batch(self, frames_per_batch, prompt=None, unique_id=None):
if unique_id is not None and prompt is not None:
requeue = prompt[unique_id]['inputs'].get('requeue', 0)
else:
requeue = 0
print(f'update_batch >> unique_id: {unique_id}; requeue: {requeue}')
if requeue == 0:
self.reset()
self.frames_per_batch = frames_per_batch
self.unique_id = unique_id
else:
num_batches = (self.total_frames+self.frames_per_batch-1)//frames_per_batch
print(f'Meta-Batch {requeue}/{num_batches}')
#onExecuted seems to not be called unless some message is sent
return (self,)
@classmethod
def IS_CHANGED(self, frames_per_batch, prompt=None, unique_id=None):
print(f"BatchManagerNode >>> IS_CHANGED : {result}")
random_bytes = os.urandom(32)
result = hashlib.sha256(random_bytes).hexdigest()
return result
|