diff --git a/color_transfer.py b/color_transfer.py index c1f2473..48a6aaa 100644 --- a/color_transfer.py +++ b/color_transfer.py @@ -4,12 +4,12 @@ import torch import ast -def EuclideanDistance(current_colors, target_colors): - return np.linalg.norm(current_colors - target_colors, axis=1) +def EuclideanDistance(detected_colors, target_colors): + return np.linalg.norm(detected_colors - target_colors, axis=1) -def ManhattanDistance(current_colors, target_colors): - return np.sum(np.abs(current_colors - target_colors), axis=1) +def ManhattanDistance(detected_colors, target_colors): + return np.sum(np.abs(detected_colors - target_colors), axis=1) def ColorClustering(image, k, cluster_method): @@ -20,14 +20,14 @@ def ColorClustering(image, k, cluster_method): "Mini batch Kmeans": MiniBatchKMeans } - kmeans = cluster_methods.get(cluster_method)(n_clusters=k) + clustering_model = cluster_methods.get(cluster_method)(n_clusters=k, n_init='auto') - kmeans.fit(img_array) - main_colors = kmeans.cluster_centers_ - return image, main_colors.astype(int), kmeans + clustering_model.fit(img_array) + main_colors = clustering_model.cluster_centers_ + return image, main_colors.astype(int), clustering_model -def SwitchColors(image, current_colors, target_colors, kmeans, distance_method): +def SwitchColors(image, detected_colors, target_colors, clustering_model, distance_method): closest_colors = [] distance_methods = { @@ -37,17 +37,18 @@ def SwitchColors(image, current_colors, target_colors, kmeans, distance_method): distance_method = distance_methods.get(distance_method) - for color in current_colors: + for color in detected_colors: distances = distance_method(color, target_colors) closest_color = target_colors[np.argmin(distances)] closest_colors.append(closest_color) + closest_colors = np.array(closest_colors) - image = closest_colors[kmeans.labels_].reshape(image.shape) + image = closest_colors[clustering_model.labels_].reshape(image.shape) image = np.array(image).astype(np.float32) / 255.0 - image = torch.from_numpy(image)[None,] + processedImage = torch.from_numpy(image)[None,] - return image + return processedImage class PaletteTransferNode: @@ -56,7 +57,7 @@ class PaletteTransferNode: data_in = { "required": { "image": ("IMAGE",), - "colors": ("COLORS",), + "target_colors": ("COLORS",), "cluster_method": (["Kmeans","Mini batch Kmeans"], {'default': 'Kmeans'}, ), "distance_method": (["Euclidean", "Manhattan"], {'default': 'Euclidean'}, ) } @@ -68,22 +69,23 @@ class PaletteTransferNode: CATEGORY = "Palette Transfer" - def color_transfer(self, image, colors, cluster_method, distance_method): + def color_transfer(self, image, target_colors, cluster_method, distance_method): - if len(colors) == 0: + if len(target_colors) == 0: return (image,) - else: - processedImages = [] + + processedImages = [] - for image in image: - img = 255. * image.cpu().numpy() + for image in image: + img = 255. * image.cpu().numpy() - img, current_colors, kmeans = ColorClustering(img, len(colors), cluster_method) - processed = SwitchColors(img, current_colors, colors, kmeans, distance_method) - processedImages.append(processed) - output = torch.cat(processedImages, dim=0) + clustered_img, detected_colors, clustering_model = ColorClustering(img, len(target_colors), cluster_method) + processed = SwitchColors(clustered_img, detected_colors, target_colors, clustering_model, distance_method) + processedImages.append(processed) + + output = torch.cat(processedImages, dim=0) - return (output, ) + return (output, ) class ColorPaletteNode: @@ -91,7 +93,7 @@ class ColorPaletteNode: def INPUT_TYPES(s): return { "required": { - "colors": ("STRING", {'default': '', 'multiline': True}) + "color_palette": ("STRING", {'default': '', 'multiline': True}) }, } @@ -99,5 +101,5 @@ class ColorPaletteNode: RETURN_NAMES = ("Color palette", ) FUNCTION = "color_list" - def color_list(self, colors): - return (ast.literal_eval(colors), ) + def color_list(self, color_palette): + return (ast.literal_eval(color_palette), ) diff --git a/color_transfer_example.png b/color_transfer_example.png index 93c0ece..6ce458d 100644 Binary files a/color_transfer_example.png and b/color_transfer_example.png differ diff --git a/workflow_examples/example_workflow.json b/workflow_examples/example_workflow.json new file mode 100644 index 0000000..f3512b1 --- /dev/null +++ b/workflow_examples/example_workflow.json @@ -0,0 +1,821 @@ +{ + "last_node_id": 71, + "last_link_id": 113, + "nodes": [ + { + "id": 58, + "type": "PreviewImage", + "pos": [ + 1580, + 520 + ], + "size": { + "0": 562.9608764648438, + "1": 307.1505432128906 + }, + "flags": {}, + "order": 13, + "mode": 0, + "inputs": [ + { + "name": "images", + "type": "IMAGE", + "link": 107 + } + ], + "properties": { + "Node name for S&R": "PreviewImage" + } + }, + { + "id": 44, + "type": "CLIPTextEncode", + "pos": [ + 271, + 432 + ], + "size": { + "0": 407.7621154785156, + "1": 86.47399139404297 + }, + "flags": {}, + "order": 5, + "mode": 0, + "inputs": [ + { + "name": "clip", + "type": "CLIP", + "link": 67, + "slot_index": 0 + } + ], + "outputs": [ + { + "name": "CONDITIONING", + "type": "CONDITIONING", + "links": [ + 60 + ], + "shape": 3, + "slot_index": 0 + } + ], + "properties": { + "Node name for S&R": "CLIPTextEncode" + }, + "widgets_values": [ + "low quality, (bad artist, watermark, text), blurry, blurry edges," + ] + }, + { + "id": 68, + "type": "PaletteTransfer", + "pos": [ + 1230, + 680 + ], + "size": { + "0": 315, + "1": 102 + }, + "flags": {}, + "order": 10, + "mode": 0, + "inputs": [ + { + "name": "image", + "type": "IMAGE", + "link": 101 + }, + { + "name": "target_colors", + "type": "COLORS", + "link": 104 + } + ], + "outputs": [ + { + "name": "IMAGE", + "type": "IMAGE", + "links": [ + 108 + ], + "shape": 3, + "slot_index": 0 + } + ], + "properties": { + "Node name for S&R": "PaletteTransfer" + }, + "widgets_values": [ + "Kmeans", + "Manhattan" + ] + }, + { + "id": 69, + "type": "PaletteTransfer", + "pos": [ + 1230, + 830 + ], + "size": { + "0": 315, + "1": 102 + }, + "flags": {}, + "order": 11, + "mode": 0, + "inputs": [ + { + "name": "image", + "type": "IMAGE", + "link": 102 + }, + { + "name": "target_colors", + "type": "COLORS", + "link": 111 + } + ], + "outputs": [ + { + "name": "IMAGE", + "type": "IMAGE", + "links": [ + 109 + ], + "shape": 3, + "slot_index": 0 + } + ], + "properties": { + "Node name for S&R": "PaletteTransfer" + }, + "widgets_values": [ + "Mini batch Kmeans", + "Euclidean" + ] + }, + { + "id": 61, + "type": "PreviewImage", + "pos": [ + 2160, + 520 + ], + "size": { + "0": 562.9608764648438, + "1": 307.1505432128906 + }, + "flags": {}, + "order": 14, + "mode": 0, + "inputs": [ + { + "name": "images", + "type": "IMAGE", + "link": 108 + } + ], + "properties": { + "Node name for S&R": "PreviewImage" + } + }, + { + "id": 63, + "type": "PreviewImage", + "pos": [ + 2158, + 888 + ], + "size": { + "0": 562.9608764648438, + "1": 307.1505432128906 + }, + "flags": {}, + "order": 16, + "mode": 0, + "inputs": [ + { + "name": "images", + "type": "IMAGE", + "link": 110 + } + ], + "properties": { + "Node name for S&R": "PreviewImage" + } + }, + { + "id": 62, + "type": "PreviewImage", + "pos": [ + 1576, + 888 + ], + "size": { + "0": 562.9608764648438, + "1": 307.1505432128906 + }, + "flags": {}, + "order": 15, + "mode": 0, + "inputs": [ + { + "name": "images", + "type": "IMAGE", + "link": 109 + } + ], + "properties": { + "Node name for S&R": "PreviewImage" + } + }, + { + "id": 40, + "type": "EmptyLatentImage", + "pos": [ + 318, + 588 + ], + "size": { + "0": 315, + "1": 106 + }, + "flags": {}, + "order": 0, + "mode": 0, + "outputs": [ + { + "name": "LATENT", + "type": "LATENT", + "links": [ + 61 + ], + "slot_index": 0 + } + ], + "properties": { + "Node name for S&R": "EmptyLatentImage" + }, + "widgets_values": [ + 1344, + 768, + 1 + ] + }, + { + "id": 41, + "type": "CLIPTextEncode", + "pos": [ + 268, + 274 + ], + "size": { + "0": 411.97210693359375, + "1": 96.00001525878906 + }, + "flags": {}, + "order": 4, + "mode": 0, + "inputs": [ + { + "name": "clip", + "type": "CLIP", + "link": 62 + } + ], + "outputs": [ + { + "name": "CONDITIONING", + "type": "CONDITIONING", + "links": [ + 59 + ], + "slot_index": 0 + } + ], + "properties": { + "Node name for S&R": "CLIPTextEncode" + }, + "widgets_values": [ + "realistic, landscape, sunset" + ] + }, + { + "id": 38, + "type": "KSampler", + "pos": [ + 710, + 270 + ], + "size": { + "0": 315, + "1": 262 + }, + "flags": {}, + "order": 6, + "mode": 0, + "inputs": [ + { + "name": "model", + "type": "MODEL", + "link": 113 + }, + { + "name": "positive", + "type": "CONDITIONING", + "link": 59 + }, + { + "name": "negative", + "type": "CONDITIONING", + "link": 60 + }, + { + "name": "latent_image", + "type": "LATENT", + "link": 61 + } + ], + "outputs": [ + { + "name": "LATENT", + "type": "LATENT", + "links": [ + 64 + ], + "slot_index": 0 + } + ], + "properties": { + "Node name for S&R": "KSampler" + }, + "widgets_values": [ + 111312, + "fixed", + 15, + 11, + "heun", + "karras", + 1 + ] + }, + { + "id": 70, + "type": "PaletteTransfer", + "pos": [ + 1230, + 1000 + ], + "size": { + "0": 315, + "1": 102 + }, + "flags": {}, + "order": 12, + "mode": 0, + "inputs": [ + { + "name": "image", + "type": "IMAGE", + "link": 103 + }, + { + "name": "target_colors", + "type": "COLORS", + "link": 112 + } + ], + "outputs": [ + { + "name": "IMAGE", + "type": "IMAGE", + "links": [ + 110 + ], + "shape": 3, + "slot_index": 0 + } + ], + "properties": { + "Node name for S&R": "PaletteTransfer" + }, + "widgets_values": [ + "Mini batch Kmeans", + "Manhattan" + ] + }, + { + "id": 45, + "type": "PreviewImage", + "pos": [ + 1579, + 162 + ], + "size": { + "0": 562.9608764648438, + "1": 307.1505432128906 + }, + "flags": {}, + "order": 8, + "mode": 0, + "inputs": [ + { + "name": "images", + "type": "IMAGE", + "link": 68 + } + ], + "properties": { + "Node name for S&R": "PreviewImage" + } + }, + { + "id": 67, + "type": "PaletteTransfer", + "pos": [ + 1230, + 520 + ], + "size": { + "0": 315, + "1": 102 + }, + "flags": {}, + "order": 9, + "mode": 0, + "inputs": [ + { + "name": "image", + "type": "IMAGE", + "link": 100 + }, + { + "name": "target_colors", + "type": "COLORS", + "link": 99 + } + ], + "outputs": [ + { + "name": "IMAGE", + "type": "IMAGE", + "links": [ + 107 + ], + "shape": 3, + "slot_index": 0 + } + ], + "properties": { + "Node name for S&R": "PaletteTransfer" + }, + "widgets_values": [ + "Kmeans", + "Euclidean" + ] + }, + { + "id": 59, + "type": "ColorPalette", + "pos": [ + 980, + 600 + ], + "size": { + "0": 210, + "1": 94.74728393554688 + }, + "flags": {}, + "order": 1, + "mode": 0, + "outputs": [ + { + "name": "Color palette", + "type": "COLORS", + "links": [ + 99, + 104 + ], + "shape": 3, + "slot_index": 0 + } + ], + "properties": { + "Node name for S&R": "ColorPalette" + }, + "widgets_values": [ + "[(68,71,96), (97,105,174), \t(241,122,143), (255,187,107)\n,(244,237,237)]" + ] + }, + { + "id": 71, + "type": "ColorPalette", + "pos": [ + 977, + 919 + ], + "size": { + "0": 210, + "1": 94.74728393554688 + }, + "flags": {}, + "order": 2, + "mode": 0, + "outputs": [ + { + "name": "Color palette", + "type": "COLORS", + "links": [ + 111, + 112 + ], + "shape": 3, + "slot_index": 0 + } + ], + "properties": { + "Node name for S&R": "ColorPalette" + }, + "widgets_values": [ + "[(31,31,31), (66,54,88), \t(92,65,95), (123,86,106)\n,(146,123,134)]" + ] + }, + { + "id": 42, + "type": "VAEDecode", + "pos": [ + 1060, + 290 + ], + "size": { + "0": 140, + "1": 46 + }, + "flags": {}, + "order": 7, + "mode": 0, + "inputs": [ + { + "name": "samples", + "type": "LATENT", + "link": 64 + }, + { + "name": "vae", + "type": "VAE", + "link": 65 + } + ], + "outputs": [ + { + "name": "IMAGE", + "type": "IMAGE", + "links": [ + 68, + 100, + 101, + 102, + 103 + ], + "slot_index": 0 + } + ], + "properties": { + "Node name for S&R": "VAEDecode" + } + }, + { + "id": 39, + "type": "CheckpointLoaderSimple", + "pos": [ + -253, + 616 + ], + "size": { + "0": 350.9955749511719, + "1": 98 + }, + "flags": {}, + "order": 3, + "mode": 0, + "outputs": [ + { + "name": "MODEL", + "type": "MODEL", + "links": [ + 113 + ], + "slot_index": 0 + }, + { + "name": "CLIP", + "type": "CLIP", + "links": [ + 62, + 67 + ], + "slot_index": 1 + }, + { + "name": "VAE", + "type": "VAE", + "links": [ + 65 + ], + "slot_index": 2 + } + ], + "properties": { + "Node name for S&R": "CheckpointLoaderSimple" + }, + "widgets_values": [ + "SDXLFaetastic_v20.safetensors" + ] + } + ], + "links": [ + [ + 59, + 41, + 0, + 38, + 1, + "CONDITIONING" + ], + [ + 60, + 44, + 0, + 38, + 2, + "CONDITIONING" + ], + [ + 61, + 40, + 0, + 38, + 3, + "LATENT" + ], + [ + 62, + 39, + 1, + 41, + 0, + "CLIP" + ], + [ + 64, + 38, + 0, + 42, + 0, + "LATENT" + ], + [ + 65, + 39, + 2, + 42, + 1, + "VAE" + ], + [ + 67, + 39, + 1, + 44, + 0, + "CLIP" + ], + [ + 68, + 42, + 0, + 45, + 0, + "IMAGE" + ], + [ + 99, + 59, + 0, + 67, + 1, + "COLORS" + ], + [ + 100, + 42, + 0, + 67, + 0, + "IMAGE" + ], + [ + 101, + 42, + 0, + 68, + 0, + "IMAGE" + ], + [ + 102, + 42, + 0, + 69, + 0, + "IMAGE" + ], + [ + 103, + 42, + 0, + 70, + 0, + "IMAGE" + ], + [ + 104, + 59, + 0, + 68, + 1, + "COLORS" + ], + [ + 107, + 67, + 0, + 58, + 0, + "IMAGE" + ], + [ + 108, + 68, + 0, + 61, + 0, + "IMAGE" + ], + [ + 109, + 69, + 0, + 62, + 0, + "IMAGE" + ], + [ + 110, + 70, + 0, + 63, + 0, + "IMAGE" + ], + [ + 111, + 71, + 0, + 69, + 1, + "COLORS" + ], + [ + 112, + 71, + 0, + 70, + 1, + "COLORS" + ], + [ + 113, + 39, + 0, + 38, + 0, + "MODEL" + ] + ], + "groups": [], + "config": {}, + "extra": { + "ds": { + "scale": 0.8390545288824196, + "offset": { + "0": -781.0907873516428, + "1": -113.29259441450375 + } + } + }, + "version": 0.4 +} \ No newline at end of file