diff --git a/nodes/psd.py b/nodes/psd.py index 1a468a9..9398eee 100644 --- a/nodes/psd.py +++ b/nodes/psd.py @@ -39,7 +39,7 @@ class PsdSave: for group, layers in groups.items(): current_group = nested_layers.Group( - group, visible=True, opacity=255, layers=layers + group, visible=True, opacity=255, layers=layers, closed=False ) out_layers.append(current_group) @@ -81,9 +81,9 @@ class PsdLayer: # 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)},) + psd = utils.tensor2pytolayer(image, layer_name, mask=mask) + # log.warning("Mask is currently ignored for PSD Layers...") + return ({group: psd},) __nodes__ = [PsdLayer, PsdSave] diff --git a/utils.py b/utils.py index 7465ea6..21130e2 100644 --- a/utils.py +++ b/utils.py @@ -4,9 +4,10 @@ import torch from pathlib import Path import sys -from typing import Union, List +from typing import List, Optional from pytoshop.user import nested_layers from pytoshop import enums +# from pytoshop.layers import LayerMask, LayerRecord from .log import log @@ -106,6 +107,7 @@ def tensor2pytolayer( metadata: dict = {}, layer_color=0, color_mode=None, + mask: Optional[torch.Tensor] = None, # Add the mask parameter with default value as None ) -> nested_layers.Image: batch_count = 1 if len(tensor.shape) > 3: @@ -115,13 +117,19 @@ def tensor2pytolayer( raise Exception( f"Only one image is supported (batch size is currently {batch_count})" ) - out_channels = tensor2pil(tensor) + out_channels = tensor2pil(tensor)[0] arr = np.array(out_channels) + + # If a mask is provided, convert it to numpy array + if mask is not None: + mask_arr = np.array(tensor2pil(mask)[0]) + else: + mask_arr = np.full_like(arr, 255, dtype=np.uint8) - # 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( + channels = [arr[:, :, 0], arr[:, :, 1], arr[:, :, 2], mask_arr[:, :, 0]] + + + image = nested_layers.Image( name=name, visible=visible, opacity=opacity, @@ -134,3 +142,6 @@ def tensor2pytolayer( layer_color=layer_color, color_mode=color_mode, ) + + + return image \ No newline at end of file