770 lines
24 KiB
Python
770 lines
24 KiB
Python
"""
|
||
ComfyUI custom node port of Sprite Fusion Pixel Snapper (original Rust).
|
||
The node snaps AI‑generated pixel art to a consistent grid and palette.
|
||
|
||
Original project (MIT) © 2025 Hugo Duprez:
|
||
https://github.com/Hugo-Dz/spritefusion-pixel-snapper
|
||
Ported to Python for ComfyUI custom_nodes.
|
||
"""
|
||
|
||
from __future__ import annotations
|
||
|
||
import math
|
||
from dataclasses import dataclass
|
||
from typing import List, Optional, Sequence, Tuple
|
||
|
||
import numpy as np
|
||
import torch
|
||
|
||
|
||
@dataclass
|
||
class Config:
|
||
k_colors: int = 16
|
||
pixel_size_override: float = 0.0
|
||
k_seed: int = 42
|
||
max_kmeans_iterations: int = 15
|
||
peak_threshold_multiplier: float = 0.2
|
||
peak_distance_filter: int = 4
|
||
walker_search_window_ratio: float = 0.35
|
||
walker_min_search_window: float = 2.0
|
||
walker_strength_threshold: float = 0.5
|
||
min_cuts_per_axis: int = 4
|
||
fallback_target_segments: int = 64
|
||
max_step_ratio: float = 1.8 # Lowered from 3.0 to catch more skew cases
|
||
|
||
|
||
class PixelSnapperError(ValueError):
|
||
"""Raised when invalid input or processing failure occurs."""
|
||
|
||
|
||
# --- Utility: tensor <-> numpy ------------------------------------------------
|
||
|
||
|
||
def tensor_to_numpy_batch(image: torch.Tensor) -> np.ndarray:
|
||
"""
|
||
Convert a ComfyUI IMAGE tensor to numpy uint8 in channel-last layout (B, H, W, 3).
|
||
Accepts both ComfyUI standard (B, H, W, C) and channel-first (B, C, H, W) tensors.
|
||
If C == 4, the alpha channel is dropped.
|
||
"""
|
||
t = image.detach().cpu().clamp(0.0, 1.0)
|
||
|
||
if t.ndim == 3:
|
||
# Single image without batch
|
||
if t.shape[0] in (3, 4): # C, H, W
|
||
t = t.unsqueeze(0).movedim(1, -1)
|
||
elif t.shape[-1] in (3, 4): # H, W, C
|
||
t = t.unsqueeze(0)
|
||
else:
|
||
raise PixelSnapperError(
|
||
f"Expected IMAGE tensor with 3 or 4 channels, got shape {tuple(t.shape)}."
|
||
)
|
||
elif t.ndim == 4:
|
||
if t.shape[1] in (3, 4): # B, C, H, W
|
||
t = t.movedim(1, -1)
|
||
elif t.shape[-1] in (3, 4): # B, H, W, C
|
||
pass
|
||
else:
|
||
raise PixelSnapperError(
|
||
f"Expected IMAGE tensor with 3 or 4 channels, got shape {tuple(t.shape)}."
|
||
)
|
||
else:
|
||
raise PixelSnapperError(
|
||
"Expected IMAGE tensor with shape (B,H,W,C) or (B,C,H,W)."
|
||
)
|
||
|
||
if t.shape[-1] == 4:
|
||
t = t[..., :3]
|
||
|
||
np_imgs = (t.numpy() * 255.0).round().astype(np.uint8)
|
||
return np_imgs
|
||
|
||
|
||
def numpy_batch_to_tensor(images: np.ndarray) -> torch.Tensor:
|
||
"""
|
||
Convert numpy batch (B, H, W, 3) uint8 to ComfyUI IMAGE tensor (B, H, W, 3) float in [0,1].
|
||
"""
|
||
if images.ndim != 4 or images.shape[-1] != 3:
|
||
raise PixelSnapperError("Expected numpy batch with shape (B, H, W, 3).")
|
||
tensor = torch.from_numpy(images.astype(np.float32) / 255.0)
|
||
return tensor
|
||
|
||
|
||
def process_batch(
|
||
np_batch: np.ndarray,
|
||
config: Config,
|
||
output_scale: int,
|
||
) -> List[np.ndarray]:
|
||
if np_batch.shape[0] == 0:
|
||
raise PixelSnapperError("Input batch is empty.")
|
||
outputs: List[np.ndarray] = []
|
||
for idx, img_np in enumerate(np_batch):
|
||
try:
|
||
processed = process_image_array(img_np, config)
|
||
if output_scale > 1:
|
||
processed = upscale_nearest(processed, output_scale)
|
||
outputs.append(processed)
|
||
except PixelSnapperError as exc:
|
||
raise PixelSnapperError(
|
||
f"PixelSnapper failed on frame {idx + 1}/{len(np_batch)}: {exc}"
|
||
) from exc
|
||
return outputs
|
||
|
||
|
||
def upscale_nearest(img: np.ndarray, scale: int) -> np.ndarray:
|
||
"""
|
||
Pixel-art-safe integer upscaling by nearest neighbor.
|
||
img: (H, W, 3) uint8
|
||
scale: int >= 1
|
||
"""
|
||
if scale <= 1:
|
||
return img
|
||
return np.repeat(np.repeat(img, scale, axis=0), scale, axis=1)
|
||
|
||
|
||
# --- Core algorithm (ported from Rust) ---------------------------------------
|
||
|
||
|
||
def validate_image_dimensions(width: int, height: int) -> None:
|
||
if width <= 0 or height <= 0:
|
||
raise PixelSnapperError("Image dimensions cannot be zero")
|
||
if width > 10_000 or height > 10_000:
|
||
raise PixelSnapperError("Image dimensions too large (max 10000x10000)")
|
||
|
||
|
||
def quantize_image(img: np.ndarray, config: Config) -> np.ndarray:
|
||
"""
|
||
Simple K-Means color quantization using the Rust defaults.
|
||
img: (H, W, 3) uint8
|
||
"""
|
||
if config.k_colors <= 0:
|
||
raise PixelSnapperError("Number of colors must be greater than 0")
|
||
|
||
pixels = img.reshape(-1, 3).astype(np.float32)
|
||
n_pixels = pixels.shape[0]
|
||
if n_pixels == 0:
|
||
raise PixelSnapperError("Image has no pixels")
|
||
|
||
k = min(config.k_colors, n_pixels)
|
||
rng = np.random.default_rng(config.k_seed)
|
||
|
||
def dist_sq(points: np.ndarray, centroid: np.ndarray) -> np.ndarray:
|
||
diff = points - centroid
|
||
return np.sum(diff * diff, axis=1)
|
||
|
||
first_idx = rng.integers(0, n_pixels)
|
||
centroids = [pixels[first_idx]]
|
||
distances = np.full(n_pixels, np.inf, dtype=np.float32)
|
||
|
||
# KMeans++ style init (mirroring Rust logic)
|
||
for _ in range(1, k):
|
||
last_c = centroids[-1]
|
||
d_sq = dist_sq(pixels, last_c)
|
||
distances = np.minimum(distances, d_sq)
|
||
sum_sq = float(np.sum(distances))
|
||
if sum_sq <= 0.0 or not math.isfinite(sum_sq):
|
||
idx = rng.integers(0, n_pixels)
|
||
else:
|
||
probs = distances / sum_sq
|
||
idx = rng.choice(n_pixels, p=probs)
|
||
centroids.append(pixels[idx])
|
||
|
||
centroids = np.stack(centroids, axis=0)
|
||
prev_centroids = centroids.copy()
|
||
|
||
for iteration in range(config.max_kmeans_iterations):
|
||
# Compute distances to centroids
|
||
diff = pixels[:, None, :] - centroids[None, :, :]
|
||
dists = np.sum(diff * diff, axis=2)
|
||
labels = np.argmin(dists, axis=1)
|
||
|
||
sums = np.zeros_like(centroids)
|
||
np.add.at(sums, labels, pixels)
|
||
counts = np.bincount(labels, minlength=k).astype(np.float32)
|
||
counts_expanded = counts[:, None]
|
||
nonzero = counts_expanded[:, 0] > 0
|
||
centroids[nonzero] = sums[nonzero] / counts_expanded[nonzero]
|
||
|
||
if iteration > 0:
|
||
movement = np.max(np.sum((centroids - prev_centroids) ** 2, axis=1))
|
||
if movement < 0.01:
|
||
break
|
||
prev_centroids = centroids.copy()
|
||
|
||
diff = pixels[:, None, :] - centroids[None, :, :]
|
||
dists = np.sum(diff * diff, axis=2)
|
||
labels = np.argmin(dists, axis=1)
|
||
|
||
best = centroids[labels].round().clip(0, 255).astype(np.uint8)
|
||
return best.reshape(img.shape)
|
||
|
||
|
||
def compute_profiles(img: np.ndarray) -> Tuple[np.ndarray, np.ndarray]:
|
||
h, w, _ = img.shape
|
||
if w < 3 or h < 3:
|
||
raise PixelSnapperError("Image too small (minimum 3x3)")
|
||
|
||
gray = (
|
||
0.299 * img[:, :, 0].astype(np.float64)
|
||
+ 0.587 * img[:, :, 1].astype(np.float64)
|
||
+ 0.114 * img[:, :, 2].astype(np.float64)
|
||
)
|
||
|
||
col_proj = np.zeros(w, dtype=np.float64)
|
||
row_proj = np.zeros(h, dtype=np.float64)
|
||
|
||
col_proj[1:-1] = np.sum(np.abs(gray[:, 2:] - gray[:, :-2]), axis=0)
|
||
row_proj[1:-1] = np.sum(np.abs(gray[2:, :] - gray[:-2, :]), axis=1)
|
||
|
||
return col_proj, row_proj
|
||
|
||
|
||
def estimate_step_size(profile: Sequence[float], config: Config) -> Optional[float]:
|
||
if len(profile) == 0:
|
||
return None
|
||
max_val = float(np.max(profile))
|
||
if max_val == 0.0:
|
||
return None
|
||
threshold = max_val * config.peak_threshold_multiplier
|
||
|
||
peaks: List[int] = []
|
||
for i in range(1, len(profile) - 1):
|
||
if (
|
||
profile[i] > threshold
|
||
and profile[i] > profile[i - 1]
|
||
and profile[i] > profile[i + 1]
|
||
):
|
||
peaks.append(i)
|
||
if len(peaks) < 2:
|
||
return None
|
||
|
||
clean_peaks = [peaks[0]]
|
||
for p in peaks[1:]:
|
||
if p - clean_peaks[-1] > (config.peak_distance_filter - 1):
|
||
clean_peaks.append(p)
|
||
if len(clean_peaks) < 2:
|
||
return None
|
||
|
||
diffs = np.diff(clean_peaks)
|
||
diffs.sort()
|
||
return float(diffs[len(diffs) // 2])
|
||
|
||
|
||
def resolve_step_sizes(
|
||
step_x_opt: Optional[float],
|
||
step_y_opt: Optional[float],
|
||
width: int,
|
||
height: int,
|
||
config: Config,
|
||
) -> Tuple[float, float]:
|
||
if config.pixel_size_override != 0.0:
|
||
return config.pixel_size_override, config.pixel_size_override
|
||
|
||
if step_x_opt is not None and step_y_opt is not None:
|
||
sx, sy = step_x_opt, step_y_opt
|
||
ratio = sx / sy if sx > sy else sy / sx
|
||
if ratio > config.max_step_ratio:
|
||
smaller = min(sx, sy)
|
||
return smaller, smaller
|
||
avg = (sx + sy) / 2.0
|
||
return avg, avg
|
||
|
||
if step_x_opt is not None:
|
||
return step_x_opt, step_x_opt
|
||
if step_y_opt is not None:
|
||
return step_y_opt, step_y_opt
|
||
|
||
fallback = (
|
||
(min(width, height) / float(config.fallback_target_segments))
|
||
if config.fallback_target_segments
|
||
else 1.0
|
||
)
|
||
return max(fallback, 1.0), max(fallback, 1.0)
|
||
|
||
|
||
def walk(
|
||
profile: Sequence[float], step_size: float, limit: int, config: Config
|
||
) -> List[int]:
|
||
if len(profile) == 0:
|
||
raise PixelSnapperError("Cannot walk on empty profile")
|
||
|
||
cuts = [0]
|
||
current_pos = 0.0
|
||
search_window = max(
|
||
step_size * config.walker_search_window_ratio, config.walker_min_search_window
|
||
)
|
||
mean_val = float(np.mean(profile))
|
||
|
||
while current_pos < limit:
|
||
target = current_pos + step_size
|
||
if target >= limit:
|
||
cuts.append(limit)
|
||
break
|
||
|
||
start_search = max(int(target - search_window), int(current_pos + 1))
|
||
end_search = min(int(target + search_window), limit)
|
||
|
||
if end_search <= start_search:
|
||
current_pos = target
|
||
continue
|
||
|
||
segment = np.asarray(profile[start_search:end_search])
|
||
max_rel = int(np.argmax(segment))
|
||
max_val = float(segment[max_rel]) if segment.size else -1.0
|
||
max_idx = start_search + max_rel
|
||
|
||
if max_val > mean_val * config.walker_strength_threshold:
|
||
cuts.append(max_idx)
|
||
current_pos = float(max_idx)
|
||
else:
|
||
cuts.append(int(target))
|
||
current_pos = target
|
||
return cuts
|
||
|
||
|
||
def sanitize_cuts(cuts: List[int], limit: int) -> List[int]:
|
||
if limit == 0:
|
||
return [0]
|
||
|
||
normalized = []
|
||
has_zero = False
|
||
has_limit = False
|
||
for v in cuts:
|
||
v = 0 if v < 0 else v
|
||
if v == 0:
|
||
has_zero = True
|
||
if v >= limit:
|
||
v = limit
|
||
has_limit = True
|
||
normalized.append(v)
|
||
if not has_zero:
|
||
normalized.append(0)
|
||
if not has_limit:
|
||
normalized.append(limit)
|
||
|
||
normalized = sorted(set(normalized))
|
||
return normalized
|
||
|
||
|
||
def snap_uniform_cuts(
|
||
profile: Sequence[float],
|
||
limit: int,
|
||
target_step: float,
|
||
config: Config,
|
||
min_required: int,
|
||
) -> List[int]:
|
||
if limit == 0:
|
||
return [0]
|
||
if limit == 1:
|
||
return [0, 1]
|
||
|
||
desired_cells = (
|
||
int(round(limit / target_step))
|
||
if target_step > 0 and math.isfinite(target_step)
|
||
else 0
|
||
)
|
||
desired_cells = max(desired_cells, min_required - 1, 1)
|
||
desired_cells = min(desired_cells, limit)
|
||
|
||
cell_width = limit / float(desired_cells)
|
||
search_window = max(
|
||
cell_width * config.walker_search_window_ratio, config.walker_min_search_window
|
||
)
|
||
mean_val = float(np.mean(profile)) if len(profile) else 0.0
|
||
|
||
cuts: List[int] = [0]
|
||
for idx in range(1, desired_cells):
|
||
target = cell_width * idx
|
||
prev = cuts[-1]
|
||
if prev + 1 >= limit:
|
||
break
|
||
|
||
start = int(math.floor(target - search_window))
|
||
start = max(start, prev + 1, 0)
|
||
end = int(math.ceil(target + search_window))
|
||
end = min(end, limit - 1)
|
||
if end < start:
|
||
start = prev + 1
|
||
end = start
|
||
|
||
segment = np.asarray(profile[start : min(end + 1, len(profile))])
|
||
if segment.size:
|
||
best_rel = int(np.argmax(segment))
|
||
best_val = float(segment[best_rel])
|
||
best_idx = start + best_rel
|
||
else:
|
||
best_val = -1.0
|
||
best_idx = start
|
||
|
||
strength_threshold = mean_val * config.walker_strength_threshold
|
||
if best_val < strength_threshold:
|
||
fallback_idx = int(round(target))
|
||
if fallback_idx <= prev:
|
||
fallback_idx = prev + 1
|
||
if fallback_idx >= limit:
|
||
fallback_idx = max(limit - 1, prev + 1)
|
||
best_idx = fallback_idx
|
||
|
||
cuts.append(best_idx)
|
||
|
||
if cuts[-1] != limit:
|
||
cuts.append(limit)
|
||
|
||
return sanitize_cuts(cuts, limit)
|
||
|
||
|
||
def stabilize_cuts(
|
||
profile: Sequence[float],
|
||
cuts: List[int],
|
||
limit: int,
|
||
sibling_cuts: Sequence[int],
|
||
sibling_limit: int,
|
||
config: Config,
|
||
) -> List[int]:
|
||
if limit == 0:
|
||
return [0]
|
||
|
||
cuts = sanitize_cuts(cuts, limit)
|
||
min_required = max(config.min_cuts_per_axis, 2)
|
||
min_required = min(min_required, limit + 1)
|
||
|
||
axis_cells = max(len(cuts) - 1, 0)
|
||
sibling_cells = max(len(sibling_cuts) - 1, 0)
|
||
sibling_has_grid = (
|
||
sibling_limit > 0 and sibling_cells >= (min_required - 1) and sibling_cells > 0
|
||
)
|
||
steps_skewed = False
|
||
if sibling_has_grid and axis_cells > 0:
|
||
axis_step = limit / float(axis_cells)
|
||
sibling_step = sibling_limit / float(sibling_cells)
|
||
step_ratio = axis_step / sibling_step
|
||
steps_skewed = (
|
||
step_ratio > config.max_step_ratio
|
||
or step_ratio < 1.0 / config.max_step_ratio
|
||
)
|
||
|
||
has_enough = len(cuts) >= min_required
|
||
if has_enough and not steps_skewed:
|
||
return cuts
|
||
|
||
if sibling_has_grid:
|
||
target_step = sibling_limit / float(sibling_cells)
|
||
elif config.fallback_target_segments > 1:
|
||
target_step = limit / float(config.fallback_target_segments)
|
||
elif axis_cells > 0:
|
||
target_step = limit / float(axis_cells)
|
||
else:
|
||
target_step = float(limit)
|
||
|
||
if not math.isfinite(target_step) or target_step <= 0.0:
|
||
target_step = 1.0
|
||
|
||
return snap_uniform_cuts(profile, limit, target_step, config, min_required)
|
||
|
||
|
||
def stabilize_both_axes(
|
||
profile_x: Sequence[float],
|
||
profile_y: Sequence[float],
|
||
raw_col_cuts: List[int],
|
||
raw_row_cuts: List[int],
|
||
width: int,
|
||
height: int,
|
||
config: Config,
|
||
) -> Tuple[List[int], List[int]]:
|
||
col_cuts_pass1 = stabilize_cuts(
|
||
profile_x, list(raw_col_cuts), width, raw_row_cuts, height, config
|
||
)
|
||
row_cuts_pass1 = stabilize_cuts(
|
||
profile_y, list(raw_row_cuts), height, raw_col_cuts, width, config
|
||
)
|
||
|
||
col_cells = max(len(col_cuts_pass1) - 1, 1)
|
||
row_cells = max(len(row_cuts_pass1) - 1, 1)
|
||
col_step = width / float(col_cells)
|
||
row_step = height / float(row_cells)
|
||
step_ratio = col_step / row_step if col_step > row_step else row_step / col_step
|
||
|
||
if step_ratio > config.max_step_ratio:
|
||
target_step = min(col_step, row_step)
|
||
if col_step > target_step * 1.2:
|
||
final_cols = snap_uniform_cuts(
|
||
profile_x, width, target_step, config, config.min_cuts_per_axis
|
||
)
|
||
else:
|
||
final_cols = col_cuts_pass1
|
||
|
||
if row_step > target_step * 1.2:
|
||
final_rows = snap_uniform_cuts(
|
||
profile_y, height, target_step, config, config.min_cuts_per_axis
|
||
)
|
||
else:
|
||
final_rows = row_cuts_pass1
|
||
return final_cols, final_rows
|
||
|
||
return col_cuts_pass1, row_cuts_pass1
|
||
|
||
|
||
def resample(img: np.ndarray, cols: Sequence[int], rows: Sequence[int]) -> np.ndarray:
|
||
if len(cols) < 2 or len(rows) < 2:
|
||
raise PixelSnapperError("Insufficient grid cuts for resampling")
|
||
|
||
out_w = max(len(cols) - 1, 1)
|
||
out_h = max(len(rows) - 1, 1)
|
||
final_img = np.zeros((out_h, out_w, 3), dtype=np.uint8)
|
||
|
||
for y_i, (ys, ye) in enumerate(zip(rows[:-1], rows[1:])):
|
||
for x_i, (xs, xe) in enumerate(zip(cols[:-1], cols[1:])):
|
||
if xe <= xs or ye <= ys:
|
||
continue
|
||
cell = img[ys:ye, xs:xe]
|
||
if cell.size == 0:
|
||
continue
|
||
pixels = cell.reshape(-1, 3)
|
||
values, counts = np.unique(pixels, axis=0, return_counts=True)
|
||
best_idx = int(np.argmax(counts))
|
||
final_img[y_i, x_i] = values[best_idx]
|
||
|
||
return final_img
|
||
|
||
|
||
def process_image_array(img: np.ndarray, config: Config) -> np.ndarray:
|
||
"""
|
||
Full pipeline on a single RGB uint8 numpy image (H, W, 3).
|
||
Returns processed uint8 numpy image.
|
||
"""
|
||
h, w, _ = img.shape
|
||
|
||
# Guard against NaN / Inf inputs (can appear with half-precision pipelines)
|
||
if not np.isfinite(img).all():
|
||
raise PixelSnapperError("Input image contains NaN or Inf values.")
|
||
|
||
validate_image_dimensions(w, h)
|
||
if config.pixel_size_override != 0.0:
|
||
max_pixel_size = min(w, h) / 2.0
|
||
px = config.pixel_size_override
|
||
if not math.isfinite(px) or px < 1.0 or px > max_pixel_size:
|
||
raise PixelSnapperError(
|
||
f"pixel_size_override {px:.1f} is out of valid range "
|
||
f"[1, {max_pixel_size:.1f}]"
|
||
)
|
||
|
||
quantized = quantize_image(img, config)
|
||
profile_x, profile_y = compute_profiles(quantized)
|
||
|
||
step_x_opt = estimate_step_size(profile_x, config)
|
||
step_y_opt = estimate_step_size(profile_y, config)
|
||
step_x, step_y = resolve_step_sizes(step_x_opt, step_y_opt, w, h, config)
|
||
|
||
raw_col_cuts = walk(profile_x, step_x, w, config)
|
||
raw_row_cuts = walk(profile_y, step_y, h, config)
|
||
|
||
col_cuts, row_cuts = stabilize_both_axes(
|
||
profile_x, profile_y, raw_col_cuts, raw_row_cuts, w, h, config
|
||
)
|
||
|
||
return resample(quantized, col_cuts, row_cuts)
|
||
|
||
|
||
# --- ComfyUI node ------------------------------------------------------------
|
||
|
||
|
||
class PixelSnapperNode:
|
||
"""
|
||
Pixel Snapper that returns a list of images to preserve per-frame sizes.
|
||
"""
|
||
|
||
@classmethod
|
||
def INPUT_TYPES(cls):
|
||
return {
|
||
"required": {
|
||
"image": (
|
||
"IMAGE",
|
||
{
|
||
"tooltip": "IMAGE tensor (batch supported). Returns list of images."
|
||
},
|
||
),
|
||
"k_colors": (
|
||
"INT",
|
||
{
|
||
"default": 16,
|
||
"min": 1,
|
||
"max": 512,
|
||
"tooltip": "Palette size for color quantization (K-Means)",
|
||
},
|
||
),
|
||
"k_seed": (
|
||
"INT",
|
||
{
|
||
"default": 42,
|
||
"min": 0,
|
||
"max": 2**31 - 1,
|
||
"tooltip": "Random seed used when seeding K-Means++ centroids",
|
||
},
|
||
),
|
||
},
|
||
"optional": {
|
||
"output_scale": (
|
||
"INT",
|
||
{
|
||
"default": 1,
|
||
"min": 1,
|
||
"max": 16,
|
||
"tooltip": "Integer upscaling factor (nearest-neighbor) applied after snapping",
|
||
},
|
||
),
|
||
"max_kmeans_iterations": (
|
||
"INT",
|
||
{
|
||
"default": 15,
|
||
"min": 1,
|
||
"max": 200,
|
||
"tooltip": "Upper limit on K-Means iterations while learning the palette",
|
||
},
|
||
),
|
||
"peak_threshold_multiplier": (
|
||
"FLOAT",
|
||
{
|
||
"default": 0.2,
|
||
"min": 0.0,
|
||
"max": 5.0,
|
||
"step": 0.01,
|
||
"tooltip": "Fraction of max gradient used as threshold to keep profile peaks",
|
||
},
|
||
),
|
||
"peak_distance_filter": (
|
||
"INT",
|
||
{
|
||
"default": 4,
|
||
"min": 1,
|
||
"max": 512,
|
||
"tooltip": "Minimum pixel spacing between retained peaks in the profile",
|
||
},
|
||
),
|
||
"walker_search_window_ratio": (
|
||
"FLOAT",
|
||
{
|
||
"default": 0.35,
|
||
"min": 0.01,
|
||
"max": 5.0,
|
||
"step": 0.01,
|
||
"tooltip": "Search window size as a fraction of estimated step when walking cuts",
|
||
},
|
||
),
|
||
"walker_min_search_window": (
|
||
"FLOAT",
|
||
{
|
||
"default": 2.0,
|
||
"min": 0.1,
|
||
"max": 32.0,
|
||
"step": 0.1,
|
||
"tooltip": "Minimum search window (pixels) around each target cut",
|
||
},
|
||
),
|
||
"walker_strength_threshold": (
|
||
"FLOAT",
|
||
{
|
||
"default": 0.5,
|
||
"min": 0.0,
|
||
"max": 5.0,
|
||
"step": 0.01,
|
||
"tooltip": "Peak must exceed mean*threshold; otherwise fallback to uniform cut",
|
||
},
|
||
),
|
||
"min_cuts_per_axis": (
|
||
"INT",
|
||
{
|
||
"default": 4,
|
||
"min": 2,
|
||
"max": 512,
|
||
"tooltip": "Lowest number of cut positions per axis (including ends)",
|
||
},
|
||
),
|
||
"fallback_target_segments": (
|
||
"INT",
|
||
{
|
||
"default": 64,
|
||
"min": 1,
|
||
"max": 2048,
|
||
"tooltip": "Target number of cells when step detection fails; derives fallback step",
|
||
},
|
||
),
|
||
"max_step_ratio": (
|
||
"FLOAT",
|
||
{
|
||
"default": 1.8,
|
||
"min": 1.0,
|
||
"max": 5.0,
|
||
"step": 0.05,
|
||
"tooltip": "Max allowed X/Y step ratio before snapping to a uniform grid",
|
||
},
|
||
),
|
||
"pixel_size": (
|
||
"FLOAT",
|
||
{
|
||
"default": 0.0,
|
||
"min": 0.0,
|
||
"max": 10000.0,
|
||
"step": 0.1,
|
||
"tooltip": "Override detected pixel size in input pixels; 0 keeps auto-detection",
|
||
},
|
||
),
|
||
},
|
||
}
|
||
|
||
RETURN_TYPES = ("IMAGE",)
|
||
OUTPUT_IS_LIST = (True,)
|
||
FUNCTION = "snap"
|
||
CATEGORY = "image/transform"
|
||
|
||
def snap(
|
||
self,
|
||
image: torch.Tensor,
|
||
k_colors: int,
|
||
k_seed: int,
|
||
output_scale: int = 1,
|
||
max_kmeans_iterations: int = 15,
|
||
peak_threshold_multiplier: float = 0.2,
|
||
peak_distance_filter: int = 4,
|
||
walker_search_window_ratio: float = 0.35,
|
||
walker_min_search_window: float = 2.0,
|
||
walker_strength_threshold: float = 0.5,
|
||
min_cuts_per_axis: int = 4,
|
||
fallback_target_segments: int = 64,
|
||
max_step_ratio: float = 1.8,
|
||
pixel_size: float = 0.0,
|
||
):
|
||
config = Config(
|
||
k_colors=k_colors,
|
||
pixel_size_override=pixel_size,
|
||
k_seed=k_seed,
|
||
max_kmeans_iterations=max_kmeans_iterations,
|
||
peak_threshold_multiplier=peak_threshold_multiplier,
|
||
peak_distance_filter=peak_distance_filter,
|
||
walker_search_window_ratio=walker_search_window_ratio,
|
||
walker_min_search_window=walker_min_search_window,
|
||
walker_strength_threshold=walker_strength_threshold,
|
||
min_cuts_per_axis=min_cuts_per_axis,
|
||
fallback_target_segments=fallback_target_segments,
|
||
max_step_ratio=max_step_ratio,
|
||
)
|
||
|
||
np_batch = tensor_to_numpy_batch(image)
|
||
outputs = process_batch(np_batch, config, output_scale)
|
||
out_list = [numpy_batch_to_tensor(img[None, ...]) for img in outputs]
|
||
return (out_list,)
|
||
|
||
|
||
NODE_CLASS_MAPPINGS = {
|
||
"PixelSnapper": PixelSnapperNode,
|
||
}
|
||
|
||
NODE_DISPLAY_NAME_MAPPINGS = {
|
||
"PixelSnapper": "Sprite Fusion Pixel Snapper",
|
||
}
|
||
|
||
# Friendly name for ComfyUI extension listing
|
||
__all__ = [
|
||
"PixelSnapperNode",
|
||
"NODE_CLASS_MAPPINGS",
|
||
"NODE_DISPLAY_NAME_MAPPINGS",
|
||
]
|