204 lines
6.9 KiB
Python
204 lines
6.9 KiB
Python
"""ComfyUI node backend for Depth to Layers.
|
|
|
|
This module provides:
|
|
- Node runtime logic for the `depth to layers` node.
|
|
- Interval parsing/validation and grayscale-to-layer color mapping.
|
|
- PromptServer event payload used by the frontend editor to initialize previews.
|
|
|
|
The node converts image intensity to discrete color layers using ordered
|
|
intervals in the normalized range [0, 1].
|
|
"""
|
|
|
|
import base64
|
|
import json
|
|
from io import BytesIO
|
|
|
|
import numpy as np
|
|
import torch
|
|
from PIL import Image
|
|
from server import PromptServer # type: ignore
|
|
|
|
from .constants import VGA_256_COLORS
|
|
|
|
|
|
def tensor_to_gray_uint8(image_tensor):
|
|
"""Convert a ComfyUI image tensor to uint8 grayscale using torch only."""
|
|
source = image_tensor.mul(255.0).clamp(0, 255).to(torch.uint8)
|
|
|
|
if source.ndim == 2:
|
|
return source
|
|
|
|
if source.ndim in (3, 4) and source.shape[-1] == 1:
|
|
return source[..., 0]
|
|
|
|
if source.ndim in (3, 4) and source.shape[-1] >= 3:
|
|
rgb = source[..., :3].to(torch.float32)
|
|
return rgb.mean(dim=-1).round().clamp(0, 255).to(torch.uint8)
|
|
|
|
raise ValueError(
|
|
f"Unsupported image shape for intensity conversion: {source.shape}"
|
|
)
|
|
|
|
|
|
def tensor_to_png_base64(gray_tensor: torch.Tensor):
|
|
"""Helper to encode a grayscale tensor as a PNG base64 string."""
|
|
img = Image.fromarray(gray_tensor.cpu().numpy(), mode="L")
|
|
buffered = BytesIO()
|
|
img.save(buffered, format="PNG")
|
|
return base64.b64encode(buffered.getvalue()).decode("utf-8")
|
|
|
|
|
|
class DepthToLayers:
|
|
"""Backend implementation of the `depth to layers` ComfyUI node.
|
|
|
|
ComfyUI contract:
|
|
- INPUT_TYPES: JSON interval definition plus optional IMAGE input.
|
|
- RETURN_TYPES: one IMAGE tensor with VGA-style colorized layers.
|
|
- Hidden `node_id`: used to push initial UI payload to the frontend.
|
|
"""
|
|
|
|
DEFAULT_INTERVALS = [
|
|
{"start": 0.0, "end": 0.3},
|
|
]
|
|
|
|
VGA_COLORS = VGA_256_COLORS
|
|
VGA_PALETTE_FLOAT_CPU = torch.tensor(VGA_COLORS, dtype=torch.float32).div(255.0)
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"intervals": (
|
|
"STRING",
|
|
{"default": json.dumps(cls.DEFAULT_INTERVALS), "multiline": True},
|
|
),
|
|
},
|
|
"optional": {
|
|
"images": ("IMAGE",),
|
|
},
|
|
"hidden": {
|
|
"node_id": "UNIQUE_ID",
|
|
},
|
|
}
|
|
|
|
RETURN_TYPES = ("IMAGE",)
|
|
RETURN_NAMES = ("images",)
|
|
FUNCTION = "transform"
|
|
CATEGORY = "Game Development/Image Adjustments"
|
|
DESCRIPTION = """
|
|
Generate layered images based on depth information.
|
|
Define intervals to group depth values and visualize
|
|
them with distinct colors. Useful for game development
|
|
workflows that require depth-based layer separation.
|
|
|
|
Click on the histogram to add intervals, and drag the
|
|
handles to adjust them. Use shift + click to remove them.
|
|
"""
|
|
|
|
@classmethod
|
|
def parse_intervals(cls, raw_intervals):
|
|
"""Parse and normalize interval JSON.
|
|
|
|
Returns a list of tuples: (start, end, layer_index), where
|
|
`layer_index` starts at 1 because index 0 is reserved as background.
|
|
"""
|
|
try:
|
|
intervals = json.loads(raw_intervals)
|
|
except (TypeError, json.JSONDecodeError):
|
|
intervals = cls.DEFAULT_INTERVALS
|
|
|
|
if not isinstance(intervals, list):
|
|
intervals = cls.DEFAULT_INTERVALS
|
|
|
|
valid_intervals = []
|
|
for p in intervals:
|
|
if not isinstance(p, dict):
|
|
continue
|
|
if not isinstance(p.get("start"), (int, float)):
|
|
continue
|
|
if not isinstance(p.get("end"), (int, float)):
|
|
continue
|
|
valid_intervals.append(p)
|
|
|
|
if not valid_intervals:
|
|
valid_intervals = cls.DEFAULT_INTERVALS
|
|
|
|
extended_intervals = []
|
|
for i, p in enumerate(valid_intervals):
|
|
start = float(np.clip(p["start"], 0.0, 1.0))
|
|
next_start = (
|
|
float(np.clip(valid_intervals[i + 1]["start"], 0.0, 1.0))
|
|
if i + 1 < len(valid_intervals)
|
|
else 1.0
|
|
)
|
|
end = float(np.clip(p["end"], start, next_start))
|
|
if end > start and start < 1.0:
|
|
extended_intervals.append((start, end, len(extended_intervals) + 1))
|
|
|
|
return extended_intervals
|
|
|
|
@classmethod
|
|
def IS_CHANGED(cls, **kwargs):
|
|
"""
|
|
Indicates that the node should reprocess whenever any input changes.
|
|
This is a common pattern hack as NaN never equals itself, forcing the node to run on every input change.
|
|
This is more efficient than using timestamps or hashing for change detection in this case, since the node needs
|
|
to update whenever intervals or images change.
|
|
"""
|
|
return float("NaN")
|
|
|
|
@classmethod
|
|
def process_images(cls, gray_tensors, intervals):
|
|
"""Map batched uint8 grayscale tensors to normalized RGB layer colors."""
|
|
norm = gray_tensors.to(torch.float32).div(255.0)
|
|
output = torch.zeros(
|
|
(*gray_tensors.shape, 3), dtype=torch.float32, device=gray_tensors.device
|
|
)
|
|
palette = cls.VGA_PALETTE_FLOAT_CPU.to(device=gray_tensors.device)
|
|
|
|
for start, end, layer_index in intervals:
|
|
if layer_index >= len(cls.VGA_COLORS):
|
|
raise RuntimeError(
|
|
f"Too many intervals defined. Max supported is {len(cls.VGA_COLORS) - 1}."
|
|
)
|
|
output[(start <= norm) & (norm <= end)] = palette[layer_index]
|
|
|
|
return output
|
|
|
|
def _send_initial_data(self, node_id, gray_tensor):
|
|
"""Send initial preview/histogram data to the frontend editor."""
|
|
|
|
histogram = torch.bincount(
|
|
gray_tensor.reshape(-1).to(torch.int64), minlength=256
|
|
).to(torch.float32)
|
|
if histogram.numel() and (max_val := float(histogram.max().item())):
|
|
histogram /= max_val
|
|
|
|
PromptServer.instance.send_sync(
|
|
"game_dev.depth_to_layers.initialise",
|
|
{
|
|
"node_id": node_id,
|
|
"preview_image": tensor_to_png_base64(gray_tensor),
|
|
"histogram_data": histogram.cpu().tolist(),
|
|
"default_intervals": self.DEFAULT_INTERVALS,
|
|
"vga_palette": self.VGA_COLORS,
|
|
},
|
|
)
|
|
|
|
def transform(self, intervals, images=None, node_id=None):
|
|
"""ComfyUI entry point.
|
|
|
|
Applies interval-based layer coloring to each image in the batch.
|
|
If `node_id` is available, sends UI bootstrap data for the JS editor.
|
|
"""
|
|
if images is None:
|
|
return (None,)
|
|
|
|
parsed_intervals = self.parse_intervals(intervals)
|
|
gray_images = tensor_to_gray_uint8(images)
|
|
|
|
if node_id:
|
|
self._send_initial_data(node_id, gray_images[0])
|
|
|
|
return (self.process_images(gray_images, parsed_intervals),)
|