362 lines
12 KiB
Python
362 lines
12 KiB
Python
"""ComfyUI node backend for nearest-neighbor resize + dithering."""
|
|
|
|
import base64
|
|
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
|
|
|
|
|
|
class ResizeDither:
|
|
"""Resize image batches with nearest-neighbor and apply dithering."""
|
|
|
|
VGA_COLORS = VGA_256_COLORS
|
|
VGA_PALETTE_FLOAT_CPU = torch.tensor(VGA_COLORS, dtype=torch.float32).div(255.0)
|
|
|
|
BAYER_8X8 = torch.tensor(
|
|
[
|
|
[0, 32, 8, 40, 2, 34, 10, 42],
|
|
[48, 16, 56, 24, 50, 18, 58, 26],
|
|
[12, 44, 4, 36, 14, 46, 6, 38],
|
|
[60, 28, 52, 20, 62, 30, 54, 22],
|
|
[3, 35, 11, 43, 1, 33, 9, 41],
|
|
[51, 19, 59, 27, 49, 17, 57, 25],
|
|
[15, 47, 7, 39, 13, 45, 5, 37],
|
|
[63, 31, 55, 23, 61, 29, 53, 21],
|
|
],
|
|
dtype=torch.float32,
|
|
).div(64.0)
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"images": ("IMAGE",),
|
|
"width": ("INT", {"default": 0, "min": 0, "max": 8192}),
|
|
"height": ("INT", {"default": 200, "min": 0, "max": 8192}),
|
|
"dither_mode": (("positional", "floyd"), {"default": "positional"}),
|
|
"palette_mode": (
|
|
("vga_256", "optimized_x_colors"),
|
|
{"default": "vga_256"},
|
|
),
|
|
"palette_colors": (
|
|
"INT",
|
|
{"default": 24, "min": 2, "max": 256},
|
|
),
|
|
},
|
|
"hidden": {
|
|
"node_id": "UNIQUE_ID",
|
|
},
|
|
}
|
|
|
|
RETURN_TYPES = ("IMAGE",)
|
|
RETURN_NAMES = ("images",)
|
|
FUNCTION = "transform"
|
|
CATEGORY = "Game Development/Image Adjustments"
|
|
DESCRIPTION = """
|
|
Resize each image in a batch using nearest-neighbor, then apply
|
|
RGB dithering using either positional (ordered) Bayer dithering
|
|
or Floyd-Steinberg error diffusion.
|
|
"""
|
|
|
|
@classmethod
|
|
def _tensor_to_rgb(cls, images):
|
|
"""Convert image batch tensor [B,H,W,C] to quantized RGB [B,H,W,3]."""
|
|
if images.ndim != 4:
|
|
raise ValueError(f"Unsupported image shape: {tuple(images.shape)}")
|
|
|
|
if images.shape[-1] == 1:
|
|
mono = (
|
|
images[..., 0]
|
|
.to(torch.float32)
|
|
.mul(255.0)
|
|
.round()
|
|
.clamp(0.0, 255.0)
|
|
.div(255.0)
|
|
)
|
|
return mono.unsqueeze(-1).repeat(1, 1, 1, 3)
|
|
|
|
if images.shape[-1] >= 3:
|
|
return (
|
|
images[..., :3]
|
|
.to(torch.float32)
|
|
.mul(255.0)
|
|
.round()
|
|
.clamp(0.0, 255.0)
|
|
.div(255.0)
|
|
)
|
|
|
|
raise ValueError(f"Unsupported channel count: {images.shape[-1]}")
|
|
|
|
@classmethod
|
|
def _resize_nearest(cls, rgb_images, width, height):
|
|
"""Resize RGB batch [B,H,W,3] to [B,height,width,3] with nearest-neighbor."""
|
|
_, src_h, src_w, _ = rgb_images.shape
|
|
y_idx = (
|
|
torch.div(
|
|
torch.arange(height, device=rgb_images.device) * src_h,
|
|
height,
|
|
rounding_mode="floor",
|
|
)
|
|
.clamp(0, src_h - 1)
|
|
.to(torch.int64)
|
|
)
|
|
x_idx = (
|
|
torch.div(
|
|
torch.arange(width, device=rgb_images.device) * src_w,
|
|
width,
|
|
rounding_mode="floor",
|
|
)
|
|
.clamp(0, src_w - 1)
|
|
.to(torch.int64)
|
|
)
|
|
return rgb_images.index_select(1, y_idx).index_select(2, x_idx)
|
|
|
|
@classmethod
|
|
def _nearest_palette_torch(cls, rgb_images, palette):
|
|
"""Nearest-color quantization in torch against provided palette."""
|
|
palette = palette.to(device=rgb_images.device, dtype=rgb_images.dtype)
|
|
flat = rgb_images.reshape(-1, 3)
|
|
out = torch.empty_like(flat)
|
|
|
|
chunk_size = 65536
|
|
for start in range(0, flat.shape[0], chunk_size):
|
|
end = min(start + chunk_size, flat.shape[0])
|
|
block = flat[start:end]
|
|
dist2 = ((block.unsqueeze(1) - palette.unsqueeze(0)) ** 2).sum(dim=-1)
|
|
idx = dist2.argmin(dim=1)
|
|
out[start:end] = palette.index_select(0, idx)
|
|
|
|
return out.reshape_as(rgb_images)
|
|
|
|
@classmethod
|
|
def _positional_dither(cls, rgb_images, palette):
|
|
"""GlitchNodes-style ordered dithering: Bayer nudge + nearest palette."""
|
|
_, h, w, _ = rgb_images.shape
|
|
bayer = cls.BAYER_8X8.to(device=rgb_images.device)
|
|
|
|
repeated = bayer.repeat((h + 7) // 8, (w + 7) // 8)[:h, :w]
|
|
threshold = repeated - 0.5
|
|
nudged = torch.clamp(
|
|
rgb_images + threshold.unsqueeze(0).unsqueeze(-1) / 64.0,
|
|
0.0,
|
|
1.0,
|
|
)
|
|
return cls._nearest_palette_torch(nudged, palette)
|
|
|
|
@classmethod
|
|
def _floyd_dither(cls, rgb_images, palette):
|
|
"""GlitchNodes-style Floyd-Steinberg error diffusion to full palette."""
|
|
original_device = rgb_images.device
|
|
|
|
working = rgb_images.detach().to("cpu", dtype=torch.float32).numpy()
|
|
output = np.zeros_like(working, dtype=np.float32)
|
|
palette = palette.detach().to("cpu", dtype=torch.float32).numpy()
|
|
|
|
for bi in range(working.shape[0]):
|
|
img = working[bi].copy()
|
|
out = output[bi]
|
|
h, w, _ = img.shape
|
|
|
|
for y in range(h):
|
|
for x in range(w):
|
|
old_pixel = img[y, x, :]
|
|
diff = palette - old_pixel[None, :]
|
|
idx = np.argmin(np.sum(diff * diff, axis=1))
|
|
new_pixel = palette[idx]
|
|
out[y, x] = new_pixel
|
|
|
|
error = old_pixel - new_pixel
|
|
if x + 1 < w:
|
|
img[y, x + 1, :] = np.clip(
|
|
img[y, x + 1, :] + error * (7.0 / 16.0), 0.0, 1.0
|
|
)
|
|
if y + 1 < h:
|
|
if x > 0:
|
|
img[y + 1, x - 1, :] = np.clip(
|
|
img[y + 1, x - 1, :] + error * (3.0 / 16.0),
|
|
0.0,
|
|
1.0,
|
|
)
|
|
img[y + 1, x, :] = np.clip(
|
|
img[y + 1, x, :] + error * (5.0 / 16.0), 0.0, 1.0
|
|
)
|
|
if x + 1 < w:
|
|
img[y + 1, x + 1, :] = np.clip(
|
|
img[y + 1, x + 1, :] + error * (1.0 / 16.0),
|
|
0.0,
|
|
1.0,
|
|
)
|
|
|
|
return torch.from_numpy(output).to(device=original_device, dtype=torch.float32)
|
|
|
|
@classmethod
|
|
def _build_optimized_palette(cls, rgb_images, k_colors):
|
|
"""Build an adaptive K-color palette from the resized images.
|
|
|
|
Uses a deterministic k-means style clustering over sampled pixels.
|
|
"""
|
|
k = int(np.clip(k_colors, 2, 256))
|
|
pixels = (
|
|
rgb_images.detach().to("cpu", dtype=torch.float32).numpy().reshape(-1, 3)
|
|
)
|
|
if pixels.shape[0] == 0:
|
|
return cls.VGA_PALETTE_FLOAT_CPU
|
|
|
|
sample_limit = 20000
|
|
if pixels.shape[0] > sample_limit:
|
|
step = max(1, pixels.shape[0] // sample_limit)
|
|
sample = pixels[::step][:sample_limit]
|
|
else:
|
|
sample = pixels
|
|
|
|
if sample.shape[0] < k:
|
|
reps = int(np.ceil(k / sample.shape[0]))
|
|
sample = np.tile(sample, (reps, 1))[:k]
|
|
|
|
init_idx = np.linspace(0, sample.shape[0] - 1, k, dtype=np.int64)
|
|
centers = sample[init_idx].copy()
|
|
|
|
for _ in range(10):
|
|
diff = sample[:, None, :] - centers[None, :, :]
|
|
dist2 = np.sum(diff * diff, axis=2)
|
|
labels = np.argmin(dist2, axis=1)
|
|
|
|
for j in range(k):
|
|
pts = sample[labels == j]
|
|
if pts.size > 0:
|
|
centers[j] = pts.mean(axis=0)
|
|
|
|
lum = 0.2126 * centers[:, 0] + 0.7152 * centers[:, 1] + 0.0722 * centers[:, 2]
|
|
centers = centers[np.argsort(lum)]
|
|
centers = np.clip(centers, 0.0, 1.0)
|
|
return torch.from_numpy(centers).to(dtype=torch.float32)
|
|
|
|
@classmethod
|
|
def _resolve_palette(cls, rgb_images, palette_mode, palette_colors):
|
|
if palette_mode == "optimized_x_colors":
|
|
return cls._build_optimized_palette(rgb_images, palette_colors)
|
|
return cls.VGA_PALETTE_FLOAT_CPU
|
|
|
|
@classmethod
|
|
def _resolve_target_size(
|
|
cls, requested_width, requested_height, src_width, src_height
|
|
):
|
|
"""Resolve width/height with auto(0) semantics.
|
|
|
|
Rules:
|
|
- (0, 0): keep original size
|
|
- (0, h): compute width proportionally
|
|
- (w, 0): compute height proportionally
|
|
"""
|
|
width = int(requested_width)
|
|
height = int(requested_height)
|
|
src_width = max(1, int(src_width))
|
|
src_height = max(1, int(src_height))
|
|
|
|
if width <= 0 and height <= 0:
|
|
width, height = src_width, src_height
|
|
elif width <= 0:
|
|
width = max(1, int(round(src_width * (height / src_height))))
|
|
elif height <= 0:
|
|
height = max(1, int(round(src_height * (width / src_width))))
|
|
|
|
width = int(np.clip(width, 1, 8192))
|
|
height = int(np.clip(height, 1, 8192))
|
|
return width, height
|
|
|
|
@classmethod
|
|
def _tensor_to_png_base64(cls, image_tensor):
|
|
"""Encode a single tensor image as PNG base64 for frontend previews."""
|
|
img = image_tensor.detach().cpu().clamp(0.0, 1.0)
|
|
|
|
if img.ndim == 2:
|
|
arr = (img.numpy() * 255.0).round().astype(np.uint8)
|
|
pil_img = Image.fromarray(arr, mode="L")
|
|
elif img.ndim == 3 and img.shape[-1] == 1:
|
|
arr = (img[..., 0].numpy() * 255.0).round().astype(np.uint8)
|
|
pil_img = Image.fromarray(arr, mode="L")
|
|
elif img.ndim == 3 and img.shape[-1] >= 3:
|
|
arr = (img[..., :3].numpy() * 255.0).round().astype(np.uint8)
|
|
pil_img = Image.fromarray(arr, mode="RGB")
|
|
else:
|
|
raise ValueError(f"Unsupported preview tensor shape: {tuple(img.shape)}")
|
|
|
|
buffer = BytesIO()
|
|
pil_img.save(buffer, format="PNG")
|
|
return base64.b64encode(buffer.getvalue()).decode("utf-8")
|
|
|
|
def _send_initial_data(
|
|
self,
|
|
node_id,
|
|
before_image,
|
|
after_image,
|
|
width,
|
|
height,
|
|
dither_mode,
|
|
palette_mode,
|
|
palette_colors,
|
|
):
|
|
"""Send first-run preview payload to frontend."""
|
|
PromptServer.instance.send_sync(
|
|
"game_dev.resize_dither.initialise",
|
|
{
|
|
"node_id": node_id,
|
|
"preview_before": self._tensor_to_png_base64(before_image),
|
|
"preview_after": self._tensor_to_png_base64(after_image),
|
|
"width": int(width),
|
|
"height": int(height),
|
|
"dither_mode": dither_mode,
|
|
"palette_mode": palette_mode,
|
|
"palette_colors": int(palette_colors),
|
|
},
|
|
)
|
|
|
|
def transform(
|
|
self,
|
|
images,
|
|
width,
|
|
height,
|
|
dither_mode,
|
|
palette_mode,
|
|
palette_colors,
|
|
node_id=None,
|
|
):
|
|
"""ComfyUI entry point for batch resize + dithering."""
|
|
requested_width = int(width)
|
|
requested_height = int(height)
|
|
palette_colors = int(palette_colors)
|
|
|
|
rgb = self._tensor_to_rgb(images)
|
|
_, src_h, src_w, _ = rgb.shape
|
|
width, height = self._resolve_target_size(
|
|
requested_width, requested_height, src_w, src_h
|
|
)
|
|
resized = self._resize_nearest(rgb, width, height)
|
|
palette = self._resolve_palette(resized, palette_mode, palette_colors)
|
|
|
|
if dither_mode == "floyd":
|
|
dithered = self._floyd_dither(resized, palette)
|
|
else:
|
|
dithered = self._positional_dither(resized, palette)
|
|
|
|
output = dithered
|
|
|
|
if node_id:
|
|
self._send_initial_data(
|
|
node_id=node_id,
|
|
before_image=images[0],
|
|
after_image=output[0],
|
|
width=requested_width,
|
|
height=requested_height,
|
|
dither_mode=dither_mode,
|
|
palette_mode=palette_mode,
|
|
palette_colors=palette_colors,
|
|
)
|
|
|
|
return (output,)
|