feat: ✨ half working POC

Supports group but create trimming mask for each group for some reason.
I also want a better logic to group batches.

nodes/psd.py

@@ -0,0 +1,89 @@
+from pytoshop.user import nested_layers
+
+# from pytoshop.image_data import ImageData
+from .. import utils
+from ..log import log
+from uuid import uuid4
+from pathlib import Path
+import folder_paths
+from importlib import reload
+
+
+class PsdSave:
+    def __init__(self):
+        pass
+
+    @classmethod
+    def INPUT_TYPES(cls):
+        return {
+            "required": {
+                "input_1": ("PSDLAYER",),
+            },
+        }
+
+    RETURN_TYPES = ()
+    FUNCTION = "psd_save"
+    CATEGORY = "psd"
+    OUTPUT_NODE = True
+
+    def psd_save(self, **kwargs):
+        groups = {
+            "main": [],
+        }
+        out_layers = []
+        for input, item in kwargs.items():
+            for group, layer in item.items():
+                if group not in groups:
+                    groups[group] = []
+                groups[group].append(layer)
+
+        for group, layers in groups.items():
+            current_group = nested_layers.Group(
+                group, visible=True, opacity=255, layers=layers
+            )
+            out_layers.append(current_group)
+
+        out_layers = nested_layers.nested_layers_to_psd(out_layers, color_mode=3)
+        output_name = f"{uuid4()}.psd"
+        output_path = Path(folder_paths.output_directory) / output_name
+
+        log.info(f"Saving PSD to {output_name}")
+
+        with open(output_path, "wb") as f:
+            out_layers.write(f)
+
+        return ()
+
+
+class PsdLayer:
+    def __init__(self):
+        pass
+
+    @classmethod
+    def INPUT_TYPES(cls):
+        return {
+            "required": {
+                "layer_name": ("STRING", {"default": "layer"}),
+                "image": ("IMAGE",),
+            },
+            "optional": {"mask": ("MASK",)},
+        }
+
+    RETURN_TYPES = ("PSDLAYER",)
+    FUNCTION = "psd_layer"
+    CATEGORY = "psd"
+
+    def psd_layer(self, layer_name, image, mask=None):
+        reload(utils)
+        group = "main"
+        if "/" in layer_name:
+            sepname = layer_name.split("/")
+            # layer_name = sepname.pop() # todo: support nesting?
+            group = sepname[0]
+            layer_name = sepname[1]
+
+        log.warning("Mask is currently ignored for PSD Layers...")
+        return ({group: utils.tensor2pytolayer(image, layer_name)},)
+
+
+__nodes__ = [PsdLayer, PsdSave]

requirements.txt

@@ -14,4 +14,5 @@ tb-nightly==2.12.0a20230126; platform_system == "Windows"
 tensorflow; platform_system != "Windows"
 # the old tf version on windows comes with a breaking protobuf version
 protobuf==3.19.6
-gdown @ git+https://github.com/melMass/gdown@main
+gdown @ git+https://github.com/melMass/gdown@main
+pytoshop

utils.py

@@ -5,6 +5,9 @@ from pathlib import Path
 import sys
 
 from typing import Union, List
+from pytoshop.user import nested_layers
+from pytoshop import enums
+from .log import log
 
 
 def add_path(path, prepend=False):
@@ -32,7 +35,7 @@ comfy_dir = here.parent.parent
 # Construct the path to the font file
 font_path = here / "font.ttf"
 
-# Add extern folder to path
+# Add exteextern folder to path
 extern_root = here / "extern"
 add_path(extern_root)
 for pth in extern_root.iterdir():
@@ -45,16 +48,21 @@ add_path(comfy_dir)
 add_path((comfy_dir / "custom_nodes"))
 
 
-def tensor2pil(image: torch.Tensor) -> Union[Image.Image, List[Image.Image]]:
+def tensor2pil(image: torch.Tensor) -> List[Image.Image]:
     batch_count = 1
     if len(image.shape) > 3:
         batch_count = image.size(0)
 
-    if batch_count == 1:
-        return Image.fromarray(
+    if batch_count > 1:
+        out = []
+        out.extend([tensor2pil(image[i]) for i in range(batch_count)])
+        return out
+
+    return [
+        Image.fromarray(
             np.clip(255.0 * image.cpu().numpy().squeeze(), 0, 255).astype(np.uint8)
         )
-    return [tensor2pil(image[i]) for i in range(batch_count)]
+    ]
 
 
 def pil2tensor(image: Image.Image | List[Image.Image]) -> torch.Tensor:
@@ -76,5 +84,51 @@ def tensor2np(tensor: torch.Tensor) -> Union[np.ndarray, List[np.ndarray]]:
     if len(tensor.shape) > 3:
         batch_count = tensor.size(0)
     if batch_count > 1:
-        return [tensor2np(tensor[i]) for i in range(batch_count)]
-    return np.clip(255.0 * tensor.cpu().numpy().squeeze(), 0, 255).astype(np.uint8)
+        out = []
+        out.extend([tensor2np(tensor[i]) for i in range(batch_count)])
+        return out
+
+    return [np.clip(255.0 * tensor.cpu().numpy().squeeze(), 0, 255).astype(np.uint8)]
+
+
+def tensor2pytolayer(
+    tensor: torch.Tensor,
+    name: str,
+    visible: bool = True,
+    opacity: int = 255,
+    group_id: int = 0,
+    blend_mode=enums.BlendMode.normal,
+    x: int = 0,
+    y: int = 0,
+    # channels: int = 3,
+    metadata: dict = {},
+    layer_color=0,
+    color_mode=None,
+) -> nested_layers.Image:
+    batch_count = 1
+    if len(tensor.shape) > 3:
+        batch_count = tensor.size(0)
+
+    if batch_count > 1:
+        raise Exception(
+            f"Only one image is supported (batch size is currently {batch_count})"
+        )
+    out_channels = tensor2pil(tensor)
+    arr = np.array(out_channels)
+
+    # the array is currently H, W, C but we want C, H, W
+    # out_channels = np.transpose(out_channels, (2, 0, 1))
+    channels = [arr[:, :, 0], arr[:, :, 1], arr[:, :, 2]]
+    return nested_layers.Image(
+        name=name,
+        visible=visible,
+        opacity=opacity,
+        group_id=group_id,
+        blend_mode=blend_mode,
+        top=y,
+        left=x,
+        channels=channels,
+        metadata=metadata,
+        layer_color=layer_color,
+        color_mode=color_mode,
+    )

web/mtb_widget_register.js

@@ -0,0 +1,40 @@
+import { app } from "/scripts/app.js";
+/**
+ * @returns {import("./types/comfy").ComfyExtension} extension
+ */
+const mtb_widgets = {
+    name: "mtb.core.register",
+
+    /**
+     * 
+     * @param {import("./types/litegraph").LGraphNode} node 
+     */
+    async nodeCreated(node, app) {
+        if (node.comfyClass === "Psd Save (mtb)") {
+            node.onConnectionsChange = function (type, index, connected, link_info) {
+
+                // remove all non connected inputs
+                if (!connected && node.inputs.length > 1) {
+                    node.removeInput(index)
+
+                    // make inputs sequential again
+                    for (let i = 0; i < node.inputs.length; i++) {
+                        node.inputs[i].name = `input_${i + 1}`
+                    }
+                }
+
+                // add an extra input
+                if (node.inputs[node.inputs.length - 1].link != undefined) {
+                    node.addInput(`input_${node.inputs.length + 1}`, "PSDLAYER")
+                }
+
+
+            }
+
+        }
+    },
+
+};
+
+
+app.registerExtension(mtb_widgets);