diff --git a/__init__.py b/__init__.py index 1bcfe74..1312f2d 100644 --- a/__init__.py +++ b/__init__.py @@ -1,8 +1,10 @@ from .control_nodes import FloatControl, IntControl +from .quick_shape_mask import QuickShapeMask NODE_CLASS_MAPPINGS = { "FloatControl": FloatControl, - "IntControl": IntControl + "IntControl": IntControl, + "QuickShapeMask": QuickShapeMask } NODE_DISPLAY_NAME_MAPPINGS = { diff --git a/control_nodes.py b/control_nodes.py index 04106f0..360a7f5 100644 --- a/control_nodes.py +++ b/control_nodes.py @@ -3,10 +3,8 @@ import time class ValueControlBase: """Base class for float and integer control nodes""" - # Class level storage for persistent values between executions instances = {} - @classmethod def INPUT_TYPES(cls): """Base input types that are common to both Float and Int controls""" @@ -58,7 +56,6 @@ class ValueControlBase: def update_value_base(self, maximum_value, minimum_value, starting_value, steps_per_cycle, movement_type, always_execute=True): instance = self.__class__.instances[self.instance_id] - # Initialize if this is the first run if instance['current_value'] is None: instance['current_value'] = starting_value diff --git a/examples/shape_resize_example_API.json b/examples/shape_resize_example_API.json new file mode 100644 index 0000000..2e03f73 --- /dev/null +++ b/examples/shape_resize_example_API.json @@ -0,0 +1,110 @@ +{ + "206": { + "inputs": { + "image": "example.png", + "upload": "image" + }, + "class_type": "LoadImage", + "_meta": { + "title": "Load Image" + } + }, + "224": { + "inputs": { + "steps_per_cycle": 30, + "movement_type": "sine_wave", + "always_execute": true, + "maximum_value": 512, + "minimum_value": 1, + "starting_value": 1 + }, + "class_type": "IntControl", + "_meta": { + "title": "Int Control (RyanOnTheInside)" + } + }, + "227": { + "inputs": { + "x": 0, + "y": 0, + "resize_source": false, + "destination": [ + "228", + 0 + ], + "source": [ + "206", + 0 + ], + "mask": [ + "234", + 0 + ] + }, + "class_type": "ImageCompositeMasked", + "_meta": { + "title": "ImageCompositeMasked" + } + }, + "228": { + "inputs": { + "width": 512, + "height": 512, + "batch_size": 1, + "color": 0 + }, + "class_type": "EmptyImage", + "_meta": { + "title": "EmptyImage" + } + }, + "231": { + "inputs": { + "steps_per_cycle": 29, + "movement_type": "sine_wave", + "always_execute": true, + "maximum_value": 502, + "minimum_value": 1, + "starting_value": 474 + }, + "class_type": "IntControl", + "_meta": { + "title": "Int Control (RyanOnTheInside)" + } + }, + "233": { + "inputs": { + "filename_prefix": "ComfyUI", + "images": [ + "227", + 0 + ] + }, + "class_type": "SaveImage", + "_meta": { + "title": "Save Image" + } + }, + "234": { + "inputs": { + "shape": "circle", + "width": [ + "224", + 0 + ], + "height": [ + "231", + 0 + ], + "x": 256, + "y": 256, + "canvas_width": 512, + "canvas_height": 512, + "batch_size": 1 + }, + "class_type": "QuickShapeMask", + "_meta": { + "title": "QuickShapeMask" + } + } +} \ No newline at end of file diff --git a/examples/shape_resize_example_WORKFLOW.json b/examples/shape_resize_example_WORKFLOW.json new file mode 100644 index 0000000..bc3c641 --- /dev/null +++ b/examples/shape_resize_example_WORKFLOW.json @@ -0,0 +1,390 @@ +{ + "last_node_id": 234, + "last_link_id": 59, + "nodes": [ + { + "id": 224, + "type": "IntControl", + "pos": { + "0": 270, + "1": 350 + }, + "size": { + "0": 315, + "1": 178 + }, + "flags": {}, + "order": 0, + "mode": 0, + "inputs": [], + "outputs": [ + { + "name": "INT", + "type": "INT", + "links": [ + 57 + ], + "slot_index": 0, + "shape": 3 + } + ], + "properties": { + "Node name for S&R": "IntControl" + }, + "widgets_values": [ + 30, + "sine_wave", + true, + 512, + 1, + 1 + ] + }, + { + "id": 231, + "type": "IntControl", + "pos": { + "0": 270, + "1": 590 + }, + "size": { + "0": 315, + "1": 178 + }, + "flags": {}, + "order": 1, + "mode": 0, + "inputs": [], + "outputs": [ + { + "name": "INT", + "type": "INT", + "links": [ + 58 + ], + "slot_index": 0, + "shape": 3 + } + ], + "properties": { + "Node name for S&R": "IntControl" + }, + "widgets_values": [ + 29, + "sine_wave", + true, + 502, + 1, + 474 + ] + }, + { + "id": 218, + "type": "FloatControl", + "pos": { + "0": 274, + "1": 919 + }, + "size": { + "0": 390.5999755859375, + "1": 178 + }, + "flags": {}, + "order": 2, + "mode": 4, + "inputs": [], + "outputs": [ + { + "name": "FLOAT", + "type": "FLOAT", + "links": [], + "slot_index": 0, + "shape": 3 + } + ], + "properties": { + "Node name for S&R": "FloatControl" + }, + "widgets_values": [ + 30, + "sine_wave", + true, + 1, + 1, + 1 + ] + }, + { + "id": 206, + "type": "LoadImage", + "pos": { + "0": 765, + "1": 262 + }, + "size": { + "0": 315, + "1": 314 + }, + "flags": {}, + "order": 3, + "mode": 0, + "inputs": [], + "outputs": [ + { + "name": "IMAGE", + "type": "IMAGE", + "links": [ + 47 + ], + "slot_index": 0, + "shape": 3 + }, + { + "name": "MASK", + "type": "MASK", + "links": null, + "shape": 3 + } + ], + "properties": { + "Node name for S&R": "LoadImage" + }, + "widgets_values": [ + "example.png", + "image" + ] + }, + { + "id": 227, + "type": "ImageCompositeMasked", + "pos": { + "0": 1278, + "1": 316 + }, + "size": { + "0": 315, + "1": 146 + }, + "flags": {}, + "order": 6, + "mode": 0, + "inputs": [ + { + "name": "destination", + "type": "IMAGE", + "link": 46 + }, + { + "name": "source", + "type": "IMAGE", + "link": 47 + }, + { + "name": "mask", + "type": "MASK", + "link": 59 + } + ], + "outputs": [ + { + "name": "IMAGE", + "type": "IMAGE", + "links": [ + 56 + ], + "slot_index": 0, + "shape": 3 + } + ], + "properties": { + "Node name for S&R": "ImageCompositeMasked" + }, + "widgets_values": [ + 0, + 0, + false + ] + }, + { + "id": 233, + "type": "SaveImage", + "pos": { + "0": 1684, + "1": 317 + }, + "size": [ + 315, + 270 + ], + "flags": {}, + "order": 7, + "mode": 0, + "inputs": [ + { + "name": "images", + "type": "IMAGE", + "link": 56 + } + ], + "outputs": [], + "properties": {}, + "widgets_values": [ + "ComfyUI" + ] + }, + { + "id": 228, + "type": "EmptyImage", + "pos": { + "0": 751, + "1": 58 + }, + "size": { + "0": 315, + "1": 130 + }, + "flags": {}, + "order": 4, + "mode": 0, + "inputs": [], + "outputs": [ + { + "name": "IMAGE", + "type": "IMAGE", + "links": [ + 46 + ], + "slot_index": 0, + "shape": 3 + } + ], + "properties": { + "Node name for S&R": "EmptyImage" + }, + "widgets_values": [ + 512, + 512, + 1, + 0 + ] + }, + { + "id": 234, + "type": "QuickShapeMask", + "pos": { + "0": 770, + "1": 642 + }, + "size": [ + 315, + 226 + ], + "flags": {}, + "order": 5, + "mode": 0, + "inputs": [ + { + "name": "width", + "type": "INT", + "link": 57, + "widget": { + "name": "width" + } + }, + { + "name": "height", + "type": "INT", + "link": 58, + "widget": { + "name": "height" + } + } + ], + "outputs": [ + { + "name": "MASK", + "type": "MASK", + "links": [ + 59 + ], + "shape": 3 + } + ], + "properties": { + "Node name for S&R": "QuickShapeMask" + }, + "widgets_values": [ + "circle", + 64, + 64, + 256, + 256, + 512, + 512, + 1 + ] + } + ], + "links": [ + [ + 46, + 228, + 0, + 227, + 0, + "IMAGE" + ], + [ + 47, + 206, + 0, + 227, + 1, + "IMAGE" + ], + [ + 56, + 227, + 0, + 233, + 0, + "IMAGE" + ], + [ + 57, + 224, + 0, + 234, + 0, + "INT" + ], + [ + 58, + 231, + 0, + 234, + 1, + "INT" + ], + [ + 59, + 234, + 0, + 227, + 2, + "MASK" + ] + ], + "groups": [], + "config": {}, + "extra": { + "ds": { + "scale": 1.0834705943389067, + "offset": [ + -150.8890623486404, + -198.28032954985966 + ] + } + }, + "version": 0.4 +} \ No newline at end of file diff --git a/quick_shape_mask.py b/quick_shape_mask.py new file mode 100644 index 0000000..b823e2e --- /dev/null +++ b/quick_shape_mask.py @@ -0,0 +1,111 @@ +import numpy as np +import torch + +class QuickShapeMask: + """A node that quickly generates basic shape masks""" + + @classmethod + def INPUT_TYPES(cls): + return { + "required": { + "shape": (["circle", "square"], { + "default": "circle", + "tooltip": "The shape of the mask to generate" + }), + "width": ("INT", { + "default": 64, + "min": 1, + "max": 4096, + "step": 1, + "tooltip": "Width of the shape in pixels" + }), + "height": ("INT", { + "default": 64, + "min": 1, + "max": 4096, + "step": 1, + "tooltip": "Height of the shape in pixels" + }), + "x": ("INT", { + "default": 256, + "min": 0, + "max": 4096, + "step": 1, + "tooltip": "X position of the shape center (0 = left edge)" + }), + "y": ("INT", { + "default": 256, + "min": 0, + "max": 4096, + "step": 1, + "tooltip": "Y position of the shape center (0 = top edge)" + }), + "canvas_width": ("INT", { + "default": 512, + "min": 1, + "max": 4096, + "step": 1, + "tooltip": "Width of the output mask" + }), + "canvas_height": ("INT", { + "default": 512, + "min": 1, + "max": 4096, + "step": 1, + "tooltip": "Height of the output mask" + }), + "batch_size": ("INT", { + "default": 1, + "min": 1, + "max": 64, + "step": 1, + "tooltip": "Number of identical masks to generate" + }) + } + } + + RETURN_TYPES = ("MASK",) + FUNCTION = "generate_mask" + CATEGORY = "mask" + + DESCRIPTION = "Generates a mask containing a basic shape (circle or square) with high performance" + + def generate_mask(self, shape, width, height, x, y, canvas_width, canvas_height, batch_size): + # Create empty mask + mask = np.zeros((canvas_height, canvas_width), dtype=np.float32) + + # Calculate boundaries + half_width = width // 2 + half_height = height // 2 + + # Calculate shape boundaries + left = max(0, x - half_width) + right = min(canvas_width, x + half_width) + top = max(0, y - half_height) + bottom = min(canvas_height, y + half_height) + + if shape == "square": + # Simple square mask + mask[top:bottom, left:right] = 1.0 + + else: # circle + # Create coordinate grids for the region of interest + Y, X = np.ogrid[top:bottom, left:right] + + # Calculate distances from center for the region + dist_x = (X - x) + dist_y = (Y - y) + + # Create circle mask using distance formula + circle_mask = (dist_x**2 / (width/2)**2 + dist_y**2 / (height/2)**2) <= 1 + + # Apply circle to the region + mask[top:bottom, left:right][circle_mask] = 1.0 + + # Convert to torch tensor and add batch dimension (BHW format) + mask_tensor = torch.from_numpy(mask) + + # Expand to requested batch size + mask_tensor = mask_tensor.unsqueeze(0).expand(batch_size, -1, -1) + + return (mask_tensor,) \ No newline at end of file