diff --git a/examples/ic-light-simple-light-source.png b/examples/ic-light-simple-light-source.png index e76efd8..e4fe567 100644 Binary files a/examples/ic-light-simple-light-source.png and b/examples/ic-light-simple-light-source.png differ diff --git a/nodes.py b/nodes.py index 3c14680..bd12058 100644 --- a/nodes.py +++ b/nodes.py @@ -162,6 +162,79 @@ To use the "opt_background" input, you also need to use the n = [t[0], d] c.append(n) out.append(c) + return (out[0], out[1], negative, out_latent) + +### Light Source +import numpy as np +from enum import Enum +from nodes import MAX_RESOLUTION + +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 generate_gradient_image(width:int, height:int, color_rgb: tuple, multiplier: float, lightPosition:LightPosition): + """ + Generate a gradient image with a light source effect. + + Parameters: + width (int): Width of the image. + height (int): Height of the image. + color_rgb: Color RGB of the image. + multiplier: weight of light. + lightPosition (str): Position of the light source. + It can be 'Left Light', 'Right Light', 'Top Light', 'Bottom Light', + 'Top Left Light', 'Top Right Light', 'Bottom Left Light', 'Bottom Right Light'. + + Returns: + np.array: 2D gradient image array. + """ + if lightPosition == LightPosition.LEFT: + gradient = np.tile(np.linspace(1, 0, width), (height, 1)) + elif lightPosition == LightPosition.RIGHT: + gradient = np.tile(np.linspace(0, 1, width), (height, 1)) + elif lightPosition == LightPosition.TOP: + gradient = np.tile(np.linspace(1, 0, height), (width, 1)).T + elif lightPosition == LightPosition.BOTTOM: + gradient = np.tile(np.linspace(0, 1, height), (width, 1)).T + elif lightPosition == LightPosition.TOP_LEFT: + x = np.linspace(1, 0, width) + y = np.linspace(1, 0, height) + x_mesh, y_mesh = np.meshgrid(x, y) + gradient = (x_mesh + y_mesh) / 2 + elif lightPosition == LightPosition.TOP_RIGHT: + x = np.linspace(0, 1, width) + y = np.linspace(1, 0, height) + x_mesh, y_mesh = np.meshgrid(x, y) + gradient = (x_mesh + y_mesh) / 2 + elif lightPosition == LightPosition.BOTTOM_LEFT: + x = np.linspace(1, 0, width) + y = np.linspace(0, 1, height) + x_mesh, y_mesh = np.meshgrid(x, y) + gradient = (x_mesh + y_mesh) / 2 + elif lightPosition == LightPosition.BOTTOM_RIGHT: + x = np.linspace(0, 1, width) + y = np.linspace(0, 1, height) + x_mesh, y_mesh = np.meshgrid(x, y) + gradient = (x_mesh + y_mesh) / 2 + else: + raise ValueError("Unsupported position. Choose from 'Left Light', 'Right Light', 'Top Light', 'Bottom Light','Top Left Light', 'Top Right Light', 'Bottom Left Light', 'Bottom Right Light'.") + + gradient = gradient * multiplier + gradient_x = gradient * color_rgb[0] + gradient_y = gradient * color_rgb[1] + gradient_z = gradient * color_rgb[2] + gradient = [gradient_x, gradient_y, gradient_z] + gradient = np.stack(gradient, axis=-1).astype(np.uint8) + + return gradient + return (out[0], out[1], {"samples": out_latent}) class LightSource: @@ -170,6 +243,8 @@ class LightSource: return { "required": { "light_position": (["Left Light", "Right Light", "Top Light", "Bottom Light",'Top Left Light', 'Top Right Light', 'Bottom Left Light', 'Bottom Right Light'],), + "multiplier": ("FLOAT", {"default": 0.5, "min": 0.0, "max": 1.0, "step": 0.001}), + "color": ("STRING", {"default": "#ffffff"}), "width": ("INT", { "default": 512, "min": 0, "max": MAX_RESOLUTION, "step": 8, }), "height": ("INT", { "default": 512, "min": 0, "max": MAX_RESOLUTION, "step": 8, }), "multiplier": ("FLOAT", { "default": 1.0, "min": 0.0, "max": 100.0, "step": 0.01, }), @@ -187,17 +262,16 @@ as a simple light source. The color can be specified in RGB or hex format. """ - def execute(self, width, height, light_position, multiplier, color): + def execute(self, light_position, multiplier, color, width, height): if color.startswith('#') and len(color) == 7: # e.g. "#RRGGBB" - r = int(color[1:3], 16) - g = int(color[3:5], 16) - b = int(color[5:7], 16) + color_hex = color.lstrip('#') + color_rgb =tuple(int(color_hex[i:i+2], 16) for i in (0, 2, 4)) else: - r, g, b = map(int, color.split(',')) + color_rgb = tuple(int(i) for i in color.split(',')) lightPosition = LightPosition(light_position) - image = generate_gradient_image(width, height, lightPosition) - image = image * [r / 255.0, g / 255.0, b / 255.0] + image = generate_gradient_image(width, height, color_rgb, multiplier, lightPosition) + # Convert a numpy array to a tensor and scale its values from 0-255 to 0-1 image = image.astype(np.float32) / 255.0 image = image * multiplier