Add optimized palette mode to Resize + Dither

2026-04-08 10:47:02 +10:00
parent b129932eeb
commit 54ff507857
3 changed files with 226 additions and 17 deletions
@@ -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`)
@@ -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);
  }
@@ -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,)