diff --git a/nodes/psd.py b/nodes/psd.py new file mode 100644 index 0000000..1a468a9 --- /dev/null +++ b/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] diff --git a/requirements.txt b/requirements.txt index abb3097..3586dad 100644 --- a/requirements.txt +++ b/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 \ No newline at end of file +gdown @ git+https://github.com/melMass/gdown@main +pytoshop \ No newline at end of file diff --git a/utils.py b/utils.py index ce94de4..6a98efb 100644 --- a/utils.py +++ b/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, + ) diff --git a/web/mtb_widget_register.js b/web/mtb_widget_register.js new file mode 100644 index 0000000..44f9704 --- /dev/null +++ b/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); \ No newline at end of file