import torch import numpy as np from PIL import Image import json import base64 from io import BytesIO from server import PromptServer from comfy_execution.graph import ExecutionBlocker class ImageColorSampler: """ This node allows clicking on an input image to sample colors, creating a color palette from the selected sample points. """ @classmethod def INPUT_TYPES(cls): return { "required": { "image": ("IMAGE",), "sample_points": ("STRING", {"default": "[]", "multiline": True}), "palette_size": ("INT", {"default": 128, "min": 32, "max": 512}), "sample_size": ("INT", {"default": 5, "min": 1, "max": 15}), "wait_for_input": ("BOOLEAN", {"default": True}) }, "hidden": { "node_id": "UNIQUE_ID", }, } RETURN_TYPES = ("IMAGE", "STRING") RETURN_NAMES = ("palette", "sampled_colors") FUNCTION = "create_palette" CATEGORY = "image/color" # Track which nodes are waiting for user input waiting_nodes = set() def tensor_to_base64_image(self, tensor): """Convert a torch tensor to a base64 encoded image string""" # Convert tensor to numpy and then to PIL image img_np = tensor.cpu().numpy().squeeze(0) img_np = (img_np * 255).astype(np.uint8) img_pil = Image.fromarray(img_np) # Save to a bytes buffer and convert to base64 buffered = BytesIO() img_pil.save(buffered, format="PNG") img_str = base64.b64encode(buffered.getvalue()).decode('utf-8') return f"data:image/png;base64,{img_str}" def create_palette(self, image, sample_points, palette_size=128, sample_size=5, wait_for_input=True, node_id=None): """ Creates a color palette from the sampled points on the image. Args: image: Input image tensor sample_points: JSON string of sample points coordinates and colors palette_size: Size of the palette image (height in pixels) sample_size: Size of sample area (radius) for color averaging wait_for_input: Whether to block execution waiting for user input node_id: Unique ID of this node instance Returns: Tuple of (palette_image, sampled_colors_json) """ # Parse the sample points try: points = json.loads(sample_points) except json.JSONDecodeError: points = [] # Check if this is the initial call or a resumption after user input is_initial_call = node_id not in self.waiting_nodes # For initial call, send image data to the UI for interactive editing if is_initial_call and wait_for_input: # Convert image tensor to a base64 string for sending to UI img_base64 = self.tensor_to_base64_image(image) # Send image and current points to the UI ui_data = { "image": img_base64, "sample_points": points, "sample_size": sample_size, "node_id": node_id } # Send message to UI to display the image for interaction PromptServer.instance.send_sync("image_sampler_init", {"node": node_id, "data": ui_data}) # Add to waiting nodes and block execution self.waiting_nodes.add(node_id) # Return ExecutionBlocker for both outputs return (ExecutionBlocker(None), ExecutionBlocker(None)) # Remove from waiting list if resuming if node_id in self.waiting_nodes: self.waiting_nodes.remove(node_id) # Convert image tensor to numpy array img_np = image.cpu().numpy().squeeze(0) # Image dimensions height, width, _ = img_np.shape # If no points, return empty palette if not points: # Create empty palette palette_img = np.zeros((palette_size, palette_size, 3), dtype=np.float32) palette_tensor = torch.from_numpy(palette_img)[None, ] return (palette_tensor, "[]") # Calculate colors for each sample point sampled_colors = [] for point in points: x = int(point["x"] * width) y = int(point["y"] * height) # Ensure coordinates are within bounds x = max(sample_size, min(width - sample_size - 1, x)) y = max(sample_size, min(height - sample_size - 1, y)) # Sample area - take average color in the sample radius sample_area = img_np[y-sample_size:y+sample_size+1, x-sample_size:x+sample_size+1] avg_color = np.mean(sample_area, axis=(0, 1)) # Convert to 8-bit RGB r, g, b = [int(c * 255) for c in avg_color] # Add to colors list with position info sampled_colors.append({ "position": {"x": point["x"], "y": point["y"]}, "color": {"r": r, "g": g, "b": b}, "hex": f"#{r:02X}{g:02X}{b:02X}" }) # Create palette image num_colors = len(sampled_colors) if num_colors == 0: # Create empty palette palette_img = np.zeros((palette_size, palette_size, 3), dtype=np.float32) palette_tensor = torch.from_numpy(palette_img)[None, ] return (palette_tensor, "[]") # Create palette image - a horizontal strip of colors stripe_height = palette_size stripe_width = palette_size // num_colors if num_colors > 0 else palette_size palette_img = np.zeros((stripe_height, palette_size, 3), dtype=np.float32) for i, color_data in enumerate(sampled_colors): color = color_data["color"] start_x = i * stripe_width end_x = (i + 1) * stripe_width if i < num_colors - 1 else palette_size # Fill the stripe with the color palette_img[:, start_x:end_x, 0] = color["r"] / 255.0 palette_img[:, start_x:end_x, 1] = color["g"] / 255.0 palette_img[:, start_x:end_x, 2] = color["b"] / 255.0 # Convert to tensor palette_tensor = torch.from_numpy(palette_img)[None, ] # Return palette image and colors data as JSON return (palette_tensor, json.dumps(sampled_colors)) # Add routes for handling the continuation of workflow @PromptServer.instance.routes.post("/image_sampler/continue") async def image_sampler_continue(request): """Handle when user is done selecting color samples and wants to continue""" data = await request.json() node_id = data.get("node_id") sample_points = data.get("sample_points", "[]") # Update the sample_points widget value to continue with new points PromptServer.instance.send_sync("image_sampler_update", { "node": node_id, "widget_name": "sample_points", "value": json.dumps(sample_points) }) return {"status": "success"}