import torch import random import numpy as np from enum import Enum from nodes import MAX_RESOLUTION # Images generated by combination of different creases need to be normalized def normalize_img(img): return (img / img.max()) * 255 # cross = True if cross crease is wanted def create_uneven_array(low, up, steps, spacing=1, cross=False): span = up - low dx = 1.0 / steps if cross: arr = np.array([low + (i*dx)**spacing*span for i in range(steps//2)]) return np.append(arr, arr[::-1]) else : arr = np.array([low + (i*dx)**spacing*span for i in range(steps)]) return arr def parabolic_crease(spacing, c, scale=100, corner=1, resolution = 1000): """ Parameters: spacing = controls how close the intermediate values will be to lower value c = higher the c more spread out the gradient will be scale = lesser the scale more concentrated is gradient towards the corner """ img = np.zeros((resolution, resolution)) # Varying the scaling parameter of create_uneven_array will give the parabolic gradient transition for i in range(resolution): img[i] = create_uneven_array(255, 0, resolution, spacing + c*i/scale) if corner == 1: return img elif corner == 2: return img[::-1] elif corner == 3: return img.T else: return img.T[::-1] # If cross=1, then cross crease else linear crease is returned def cross_crease(spacing, cross=1, resolution = 1000): a = create_uneven_array(255, 0, resolution, spacing, cross=True) img = np.tile(a, (resolution, 1)) return normalize_img(img*img.T) if cross else img # Final function to return some random crease from 8 different types def custom_crease(): spacing = random.uniform(1, 1.5) scale = random.randint(100, 300) corner = random.randint(1, 4) # constant determines the type of crease and also is used to scale spacing in parabolic_crease constant = random.randint(1, 10) # Returning those creases which are based on parabolic parabolic = parabolic_crease(spacing, constant, scale, corner) if constant == 1: return parabolic elif constant == 2: return normalize_img(parabolic*parabolic.T*parabolic[::-1]*parabolic.T[::-1]) # Returning those creases which are based on parabolic and cross cross = cross_crease(spacing) if constant == 3: return cross elif constant == 4: return normalize_img(parabolic * cross) elif constant == 5: return normalize_img(cross * parabolic * parabolic.T) # Returning those creases which are based on parabolic and linear linear = cross_crease(spacing, 0) if constant == 6: return linear elif constant == 7: return linear.T else: return normalize_img(linear * parabolic) class LightPosition(Enum): LEFT = "Left Light" RIGHT = "Right Light" TOP = "Top Light" BOTTOM = "Bottom Light" TOP_LEFT = "Top Left Light" TOP_RIGHT = "Top Right Light" BOTTOM_LEFT = "Bottom Left Light" BOTTOM_RIGHT = "Bottom Right Light" def toRgb(color): if color.startswith('#') and len(color) == 7: # e.g. "#RRGGBB" color_rgb =tuple(int(color[i:i+2], 16) for i in (1, 3, 5)) else: # e.g. "255,255,255" color_rgb = tuple(int(i) for i in color.split(',')) return color_rgb def rgb_to_int(rgb): r, g, b = rgb return (r << 16) + (g << 8) + b def generate_gradient_image(width:int, height:int, start_color:tuple=(255,255,255), end_color:tuple=(0,0,0), multiplier:float=1.0, lightPosition:LightPosition=LightPosition.LEFT): """ Generate a gradient image with a light source effect. Parameters: width (int): Width of the image. height (int): Height of the image. start_color: Starting color RGB of the gradient. end_color: Ending color RGB of the gradient. multiplier: Weight of light. lightPosition (LightPosition): Position of the light source. Returns: np.array: 2D gradient image array. """ # Create a gradient from 0 to 1 and apply multiplier if lightPosition == LightPosition.LEFT: gradient = np.tile(np.linspace(0, 1, width)**multiplier, (height, 1)) elif lightPosition == LightPosition.RIGHT: gradient = np.tile(np.linspace(1, 0, width)**multiplier, (height, 1)) elif lightPosition == LightPosition.TOP: gradient = np.tile(np.linspace(0, 1, height)**multiplier, (width, 1)).T elif lightPosition == LightPosition.BOTTOM: gradient = np.tile(np.linspace(1, 0, height)**multiplier, (width, 1)).T elif lightPosition == LightPosition.BOTTOM_RIGHT: x = np.linspace(1, 0, width)**multiplier y = np.linspace(1, 0, height)**multiplier x_mesh, y_mesh = np.meshgrid(x, y) gradient = np.sqrt(x_mesh**2 + y_mesh**2) / np.sqrt(2.0) elif lightPosition == LightPosition.BOTTOM_LEFT: x = np.linspace(0, 1, width)**multiplier y = np.linspace(1, 0, height)**multiplier x_mesh, y_mesh = np.meshgrid(x, y) gradient = np.sqrt(x_mesh**2 + y_mesh**2) / np.sqrt(2.0) elif lightPosition == LightPosition.TOP_RIGHT: x = np.linspace(1, 0, width)**multiplier y = np.linspace(0, 1, height)**multiplier x_mesh, y_mesh = np.meshgrid(x, y) gradient = np.sqrt(x_mesh**2 + y_mesh**2) / np.sqrt(2.0) elif lightPosition == LightPosition.TOP_LEFT: x = np.linspace(0, 1, width)**multiplier y = np.linspace(0, 1, height)**multiplier x_mesh, y_mesh = np.meshgrid(x, y) gradient = np.sqrt(x_mesh**2 + y_mesh**2) / np.sqrt(2.0) else: raise ValueError(f"Unsupported position. Choose from {', '.join([member.value for member in LightPosition])}.") # Interpolate between start_color and end_color based on the gradient gradient_img = np.zeros((height, width, 3), dtype=np.float32) for i in range(3): gradient_img[..., i] = start_color[i] + (end_color[i] - start_color[i]) * gradient gradient_img = np.clip(gradient_img, 0, 255).astype(np.uint8) return gradient_img def numpy_to_tensor(array: np.ndarray) -> torch.Tensor: """Convert a numpy array to a tensor and scale its values from 0-255 to 0-1.""" array = array.astype(np.float32) / 255.0 return torch.from_numpy(array)[None,] class ImageGradient: @classmethod def INPUT_TYPES(s): return { "required": { "light_position": ([member.value for member in LightPosition],), "multiplier": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 1.0, "step": 0.001}), "start_color": ("STRING", {"default": "#FFFFFF"}), "end_color": ("STRING", {"default": "#000000"}), "width": ("INT", { "default": 512, "min": 0, "max": MAX_RESOLUTION, "step": 8, }), "height": ("INT", { "default": 512, "min": 0, "max": MAX_RESOLUTION, "step": 8, }) } } RETURN_TYPES = ("IMAGE",) FUNCTION = "execute" CATEGORY = "LightGradient" DESCRIPTION = """Simple Light Gradient""" def execute(self, light_position, multiplier, start_color, end_color, width, height): lightPosition = LightPosition(light_position) start_color_rgb = toRgb(start_color) end_color_rgb = toRgb(end_color) image = generate_gradient_image(width, height, start_color_rgb, end_color_rgb, multiplier, lightPosition) image = numpy_to_tensor(image) mask = generate_gradient_image(width, height, multiplier=multiplier, lightPosition=lightPosition) mask = numpy_to_tensor(mask) mask = mask[:, :, :, 0] return (image,mask,) class MaskGradient: @classmethod def INPUT_TYPES(s): return { "required": { "light_position": ([member.value for member in LightPosition],), "multiplier": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 1.0, "step": 0.001}), "width": ("INT", { "default": 512, "min": 0, "max": MAX_RESOLUTION, "step": 8, }), "height": ("INT", { "default": 512, "min": 0, "max": MAX_RESOLUTION, "step": 8, }) } } RETURN_TYPES = ("MASK",) FUNCTION = "execute" CATEGORY = "LightGradient" DESCRIPTION = """Mask Gradient""" def execute(self, light_position, multiplier, width, height): lightPosition = LightPosition(light_position) mask = generate_gradient_image(width, height, multiplier=multiplier, lightPosition=lightPosition) mask = numpy_to_tensor(mask) mask = mask[:, :, :, 0] return (mask,) NODE_CLASS_MAPPINGS = { "ImageGradient": ImageGradient, "MaskGradient": MaskGradient } NODE_DISPLAY_NAME_MAPPINGS = { "ImageGradient": "Image Gradient", "MaskGradient": "Mask Gradient" }