diff --git a/displacement_map_node.py b/displacement_map_node.py new file mode 100644 index 0000000..a90c453 --- /dev/null +++ b/displacement_map_node.py @@ -0,0 +1,55 @@ +import torch +import numpy as np +from PIL import Image +import torchvision.transforms as transforms + +class ExtractDisplacementMap: + @classmethod + def INPUT_TYPES(cls): + return { + "required": { + "image": ("IMAGE",), # Input image (e.g., normal or bump map) + "intensity": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 10.0, "step": 0.1}), # Displacement intensity + } + } + + RETURN_TYPES = ("IMAGE",) + RETURN_NAMES = ("displacement_map",) + FUNCTION = "extract_displacement_map" + CATEGORY = "image/processing" + OUTPUT_NODE = False + + def extract_displacement_map(self, image, intensity): + # Convert ComfyUI image tensor to numpy array + if isinstance(image, torch.Tensor): + image = image.cpu().numpy() + if image.ndim == 4: + image = image[0] # Remove batch dimension if present + if image.shape[0] in [1, 3]: + image = np.transpose(image, (1, 2, 0)) # Convert (C, H, W) to (H, W, C) + image = (image * 255).astype(np.uint8) # Scale to 0-255 + + # Convert to grayscale + if image.shape[-1] == 3: + image = np.dot(image[..., :3], [0.2989, 0.5870, 0.1140]) # RGB to grayscale + else: + image = image[..., 0] # Use the first channel if already grayscale + + # Normalize and scale by intensity + displacement_map = image.astype(np.float32) / 255.0 + displacement_map = displacement_map * intensity + + # Convert to tensor and add batch dimension + displacement_map = torch.from_numpy(displacement_map).unsqueeze(0).unsqueeze(0) + displacement_map = displacement_map.to(torch.float32) + + return (displacement_map,) + +# Node mappings +NODE_CLASS_MAPPINGS = { + "ExtractDisplacementMap": ExtractDisplacementMap +} + +NODE_DISPLAY_NAME_MAPPINGS = { + "ExtractDisplacementMap": "Extract Displacement Map" +}