From 7c4ed7ec494fd61847fd59199db8a005e67cfa12 Mon Sep 17 00:00:00 2001
From: cerspense <cerspense@gmail.com>
Date: Wed, 31 Jul 2024 23:25:44 -0700
Subject: [PATCH] added depth to normal node

---
 README.md   | 21 +++++++++++++++++-
 cspnodes.py | 64 +++++++++++++++++++++++++++++++++++++++++++++++++++--
 2 files changed, 82 insertions(+), 3 deletions(-)

diff --git a/README.md b/README.md
index 5c13b0e..7da3b8e 100644
--- a/README.md
+++ b/README.md
@@ -14,6 +14,7 @@ This package contains a collection of custom nodes for ComfyUI, designed to enha
 7. [RemapRange](#remaprange)
 8. [ResizeByImage](#resizebyimage)
 9. [IncrementEveryN](#incrementeveryn)
+10. [DepthToNormalMap](#depthtormalmap)
 
 ## TextFileLineIterator
 
@@ -159,4 +160,22 @@ This node divides the input value by the step size, adds the offset, and returns
   - Input values 12-17 will output 12
   - And so on...
 
-This node is useful for creating slower-changing values from rapidly incrementing inputs, which can be helpful in various animation and procedural generation scenarios. The offset parameter allows for further customization of the output range.
\ No newline at end of file
+This node is useful for creating slower-changing values from rapidly incrementing inputs, which can be helpful in various animation and procedural generation scenarios. The offset parameter allows for further customization of the output range.
+
+## DepthToNormalMap
+
+This node converts depth maps to normal maps, with options to control intensity and axis flipping.
+
+### Parameters:
+- `depth_maps` (IMAGE): Input depth map image(s).
+- `normal_intensity` (FLOAT, default: 14.0, range: 0.01 to 100.0): Intensity of the normal map effect. Note that the actual intensity is 10 times this value internally.
+- `flip_x` (BOOLEAN, default: True): Whether to flip the X-axis of the normal map.
+- `flip_y` (BOOLEAN, default: False): Whether to flip the Y-axis of the normal map.
+
+### Output:
+- (IMAGE): The generated normal map(s).
+
+### Behavior:
+This node takes depth map images as input and converts them to normal maps. The conversion process involves calculating gradients in the X and Y directions, then using these to create a 3D normal vector for each pixel. The `normal_intensity` parameter controls the strength of the effect, with higher values resulting in more pronounced normal maps. The `flip_x` and `flip_y` options allow for adjusting the orientation of the normal map to match different coordinate systems or depth map conventions.
+
+Note that the actual intensity applied is 10 times the input `normal_intensity` value, allowing for a wide range of effect strengths. The node can handle batches of images, processing multiple depth maps in a single operation.
\ No newline at end of file
diff --git a/cspnodes.py b/cspnodes.py
index 808a66e..e64cd72 100644
--- a/cspnodes.py
+++ b/cspnodes.py
@@ -1,14 +1,12 @@
 import os
 import torch
 from diffusers import DiffusionPipeline, DPMSolverMultistepScheduler
-from diffusers.utils import export_to_video
 from PIL import Image, ImageOps
 import numpy as np
 import random
 import torch.nn.functional as F
 import glob
 
-
 class TextFileLineIterator:
     @classmethod
     def INPUT_TYPES(cls):
@@ -418,9 +416,70 @@ class IncrementEveryN:
     def increment_every_n(self, input_value, step_size, offset):
         output_value = (input_value // step_size) + offset
         return (output_value,)
+    
+import torch
+import torch.nn.functional as F
+
+class DepthToNormalMap:
+    @classmethod
+    def INPUT_TYPES(cls):
+        return {
+            "required": {
+                "depth_maps": ("IMAGE",),
+                "normal_intensity": ("FLOAT", {"default": 14.0, "min": 0.01, "max": 100.0, "step": 0.01}),
+                "flip_x": ("BOOLEAN", {"default": True}),
+                "flip_y": ("BOOLEAN", {"default": False}),
+            }
+        }
+
+    RETURN_TYPES = ("IMAGE",)
+    FUNCTION = "convert_depth_to_normal"
+    CATEGORY = "cspnodes"
+
+    def convert_depth_to_normal(self, depth_maps, normal_intensity, flip_x, flip_y):
+        # Ensure depth_maps is a float tensor and normalize to [0, 1]
+        depth_maps = depth_maps.float()
+        if depth_maps.max() > 1.0:
+            depth_maps = depth_maps / 255.0
+
+        # Extract only the first channel if the input has multiple channels
+        if depth_maps.shape[-1] > 1:
+            depth_maps = depth_maps[..., 0].unsqueeze(-1)
+
+        # Compute gradients
+        grad_y, grad_x = torch.gradient(depth_maps[..., 0], dim=(1, 2))
+
+        # Reshape gradients to match input shape
+        grad_x = grad_x.unsqueeze(-1)
+        grad_y = grad_y.unsqueeze(-1)
+
+        # Apply 10x stronger intensity
+        intensity = normal_intensity * 10
+
+        # Create normal map
+        normal_maps = torch.cat([grad_x * intensity, 
+                                 grad_y * intensity, 
+                                 torch.ones_like(grad_x)], dim=-1)
+
+        # Normalize
+        normal_maps = F.normalize(normal_maps, p=2, dim=-1)
+
+        # Flip X axis if requested
+        if flip_x:
+            normal_maps[..., 0] *= -1
+
+        # Invert Y axis if requested
+        if flip_y:
+            normal_maps[..., 1] *= -1
+
+        # Scale to [0, 1] range
+        normal_maps = (normal_maps + 1) / 2
+
+        return (normal_maps,)
 
     
 NODE_CLASS_MAPPINGS = {
+    "DepthToNormalMap": DepthToNormalMap,
     "IncrementEveryN": IncrementEveryN,
     "ResizeByImage": ResizeByImage,
     "SplitImageChannels": SplitImageChannels,
@@ -433,6 +492,7 @@ NODE_CLASS_MAPPINGS = {
 }
 
 NODE_DISPLAY_NAME_MAPPINGS = {
+    "DepthToNormalMap": "Depth to Normal Map",
     "IncrementEveryN": "Increment Every N",
     "ResizeByImage": "Resize By Image",
     "SplitImageChannels": "Split Image Channels",