diff --git a/node_list.json b/node_list.json index 2d416b5..014dcd0 100644 --- a/node_list.json +++ b/node_list.json @@ -2,6 +2,9 @@ "Animation Builder (mtb)": "Convenient way to manage basic animation maths at the core of many of my workflows", "Any To String (mtb)": "Tries to take any input and convert it to a string", "Batch Float (mtb)": "Generates a batch of float values with interpolation", + "Batch Float Assemble (mtb)": "Assembles mutiple batches of floats into a single stream (batch)", + "Batch Float Fill (mtb)": "Fills a batch float with a single value until it reaches the target length", + "Batch Make (mtb)": "Simply duplicates the input frame as a batch", "Batch Shape (mtb)": "Generates a batch of 2D shapes with optional shading (experimental)", "Batch Transform (mtb)": "Transform a batch of images using a batch of keyframes", "Bbox (mtb)": "The bounding box (BBOX) custom type used by other nodes", diff --git a/nodes/batch.py b/nodes/batch.py index 7bbdb9e..ade068f 100644 --- a/nodes/batch.py +++ b/nodes/batch.py @@ -19,6 +19,29 @@ def hex_to_rgb(hex_color, bgr=False): return tuple(int(hex_color[i : i + 2], 16) for i in (0, 2, 4)) +class BatchMake: + """Simply duplicates the input frame as a batch""" + + @classmethod + def INPUT_TYPES(cls): + return { + "required": { + "image": ("IMAGE",), + "count": ("INT", {"default": 1}), + } + } + + RETURN_TYPES = ("IMAGE",) + FUNCTION = "generate_batch" + CATEGORY = "mtb/batch" + + def generate_batch(self, image: torch.Tensor, count): + if len(image.shape) == 3: + image = image.unsqueeze(0) + + return (image.repeat(count, 1, 1, 1),) + + class BatchShape: """Generates a batch of 2D shapes with optional shading (experimental)""" @@ -112,6 +135,60 @@ class BatchShape: return (pil2tensor(res),) +class BatchFloatFill: + """Fills a batch float with a single value until it reaches the target length""" + + @classmethod + def INPUT_TYPES(cls): + return { + "required": { + "floats": ("FLOATS",), + "direction": (["head", "tail"], {"default": "tail"}), + "value": ("FLOAT", {"default": 0.0}), + "count": ("INT", {"default": 1}), + } + } + + FUNCTION = "fill_floats" + RETURN_TYPES = ("FLOATS",) + CATEGORY = "mtb/batch" + + def fill_floats(self, floats, direction, value, count): + size = len(floats) + if size > count: + raise ValueError(f"Size ({size}) is less then target count ({count})") + + rem = count - size + if direction == "tail": + floats = floats + [value] * rem + else: + floats = [value] * rem + floats + return (floats,) + + +class BatchFloatAssemble: + """Assembles mutiple batches of floats into a single stream (batch)""" + + @classmethod + def INPUT_TYPES(cls): + return {"required": {"reverse": ("BOOLEAN", {"default": False})}} + + FUNCTION = "assemble_floats" + RETURN_TYPES = ("FLOATS",) + CATEGORY = "mtb/batch" + + def assemble_floats(self, reverse, **kwargs): + res = [] + if reverse: + for x in reversed(kwargs.values()): + res += x + else: + for x in kwargs.values(): + res += x + + return (res,) + + class BatchFloat: """Generates a batch of float values with interpolation""" @@ -257,4 +334,11 @@ class Batch2dTransform: return (torch.cat(res, dim=0),) -__nodes__ = [BatchFloat, Batch2dTransform, BatchShape] +__nodes__ = [ + BatchFloat, + Batch2dTransform, + BatchShape, + BatchMake, + BatchFloatAssemble, + BatchFloatFill, +] diff --git a/web/comfy_shared.js b/web/comfy_shared.js index e784a14..c034c87 100644 --- a/web/comfy_shared.js +++ b/web/comfy_shared.js @@ -92,7 +92,27 @@ export function getWidgetType(config) { } return { type, linkType } } +export const setupDynamicConnections = (nodeType, prefix, inputType) => { + const onNodeCreated = nodeType.prototype.onNodeCreated + nodeType.prototype.onNodeCreated = function () { + const r = onNodeCreated ? onNodeCreated.apply(this, arguments) : undefined + this.addInput(`${prefix}_1`, inputType) + return r + } + const onConnectionsChange = nodeType.prototype.onConnectionsChange + nodeType.prototype.onConnectionsChange = function ( + type, + index, + connected, + link_info + ) { + const r = onConnectionsChange + ? onConnectionsChange.apply(this, arguments) + : undefined + dynamic_connection(this, index, connected, `${prefix}_`, inputType) + } +} export const dynamic_connection = ( node, index, diff --git a/web/mtb_widgets.js b/web/mtb_widgets.js index 287ab4f..3f147a2 100644 --- a/web/mtb_widgets.js +++ b/web/mtb_widgets.js @@ -878,27 +878,17 @@ const mtb_widgets = { break } case 'Stack Images (mtb)': { - const onNodeCreated = nodeType.prototype.onNodeCreated - nodeType.prototype.onNodeCreated = function () { - const r = onNodeCreated - ? onNodeCreated.apply(this, arguments) - : undefined - this.addInput(`image_1`, 'IMAGE') - return r - } + shared.setupDynamicConnections(nodeType, 'image', 'IMAGE') + + break + } + case 'Batch Float Assemble (mtb)': { + shared.setupDynamicConnections(nodeType, 'floats', 'FLOATS') + break + } + case 'Batch Merge (mtb)': { + shared.setupDynamicConnections(nodeType, 'batches', 'IMAGE') - const onConnectionsChange = nodeType.prototype.onConnectionsChange - nodeType.prototype.onConnectionsChange = function ( - type, - index, - connected, - link_info - ) { - const r = onConnectionsChange - ? onConnectionsChange.apply(this, arguments) - : undefined - shared.dynamic_connection(this, index, connected, 'image_', 'IMAGE') - } break } case 'Save Tensors (mtb)': {