Merge pull request #5 from FuouM/extra

add coral color transfer extra node

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
 
 ```

__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",
 }
 
 

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())

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

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,)

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
+}