From 331a07df78681b935da4ebd7a6018927e1546cea Mon Sep 17 00:00:00 2001 From: hakoniwa <137879019+852wa@users.noreply.github.com> Date: Wed, 29 Jan 2025 21:13:40 +0900 Subject: [PATCH] Add files via upload --- nodes.py | 239 ++++++++++++++++++++++++++++++++----------------------- 1 file changed, 140 insertions(+), 99 deletions(-) diff --git a/nodes.py b/nodes.py index b074036..05eec51 100644 --- a/nodes.py +++ b/nodes.py @@ -1,7 +1,7 @@ # nodes.py import torch import numpy as np -from sklearn.cluster import KMeans +from sklearn.cluster import KMeans # GPUでクラスタリングが難しいので、CPUで実行する import colorsys import json import random @@ -14,67 +14,79 @@ class ColorshiftColorNode: def INPUT_TYPES(cls): return { "required": { - "image": ("IMAGE",), + "images": ("IMAGE",), "color_count": ("INT", {"default": 8, "min": 2, "max": 32}), }, "optional": { - "lock_mask": ("MASK", {"default": None}), + "lock_masks": ("MASK", {"default": None}), "palette_override": ("PALETTE", {"default": None}), "font_size": ("INT", {"default": 20, "min": 10, "max": 50}), } } RETURN_TYPES = ("IMAGE", "PALETTE", "MASK", "IMAGE") - RETURN_NAMES = ("image", "palette", "index_map", "palette_preview") + RETURN_NAMES = ("images", "palette", "index_maps", "palette_preview") FUNCTION = "process" CATEGORY = "Image/Color" - def process(self, image, color_count, lock_mask=None, palette_override=None, font_size=20): - # Convert torch tensor to numpy array - img_np = image.numpy().squeeze() - height, width = img_np.shape[:2] - pixels = img_np.reshape(-1, 3) + def process(self, images, color_count, lock_masks=None, palette_override=None, font_size=20): + # デバイスを取得 + device = images.device + + # バッチ処理のために画像をスタック + batch_size, height, width, channels = images.shape + pixels = images.reshape(batch_size, -1, 3) if palette_override is not None: - if hasattr(palette_override, 'numpy'): - palette = palette_override.numpy() - else: - palette = palette_override - labels = self._match_palette(pixels, palette) + palette = palette_override.to(device) + # バッチ全体に対して一回の距離計算でラベルを取得 + labels = self._match_palette_batch(pixels, palette) else: - kmeans = KMeans(n_clusters=color_count, random_state=42) - labels = kmeans.fit_predict(pixels) - palette = kmeans.cluster_centers_ + # 複数の画像に対して、それぞれでクラスタリングをすると、計算量も増えるし、パレットがずれる。 + # 解決策として、複数の画像を混ぜてクラスタリングを行う + # 混ぜてクラスタリング + pixels_for_clustering = pixels.cpu().numpy().copy().reshape(-1,3) + random_indices = np.random.choice(pixels_for_clustering.shape[0], size=int(pixels_for_clustering.shape[0] * 0.5), replace=False) + pixels_for_clustering = pixels_for_clustering[random_indices] + kmeans = KMeans(n_clusters=color_count, random_state=42, n_init=10) + kmeans.fit(pixels_for_clustering) - # Generate palette preview - preview_img = self.generate_palette_preview(palette, font_size) + # バッチ内の各画像に対してラベルを取得 + labels = torch.from_numpy(kmeans.predict(pixels.cpu().numpy().reshape(-1,3)).reshape(batch_size, -1)).long().to(device) + palette = torch.from_numpy(kmeans.cluster_centers_).float().to(device) - # Get the unique index map first - index_map = labels.reshape(height, width) - index_map_tensor = torch.from_numpy(index_map.astype(np.float32)) + # Generate palette preview (最初の画像のみ, CPUで処理) + preview_img = self.generate_palette_preview(palette.cpu().numpy(), font_size) - if lock_mask is not None: - if hasattr(lock_mask, 'numpy'): - mask = lock_mask.numpy() - else: - mask = lock_mask - - pixel_mask = np.zeros_like(labels, dtype=bool) - for i, is_locked in enumerate(mask): - if is_locked: - pixel_mask[labels == i] = True - - original_pixels = pixels.copy() + # バッチ全体のインデックスマップを取得 + index_maps = labels.reshape(batch_size, height, width) + + # GPUで一括処理 + if lock_masks is not None: + lock_masks_batch = [] + for i in range(batch_size): + lock_mask = lock_masks[i] if i < lock_masks.shape[0] else lock_masks[-1] + lock_masks_batch.append(lock_mask) + lock_masks = torch.stack(lock_masks_batch).to(device) + + pixel_masks = torch.zeros_like(index_maps, dtype=torch.bool) + for i in range(batch_size): + for j, is_locked in enumerate(lock_masks[i]): + if is_locked: + pixel_masks[i][index_maps[i] == j] = True + + original_pixels = pixels.clone() # ディープコピー new_pixels = palette[labels] - pixels[~pixel_mask] = new_pixels[~pixel_mask] - pixels[pixel_mask] = original_pixels[pixel_mask] + pixels[~pixel_masks] = new_pixels[~pixel_masks] + pixels[pixel_masks] = original_pixels[pixel_masks] else: pixels = palette[labels] - processed_image = torch.from_numpy(pixels.reshape((1, height, width, 3))).float() - return (processed_image, torch.from_numpy(palette), index_map_tensor, preview_img) + processed_images_tensor = pixels.reshape(batch_size, height, width, 3) + return (processed_images_tensor, palette, index_maps, preview_img) def generate_palette_preview(self, palette, font_size): + # 既存のコード (変更なし) if isinstance(palette, torch.Tensor): colors = palette.numpy() else: @@ -109,18 +121,19 @@ class ColorshiftColorNode: img_tensor = torch.from_numpy(np.array(img).astype(np.float32) / 255.0).unsqueeze(0) return img_tensor - def _match_palette(self, pixels, palette): - distances = np.linalg.norm(pixels[:, np.newaxis] - palette, axis=2) - return np.argmin(distances, axis=1) + def _match_palette_batch(self, pixels, palette): + # バッチ処理のために距離計算を最適化（PyTorchで書き換え） + distances = torch.cdist(pixels, palette) + return torch.argmin(distances, dim=2) class PaletteEditorNode: @classmethod def INPUT_TYPES(cls): return { "required": { - "image": ("IMAGE",), + "images": ("IMAGE",), "palette": ("PALETTE",), - "index_map": ("MASK",), + "index_maps": ("MASK",), }, "optional": { "seed": ("INT", {"default": 0, "min": 0, "max": 0xffffffffffffffff}), @@ -138,37 +151,80 @@ class PaletteEditorNode: } RETURN_TYPES = ("IMAGE", "PALETTE", "MASK") - RETURN_NAMES = ("image", "palette", "mask") + RETURN_NAMES = ("images", "palette", "mask") FUNCTION = "process_palette" CATEGORY = "Image/Color" def rgb_to_hsv(self, rgb): - return np.array([colorsys.rgb_to_hsv(*c) for c in rgb]) + # PyTorchで書き換え + r, g, b = rgb.unbind(-1) + max_rgb, argmax_rgb = rgb.max(-1) + min_rgb, _ = rgb.min(-1) + delta = max_rgb - min_rgb + + h = torch.zeros_like(r) + s = torch.zeros_like(r) + v = max_rgb + + non_zero_delta = delta != 0 + h[non_zero_delta] = torch.where(argmax_rgb[non_zero_delta] == 0, + (g[non_zero_delta] - b[non_zero_delta]) / delta[non_zero_delta], + torch.where(argmax_rgb[non_zero_delta] == 1, + 2 + (b[non_zero_delta] - r[non_zero_delta]) / delta[non_zero_delta], + 4 + (r[non_zero_delta] - g[non_zero_delta]) / delta[non_zero_delta])) + h = (h / 6.0) % 1.0 + s[non_zero_delta] = delta[non_zero_delta] / max_rgb[non_zero_delta] + + return torch.stack((h, s, v), dim=-1) def hsv_to_rgb(self, hsv): - return np.array([colorsys.hsv_to_rgb(*c) for c in hsv]) + # PyTorchで書き換え + h, s, v = hsv.unbind(-1) - def process_palette(self, image, palette, index_map, + c = v * s + x = c * (1 - torch.abs((h * 6) % 2 - 1)) + m = v - c + + rgb_prime = torch.zeros_like(hsv) + + h_category = (h * 6).long() % 6 + + rgb_prime[..., 0] = torch.where((h_category == 0) | (h_category == 5), c, torch.where((h_category == 1) | (h_category == 4), x, 0)) + rgb_prime[..., 1] = torch.where((h_category == 1) | (h_category == 2), c, torch.where((h_category == 0) | (h_category == 3), x, 0)) + rgb_prime[..., 2] = torch.where((h_category == 3) | (h_category == 4), c, torch.where((h_category == 2) | (h_category == 5), x, 0)) + + rgb = (rgb_prime + m[..., None]) + return rgb + + def process_palette(self, images, palette, index_maps, seed, hue_random_enable, hue_shift, saturation_random_enable, saturation_scale, value_random_enable, value_scale, lock_color_num="0,", mask_enable=False, invert_mask=False, mask_type="lock", operations="[]"): + + # デバイスを取得 + device = images.device + palette = palette.to(device) + index_maps = index_maps.to(device) + + batch_size = images.shape[0] + # If mask is not enabled, apply changes to the entire palette if not mask_enable: - mask = np.zeros(len(palette), dtype=bool) + mask = torch.zeros(len(palette), dtype=torch.bool, device=device) else: try: indices = [int(idx.strip()) for idx in lock_color_num.split(",") if idx.strip()] - mask = np.zeros(len(palette), dtype=bool) + mask = torch.zeros(len(palette), dtype=torch.bool, device=device) for idx in indices: if 0 <= idx < len(palette): mask[idx] = True if invert_mask: mask = ~mask except: - mask = np.zeros(len(palette), dtype=bool) + mask = torch.zeros(len(palette), dtype=torch.bool, device=device) current_hue = random.uniform(-180, 180) if hue_random_enable else hue_shift current_saturation = random.uniform(0, 5) if saturation_random_enable else saturation_scale @@ -177,21 +233,42 @@ class PaletteEditorNode: modified_palette = self.edit_palette(palette, operations, mask, current_hue, current_saturation, current_value) - # Apply the modified palette to the image - processed_image = self.apply_palette(image, modified_palette, index_map, mask, mask_type) - + # バッチ処理 + pixels = images.reshape(batch_size, -1, 3) + labels = index_maps.reshape(batch_size, -1) + new_pixels = modified_palette[labels.long()] + original_pixels = pixels.clone() + + if mask_type == "change": + # Create pixel_mask based on color_mask + pixel_mask = torch.zeros_like(labels, dtype=torch.bool) + for i, is_masked in enumerate(mask): + if is_masked: + pixel_mask[labels == i] = True + # Apply new_pixels to masked areas + pixels = torch.where(pixel_mask.unsqueeze(-1), new_pixels, original_pixels) + else: + # Create pixel_mask based on color_mask + pixel_mask = torch.zeros_like(labels, dtype=torch.bool) + for i, is_masked in enumerate(mask): + if is_masked: + pixel_mask[labels == i] = True + # Apply new_pixels to non-masked areas + pixels = torch.where(pixel_mask.unsqueeze(-1), original_pixels, new_pixels) + + processed_images = pixels.reshape(images.shape) + random.seed(None) - return (processed_image, modified_palette, mask) + return (processed_images, modified_palette, mask) def edit_palette(self, palette, operations, mask, hue_shift, saturation_scale, value_scale): - palette_np = palette.numpy() - hsv_palette = self.rgb_to_hsv(palette_np) + hsv_palette = self.rgb_to_hsv(palette) # Apply HSV adjustments to unmasked colors hsv_palette[~mask, 0] = (hsv_palette[~mask, 0] + hue_shift/360.0) % 1.0 - hsv_palette[~mask, 1] = np.clip(hsv_palette[~mask, 1] * saturation_scale, 0, 1) - hsv_palette[~mask, 2] = np.clip(hsv_palette[~mask, 2] * value_scale, 0, 1) + hsv_palette[~mask, 1] = torch.clamp(hsv_palette[~mask, 1] * saturation_scale, 0, 1) + hsv_palette[~mask, 2] = torch.clamp(hsv_palette[~mask, 2] * value_scale, 0, 1) try: ops = json.loads(operations) @@ -199,54 +276,18 @@ class PaletteEditorNode: for op in ops: if isinstance(op, dict) and "index" in op: idx = op["index"] - if 0 <= idx < len(palette_np): + if 0 <= idx < len(palette): if "color" in op and len(op["color"]) == 3: - new_color = np.array(op["color"], dtype=np.float32) - palette_np[idx] = new_color + new_color = torch.tensor(op["color"], dtype=torch.float32, device=palette.device) + palette[idx] = new_color hsv_palette[idx] = self.rgb_to_hsv(new_color.reshape(1,3))[0] if "hsv" in op and len(op["hsv"]) == 3: - hsv_palette[idx] = np.clip(op["hsv"], [0,0,0], [1,1,1]) + hsv_palette[idx] = torch.clamp(torch.tensor(op["hsv"], device=palette.device), 0, 1) except Exception as e: print(f"Operation Error: {str(e)}") modified_palette = self.hsv_to_rgb(hsv_palette) - return torch.from_numpy(modified_palette).float() - - def apply_palette(self, image, palette, index_map, mask, mask_type="lock"): - img_np = image.numpy().squeeze() - height, width = img_np.shape[:2] - pixels = img_np.reshape(-1, 3) - - if hasattr(palette, 'numpy'): - palette = palette.numpy() - - index_map = index_map.numpy() if hasattr(index_map, 'numpy') else index_map - color_mask = mask.numpy() if hasattr(mask, 'numpy') else mask - - labels = index_map.reshape(-1) - new_pixels = palette[labels.astype(int)] - original_pixels = pixels.copy() - - # Check if mask is entirely False (mask disabled) - if np.all(~color_mask): - # Apply new_pixels to all pixels - pixels = new_pixels - else: - # Create pixel_mask based on color_mask - pixel_mask = np.zeros_like(labels, dtype=bool) - for i, is_masked in enumerate(color_mask): - if is_masked: - pixel_mask[labels == i] = True - - if mask_type == "change": - # Apply new_pixels to masked areas - pixels = np.where(pixel_mask[:, np.newaxis], new_pixels, original_pixels) - else: - # Apply new_pixels to non-masked areas - pixels = np.where(pixel_mask[:, np.newaxis], original_pixels, new_pixels) - - processed_image = torch.from_numpy(pixels.reshape((1, height, width, 3))).float() - return processed_image + return modified_palette NODE_CLASS_MAPPINGS = { "ColorshiftColor": ColorshiftColorNode,