Merge pull request #1 from huagetai/master

first commit
2024-05-11 02:17:37 +08:00
parent fb214d4911 67cf0dd79b
commit 8cacc514a2
5 changed files with 302 additions and 0 deletions
@@ -0,0 +1,18 @@
+__pycache__/
+*.py[cod]
+/output/
+/input/
+!/input/example.png
+/models/
+/temp/
+/custom_nodes/
+!custom_nodes/example_node.py.example
+extra_model_paths.yaml
+/.vs
+.idea/
+venv/
+/web/extensions/*
+!/web/extensions/logging.js.example
+!/web/extensions/core/
+/tests-ui/data/object_info.json
+/user/
@@ -0,0 +1,44 @@
+# ComfyUI Light Gradient
+
+A set of nodes for ComfyUI Light Gradient.  
+
+
+![image](example/workflow.png)  
+<font size="1">*this workflow (workflow.png) is in the example directory.  </font><br /> 
+
+
+## How to install 
+
+* Recommended use ComfyUI Manager for installation.
+
+* Or open the cmd window in the plugin directory of ComfyUI, like ```ComfyUI\custom_nodes```，type    
+```
+git clone https://github.com/huagetai/ComfyUI_LightGradient.git
+```
+* Or download the zip file and extracted, copy the resulting folder to ```ComfyUI\custom_ Nodes```    
+
+* Restart ComfyUI.
+
+
+## Nodes
+
+1. **Image Gradient Node**
+* input
+  * width - width of Mask
+  * height - height of Mask
+  * light_position - position of Light
+  * multiplier - strength of Light
+  * start_color - start color of gradient
+  * end_color - end color of gradient
+* output
+  * IMAGE
+
+2. **Mask Gradient Node**
+* input
+  * width - width of Mask
+  * height - height of Mask
+  * light_position - position of Light
+  * multiplier - strength of Light
+
+* output
+  * Mask
@@ -0,0 +1,2 @@
+from .nodes import NODE_CLASS_MAPPINGS, NODE_DISPLAY_NAME_MAPPINGS
+__all__ = ["NODE_CLASS_MAPPINGS", "NODE_DISPLAY_NAME_MAPPINGS"]
@@ -0,0 +1,238 @@
+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"
+}