color_transfer.py

import gc
from itertools import combinations
import ast
import numpy as np
import torch
import cv2
import ot

from sklearn.cluster import KMeans, MiniBatchKMeans
from scipy.spatial import Delaunay
from skimage import color
from scipy.interpolate import Rbf

from comfy.comfy_types import IO, ComfyNodeABC

from .utils import (
    EuclideanDistance,
    ManhattanDistance,
    CosineSimilarity,
    HSVColorSimilarity,
    RGBWeightedDistance,
    RGBWeightedSimilarity,
    Blur,
)


class PaletteExtension:
    MAX_PALETTE_SIZE = 512

    @staticmethod
    def dense_palette(
        base_palette: list[tuple[int, int, int]],
        points: int = 5,
        iterations: int = 2,
        extend_bw: bool = True,
    ) -> list[tuple[int, int, int]]:
        """
        Interpolate N points between each pair of colors in the base_palette.
        """
        if extend_bw:
            base_palette = set(base_palette)
            base_palette.add((0, 0, 0))
            base_palette.add((255, 255, 255))

        palette = np.array(list(base_palette), dtype=np.float32)

        for _ in range(iterations):
            if len(palette) > PaletteExtension.MAX_PALETTE_SIZE:
                break

            idx_pairs = list(combinations(range(len(palette)), 2))
            t = np.linspace(0, 1, points + 2)[1:-1].astype(np.float32)

            new_colors = []
            for i, j in idx_pairs:
                c1, c2 = palette[i], palette[j]
                inter = c1[None, :] + (c2 - c1)[None, :] * t[:, None]
                new_colors.append(inter)
            if new_colors:
                new_colors = np.vstack(new_colors)
                palette = np.vstack([palette, new_colors])

            palette = np.unique(np.rint(palette).astype(int), axis=0).astype(np.float32)

        result = [tuple(rgb.astype(int)) for rgb in palette]
        return result

    @staticmethod
    def edge_based_palette(
        base_palette: list[tuple[int, int, int]], points: int = 5, iterations: int = 2
    ) -> list[tuple[int, int, int]]:
        """
        Use Delaunay triangulation to find edges between colors.
        """
        base_palette = set(base_palette)
        base_palette.add((0, 0, 0))
        base_palette.add((255, 255, 255))

        palette = np.array(list(base_palette), dtype=np.float32)

        for _ in range(iterations):
            if len(palette) > PaletteExtension.MAX_PALETTE_SIZE:
                break

            if len(palette) >= palette.shape[1] + 1:
                try:
                    tri = Delaunay(palette)
                    edges = set()
                    for simplex in tri.simplices:
                        for i in range(len(simplex)):
                            for j in range(i + 1, len(simplex)):
                                a, b = simplex[i], simplex[j]
                                edges.add(tuple(sorted((a, b))))
                except Exception:
                    # Fallback if Delaunay fails (e.g. coplanar points)
                    idxs = range(len(palette))
                    edges = set(tuple(sorted((i, j))) for i in idxs for j in idxs if i < j)
            else:
                idxs = range(len(palette))
                edges = set(tuple(sorted((i, j))) for i in idxs for j in idxs if i < j)

            new_colors = []
            t = np.linspace(0, 1, points + 2)[1:-1][:, None].astype(np.float32)

            for i, j in edges:
                c1, c2 = palette[i], palette[j]
                inters = c1 + (c2 - c1) * t
                new_colors.append(inters)

            if new_colors:
                new_stack = np.vstack(new_colors)
                palette = np.vstack([palette, new_stack])
                palette = np.unique(np.rint(palette).astype(int), axis=0).astype(np.float32)

        return [tuple(c.astype(int)) for c in palette]


class ColorSpaceConvert:
    @staticmethod
    def convert_to_target_space(
        image: np.ndarray, target_colors: list[tuple[int, int, int]], color_space: str
    ) -> tuple[np.ndarray, list[tuple[int, int, int]]]:
        if color_space == "RGB":
            return image, target_colors

        conversion_map = {
            "HSV": (cv2.COLOR_RGB2HSV, cv2.COLOR_HSV2RGB),
            "LAB": (cv2.COLOR_RGB2LAB, cv2.COLOR_LAB2RGB),
        }

        forward_conversion, _ = conversion_map[color_space]
        
        # Ensure image is float32 compliant for OpenCV conversions if needed, 
        # though cv2 usually handles standard types. keeping input type is safer.
        converted_image = cv2.cvtColor(image, forward_conversion)

        target_colors_array = np.array(target_colors, dtype=np.uint8).reshape(-1, 1, 3)
        converted_colors = cv2.cvtColor(target_colors_array, forward_conversion)
        converted_colors = [tuple(color[0]) for color in converted_colors]

        return converted_image, converted_colors

    @staticmethod
    def convert_to_rgb(image: np.ndarray, color_space: str) -> np.ndarray:
        if color_space == "RGB":
            return image

        conversion_map = {"HSV": cv2.COLOR_HSV2RGB, "LAB": cv2.COLOR_LAB2RGB}
        # Ensure uint8 for consistent cvtColor behavior across color spaces
        return cv2.cvtColor(image.clip(0, 255).astype(np.uint8), conversion_map[color_space])


class ColorClustering:
    def __init__(self, cluster_method: str):
        self.clustering_methods = {
            "Kmeans": KMeans,
            "Mini batch Kmeans": MiniBatchKMeans,
        }
        self.method_name = cluster_method
        self.method = self.clustering_methods[cluster_method]

    def cluster_colors(self, image: np.ndarray, k: int) -> dict:
        """Perform color clustering on the image."""
        # Ensure we are working with float32 to save memory
        img_array = image.reshape((-1, 3)).astype(np.float32)
        
        # Add explicit parameters to control memory usage
        if self.method_name == "Mini batch Kmeans":
            clustering_model = self.method(
                n_clusters=k, 
                n_init="auto", 
                batch_size=4096, # Process in chunks
                reassignment_ratio=0 # Reduce bookkeeping
            )
        else:
            clustering_model = self.method(n_clusters=k, n_init="auto")
            
        clustering_model.fit(img_array)

        return {
            "image": image,
            "main_colors": clustering_model.cluster_centers_.astype(int),
            "model": clustering_model,
        }


class ColorMatcher:
    def __init__(self, distance_method: str):
        self.distance_methods = {
            "Euclidean": EuclideanDistance,
            "Manhattan": ManhattanDistance,
            "Cosine Similarity": CosineSimilarity,
            "HSV Distance": HSVColorSimilarity,
            "RGB Weighted Distance": RGBWeightedDistance,
            "RGB Weighted Similarity": RGBWeightedSimilarity,
        }
        self.distance_func = self.distance_methods[distance_method]

    def match_colors(
        self,
        detected_colors: np.ndarray,
        target_colors: list[tuple[int, int, int]],
        clustering_model: KMeans | MiniBatchKMeans,
        image_shape: tuple[int, int, int],
    ) -> np.ndarray:
        """Match detected colors with target colors."""
        closest_colors = []

        for color in detected_colors:
            distances = self.distance_func(color, target_colors)
            closest_color = target_colors[np.argmin(distances)]
            closest_colors.append(closest_color)

        closest_colors = np.array(closest_colors, dtype=np.float32)
        # Reconstruct image from labels
        return closest_colors[clustering_model.labels_].reshape(image_shape)


class ImagePostProcessor:
    def __init__(self, gaussian_blur: int = 0):
        self.gaussian_blur = gaussian_blur

    def process_image(self, image: np.ndarray) -> np.ndarray:
        processed = np.array(image, dtype=np.float32)

        if self.gaussian_blur:
            processed = Blur(processed, self.gaussian_blur)

        return processed / 255.0


def process_image_with_palette(
    image: list[torch.Tensor],
    target_colors: list[tuple[int, int, int]],
    color_space: str,
    cluster_method: str,
    distance_method: str,
    gaussian_blur: int,
) -> torch.Tensor:
    """
    Shared function to process an image with the given parameters.
    """
    processedImages = []

    converter = ColorSpaceConvert()
    clustering_engine = ColorClustering(cluster_method)
    color_matcher = ColorMatcher(distance_method)
    image_processor = ImagePostProcessor(gaussian_blur)

    for img_tensor in image:
        # Convert to uint8 0-255 so cv2.cvtColor handles all color spaces correctly
        img = (255.0 * img_tensor.cpu().numpy()).clip(0, 255).astype(np.uint8)

        converted_img, converted_colors = converter.convert_to_target_space(
            img, target_colors, color_space
        )

        clustering_result = clustering_engine.cluster_colors(
            converted_img, len(target_colors)
        )

        processed = color_matcher.match_colors(
            clustering_result["main_colors"],
            converted_colors,
            clustering_result["model"],
            converted_img.shape,
        )

        processed = converter.convert_to_rgb(processed, color_space)
        processed = image_processor.process_image(processed)
        
        # Ensure contiguous memory before converting to tensor
        processed = np.ascontiguousarray(processed)
        processed_tensor = torch.from_numpy(processed)[None,]
        processedImages.append(processed_tensor)

        # CLEANUP: Explicitly delete large temporary arrays and force GC
        del img
        del converted_img
        del clustering_result
        del processed
        
        # Only collect garbage if batch size is large to avoid perf hit on single images
        if len(image) > 1:
            gc.collect()

    return torch.cat(processedImages, dim=0)


class ReferenceTransferReinhard(ComfyNodeABC):
    @classmethod
    def INPUT_TYPES(cls) -> dict:
        return {
            "required": {
                "image": (IO.IMAGE,),
                "image_reference": (IO.IMAGE,),
            }
        }

    RETURN_TYPES = (IO.IMAGE,)
    FUNCTION = "color_transfer"
    CATEGORY = "Color Transfer/Reference Transfer"

    def color_transfer(
        self, image: torch.Tensor, image_reference: torch.Tensor
    ) -> tuple[torch.Tensor]:

        processed_images = []

        target = image_reference.cpu().numpy().astype(np.float32)
        if len(target.shape) == 4:
            target = np.concatenate([target[i] for i in range(target.shape[0])], axis=0)

        for img_tensor in image:
            source = img_tensor.cpu().numpy().astype(np.float32)

            source_lab = color.rgb2lab(source)
            target_lab = color.rgb2lab(target)

            s_mean, s_std = source_lab.mean(axis=(0, 1)), source_lab.std(axis=(0, 1))
            t_mean, t_std = target_lab.mean(axis=(0, 1)), target_lab.std(axis=(0, 1))

            # Avoid division by zero for uniform channels
            s_std = np.maximum(s_std, 1e-6)

            result_lab = (source_lab - s_mean) / s_std * t_std + t_mean
            result_rgb = np.clip(color.lab2rgb(result_lab), 0, 1)

            result_tensor = torch.from_numpy(result_rgb.astype(np.float32)).unsqueeze(0)
            processed_images.append(result_tensor)

            del source
            del source_lab
            del result_lab
            del result_rgb

        return (torch.cat(processed_images, dim=0),)


class PaletteOptimalTransportTransfer(ComfyNodeABC):
    @classmethod
    def INPUT_TYPES(cls) -> dict:
        return {
            "required": {
                "image": (IO.IMAGE,),
                "target_colors": ("COLOR_LIST",),
                "palette_extension_method": (["Dense", "Edge", "None"], {"default": "None"}),
                "palette_extension_points": ("INT", {"min": 2, "max": 20, "step": 1, "default": 5}),
                "blend_mode": (["Original", "Grayscale"], {"default": "Original"}),
                "blend_ratio": ("FLOAT", {"min": 0, "max": 1, "step": 0.1, "default": 0.5}),
            }
        }

    RETURN_TYPES = (IO.IMAGE,)
    FUNCTION = "color_transfer"
    CATEGORY = "Color Transfer/Palette Transfer"

    def color_transfer(
        self,
        image: torch.Tensor,
        target_colors: list[tuple[int, int, int]],
        palette_extension_method: str,
        palette_extension_points: int,
        blend_mode: str,
        blend_ratio: float,
    ) -> tuple[torch.Tensor]:

        if palette_extension_method == "Dense":
            target_colors = PaletteExtension.dense_palette(target_colors, points=palette_extension_points)
        elif palette_extension_method == "Edge":
            target_colors = PaletteExtension.edge_based_palette(target_colors, points=palette_extension_points)

        palette = np.array(target_colors, dtype=np.float32) / 255.0
        n_palette = palette.shape[0]
        palette_weights = np.ones((n_palette,)) / n_palette

        processed_images = []

        for img_tensor in image:
            source = img_tensor.cpu().numpy().astype(np.float32)
            h, w, _ = source.shape
            pixels = source.reshape(-1, 3)

            n_source_colors = 1000
            # MEMORY FIX: Reduce batch size and disable reassignment to save memory
            kmeans = MiniBatchKMeans(
                n_clusters=n_source_colors, 
                n_init="auto",
                batch_size=2048,
                reassignment_ratio=0
            )
            kmeans.fit(pixels)
            source_centroids = kmeans.cluster_centers_
            pixel_labels = kmeans.labels_

            source_weights = np.bincount(pixel_labels) / len(pixel_labels)

            cost_matrix = ot.dist(source_centroids, palette, metric="euclidean") ** 2

            transport_plan = ot.sinkhorn(
                source_weights,
                palette_weights,
                cost_matrix,
                reg=1e-2,
                numItermax=100000,
            )

            mapped_centroids = np.dot(transport_plan, palette) / np.sum(
                transport_plan, axis=1, keepdims=True
            )

            if blend_mode == "Original":
                recolored_pixels = (1 - blend_ratio) * pixels + blend_ratio * mapped_centroids[pixel_labels]
            elif blend_mode == "Grayscale":
                gray = color.rgb2gray(source)
                gray_rgb = np.stack([gray] * 3, axis=-1)
                gray_pixels = gray_rgb.reshape(-1, 3)
                recolored_pixels = (1 - blend_ratio) * gray_pixels + blend_ratio * mapped_centroids[pixel_labels]

            recolored_image = recolored_pixels.reshape(h, w, 3).astype(np.float32)

            result_tensor = torch.from_numpy(recolored_image).unsqueeze(0)
            processed_images.append(result_tensor)

            # CLEANUP
            del kmeans
            del source
            del pixels
            del cost_matrix
            del transport_plan
            
            if len(image) > 1:
                gc.collect()

        return (torch.cat(processed_images, dim=0),)


class PaletteRbfTransfer(ComfyNodeABC):
    @classmethod
    def INPUT_TYPES(cls) -> dict:
        return {
            "required": {
                "image": (IO.IMAGE,),
                "target_colors": ("COLOR_LIST",),
                "rbf_function": (["thin_plate", "multiquadric", "inverse", "gaussian"], {"default": "gaussian"}),
                "epsilon": ("FLOAT", {"min": 0.01, "max": 100, "step": 0.1, "default": 1.0}),
            }
        }

    RETURN_TYPES = (IO.IMAGE,)
    FUNCTION = "color_transfer"
    CATEGORY = "Color Transfer/Palette Transfer"

    def color_transfer(
        self,
        image: torch.Tensor,
        target_colors: list[tuple[int, int, int]],
        rbf_function: str,
        epsilon: float,
    ) -> tuple[torch.Tensor]:
        
        palette = np.array(target_colors, dtype=np.float32) / 255
        r, g, b = palette[:, 0], palette[:, 1], palette[:, 2]

        rbf_r = Rbf(r, g, b, r, function=rbf_function, epsilon=epsilon)
        rbf_g = Rbf(r, g, b, g, function=rbf_function, epsilon=epsilon)
        rbf_b = Rbf(r, g, b, b, function=rbf_function, epsilon=epsilon)

        processed_images = []

        for img_tensor in image:
            source = img_tensor.cpu().numpy().astype(np.float32)
            h, w, _ = source.shape
            pixels = source.reshape(-1, 3)

            mapped_r = rbf_r(pixels[:, 0], pixels[:, 1], pixels[:, 2])
            mapped_g = rbf_g(pixels[:, 0], pixels[:, 1], pixels[:, 2])
            mapped_b = rbf_b(pixels[:, 0], pixels[:, 1], pixels[:, 2])

            mapped_pixels = np.stack((mapped_r, mapped_g, mapped_b), axis=-1)
            mapped_pixels = np.clip(mapped_pixels, 0, 1).astype(np.float32)
            recolored_image = mapped_pixels.reshape(h, w, 3)

            result_tensor = torch.from_numpy(recolored_image).unsqueeze(0)
            processed_images.append(result_tensor)
            
            del source
            del pixels
            del mapped_pixels

        return (torch.cat(processed_images, dim=0),)


class PaletteSoftTransfer(ComfyNodeABC):
    @classmethod
    def INPUT_TYPES(cls) -> dict:
        return {
            "required": {
                "image": (IO.IMAGE,),
                "target_colors": ("COLOR_LIST",),
                "blend_mode": (["Original", "Grayscale"], {"default": "Original"}),
                "blend_ratio": ("FLOAT", {"min": 0, "max": 1, "step": 0.1, "default": 0.5}),
                "softness": ("FLOAT", {"min": 0, "max": 20, "step": 0.1, "default": 1}),
            }
        }

    RETURN_TYPES = (IO.IMAGE,)
    FUNCTION = "color_transfer"
    CATEGORY = "Color Transfer/Palette Transfer"

    def color_transfer(
        self,
        image: torch.Tensor,
        target_colors: list[tuple[int, int, int]],
        blend_mode: str,
        blend_ratio: float,
        softness: float,
    ) -> tuple[torch.Tensor]:
        
        if len(target_colors) < 2:
            raise ValueError("Palette must contain at least 2 colors")

        processed_images = []

        for img_tensor in image:
            source = img_tensor.cpu().numpy().astype(np.float32)

            img_lab = color.rgb2lab(source)
            palette_lab = np.array(
                [color.rgb2lab(np.array([[c]]) / 255.0)[0, 0] for c in target_colors]
            )

            pixels = img_lab.reshape(-1, 3).astype(np.float32)
            n_pixels = pixels.shape[0]
            n_colors = len(palette_lab)
            chunk_size = max(1, 500_000 // n_colors)

            projected = np.empty_like(pixels)
            for start in range(0, n_pixels, chunk_size):
                end = min(start + chunk_size, n_pixels)
                chunk = pixels[start:end]
                # (n_colors, chunk_len)
                dists = np.array([np.linalg.norm(chunk - p, axis=1) for p in palette_lab])
                weights = np.exp(-softness * dists)
                weights /= weights.sum(axis=0)
                projected[start:end] = np.dot(weights.T, palette_lab)

            if blend_mode == "Original":
                blended = (1 - blend_ratio) * pixels + blend_ratio * projected
            elif blend_mode == "Grayscale":
                gray = color.rgb2gray(source)
                gray_rgb = np.stack([gray] * 3, axis=-1)
                gray_lab = color.rgb2lab(gray_rgb)
                gray_pixels = gray_lab.reshape(-1, 3)
                blended = (1 - blend_ratio) * gray_pixels + blend_ratio * projected

            blended_lab = blended.reshape(img_lab.shape)
            blended_rgb = np.clip(color.lab2rgb(blended_lab), 0, 1).astype(np.float32)

            result_tensor = torch.from_numpy(blended_rgb).unsqueeze(0)
            processed_images.append(result_tensor)
            
            del source
            del img_lab
            del pixels

        return (torch.cat(processed_images, dim=0),)


class PaletteTransferReinhard(ComfyNodeABC):
    @classmethod
    def INPUT_TYPES(cls) -> dict:
        return {
            "required": {
                "image": (IO.IMAGE,),
                "target_colors": ("COLOR_LIST",),
            }
        }

    RETURN_TYPES = (IO.IMAGE,)
    FUNCTION = "color_transfer"
    CATEGORY = "Color Transfer/Palette Transfer"

    def color_transfer(
        self, image: torch.Tensor, target_colors: list[tuple[int, int, int]]
    ) -> tuple[torch.Tensor]:
        if len(target_colors) == 0:
            return (image,)

        target_colors = PaletteExtension.dense_palette(target_colors, points=3)

        # Compute target LAB stats directly from palette instead of creating a large image
        palette_rgb = np.array(target_colors, dtype=np.float32).reshape(-1, 1, 3) / 255.0
        palette_lab = color.rgb2lab(palette_rgb).reshape(-1, 3)
        t_mean = palette_lab.mean(axis=0)
        t_std = np.maximum(palette_lab.std(axis=0), 1e-6)

        processed_images = []

        for img_tensor in image:
            source = img_tensor.cpu().numpy().astype(np.float32)

            source_lab = color.rgb2lab(source)

            s_mean, s_std = source_lab.mean(axis=(0, 1)), source_lab.std(axis=(0, 1))

            # Avoid division by zero for uniform channels
            s_std = np.maximum(s_std, 1e-6)

            result_lab = (source_lab - s_mean) / s_std * t_std + t_mean
            result_rgb = np.clip(color.lab2rgb(result_lab), 0, 1).astype(np.float32)

            result_tensor = torch.from_numpy(result_rgb).unsqueeze(0)
            processed_images.append(result_tensor)

            del source
            del source_lab

        return (torch.cat(processed_images, dim=0),)


class PalleteTransferClustering(ComfyNodeABC):
    @classmethod
    def INPUT_TYPES(cls) -> dict:
        return {
            "required": {
                "image": (IO.IMAGE,),
                "target_colors": ("COLOR_LIST",),
                "palette_extension_method": (["Dense", "Edge", "None"], {"default": "None"}),
                "palette_extension_points": ("INT", {"min": 2, "max": 20, "step": 1, "default": 5}),
                "gaussian_blur": ("INT", {"default": 3, "min": 0, "max": 27, "step": 1}),
            }
        }

    RETURN_TYPES = (IO.IMAGE,)
    FUNCTION = "color_transfer"
    CATEGORY = "Color Transfer/Palette Transfer"

    def color_transfer(
        self,
        image: torch.Tensor,
        target_colors: list[tuple[int, int, int]],
        palette_extension_method: str,
        palette_extension_points: int,
        gaussian_blur: int,
    ) -> tuple[torch.Tensor]:
        if len(target_colors) == 0:
            return (image,)

        if palette_extension_method == "Dense":
            target_colors = PaletteExtension.dense_palette(target_colors, points=palette_extension_points)
        elif palette_extension_method == "Edge":
            target_colors = PaletteExtension.edge_based_palette(target_colors, points=palette_extension_points)

        output = process_image_with_palette(
            image=image,
            target_colors=target_colors,
            color_space="RGB",
            cluster_method="Mini batch Kmeans",
            distance_method="Euclidean",
            gaussian_blur=gaussian_blur,
        )

        return (output,)


class PaletteTransferNode(ComfyNodeABC):
    @classmethod
    def INPUT_TYPES(cls) -> dict:
        return {
            "required": {
                "image": (IO.IMAGE,),
                "target_colors": ("COLOR_LIST",),
                "color_space": (["RGB", "HSV", "LAB"], {"default": "RGB"}),
                "cluster_method": (["Kmeans", "Mini batch Kmeans"], {"default": "Kmeans"}),
                "distance_method": ([
                        "Euclidean",
                        "Manhattan",
                        "Cosine Similarity",
                        "HSV Distance",
                        "RGB Weighted Distance",
                        "RGB Weighted Similarity",
                    ], {"default": "Euclidean"}),
                "gaussian_blur": ("INT", {"default": 3, "min": 0, "max": 27, "step": 1}),
            }
        }

    RETURN_TYPES = (IO.IMAGE,)
    FUNCTION = "color_transfer"
    CATEGORY = "Color Transfer/Palette Transfer"

    def color_transfer(
        self,
        image: torch.Tensor,
        target_colors: list[tuple[int, int, int]],
        color_space: str,
        cluster_method: str,
        distance_method: str,
        gaussian_blur: int,
    ) -> tuple[torch.Tensor]:
        if len(target_colors) == 0:
            return (image,)

        output = process_image_with_palette(
            image=image,
            target_colors=target_colors,
            color_space=color_space,
            cluster_method=cluster_method,
            distance_method=distance_method,
            gaussian_blur=gaussian_blur,
        )

        return (output,)


class ColorPaletteNode(ComfyNodeABC):
    @classmethod
    def INPUT_TYPES(s) -> dict:
        return {
            "required": {
                "color_palette": (
                    IO.STRING,
                    {
                        "default": "[(30, 32, 30), (60, 61, 55), (105, 117, 101), (236, 223, 204)]",
                        "multiline": True,
                    },
                ),
            },
        }

    RETURN_TYPES = ("COLOR_LIST",)
    RETURN_NAMES = ("Color palette",)
    FUNCTION = "color_list"
    CATEGORY = "Color Transfer/Palette Transfer"

    def color_list(self, color_palette: str) -> tuple[list[tuple[int, int, int]]]:
        parsed = ast.literal_eval(color_palette)
        if not isinstance(parsed, list) or not all(
            isinstance(c, tuple) and len(c) == 3 and all(isinstance(v, int) and 0 <= v <= 255 for v in c)
            for c in parsed
        ):
            raise ValueError("Color palette must be a list of (R, G, B) tuples with values 0-255")
        return (parsed,)


class ExtractPaletteNode(ComfyNodeABC):
    @classmethod
    def INPUT_TYPES(cls) -> dict:
        return {
            "required": {
                "image": (IO.IMAGE,),
                "num_colors": ("INT", {"default": 5, "min": 1, "max": 50, "step": 1}),
                "cluster_method": (["Kmeans", "Mini batch Kmeans"], {"default": "Mini batch Kmeans"}),
            },
        }

    RETURN_TYPES = ("COLOR_LIST",)
    RETURN_NAMES = ("Color palette",)
    FUNCTION = "extract_palette"
    CATEGORY = "Color Transfer/Palette Transfer"

    def extract_palette(
        self, image: torch.Tensor, num_colors: int, cluster_method: str
    ) -> tuple[list[tuple[int, int, int]]]:
        
        img_tensor = image[0] if len(image.shape) == 4 else image
        
        # MEMORY FIX: Ensure float32 for clustering
        img = (255.0 * img_tensor.cpu().numpy()).astype(np.float32)
        
        clustering_engine = ColorClustering(cluster_method)
        clustering_result = clustering_engine.cluster_colors(img, num_colors)
        
        colors = clustering_result["main_colors"]
        color_list = [tuple(map(int, color)) for color in colors]
        
        # Cleanup
        del img
        del clustering_result
        
        return (color_list,)