diff --git a/comfy/ldm/colormap.py b/comfy/ldm/colormap.py new file mode 100644 index 000000000..1f4d88bd9 --- /dev/null +++ b/comfy/ldm/colormap.py @@ -0,0 +1,25 @@ +"""Colormap utilities for depth and geometry visualisation.""" + +from __future__ import annotations + +import torch + + +def turbo(x: torch.Tensor) -> torch.Tensor: + """Anton Mikhailov polynomial approximation of the Turbo colormap. + + Args: + x: Float tensor with values in [0, 1]. + + Returns: + RGB tensor of the same shape as ``x`` with a trailing size-3 dimension. + """ + x = x.clamp(0.0, 1.0) + x2 = x * x + x3 = x2 * x + x4 = x2 * x2 + x5 = x4 * x + r = 0.13572138 + 4.61539260*x - 42.66032258*x2 + 132.13108234*x3 - 152.94239396*x4 + 59.28637943*x5 + g = 0.09140261 + 2.19418839*x + 4.84296658*x2 - 14.18503333*x3 + 4.27729857*x4 + 2.82956604*x5 + b = 0.10667330 + 12.64194608*x - 60.58204836*x2 + 110.36276771*x3 - 89.90310912*x4 + 27.34824973*x5 + return torch.stack([r, g, b], dim=-1).clamp(0.0, 1.0) diff --git a/comfy_extras/nodes_depth_anything_3.py b/comfy_extras/nodes_depth_anything_3.py index 5b2929fdf..de258a3c2 100644 --- a/comfy_extras/nodes_depth_anything_3.py +++ b/comfy_extras/nodes_depth_anything_3.py @@ -33,6 +33,7 @@ import torch import comfy.model_management as mm import comfy.sd import folder_paths +from comfy.ldm.colormap import turbo as _turbo from comfy.ldm.depth_anything_3 import preprocess as da3_preprocess from comfy_api.latest import ComfyExtension, io @@ -150,7 +151,7 @@ class DepthAnything3Inference(io.ComfyNode): @classmethod def define_schema(cls): return io.Schema( - node_id="DepthAnything3", + node_id="DepthAnything3Inference", search_aliases=["depth", "geometry", "da3", "depth anything", "monocular", "pointmap", "sky", "3d", "metric depth", "disparity"], display_name="Run Depth Anything 3", category="image/geometry_estimation", @@ -321,6 +322,8 @@ class DepthAnything3Inference(io.ComfyNode): return io.NodeOutput(geometry) + + class DepthAnything3Render(io.ComfyNode): """Visualise a DA3_GEOMETRY packet as a single image. @@ -328,6 +331,19 @@ class DepthAnything3Render(io.ComfyNode): Use multiple nodes in parallel to get depth + sky + confidence simultaneously. """ + _DEPTH_RENDER_INPUTS = [ + io.Combo.Input("normalization", + options=["v2_style", "min_max", "raw"], + default="v2_style", + tooltip="'v2_style': mean/std normalisation for perceptually balanced results (default). " + "'min_max': stretches the full depth range to [0, 1] for maximum contrast. " + "'raw': no scaling — preserves metric units for DA3-Metric-Large."), + io.Boolean.Input("apply_sky_clip", default=False, + tooltip="Clip sky-region depth to the 99th percentile of foreground depth before " + "normalisation. Requires a 'sky' tensor in the geometry " + "(DA3-Mono-Large or DA3-Metric-Large); raises an error otherwise."), + ] + @classmethod def define_schema(cls): return io.Schema( @@ -338,22 +354,13 @@ class DepthAnything3Render(io.ComfyNode): inputs=[ DA3Geometry.Input("geometry"), io.DynamicCombo.Input("output", - tooltip="depth: normalised depth image. " + tooltip="depth: normalised greyscale depth image. " + "depth_colored: depth mapped through the Turbo colormap. " "sky_mask: sky probability in [0, 1] (Mono/Metric variants only). " "confidence: normalised depth confidence (Small/Base variants only).", options=[ - io.DynamicCombo.Option("depth", [ - io.Combo.Input("normalization", - options=["v2_style", "min_max", "raw"], - default="v2_style", - tooltip="'v2_style': mean/std normalisation for perceptually balanced results (default). " - "'min_max': stretches the full depth range to [0, 1] for maximum contrast. " - "'raw': no scaling — preserves metric units for DA3-Metric-Large."), - io.Boolean.Input("apply_sky_clip", default=False, - tooltip="Clip sky-region depth to the 99th percentile of foreground depth before " - "normalisation. Requires a 'sky' tensor in the geometry " - "(DA3-Mono-Large or DA3-Metric-Large); raises an error otherwise."), - ]), + io.DynamicCombo.Option("depth", cls._DEPTH_RENDER_INPUTS), + io.DynamicCombo.Option("depth_colored", cls._DEPTH_RENDER_INPUTS), io.DynamicCombo.Option("sky_mask", []), io.DynamicCombo.Option("confidence", []), ]), @@ -365,7 +372,7 @@ class DepthAnything3Render(io.ComfyNode): def execute(cls, geometry, output) -> io.NodeOutput: output_val = output["output"] - if output_val == "depth": + if output_val in ("depth", "depth_colored"): normalization = output["normalization"] apply_sky_clip = output["apply_sky_clip"] if apply_sky_clip and "sky" not in geometry: @@ -380,7 +387,8 @@ class DepthAnything3Render(io.ComfyNode): da3_preprocess.apply_sky_aware_clip(depth[i], sky[i]) for i in range(depth.shape[0]) ], dim=0) - result = cls._depth_to_image(depth, sky, normalization) + grey = cls._depth_to_image(depth, sky, normalization) # (B,H,W,3) greyscale + result = _turbo(grey[..., 0]) if output_val == "depth_colored" else grey elif output_val == "sky_mask": if "sky" not in geometry: diff --git a/comfy_extras/nodes_moge.py b/comfy_extras/nodes_moge.py index 0f5d4ffd6..3c6b41822 100644 --- a/comfy_extras/nodes_moge.py +++ b/comfy_extras/nodes_moge.py @@ -9,6 +9,7 @@ import folder_paths from comfy_api.latest import ComfyExtension, Types, io from typing_extensions import override +from comfy.ldm.colormap import turbo as _turbo from comfy.ldm.moge.model import MoGeModel from comfy.ldm.moge.geometry import triangulate_grid_mesh from comfy.ldm.moge.panorama import get_panorama_cameras, split_panorama_image, merge_panorama_depth, spherical_uv_to_directions, _uv_grid @@ -31,18 +32,6 @@ DA3Geometry = io.Custom("DA3_GEOMETRY") # "image": torch.Tensor (B, H, W, 3) in [0, 1], CPU (always present) -def _turbo(x: torch.Tensor) -> torch.Tensor: - """Anton Mikhailov polynomial approximation of the turbo colormap.""" - x = x.clamp(0.0, 1.0) - x2 = x * x - x3 = x2 * x - x4 = x2 * x2 - x5 = x4 * x - r = 0.13572138 + 4.61539260*x - 42.66032258*x2 + 132.13108234*x3 - 152.94239396*x4 + 59.28637943*x5 - g = 0.09140261 + 2.19418839*x + 4.84296658*x2 - 14.18503333*x3 + 4.27729857*x4 + 2.82956604*x5 - b = 0.10667330 + 12.64194608*x - 60.58204836*x2 + 110.36276771*x3 - 89.90310912*x4 + 27.34824973*x5 - return torch.stack([r, g, b], dim=-1).clamp(0.0, 1.0) - def _normals_from_points(points: torch.Tensor) -> torch.Tensor: """Camera-space surface normals from a (B, H, W, 3) point map (v1 fallback)."""