fix: 🐛 fix from merge

Still need testing but works

nodes/conditions.py

@@ -6,6 +6,88 @@ import shutil
 import csv
 
 
+class InterpolateClipSequential:
+    @classmethod
+    def INPUT_TYPES(cls):
+        return {
+            "required": {
+                "base_text": ("STRING", {"multiline": True}),
+                "text_to_replace": ("STRING", {"default": ""}),
+                "clip": ("CLIP",),
+                "interpolation_strength": (
+                    "FLOAT",
+                    {"default": 1.0, "min": 0.0, "max": 1.0, "step": 0.01},
+                ),
+            }
+        }
+
+    RETURN_TYPES = ("CONDITIONING",)
+    FUNCTION = "interpolate_encodings_sequential"
+
+    CATEGORY = "mtb/conditioning"
+
+    def interpolate_encodings_sequential(
+        self, base_text, text_to_replace, clip, interpolation_strength, **replacements
+    ):
+        log.debug(f"Received interpolation_strength: {interpolation_strength}")
+
+        # - Ensure interpolation strength is within [0, 1]
+        interpolation_strength = max(0.0, min(1.0, interpolation_strength))
+
+        # - Check if replacements were provided
+        if not replacements:
+            raise ValueError("At least one replacement should be provided.")
+
+        num_replacements = len(replacements)
+        log.debug(f"Number of replacements: {num_replacements}")
+
+        segment_length = 1.0 / num_replacements
+        log.debug(f"Calculated segment_length: {segment_length}")
+
+        # - Find the segment that the interpolation_strength falls into
+        segment_index = min(
+            int(interpolation_strength // segment_length), num_replacements - 1
+        )
+        log.debug(f"Segment index: {segment_index}")
+
+        # - Calculate the local strength within the segment
+        local_strength = (
+            interpolation_strength - (segment_index * segment_length)
+        ) / segment_length
+        log.debug(f"Local strength: {local_strength}")
+
+        # - If it's the first segment, interpolate between base_text and the first replacement
+        if segment_index == 0:
+            replacement_text = list(replacements.values())[0]
+            log.debug("Using the base text a the base blend")
+            # -  Start with the base_text condition
+            tokens = clip.tokenize(base_text)
+            cond_from, pooled_from = clip.encode_from_tokens(tokens, return_pooled=True)
+        else:
+            base_replace = list(replacements.values())[segment_index - 1]
+            log.debug(f"Using {base_replace} a the base blend")
+
+            # - Start with the base_text condition replaced by the closest replacement
+            tokens = clip.tokenize(base_text.replace(text_to_replace, base_replace))
+            cond_from, pooled_from = clip.encode_from_tokens(tokens, return_pooled=True)
+
+            replacement_text = list(replacements.values())[segment_index]
+
+        interpolated_text = base_text.replace(text_to_replace, replacement_text)
+        tokens = clip.tokenize(interpolated_text)
+        cond_to, pooled_to = clip.encode_from_tokens(tokens, return_pooled=True)
+
+        # - Linearly interpolate between the two conditions
+        interpolated_condition = (
+            1.0 - local_strength
+        ) * cond_from + local_strength * cond_to
+        interpolated_pooled = (
+            1.0 - local_strength
+        ) * pooled_from + local_strength * pooled_to
+
+        return ([[interpolated_condition, {"pooled_output": interpolated_pooled}]],)
+
+
 class SmartStep:
     """Utils to control the steps start/stop of the KAdvancedSampler in percentage"""
 
@@ -96,4 +178,4 @@ class StylesLoader:
         return (self.options[style_name][0], self.options[style_name][1])
 
 
-__nodes__ = [SmartStep, StylesLoader]
+__nodes__ = [SmartStep, StylesLoader, InterpolateClipSequential]

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, closed=False
+            )
+            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]
+        psd = utils.tensor2pytolayer(image, layer_name, mask=mask)
+        # log.warning("Mask is currently ignored for PSD Layers...")
+        return ({group: psd},)
+
+
+__nodes__ = [PsdLayer, PsdSave]

reqs.txt

@@ -5,3 +5,4 @@ tensorflow
 facexlib==0.3.0
 insightface==0.7.3
 basicsr==1.4.2
+pytoshop

reqs_windows.txt

@@ -15,4 +15,5 @@ tensorflow==2.10.1;
 tb-nightly==2.12.0a20230126; platform_system == "Windows"
 facexlib==0.3.0
 # the old tf version on windows comes with a breaking protobuf version
-protobuf==3.19.6
+protobuf==3.19.6
+pytoshop

utils.py

@@ -3,6 +3,13 @@ import numpy as np
 import torch
 from pathlib import Path
 import sys
+
+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
 from typing import List
 import signal
 from contextlib import suppress
@@ -467,3 +474,53 @@ def apply_easing(value, easing_type):
 
 
 # endregion
+
+
+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,
+    mask: Optional[
+        torch.Tensor
+    ] = None,  # Add the mask parameter with default value as None
+) -> nested_layers.Image:
+    batch_count = tensor.size(0) if len(tensor.shape) > 3 else 1
+    if batch_count > 1:
+        raise ValueError(
+            f"Only one image is supported (batch size is currently {batch_count})"
+        )
+    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)
+
+    channels = [arr[:, :, 0], arr[:, :, 1], arr[:, :, 2], mask_arr[:, :, 0]]
+
+    image = 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,
+    )
+
+    return image