From 54ff50785779c895665cfda9569d91e753c32fef Mon Sep 17 00:00:00 2001 From: Enrique Calot Date: Wed, 8 Apr 2026 10:47:02 +1000 Subject: [PATCH] Add optimized palette mode to Resize + Dither --- README.md | 4 ++ js/resize_dither.js | 132 ++++++++++++++++++++++++++++++++++++++++++-- resize_dither.py | 107 ++++++++++++++++++++++++++++++----- 3 files changed, 226 insertions(+), 17 deletions(-) diff --git a/README.md b/README.md index 28d6d0b..46e1806 100644 --- a/README.md +++ b/README.md @@ -80,6 +80,8 @@ pip install . - `width` (`INT`) - `height` (`INT`) - `dither_mode` (`positional` or `floyd`) + - `palette_mode` (`vga_256` or `optimized_x_colors`) + - `palette_colors` (`INT`, used when `palette_mode=optimized_x_colors`) - `node_id` (hidden `UNIQUE_ID`, used for UI preview messaging) - **Output:** - `IMAGE` batch (resized + dithered RGB) @@ -113,6 +115,8 @@ pip install . 4. Use in-node buttons to: - switch dithering mode (`Positional` or `Floyd`) +- choose palette mode (`vga_256` or `optimized_x_colors`) +- set optimized palette size (`palette_colors`, usually 8-64) - apply output scale presets (`1x`, `2x`, `4x`) - nudge output size (`W-`, `W+`, `H-`, `H+`) - reset to source-size positional defaults (`Reset`) diff --git a/js/resize_dither.js b/js/resize_dither.js index 05a09b6..bd69b13 100644 --- a/js/resize_dither.js +++ b/js/resize_dither.js @@ -32,6 +32,9 @@ const BUTTON_BORDER_COLOR = "#9a9a9a"; const PREVIEW_PLACEHOLDER_TEXT_COLOR = "#ebebeb"; const DITHER_MODES = ["positional", "floyd"]; +const PALETTE_MODES = ["vga_256", "optimized_x_colors"]; +const MIN_PALETTE_COLORS = 2; +const MAX_PALETTE_COLORS = 256; const buildVga256Palette = () => { const palette = []; @@ -185,6 +188,18 @@ app.registerExtension({ return DITHER_MODES.includes(mode) ? mode : "positional"; }, + getPaletteModeValue: function () { + const widget = this.findWidget("palette_mode"); + const mode = String(widget?.value || "vga_256"); + return PALETTE_MODES.includes(mode) ? mode : "vga_256"; + }, + + getPaletteColorsValue: function () { + const widget = this.findWidget("palette_colors"); + const raw = toInt(widget?.value, 24); + return clamp(raw, MIN_PALETTE_COLORS, MAX_PALETTE_COLORS); + }, + setWidgetValue: function (name, value) { const widget = this.findWidget(name); if (!widget) return; @@ -207,11 +222,97 @@ app.registerExtension({ }, rebuildBinaryPaletteLut: function () { - const palette = this.vgaPalette?.length ? this.vgaPalette : FALLBACK_PALETTE; - this.paletteInt = palette.map((c) => [c[0], c[1], c[2]]); + this.setActivePalette(this.vgaPalette?.length ? this.vgaPalette : FALLBACK_PALETTE); + }, + + setActivePalette: function (palette) { + const normalized = normalizePalette(palette); + const resolved = normalized && normalized.length > 0 ? normalized : FALLBACK_PALETTE; + this.paletteInt = resolved.map((c) => [c[0], c[1], c[2]]); this.paletteFloat = this.paletteInt.map((c) => [c[0] / 255, c[1] / 255, c[2] / 255]); }, + buildOptimizedPalette: function (rgb, width, height, kColors) { + const pixelCount = width * height; + if (pixelCount <= 0) { + return this.vgaPalette; + } + + const k = clamp(toInt(kColors, 24), MIN_PALETTE_COLORS, MAX_PALETTE_COLORS); + const sampleLimit = 20000; + const step = Math.max(1, Math.floor(pixelCount / sampleLimit)); + const sample = []; + + for (let i = 0; i < pixelCount; i += step) { + const base = i * 3; + sample.push([rgb[base], rgb[base + 1], rgb[base + 2]]); + if (sample.length >= sampleLimit) break; + } + + if (sample.length === 0) { + return this.vgaPalette; + } + + const originalSampleLength = sample.length; + while (sample.length < k) { + sample.push(sample[sample.length % originalSampleLength].slice()); + } + + const centers = []; + for (let j = 0; j < k; j++) { + const index = Math.round((j * (sample.length - 1)) / Math.max(1, k - 1)); + centers.push(sample[index].slice()); + } + + for (let iter = 0; iter < 10; iter++) { + const sumR = new Float32Array(k); + const sumG = new Float32Array(k); + const sumB = new Float32Array(k); + const counts = new Uint32Array(k); + + for (let i = 0; i < sample.length; i++) { + const point = sample[i]; + let bestIdx = 0; + let bestDist = Number.POSITIVE_INFINITY; + + for (let j = 0; j < k; j++) { + const center = centers[j]; + const dr = point[0] - center[0]; + const dg = point[1] - center[1]; + const db = point[2] - center[2]; + const dist = dr * dr + dg * dg + db * db; + if (dist < bestDist) { + bestDist = dist; + bestIdx = j; + } + } + + sumR[bestIdx] += point[0]; + sumG[bestIdx] += point[1]; + sumB[bestIdx] += point[2]; + counts[bestIdx] += 1; + } + + for (let j = 0; j < k; j++) { + if (counts[j] > 0) { + centers[j][0] = sumR[j] / counts[j]; + centers[j][1] = sumG[j] / counts[j]; + centers[j][2] = sumB[j] / counts[j]; + } + } + } + + centers.sort( + (a, b) => 0.2126 * a[0] + 0.7152 * a[1] + 0.0722 * a[2] - (0.2126 * b[0] + 0.7152 * b[1] + 0.0722 * b[2]), + ); + + return centers.map((c) => [ + clamp(Math.round(c[0] * 255), 0, 255), + clamp(Math.round(c[1] * 255), 0, 255), + clamp(Math.round(c[2] * 255), 0, 255), + ]); + }, + nearestPaletteIndex: function (r, g, b) { const palette = this.paletteFloat?.length ? this.paletteFloat @@ -521,6 +622,8 @@ app.registerExtension({ const width = target.width; const height = target.height; const mode = this.getModeValue(); + const paletteMode = this.getPaletteModeValue(); + const paletteColors = this.getPaletteColorsValue(); const canvas = this.ensurePreviewCanvas(width, height); if (!canvas || !this.localPreviewCtx) { @@ -543,6 +646,13 @@ app.registerExtension({ } } + if (paletteMode === "optimized_x_colors") { + const optimizedPalette = this.buildOptimizedPalette(rgb, width, height, paletteColors); + this.setActivePalette(optimizedPalette); + } else { + this.setActivePalette(this.vgaPalette); + } + const dithered = mode === "floyd" ? this.applyFloydDither(rgb, width, height) : this.applyPositionalDither(rgb, width, height); @@ -572,8 +682,12 @@ app.registerExtension({ updateBackend: function () { const width = this.getIntWidgetValue("width", 0); const height = this.getIntWidgetValue("height", 0); + const paletteMode = this.getPaletteModeValue(); + const paletteColors = this.getPaletteColorsValue(); this.setWidgetValue("width", width); this.setWidgetValue("height", height); + this.setWidgetValue("palette_mode", paletteMode); + this.setWidgetValue("palette_colors", paletteColors); this.renderLocalPreview(); this.setDirtyCanvas(true, true); @@ -738,6 +852,8 @@ app.registerExtension({ this.setWidgetValue("width", width); this.setWidgetValue("height", height); this.setWidgetValue("dither_mode", "positional"); + this.setWidgetValue("palette_mode", "vga_256"); + this.setWidgetValue("palette_colors", 24); this.updateBackend(); return true; } @@ -793,9 +909,17 @@ api.addEventListener("game_dev.resize_dither.initialise", ({ detail }) => { if (typeof detail.dither_mode === "string" && DITHER_MODES.includes(detail.dither_mode)) { node.setWidgetValue("dither_mode", detail.dither_mode); } + if (typeof detail.palette_mode === "string" && PALETTE_MODES.includes(detail.palette_mode)) { + node.setWidgetValue("palette_mode", detail.palette_mode); + } + if (Number.isFinite(detail.palette_colors)) { + node.setWidgetValue( + "palette_colors", + clamp(toInt(detail.palette_colors, 24), MIN_PALETTE_COLORS, MAX_PALETTE_COLORS), + ); + } - const incomingPalette = - normalizePalette(detail.vga_palette) || normalizePalette(detail.vgaPalette) || normalizePalette(detail.palette); + const incomingPalette = normalizePalette(detail.vga_palette); if (incomingPalette) { node.setPalette(incomingPalette); } diff --git a/resize_dither.py b/resize_dither.py index b599214..0a794bc 100644 --- a/resize_dither.py +++ b/resize_dither.py @@ -40,6 +40,14 @@ class ResizeDither: "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", @@ -110,11 +118,9 @@ or Floyd-Steinberg error diffusion. return rgb_images.index_select(1, y_idx).index_select(2, x_idx) @classmethod - def _nearest_palette_torch(cls, rgb_images): - """Nearest-color quantization in torch against shared VGA palette.""" - palette = cls.VGA_PALETTE_FLOAT_CPU.to( - device=rgb_images.device, dtype=rgb_images.dtype - ) + 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) @@ -129,7 +135,7 @@ or Floyd-Steinberg error diffusion. return out.reshape_as(rgb_images) @classmethod - def _positional_dither(cls, rgb_images): + 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) @@ -141,16 +147,16 @@ or Floyd-Steinberg error diffusion. 0.0, 1.0, ) - return cls._nearest_palette_torch(nudged) + return cls._nearest_palette_torch(nudged, palette) @classmethod - def _floyd_dither(cls, rgb_images): + 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 = cls.VGA_PALETTE_NUMPY + palette = palette.detach().to("cpu", dtype=torch.float32).numpy() for bi in range(working.shape[0]): img = working[bi].copy() @@ -189,6 +195,54 @@ or Floyd-Steinberg error diffusion. 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 @@ -245,6 +299,9 @@ or Floyd-Steinberg error diffusion. width, height, dither_mode, + palette_mode, + palette_colors, + palette, ): """Send first-run preview payload to frontend.""" PromptServer.instance.send_sync( @@ -256,14 +313,34 @@ or Floyd-Steinberg error diffusion. "width": int(width), "height": int(height), "dither_mode": dither_mode, - "vga_palette": self.VGA_COLORS, + "palette_mode": palette_mode, + "palette_colors": int(palette_colors), + "vga_palette": ( + palette.detach() + .cpu() + .clamp(0.0, 1.0) + .mul(255.0) + .round() + .to(torch.int64) + .tolist() + ), }, ) - def transform(self, images, width, height, dither_mode, node_id=None): + 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 @@ -271,11 +348,12 @@ or Floyd-Steinberg error diffusion. 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) + dithered = self._floyd_dither(resized, palette) else: - dithered = self._positional_dither(resized) + dithered = self._positional_dither(resized, palette) output = dithered @@ -287,6 +365,9 @@ or Floyd-Steinberg error diffusion. width=requested_width, height=requested_height, dither_mode=dither_mode, + palette_mode=palette_mode, + palette_colors=palette_colors, + palette=palette, ) return (output,)