From bca4306bcb7cab5ebdccef438b515bd763bdfbf9 Mon Sep 17 00:00:00 2001 From: Alexopus Date: Tue, 23 Dec 2025 17:32:49 +0100 Subject: [PATCH] Add RandomShapeGenerator --- __init__.py | 3 +- nodes_literals.py | 204 ++++++++++++++++++++++++++++++++++++++++++++++ 2 files changed, 206 insertions(+), 1 deletion(-) diff --git a/__init__.py b/__init__.py index e93d49a..5b6b2c0 100644 --- a/__init__.py +++ b/__init__.py @@ -1,7 +1,7 @@ from typing import Any from .nodes import ImageSaver, ImageSaverSimple, ImageSaverMetadata -from .nodes_literals import SeedGenerator, StringLiteral, SizeLiteral, IntLiteral, FloatLiteral, CfgLiteral, ConditioningConcatOptional +from .nodes_literals import SeedGenerator, StringLiteral, SizeLiteral, IntLiteral, FloatLiteral, CfgLiteral, ConditioningConcatOptional, RandomShapeGenerator from .nodes_loaders import CheckpointLoaderWithName, UNETLoaderWithName from .nodes_selectors import SamplerSelector, SchedulerSelector, SchedulerSelectorInspire, SchedulerSelectorEfficiency, InputParameters, AnyToString, WorkflowInputValue from .civitai_nodes import CivitaiHashFetcher @@ -26,6 +26,7 @@ NODE_CLASS_MAPPINGS: dict[str, Any] = { "Int Literal (Image Saver)": IntLiteral, "Float Literal (Image Saver)": FloatLiteral, "Conditioning Concat Optional (Image Saver)": ConditioningConcatOptional, + "RandomShapeGenerator": RandomShapeGenerator, "Civitai Hash Fetcher (Image Saver)": CivitaiHashFetcher, } diff --git a/nodes_literals.py b/nodes_literals.py index 2d2cc8d..6ab02a4 100644 --- a/nodes_literals.py +++ b/nodes_literals.py @@ -3,6 +3,11 @@ from typing import Any from nodes import MAX_RESOLUTION import torch +import numpy as np +from PIL import Image, ImageDraw +import random +import math + class SeedGenerator: RETURN_TYPES = ("INT",) OUTPUT_TOOLTIPS = ("seed (INT)",) @@ -152,3 +157,202 @@ class ConditioningConcatOptional: out.append(n) return (out,) + +class RandomShapeGenerator: + """ + A ComfyUI node that generates images with random shapes. + """ + + @classmethod + def INPUT_TYPES(cls) -> dict[str, Any]: + return { + "required": { + "width": ("INT", { "default": 512, "min": 64, "max": 4096, "step": 64, "tooltip": "Width of the generated image in pixels" }), + "height": ("INT", { "default": 512, "min": 64, "max": 4096, "step": 64, "tooltip": "Height of the generated image in pixels" }), + "bg_color": (["random", "white", "black", "red", "green", "blue", "yellow", "cyan", "magenta"], { "tooltip": "Background color preset or random" }), + "fg_color": (["random", "black", "white", "red", "green", "blue", "yellow", "cyan", "magenta"], { "tooltip": "Foreground shape color preset or random" }), + "shape_type": (["random", "circle", "oval", "triangle", "square", "rectangle", "rhombus", "pentagon", "hexagon"], { "tooltip": "Type of shape to generate or random" }), + "seed": ("INT", { "default": 0, "min": 0, "max": 0xffffffffffffffff, "control_after_generate": True, "tooltip": "Random seed for reproducible shape generation" }), + }, + "optional": { + "bg_color_override": ("STRING", { "default": "", "multiline": False, "tooltip": "Override background color with hex (#AABBCC) or RGB(r, g, b) format" }), + "fg_color_override": ("STRING", { "default": "", "multiline": False, "tooltip": "Override foreground color with hex (#AABBCC) or RGB(r, g, b) format" }), + }, + } + + RETURN_TYPES = ("IMAGE", "STRING", "STRING") + RETURN_NAMES = ("image", "bg_rgb", "fg_rgb") + OUTPUT_TOOLTIPS = ("Generated image with random shape", "Background color as RGB/hex", "Foreground color as RGB/hex") + FUNCTION = "generate_shape" + CATEGORY = "image/generators" + DESCRIPTION = "Generates images with random shapes for testing and prototyping" + + def __init__(self): + self.color_map = { + "white": (255, 255, 255), + "black": (0, 0, 0), + "red": (255, 0, 0), + "green": (0, 255, 0), + "blue": (0, 0, 255), + "yellow": (255, 255, 0), + "cyan": (0, 255, 255), + "magenta": (255, 0, 255), + } + + def parse_rgb_string(self, rgb_str: str) -> tuple[int, int, int] | None: + """Parse RGB string like 'RGB(123, 45, 67)' or '#AABBCC' into tuple (123, 45, 67)""" + if not rgb_str or rgb_str.strip() == "": + return None + + rgb_str = rgb_str.strip() + + try: + # Try hex format first (#AABBCC or AABBCC) + if rgb_str.startswith("#"): + hex_str = rgb_str[1:] + else: + hex_str = rgb_str + + # Check if it's a valid hex string (6 characters) + if len(hex_str) == 6 and all(c in '0123456789ABCDEFabcdef' for c in hex_str): + r = int(hex_str[0:2], 16) + g = int(hex_str[2:4], 16) + b = int(hex_str[4:6], 16) + return (r, g, b) + + # Try RGB(r, g, b) format + rgb_str_upper = rgb_str.upper() + if rgb_str_upper.startswith("RGB(") and rgb_str_upper.endswith(")"): + values = rgb_str[4:-1].split(",") + r, g, b = [int(v.strip()) for v in values] + # Validate range + if all(0 <= val <= 255 for val in [r, g, b]): + return (r, g, b) + except (ValueError, IndexError): + return None + + return None + + def draw_shape(self, draw: ImageDraw.ImageDraw, img_width: int, img_height: int, shape_type: str, shape_color: tuple[int, int, int]) -> None: + """Draw a random shape on the image.""" + + # Random size - prefer larger sizes (40-70% of image dimensions) + size_factor = random.uniform(0.4, 0.7) + shape_width = int(img_width * size_factor) + shape_height = int(img_height * size_factor) + + # Random position (ensure shape stays fully within bounds) + x = random.randint(0, max(0, img_width - shape_width)) + y = random.randint(0, max(0, img_height - shape_height)) + + # Draw the shape based on type + if shape_type == 'circle': + # Make it a perfect circle using the minimum dimension + radius = min(shape_width, shape_height) // 2 + draw.ellipse([x, y, x + radius * 2, y + radius * 2], fill=shape_color) + + elif shape_type == 'oval': + draw.ellipse([x, y, x + shape_width, y + shape_height], fill=shape_color) + + elif shape_type == 'square': + # Make it a perfect square + side = min(shape_width, shape_height) + draw.rectangle([x, y, x + side, y + side], fill=shape_color) + + elif shape_type == 'rectangle': + draw.rectangle([x, y, x + shape_width, y + shape_height], fill=shape_color) + + elif shape_type == 'triangle': + # Equilateral-ish triangle + points = [ + (x + shape_width // 2, y), # top + (x, y + shape_height), # bottom left + (x + shape_width, y + shape_height) # bottom right + ] + draw.polygon(points, fill=shape_color) + + elif shape_type == 'rhombus': + # Diamond shape + points = [ + (x + shape_width // 2, y), # top + (x + shape_width, y + shape_height // 2), # right + (x + shape_width // 2, y + shape_height), # bottom + (x, y + shape_height // 2) # left + ] + draw.polygon(points, fill=shape_color) + + elif shape_type == 'pentagon': + # Regular pentagon + cx, cy = x + shape_width // 2, y + shape_height // 2 + radius = min(shape_width, shape_height) // 2 + points = [] + for i in range(5): + angle = i * 2 * math.pi / 5 - math.pi / 2 + px = cx + radius * math.cos(angle) + py = cy + radius * math.sin(angle) + points.append((px, py)) + draw.polygon(points, fill=shape_color) + + elif shape_type == 'hexagon': + # Regular hexagon + cx, cy = x + shape_width // 2, y + shape_height // 2 + radius = min(shape_width, shape_height) // 2 + points = [] + for i in range(6): + angle = i * 2 * math.pi / 6 + px = cx + radius * math.cos(angle) + py = cy + radius * math.sin(angle) + points.append((px, py)) + draw.polygon(points, fill=shape_color) + + def generate_shape(self, width: int, height: int, bg_color: str, fg_color: str, shape_type: str, seed: int, bg_color_override: str = "", fg_color_override: str = "") -> tuple[torch.Tensor, str, str]: + """Generate an image with a random shape.""" + + # Set random seed for reproducibility + random.seed(seed) + + # Get colors from map or generate random RGB values + # Check for override first + bg_override = self.parse_rgb_string(bg_color_override) + if bg_override is not None: + bg_rgb = bg_override + elif bg_color == "random": + bg_rgb = (random.randint(0, 255), random.randint(0, 255), random.randint(0, 255)) + else: + bg_rgb = self.color_map.get(bg_color, (255, 255, 255)) + + fg_override = self.parse_rgb_string(fg_color_override) + if fg_override is not None: + fg_rgb = fg_override + elif fg_color == "random": + fg_rgb = (random.randint(0, 255), random.randint(0, 255), random.randint(0, 255)) + else: + fg_rgb = self.color_map.get(fg_color, (0, 0, 0)) + + # Create image + img = Image.new('RGB', (width, height), bg_rgb) + draw = ImageDraw.Draw(img) + + # Select shape type + if shape_type == "random": + shapes = ['circle', 'oval', 'triangle', 'square', 'rectangle', 'rhombus', 'pentagon', 'hexagon'] + selected_shape = random.choice(shapes) + else: + selected_shape = shape_type + + # Draw the shape + self.draw_shape(draw, width, height, selected_shape, fg_rgb) + + # Convert PIL Image to torch tensor (ComfyUI format) + # ComfyUI expects images in format [batch, height, width, channels] with values 0-1 + img_array = np.array(img).astype(np.float32) / 255.0 + img_tensor = torch.from_numpy(img_array)[None,] + + # Format RGB values as strings for output (both formats) + bg_hex = f"#{bg_rgb[0]:02X}{bg_rgb[1]:02X}{bg_rgb[2]:02X}" + fg_hex = f"#{fg_rgb[0]:02X}{fg_rgb[1]:02X}{fg_rgb[2]:02X}" + bg_rgb_str = f"RGB({bg_rgb[0]}, {bg_rgb[1]}, {bg_rgb[2]}) / {bg_hex}" + fg_rgb_str = f"RGB({fg_rgb[0]}, {fg_rgb[1]}, {fg_rgb[2]}) / {fg_hex}" + + return (img_tensor, bg_rgb_str, fg_rgb_str) +