diff --git a/README.md b/README.md index 9d6b658..3cd1210 100644 --- a/README.md +++ b/README.md @@ -8,6 +8,8 @@ Click name to jump to workflow 3. [**EFDM**](#efdm). Paper: [Exact Feature Distribution Matching for Arbitrary Style Transfer and Domain Generalization](https://github.com/YBZh/EFDM) 4. [**MicroAST**](#microast). Paper: [Towards Super-Fast Ultra-Resolution Arbitrary Style Transfer](https://github.com/EndyWon/MicroAST) +[**Extra nodes**](#extra-nodes) + More will come soon ## Workflows @@ -103,6 +105,13 @@ Arguments: Similar to [EFDM](#efdm) ![wf_microast](workflows/wf_microast.png) + +## Extra nodes + +### Coral Color Transfer + +![extra_coral_color_transfer](workflows/extra/extra_coral_color_transfer.png) + ## Credits ``` diff --git a/__init__.py b/__init__.py index 887df5b..0ac2141 100644 --- a/__init__.py +++ b/__init__.py @@ -1,10 +1,12 @@ from .run import NeuralNeighbor, CAST, EFDM, MicroAST +from .run_extra import CoralColorTransfer NODE_CLASS_MAPPINGS = { "NeuralNeighbor": NeuralNeighbor, "CAST": CAST, "EFDM": EFDM, "MicroAST": MicroAST, + "CoralColorTransfer": CoralColorTransfer, } NODE_DISPLAY_NAME_MAPPINGS = { @@ -12,6 +14,7 @@ NODE_DISPLAY_NAME_MAPPINGS = { "CAST": "CAST", "EFDM": "EFDM", "MicroAST": "MicroAST", + "CoralColorTransfer": "Coral Color Transfer", } diff --git a/module_extra/function.py b/module_extra/function.py new file mode 100644 index 0000000..8c0a89f --- /dev/null +++ b/module_extra/function.py @@ -0,0 +1,41 @@ +import torch + +def coral(source: torch.Tensor, target: torch.Tensor): + # assume both source and target are 3D array (C, H, W) + # Note: flatten -> f + + source_f, source_f_mean, source_f_std = _calc_feat_flatten_mean_std(source) + source_f_norm = ( + source_f - source_f_mean.expand_as(source_f) + ) / source_f_std.expand_as(source_f) + source_f_cov_eye = torch.mm(source_f_norm, source_f_norm.t()) + torch.eye(3) + + target_f, target_f_mean, target_f_std = _calc_feat_flatten_mean_std(target) + target_f_norm = ( + target_f - target_f_mean.expand_as(target_f) + ) / target_f_std.expand_as(target_f) + target_f_cov_eye = torch.mm(target_f_norm, target_f_norm.t()) + torch.eye(3) + + source_f_norm_transfer = torch.mm( + _mat_sqrt(target_f_cov_eye), + torch.mm(torch.inverse(_mat_sqrt(source_f_cov_eye)), source_f_norm), + ) + + source_f_transfer = source_f_norm_transfer * target_f_std.expand_as( + source_f_norm + ) + target_f_mean.expand_as(source_f_norm) + + return source_f_transfer.view(source.size()) + +def _calc_feat_flatten_mean_std(feat): + # takes 3D feat (C, H, W), return mean and std of array within channels + assert feat.size()[0] == 3 + assert isinstance(feat, torch.FloatTensor) + feat_flatten = feat.view(3, -1) + mean = feat_flatten.mean(dim=-1, keepdim=True) + std = feat_flatten.std(dim=-1, keepdim=True) + return feat_flatten, mean, std + +def _mat_sqrt(x): + U, D, V = torch.svd(x) + return torch.mm(torch.mm(U, D.pow(0.5).diag()), V.t()) \ No newline at end of file diff --git a/run.py b/run.py index dcab48c..99b2d0b 100644 --- a/run.py +++ b/run.py @@ -2,7 +2,7 @@ @author: Fuou Marinas @title: ComfyUI-StyleTransferPlus @nickname: StyleTransferPlus -@description: A collection of style trasnfer nodes. +@description: A collection of style transfer nodes. """ from pathlib import Path diff --git a/run_extra.py b/run_extra.py new file mode 100644 index 0000000..269c62c --- /dev/null +++ b/run_extra.py @@ -0,0 +1,63 @@ +from pathlib import Path + +import torch + +from .module_extra.function import coral + +base_dir = Path(__file__).resolve().parent + + +class CoralColorTransfer: + @classmethod + def INPUT_TYPES(s): + return { + "required": { + "src_img": ("IMAGE",), + "style_img": ("IMAGE",), + }, + } + + RETURN_TYPES = ("IMAGE",) + RETURN_NAMES = ("res_img",) + FUNCTION = "todo" + CATEGORY = "StyleTransferPlus/Extra" + + def todo( + self, + src_img: torch.Tensor, + style_img: torch.Tensor, + ): + print(f"{src_img.shape=}") + print(f"{style_img.shape=}") + + device = torch.device("cpu") + + src_img = src_img.permute(0, 3, 1, 2).to(device) # [B, C, H, W] + style_img = style_img.permute(0, 3, 1, 2).to(device) + + result = [] + if src_img.shape[0] == 1 and style_img.shape[0] == 1: + # Case 1: Single source image, single style image + result = [coral(src_img[0], style_img[0]).unsqueeze(0)] + elif src_img.shape[0] == 1: + # Case 2: Single source image, multiple style images + num_frames = style_img.shape[0] + for i in range(num_frames): + transferred = coral(src_img[0], style_img[i]) + result.append(transferred.unsqueeze(0)) + elif style_img.shape[0] == 1: + # Case 3: Multiple source images, single style image + num_frames = src_img.shape[0] + for i in range(num_frames): + transferred = coral(src_img[i], style_img[0]) + result.append(transferred.unsqueeze(0)) + else: + # Case 4: Multiple source images, multiple style images + num_frames = min(src_img.shape[0], style_img.shape[0]) + for i in range(num_frames): + transferred = coral(src_img[i], style_img[i]) + result.append(transferred.unsqueeze(0)) + + res_tensor = torch.cat(result, dim=0).permute(0, 2, 3, 1) + + return (res_tensor,) diff --git a/workflows/extra/extra_coral_color_transfer.json b/workflows/extra/extra_coral_color_transfer.json new file mode 100644 index 0000000..168d79d --- /dev/null +++ b/workflows/extra/extra_coral_color_transfer.json @@ -0,0 +1,189 @@ +{ + "last_node_id": 8, + "last_link_id": 23, + "nodes": [ + { + "id": 5, + "type": "CoralColorTransfer", + "pos": [ + 843, + 329 + ], + "size": { + "0": 216.59999084472656, + "1": 46 + }, + "flags": {}, + "order": 2, + "mode": 0, + "inputs": [ + { + "name": "src_img", + "type": "IMAGE", + "link": 23 + }, + { + "name": "style_img", + "type": "IMAGE", + "link": 22 + } + ], + "outputs": [ + { + "name": "res_img", + "type": "IMAGE", + "links": [ + 21 + ], + "shape": 3, + "slot_index": 0 + } + ], + "properties": { + "Node name for S&R": "CoralColorTransfer" + } + }, + { + "id": 2, + "type": "LoadImage", + "pos": [ + 319, + 222 + ], + "size": { + "0": 315, + "1": 314 + }, + "flags": {}, + "order": 0, + "mode": 0, + "outputs": [ + { + "name": "IMAGE", + "type": "IMAGE", + "links": [ + 23 + ], + "shape": 3, + "slot_index": 0 + }, + { + "name": "MASK", + "type": "MASK", + "links": null, + "shape": 3 + } + ], + "properties": { + "Node name for S&R": "LoadImage" + }, + "widgets_values": [ + "006.jpg", + "image" + ] + }, + { + "id": 3, + "type": "LoadImage", + "pos": [ + 315, + 583 + ], + "size": { + "0": 315, + "1": 314 + }, + "flags": {}, + "order": 1, + "mode": 0, + "outputs": [ + { + "name": "IMAGE", + "type": "IMAGE", + "links": [ + 22 + ], + "shape": 3, + "slot_index": 0 + }, + { + "name": "MASK", + "type": "MASK", + "links": null, + "shape": 3 + } + ], + "properties": { + "Node name for S&R": "LoadImage" + }, + "widgets_values": [ + "S4.jpg", + "image" + ] + }, + { + "id": 4, + "type": "SaveImage", + "pos": [ + 845, + 430 + ], + "size": [ + 435.2039984095661, + 302.78571714680663 + ], + "flags": {}, + "order": 3, + "mode": 0, + "inputs": [ + { + "name": "images", + "type": "IMAGE", + "link": 21 + } + ], + "properties": {}, + "widgets_values": [ + "StyleTransfer" + ] + } + ], + "links": [ + [ + 21, + 5, + 0, + 4, + 0, + "IMAGE" + ], + [ + 22, + 3, + 0, + 5, + 1, + "IMAGE" + ], + [ + 23, + 2, + 0, + 5, + 0, + "IMAGE" + ] + ], + "groups": [], + "config": {}, + "extra": { + "ds": { + "scale": 1.1, + "offset": [ + -100.20813064097217, + -119.37827913027836 + ] + } + }, + "version": 0.4 +} \ No newline at end of file diff --git a/workflows/extra/extra_coral_color_transfer.png b/workflows/extra/extra_coral_color_transfer.png new file mode 100644 index 0000000..c554aa0 Binary files /dev/null and b/workflows/extra/extra_coral_color_transfer.png differ