diff --git a/examples/ic-light-simple-light-source.png b/examples/ic-light-simple-light-source.png new file mode 100644 index 0000000..e76efd8 Binary files /dev/null and b/examples/ic-light-simple-light-source.png differ diff --git a/nodes.py b/nodes.py index f918320..2e8adef 100644 --- a/nodes.py +++ b/nodes.py @@ -161,11 +161,104 @@ To use the "opt_background" input, you also need to use the 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, lightPosition:LightPosition): + """ + Generate a gradient image with a light source effect. + + Parameters: + width (int): Width of the image. + height (int): Height of the image. + 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(255, 0, width), (height, 1)) + elif lightPosition == LightPosition.RIGHT: + gradient = np.tile(np.linspace(0, 255, width), (height, 1)) + elif lightPosition == LightPosition.TOP: + gradient = np.tile(np.linspace(255, 0, height), (width, 1)).T + elif lightPosition == LightPosition.BOTTOM: + gradient = np.tile(np.linspace(0, 255, height), (width, 1)).T + elif lightPosition == LightPosition.TOP_LEFT: + x = np.linspace(255, 0, width) + y = np.linspace(255, 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, 255, width) + y = np.linspace(255, 0, height) + x_mesh, y_mesh = np.meshgrid(x, y) + gradient = (x_mesh + y_mesh) / 2 + elif lightPosition == LightPosition.BOTTOM_LEFT: + x = np.linspace(255, 0, width) + y = np.linspace(0, 255, height) + x_mesh, y_mesh = np.meshgrid(x, y) + gradient = (x_mesh + y_mesh) / 2 + elif lightPosition == LightPosition.BOTTOM_RIGHT: + x = np.linspace(0, 255, width) + y = np.linspace(0, 255, 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 = np.stack((gradient,) * 3, axis=-1).astype(np.uint8) + + return gradient + + +class LightSource: + @classmethod + def INPUT_TYPES(s): + return { + "required": { + "light_position": (["Left Light", "Right Light", "Top Light", "Bottom Light",'Top Left Light', 'Top Right Light', 'Bottom Left Light', 'Bottom Right Light'],), + "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",) + RETURN_NAMES = ("IMAGE",) + FUNCTION = "execute" + CATEGORY = "IC-Light" + DESCRIPTION = """Simple Light Source""" + + def execute(self, width, height, light_position): + lightPosition = LightPosition(light_position) + image = generate_gradient_image(width, height, 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 = torch.from_numpy(image)[None,] + return (image,) + + NODE_CLASS_MAPPINGS = { "LoadAndApplyICLightUnet": LoadAndApplyICLightUnet, - "ICLightConditioning": ICLightConditioning + "ICLightConditioning": ICLightConditioning, + "LightSource": LightSource } NODE_DISPLAY_NAME_MAPPINGS = { "LoadAndApplyICLightUnet": "Load And Apply IC-Light", - "ICLightConditioning": "IC-Light Conditioning" + "ICLightConditioning": "IC-Light Conditioning", + "LightSource": "Simple Light Source" } \ No newline at end of file