add simple light source

2024-05-10 17:53:14 +08:00
parent 3e2e30d6d9
commit 9263225e27
2 changed files with 95 additions and 2 deletions
@@ -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"
 }