First Workflow push

.github/workflows/test_node.yml

@@ -0,0 +1,10 @@
+name: Test
+on:
+  workflow_dispatch:
+jobs:
+  build:
+    runs-on: ubuntu-latest
+    steps:
+      - uses: actions/checkout@v2
+      - name: Setup upterm session
+        uses: lhotari/action-upterm@v1

README.md

@@ -1 +1,81 @@
 # ClipStuff
+## Basic Instructions.
+Clone repo into custom_nodes folder.
+Install the requirements.txt file via pip.
+
+Pass CLIP output from Load Checkpoint into SpecialClipLoader node, then use the outputted clip with standard Clip Text Encode.
+
+See example workflow in examples folder.
+
+### Example Photo
+![Example Photo](assets/first_example.png)
+
+## Functions
+
+## Syntax Elements
+
+1. **Embedding**:
+    - Syntax: `embedding:WORD`
+    - Example: `embedding:face_vector`
+    - Represents a named vector embedding(Textual Inversion).
+
+2. **Word**:
+    - Syntax: Any alphanumeric word including characters such as `,`, `_`, and `-`.
+    - Example: `cat, dog_face, id_123`
+    - Represents simple words or identifiers.
+
+3. **Quoted String**:
+    - Syntax: A string enclosed within double or single quotes. You can escape quotes inside the string using a backslash (`\`).
+    - Example: `"Hello World"`, `'It\'s a sunny day'`
+    - Represents string literals.
+
+## Functions
+
+Here are the available functions and their usage:
+
+1. **Sum Function**:
+    - Syntax: `sum(arg1 | arg2 | ... | argN)`
+    - Adds together multiple embeddings.
+    - Example: `sum(embedding:face1 | dog)`
+
+2. **Negation Function**:
+    - Syntax: `neg(arg)`
+    - Negates the output.
+    - Example: `neg(A embedding:happycats outside)`
+
+3. **Normalization Function**:
+    - Syntax: `norm(arg)`
+    - Normalizes the given vector embedding.
+    - Example: `norm(sum(embedding:face1 | embedding:face2))`
+
+4. **Difference Function**:
+    - Syntax: `diff(arg1 | arg2 | ... | argN)`
+    - Computes the difference between multiple vector embeddings.
+    - Example: `diff(embedding:face1 | embedding:face2)`
+
+### Notes on Arguments:
+- Each function takes one or more arguments.
+- An argument (`arg`) can be an embedding, a word, another function, or a quoted string.
+- For functions that accept multiple arguments, they are separated by the `|` symbol.
+
+## Examples
+
+1. Add two embeddings and normalize the result:
+   ```
+   norm(sum(cat | dog | horse | parrot))
+   ```
+
+2. Negate an embedding:
+   ```
+   neg(embedding:body_vector)
+   ```
+
+3. King - Man + Woman = Queen:
+   ```
+   sum(diff(king|man)|woman)
+   ```
+   or
+   ```
+   sum(king|neg(man)|woman)
+   ```
+   ```

__init__.py

@@ -1,11 +1,9 @@
 from .nodes import (
-    KepAdvTextEncode,
     BuildGif,
     SpecialClipLoader,
 )
 
 NODE_CLASS_MAPPINGS = {
-    "Kep Adv Text Encode": KepAdvTextEncode,
     "Build Gif": BuildGif,
     "Special CLIP Loader": SpecialClipLoader,
 }

lib/action/base.py

@@ -0,0 +1,100 @@
+from abc import ABC, abstractmethod
+from typing import Union
+
+from torch import Tensor
+from torch.nn import Embedding
+
+from custom_nodes.ClipStuff.lib.parser.prompt_segment import PromptSegment
+
+
+class Action(ABC):
+    @property
+    @abstractmethod
+    def chars(self) -> list[str] | None:
+        pass
+
+    @property
+    @abstractmethod
+    def name(self) -> str:
+        pass
+
+    @property
+    @abstractmethod
+    def grammar(self) -> str:
+        """
+        The grammar for this action. This is used to parse the action from the prompt.
+        :return:
+        """
+        pass
+
+    @abstractmethod
+    def token_length(self) -> int:
+        """
+        The length of the tokens that this action will add to the prompt.
+        :return:
+        """
+        pass
+
+    @abstractmethod
+    def get_all_segments(self) -> list[PromptSegment]:
+        """
+        Get all segments, including nested segments.
+        :return:
+        """
+        pass
+
+    @abstractmethod
+    def get_result(self, embedding_module: Embedding) -> Tensor:
+        """
+        Get the result of this action. This is called when the embeddings are being calculated.
+        :param embedding_module: The embedding module to use to get the base embeddings for tokens.
+        :return:
+        """
+        pass
+
+    def depth_repr(self, depth: int = 1) -> str:
+        raise NotImplementedError()
+
+
+class SingleArgAction(Action, ABC):
+    def get_all_segments(self) -> list[PromptSegment]:
+        segments = []
+        for seg_or_action in self.arg:
+            if isinstance(seg_or_action, Action):
+                segments.extend(seg_or_action.get_all_segments())
+            else:
+                segments.append(seg_or_action)
+        return segments
+
+    def __init__(self, arg: list[PromptSegment | Action]):
+        # TODO: Target is a list now... what does this mean for us..
+        self.arg = arg
+
+    def __repr__(self) -> str:
+        return f"{self.name}({self.arg})"
+
+class MultiArgAction(Action, ABC):
+    def get_all_segments(self) -> list[PromptSegment]:
+        segments = []
+        for seg_or_action in self.base_segment:
+            if isinstance(seg_or_action, Action):
+                segments.extend(seg_or_action.get_all_segments())
+            else:
+                segments.append(seg_or_action)
+
+        for arg in self.args:
+            for seg_or_action in arg:
+                if isinstance(seg_or_action, Action):
+                    segments.extend(seg_or_action.get_all_segments())
+                else:
+                    segments.append(seg_or_action)
+
+        return segments
+
+    def __init__(
+            self,
+            base_segment: list[PromptSegment | Action],
+            args: list[list[Union[PromptSegment, Action]]],
+    ):
+        self.base_segment = base_segment
+        self.args = args

lib/actions/__init__.py

@@ -1,6 +0,0 @@
-from custom_nodes.ClipStuff.lib.actions.arith import ArithAction
-from custom_nodes.ClipStuff.lib.actions.nudge import NudgeAction
-
-ALL_ACTIONS = [NudgeAction, ArithAction]
-ALL_START_CHARS = [action.START_CHAR for action in ALL_ACTIONS]
-ALL_END_CHARS = [action.END_CHAR for action in ALL_ACTIONS]

lib/actions/action_utils.py

@@ -0,0 +1,11 @@
+from torch import Tensor
+from torch.nn import Embedding
+
+from custom_nodes.ClipStuff.lib.action.base import Action
+from custom_nodes.ClipStuff.lib.actions.types import SegOrAction
+
+
+def get_embedding(seg_or_action: SegOrAction, embedding_module: Embedding) -> Tensor:
+    if isinstance(seg_or_action, Action):
+        return seg_or_action.get_result(embedding_module)
+    return seg_or_action.get_embeddings(embedding_module)

lib/actions/arith.py

@@ -1,131 +0,0 @@
-from typing import Callable, Union
-
-import torch
-from torch.nn import Embedding
-
-from comfy.sd1_clip import SD1Tokenizer
-from custom_nodes.ClipStuff.lib.actions.base import Action, PromptSegment
-
-
-class ArithAction(Action):
-    START_CHAR = "<"
-    END_CHAR = ">"
-
-    def __init__(self, base_segment: PromptSegment | Action, ops: dict[str, list[PromptSegment | Action]]):
-        self.base_segment = base_segment
-        self.ops = ops
-
-    def __repr__(self):
-        return f"ArithAction(\n\tbase_segment={self.base_segment},\n\tops={self.ops}\n)"
-
-    def depth_repr(self, depth=1):
-        out = "ArithAction(\n"
-        if isinstance(self.base_segment, Action):
-            base_segment_repr = self.base_segment.depth_repr(depth + 1)
-            out += "\t" * depth + f"base_segment={base_segment_repr}\n"
-        elif isinstance(self.base_segment, PromptSegment):
-            out += "\t" * depth + f'base_segment={self.base_segment.depth_repr(depth)}'
-        else:
-            out += "\t" * depth + f'base_segment="{self.base_segment}",'
-
-        for op_key, ops in self.ops.items():
-            for op in ops:
-                out += "\n" + "\t" * depth + f'"{op_key}":[\n'
-                if isinstance(op, Action):
-                    op_repr = op.depth_repr(depth + 2)
-                    out += "\t" * (depth + 1) + f"{op_repr}\n"
-                else:
-                    out += "\t" * (depth + 1) + f'{op.depth_repr()},\n'
-                out += "\t" * depth + "],"
-        out += "\n" + "\t" * (depth - 1) + ")"
-        return out
-
-    def token_length(self):
-        # ArithAction modifies the embeddings of the base segment, so the length is the length of the base segment
-        if isinstance(self.base_segment, Action):
-            return self.base_segment.token_length()
-
-        return len(self.base_segment.tokens)
-
-
-    def get_all_segments(self):
-        segments = []
-        if isinstance(self.base_segment, Action):
-            segments += self.base_segment.get_all_segments()
-        else:
-            segments.append(self.base_segment)
-
-        for op_key, ops in self.ops.items():
-            for op in ops:
-                if isinstance(op, Action):
-                    segments += op.get_all_segments()
-                else:
-                    segments.append(op)
-
-        return segments
-
-    def get_result(self, embedding_module: Embedding):
-        if isinstance(self.base_segment, Action):
-            base_segment_result = self.base_segment.get_result(embedding_module)
-        else:
-            base_segment_result = self.base_segment.get_embeddings(embedding_module)
-
-        for op_key, ops in self.ops.items():
-            for op in ops:
-                if isinstance(op, Action):
-                    op_result = op.get_result(embedding_module)
-                else:
-                    op_result = op.get_embeddings(embedding_module)
-
-
-                if op_result.shape[1] > base_segment_result.shape[1]:
-                    print('[WARN] ArithAction: op_result.shape[1] > base_segment_result.shape[1] - averaging op_result')
-                    op_result = torch.mean(op_result, dim=1, keepdim=True)
-
-                if op_key == "+":
-                    base_segment_result.add(op_result)
-                elif op_key == "-":
-                    base_segment_result.subtract(op_result)
-
-        return base_segment_result
-
-    @classmethod
-    def parse_segment(
-        cls,
-        tokens: list[str],
-        start_chars: list[str],
-        end_chars: list[str],
-        parent_parser: Callable[[list[str], SD1Tokenizer], Union[PromptSegment, 'Action']],
-        tokenizer: SD1Tokenizer,
-    ) -> Action:
-        """
-        Parse an arithmetic action from a list of tokens
-        Supported formats:
-        <base_segment:+op1-op2-op3>
-
-        :param tokens: List of tokens, will be modified
-        :param start_chars: List of start chars for all actions
-        :param end_chars: List of end chars for all actions
-        :param parent_parser: Function to parse segments to allow for nested actions
-        :return:
-        """
-        token = tokens.pop(0)
-        assert token == cls.START_CHAR, "ArithAction must start with " + cls.START_CHAR + " but got " + token
-
-        # Parse base segment
-        base_segment = parent_parser(tokens, tokenizer)
-
-        token = tokens.pop(0)
-        assert token == ":", "ArithAction must have a ':' after the base segment" + " but got " + token
-
-        # Parse ops string
-        ops = {'+': [], '-': []}
-        while tokens[0] != cls.END_CHAR:
-            op_char = tokens.pop(0)
-            assert op_char in ["+", "-"], "ArithAction must have a '+' or '-' as an op char but got " + op_char
-            ops[op_char].append(parent_parser(tokens, tokenizer))
-
-        token = tokens.pop(0)
-        assert token == cls.END_CHAR, "ArithAction must end with " + cls.END_CHAR + " but got " + token
-
-        return cls(base_segment, ops)

lib/actions/base.py

@@ -1,95 +0,0 @@
-from abc import ABC, abstractmethod
-from typing import Callable, Union
-
-import torch
-from torch import Tensor
-from torch.nn import Embedding
-
-from comfy.sd1_clip import SD1Tokenizer
-
-
-class Action(ABC):
-    @property
-    @abstractmethod
-    def START_CHAR(self):
-        pass
-
-    @property
-    @abstractmethod
-    def END_CHAR(self):
-        pass
-
-    @abstractmethod
-    def token_length(self):
-        pass
-
-    @abstractmethod
-    def get_all_segments(self):
-        pass
-
-    @abstractmethod
-    def get_result(self, embedding_module: Embedding):
-        pass
-
-    @classmethod
-    @abstractmethod
-    def parse_segment(
-        cls,
-        tokens: list[str],
-        start_chars: list[str],
-        end_chars: list[str],
-        parent_parser: Callable[[list[str], SD1Tokenizer], Union[str, 'Action']],
-        tokenizer: SD1Tokenizer,
-    ) -> 'Action':
-        pass
-
-    def depth_repr(self, depth=1):
-        raise NotImplementedError()
-
-class PromptSegment:
-    def __init__(self, text: str, tokens: list[Union[int, Tensor]]):
-        self.text = text
-        self.tokens = tokens
-
-    def token_length(self):
-        return len(self.tokens)
-
-    def get_embeddings(self, embedding_module: Embedding):
-        tensors = torch.LongTensor(self.tokens).to(torch.device('cpu'))
-        unsqueezed_tensors = tensors.unsqueeze(0)
-        return embedding_module(unsqueezed_tensors)
-
-    def depth_repr(self, depth=1):
-        out = f'"{self.text}"('
-
-        cleaned_tokens = list(map(lambda x: str(x) if isinstance(x, int) else "EMBD", self.tokens))
-        out += ", ".join(cleaned_tokens)
-
-        out += ")"
-        return out
-
-def build_prompt_segment(text: str, tokenizer: SD1Tokenizer) -> PromptSegment:
-    split_text = text.split(" ")
-    tokens = []
-    for word in split_text:
-        if word.startswith(tokenizer.embedding_identifier) and tokenizer.embedding_directory is not None:
-            embedding_name = word[len(tokenizer.embedding_identifier):].strip('\n')
-
-            get_embed_ret = tokenizer._try_get_embedding(embedding_name)
-            embedding = get_embed_ret[0]
-            leftover = get_embed_ret[1]
-            if embedding is None:
-                print(f"warning, embedding:{embedding_name} does not exist, ignoring")
-            else:
-                if len(embedding.shape) == 1:
-                    tokens.append(embedding)
-                else:
-                    tokens.extend(embedding)
-
-            if leftover != "":
-                word = leftover
-            else:
-                continue
-        tokens.extend(tokenizer.tokenizer(word)["input_ids"][1:-1])
-
-    return PromptSegment(text, tokens)

lib/actions/diff.py

@@ -0,0 +1,66 @@
+import torch
+from torch.nn import Embedding
+
+from custom_nodes.ClipStuff.lib.action.base import Action, MultiArgAction
+from custom_nodes.ClipStuff.lib.actions.action_utils import get_embedding
+
+
+class DiffAction(MultiArgAction):
+    grammar = 'diff(" arg ("|" arg)* ")"'
+    name = "diff"
+    chars = ["-", "-"]
+
+    def token_length(self) -> int:
+        # Sum adds to the embeddings of the base segment, so the length is the length of the base segment
+        return sum(seg_or_action.token_length() for seg_or_action in self.base_segment)
+
+    def get_result(self, embedding_module: Embedding) -> torch.Tensor:
+        # Calculate the embeddings for the base segment
+        all_base_embeddings = [
+            get_embedding(seg_or_action, embedding_module)
+            for seg_or_action in self.base_segment
+        ]
+
+        result = torch.cat(all_base_embeddings, dim=1)
+
+        for arg in self.args:
+            all_arg_embeddings = [
+                get_embedding(seg_or_action, embedding_module) for seg_or_action in arg
+            ]
+
+            arg_embedding = torch.cat(all_arg_embeddings, dim=1)
+
+            if (
+                arg_embedding.shape[-2] == 1
+                or result.shape[-2] == arg_embedding.shape[-2]
+            ):
+                result = result.sub(arg_embedding)
+            else:
+                print(
+                    "WARNING: shape mismatch when trying to apply sum, arg will be averaged"
+                )
+                result = result.sub(torch.mean(arg_embedding, dim=1, keepdim=True))
+
+        return result
+
+    # def __repr__(self):
+    #     return f"sum(\n\tbase_segment={self.base_segment},\n\targs={self.args}\n)"
+    def __repr__(self) -> str:
+        return f"sum({', '.join(map(str, self.args))})"
+
+    def depth_repr(self, depth=1):
+        out = "NudgeAction(\n"
+        if isinstance(self.base_segment, Action):
+            base_segment_repr = self.base_segment.depth_repr(depth + 1)
+            out += "\t" * depth + f"base_segment={base_segment_repr}\n"
+        else:
+            out += "\t" * depth + f"base_segment={self.base_segment.depth_repr()},\n"
+
+        if isinstance(self.args, Action):
+            target_repr = self.args.depth_repr(depth + 1)
+            out += "\t" * depth + f"target={target_repr},\n"
+        else:
+            out += "\t" * depth + f"target={self.args.depth_repr()},\n"
+        out += "\t" * depth + f"weight={self.weight},\n"
+        out += "\t" * (depth - 1) + ")"
+        return out

lib/actions/lib.py

@@ -1,17 +0,0 @@
-from custom_nodes.ClipStuff.lib.actions import ALL_START_CHARS, ALL_END_CHARS
-from custom_nodes.ClipStuff.lib.actions.base import Action
-
-
-def is_action_segment(action_class: Action.__class__, segment: str):
-    if not issubclass(action_class, Action):
-        raise Exception(
-            f"action_class must be a subclass of Action, got {action_class}"
-        )
-
-    return (
-        segment[0] == action_class.START_CHAR and segment[-1] == action_class.END_CHAR
-    )
-
-
-def is_any_action_segment(segment: str):
-    return segment[0] in ALL_START_CHARS and segment[-1] in ALL_END_CHARS

lib/actions/neg.py

@@ -0,0 +1,32 @@
+import torch
+from torch.nn import Embedding
+
+from custom_nodes.ClipStuff.lib.action.base import Action, SingleArgAction
+
+
+class NegAction(SingleArgAction):
+    grammar = 'neg(" arg+ ")"'
+    name = "neg"
+    chars = ["[", "]"]
+
+    def token_length(self) -> int:
+        """
+        Neg negates the embeddings of the base segment, so the length is the length of the base segment
+        :return:
+        """
+        total_length = 0
+        for seg_or_action in self.arg:
+            total_length += seg_or_action.token_length()
+
+        return total_length
+
+    def get_result(self, embedding_module: Embedding) -> torch.Tensor:
+        all_embeddings = []
+        for seg_or_action in self.arg:
+            if isinstance(seg_or_action, Action):
+                all_embeddings.append(seg_or_action.get_result(embedding_module))
+            else:
+                all_embeddings.append(seg_or_action.get_embeddings(embedding_module))
+
+        target_embeddings = torch.cat(all_embeddings, dim=1)
+        return target_embeddings * -1

lib/actions/norm.py

@@ -0,0 +1,37 @@
+import torch
+from torch.nn import Embedding
+
+from custom_nodes.ClipStuff.lib.action.base import (
+    Action,
+    SingleArgAction,
+)
+
+
+class NormAction(SingleArgAction):
+    grammar = 'norm(" arg+ ")"'
+    name = "norm"
+    chars = None
+
+    def token_length(self) -> int:
+        """
+        Norm normalizes the embeddings of the base segment, so the length is the length of the base segment
+        :return:
+        """
+        total_length = 0
+        for seg_or_action in self.arg:
+            total_length += seg_or_action.token_length()
+
+        return total_length
+
+    def get_result(self, embedding_module: Embedding) -> torch.Tensor:
+        all_embeddings = []
+        for seg_or_action in self.arg:
+            if isinstance(seg_or_action, Action):
+                target_embeddings = seg_or_action.get_result(embedding_module)
+            else:
+                target_embeddings = seg_or_action.get_embeddings(embedding_module)
+            all_embeddings.append(target_embeddings)
+
+        target_embeddings = torch.cat(all_embeddings, dim=1)
+        return torch.div(target_embeddings, torch.norm(target_embeddings, dim=-1, keepdim=True))
+

lib/actions/nudge.py

@@ -1,128 +0,0 @@
-from typing import Optional, Union, Callable
-
-import torch
-from torch.nn import Embedding
-
-from comfy.sd1_clip import SD1Tokenizer
-from custom_nodes.ClipStuff.lib.actions.base import Action, PromptSegment
-
-
-class NudgeAction(Action):
-    def token_length(self):
-        # Nudge nudges the embeddings of the base segment, so the length is the length of the base segment
-        if isinstance(self.base_segment, Action):
-            return self.base_segment.token_length()
-
-        return len(self.base_segment.tokens)
-
-    def get_all_segments(self):
-        segments = []
-        if isinstance(self.base_segment, Action):
-            segments += self.base_segment.get_all_segments()
-        else:
-            segments.append(self.base_segment)
-
-        if isinstance(self.target, Action):
-            segments += self.target.get_all_segments()
-        else:
-            segments.append(self.target)
-
-        return segments
-
-    def get_result(self, embedding_module: Embedding):
-        if isinstance(self.base_segment, Action):
-            base_segment_result = self.base_segment.get_result(embedding_module)
-        else:
-            base_segment_result = self.base_segment.get_embeddings(embedding_module)
-
-        if isinstance(self.target, Action):
-            target_segment_result = self.target.get_result(embedding_module)
-        else:
-            target_segment_result = self.target.get_embeddings(embedding_module)
-
-        base_mean = torch.mean(base_segment_result, dim=1, keepdim=True)
-        if target_segment_result.shape[1] == 1:
-            translation_vector = target_segment_result - base_mean
-        else:
-            translation_vector = torch.mean(target_segment_result, dim=1, keepdim=True) - base_mean
-
-        return base_segment_result.add(translation_vector, alpha=self.weight)
-
-    START_CHAR = "["
-    END_CHAR = "]"
-
-    def __init__(
-        self,
-        base_segment: PromptSegment | Action,
-        target: Union[PromptSegment, Action],
-        weight: Optional[float] = None,
-    ):
-        self.base_segment = base_segment
-        self.weight = weight
-        self.target = target
-
-    def __repr__(self):
-        return f"NudgeAction(\n\tbase_segment={self.base_segment},\n\ttarget={self.target},\n\tweight={self.weight}\n)"
-
-    def depth_repr(self, depth=1):
-        out = "NudgeAction(\n"
-        if isinstance(self.base_segment, Action):
-            base_segment_repr = self.base_segment.depth_repr(depth + 1)
-            out += "\t" * depth + f"base_segment={base_segment_repr}\n"
-        else:
-            out += "\t" * depth + f'base_segment={self.base_segment.depth_repr()},\n'
-
-        if isinstance(self.target, Action):
-            target_repr = self.target.depth_repr(depth + 1)
-            out += "\t" * depth + f"target={target_repr},\n"
-        else:
-            out += "\t" * depth + f"target={self.target.depth_repr()},\n"
-        out += "\t" * depth + f"weight={self.weight},\n"
-        out += "\t" * (depth - 1) + ")"
-        return out
-
-    @classmethod
-    def parse_segment(
-        cls,
-        tokens: list[str],
-        start_chars: list[str],
-        end_chars: list[str],
-        parent_parser: Callable[[list[str], SD1Tokenizer], PromptSegment | Action],
-        tokenizer: SD1Tokenizer,
-    ) -> Action:
-        """
-        Parse a nudge action from a list of tokens
-        Supported formats:
-        [base_segment:target_segment]
-        [base_segment:target_segment:weight]
-
-        Weight is optional, if not provided it will be None
-        :param tokens: List of tokens, will be modified
-        :param start_chars: List of start chars for all actions
-        :param end_chars: List of end chars for all actions
-        :param parent_parser: Function to parse segments to allow for nested actions
-        :return:
-        """
-        token = tokens.pop(0)
-        assert token == cls.START_CHAR, "NudgeAction must start with " + cls.START_CHAR + " got " + token
-
-        # Parse base segment
-        base_segment = parent_parser(tokens, tokenizer)
-
-        token = tokens.pop(0)
-        assert token == ":", "NudgeAction must have a ':' after the base segment" + " but got " + token
-
-        # Parse target segment
-        target_segment = parent_parser(tokens, tokenizer)
-
-        # Parse weight if it exists
-        weight = None
-        if tokens[0] == ":":
-            # Parse weight
-            tokens.pop(0)
-            weight = float(tokens.pop(0))
-
-        token = tokens.pop(0)
-        assert token == cls.END_CHAR, "NudgeAction must end with " + cls.END_CHAR + " got " + token
-
-        return cls(base_segment, target_segment, weight)

lib/actions/sum.py

@@ -0,0 +1,96 @@
+from typing import Union
+
+import torch
+from torch.nn import Embedding
+
+from custom_nodes.ClipStuff.lib.actions.action_utils import get_embedding
+from custom_nodes.ClipStuff.lib.parser.prompt_segment import PromptSegment
+from custom_nodes.ClipStuff.lib.action.base import Action
+
+
+class SumAction(Action):
+    grammar = 'sum(" arg ("|" arg)* ")"'
+    name = "sum"
+    chars = ["+", "+"]
+
+    def __init__(
+        self,
+        base_segment: list[PromptSegment | Action],
+        args: list[list[Union[PromptSegment, Action]]],
+    ):
+        self.base_segment = base_segment
+        self.args = args
+
+    def token_length(self) -> int:
+        # Sum adds to the embeddings of the base segment, so the length is the length of the base segment
+        return sum(seg_or_action.token_length() for seg_or_action in self.base_segment)
+
+    def get_all_segments(self) -> list[PromptSegment]:
+        segments = []
+        for seg_or_action in self.base_segment:
+            if isinstance(seg_or_action, Action):
+                segments.extend(seg_or_action.get_all_segments())
+            else:
+                segments.append(seg_or_action)
+
+        for arg in self.args:
+            for seg_or_action in arg:
+                if isinstance(seg_or_action, Action):
+                    segments.extend(seg_or_action.get_all_segments())
+                else:
+                    segments.append(seg_or_action)
+
+        return segments
+
+    def get_result(self, embedding_module: Embedding) -> torch.Tensor:
+        # Calculate the embeddings for the base segment
+        all_base_embeddings = [
+            get_embedding(seg_or_action, embedding_module)
+            for seg_or_action in self.base_segment
+        ]
+
+        result = torch.cat(all_base_embeddings, dim=1)
+
+        for arg in self.args:
+            all_arg_embeddings = [
+                get_embedding(seg_or_action, embedding_module) for seg_or_action in arg
+            ]
+
+            arg_embedding = torch.cat(all_arg_embeddings, dim=1)
+
+            if (
+                arg_embedding.shape[-2] == 1
+                or result.shape[-2] == arg_embedding.shape[-2]
+            ):
+                result = result.add(arg_embedding)
+            else:
+                print(
+                    "WARNING: shape mismatch when trying to apply sum, arg will be averaged"
+                )
+                result = result.add(
+                    torch.mean(arg_embedding, dim=1, keepdim=True)
+                )
+
+        return result
+
+    # def __repr__(self):
+    #     return f"sum(\n\tbase_segment={self.base_segment},\n\targs={self.args}\n)"
+    def __repr__(self) -> str:
+        return f"sum({', '.join(map(str, self.args))})"
+
+    def depth_repr(self, depth=1):
+        out = "NudgeAction(\n"
+        if isinstance(self.base_segment, Action):
+            base_segment_repr = self.base_segment.depth_repr(depth + 1)
+            out += "\t" * depth + f"base_segment={base_segment_repr}\n"
+        else:
+            out += "\t" * depth + f"base_segment={self.base_segment.depth_repr()},\n"
+
+        if isinstance(self.args, Action):
+            target_repr = self.args.depth_repr(depth + 1)
+            out += "\t" * depth + f"target={target_repr},\n"
+        else:
+            out += "\t" * depth + f"target={self.args.depth_repr()},\n"
+        out += "\t" * depth + f"weight={self.weight},\n"
+        out += "\t" * (depth - 1) + ")"
+        return out

lib/actions/types.py

@@ -0,0 +1,6 @@
+from typing import Union
+
+from custom_nodes.ClipStuff.lib.parser.prompt_segment import PromptSegment
+from custom_nodes.ClipStuff.lib.action.base import Action
+
+SegOrAction = Union[PromptSegment, Action]

lib/actions/utils.py

@@ -0,0 +1,6 @@
+from custom_nodes.ClipStuff.lib.actions.types import SegOrAction
+
+def batch_size_info(batch: list[SegOrAction]):
+    for segment in batch:
+        print("Token Len: " + str(segment.token_length()))
+        print(segment.depth_repr())

lib/clip_model.py

@@ -1,18 +1,19 @@
 import contextlib
 import os
-from typing import Union
 
 import torch
 from transformers import CLIPTextConfig, modeling_utils
 
 from comfy import model_management
 import comfy.ops
-from comfy.sd import CLIP
-from custom_nodes.ClipStuff.lib.actions.base import PromptSegment, Action
-from custom_nodes.ClipStuff.lib.fun_clip_stuff import MyCLIPTextModel
-from custom_nodes.ClipStuff.lib.tokenizer import TokenDict
+from custom_nodes.ClipStuff.lib.action.base import Action
+from custom_nodes.ClipStuff.lib.actions.types import SegOrAction
+from custom_nodes.ClipStuff.lib.fun_clip_stuff import PromptLangTextModel
+from custom_nodes.ClipStuff.lib.parser.prompt_segment import PromptSegment
 
-class SD1FunClipModel(torch.nn.Module):
+
+# Methods with no comment can be assumed to be the same as comfy.sd1_clip.SD1ClipModel
+class PromptLangClipModel(torch.nn.Module):
     """Uses the CLIP transformer encoder for text (from huggingface)"""
     LAYERS = [
         "last",
@@ -27,15 +28,20 @@ class SD1FunClipModel(torch.nn.Module):
         assert layer in self.LAYERS
         self.num_layers = 12
         if textmodel_path is not None:
-            self.transformer = MyCLIPTextModel.from_pretrained(textmodel_path)
+            # Our transformer
+            self.transformer = PromptLangTextModel.from_pretrained(textmodel_path)
         else:
             if textmodel_json_config is None:
+                # TODO: Maybe re-use clip config?
+                # Config could come from cond_stage_model.transformer.config
+                # Copied clip_config
                 textmodel_json_config = os.path.join(os.path.dirname(os.path.realpath(__file__)), "clip_config.json")
             config = CLIPTextConfig.from_json_file(textmodel_json_config)
             self.num_layers = config.num_hidden_layers
             with comfy.ops.use_comfy_ops():
                 with modeling_utils.no_init_weights():
-                    self.transformer = MyCLIPTextModel(config)
+                    # Our transformer
+                    self.transformer = PromptLangTextModel(config)
 
         self.max_length = max_length
         if freeze:
@@ -68,7 +74,8 @@ class SD1FunClipModel(torch.nn.Module):
         self.layer = self.layer_default[0]
         self.layer_idx = self.layer_default[1]
 
-    def set_up_textual_embeddings(self, tokens: list[list[PromptSegment | Action]], current_embeds):
+    # Completely changed to support Segments and actions
+    def set_up_textual_embeddings(self, tokens: list[list[SegOrAction]], current_embeds):
         next_new_token = token_dict_size = current_embeds.weight.shape[0] - 1
         embedding_weights = []
 
@@ -131,7 +138,8 @@ class SD1FunClipModel(torch.nn.Module):
                         if segment.tokens[tokenIdx] == -1:
                             segment.tokens[tokenIdx] = n
 
-    def forward(self, tokens, **kwargs):
+    # Support our set_up_textual_embeddings which modifies the input embeddings
+    def forward(self, tokens):
         backup_embeds = self.transformer.get_input_embeddings()
         device = backup_embeds.weight.device
         self.set_up_textual_embeddings(tokens, backup_embeds)
@@ -142,16 +150,8 @@ class SD1FunClipModel(torch.nn.Module):
         else:
             precision_scope = contextlib.nullcontext
 
-
-        if (kwargs.get("position_ids", None) is not None):
-            position_ids = torch.LongTensor(kwargs["position_ids"]).to(device)
-        else:
-            position_ids = None
-
-
         with precision_scope(model_management.get_autocast_device(device)):
-            outputs = self.transformer(input_ids=tokens, output_hidden_states=self.layer == "hidden",
-                                       position_ids=position_ids)
+            outputs = self.transformer(input_ids=tokens, output_hidden_states=self.layer == "hidden")
             self.transformer.set_input_embeddings(backup_embeds)
 
             if self.layer == "last":
@@ -168,18 +168,20 @@ class SD1FunClipModel(torch.nn.Module):
                 pooled_output = pooled_output.to(self.text_projection.device) @ self.text_projection
         return z.float(), pooled_output.float()
 
-    def encode(self, tokens, **kwargs):
-        return self(tokens, **kwargs)
+    def encode(self, tokens):
+        return self(tokens)
 
     def load_sd(self, sd):
         return self.transformer.load_state_dict(sd, strict=False)
 
-    def encode_token_weights(self, prompt_segments: list[list[Union[PromptSegment | Action]]], **kwargs):
+    # Changed from comfy.sd1_clip.ClipTokenWeightEncoder
+    # Changed to use PromptSegments
+    def encode_token_weights(self, prompt_segments: list[list[SegOrAction]]):
         to_encode = [[PromptSegment(text="_Empty Batch_", tokens=self.empty_tokens[0])]]
         for batch in prompt_segments:
             to_encode.append(batch)
 
-        out, pooled = self.encode(to_encode, **kwargs)
+        out, pooled = self.encode(to_encode)
         z_empty = out[0:1]
         if pooled.shape[0] > 1:
             first_pooled = pooled[1:2]

lib/fun_clip_stuff.py

@@ -1,14 +1,13 @@
 from typing import Optional, Tuple, Union
 
 import torch
-from torch import device
 from transformers import CLIPTextConfig
 from transformers.modeling_outputs import BaseModelOutputWithPooling
 from transformers.models.clip.modeling_clip import _expand_mask, CLIPTextEmbeddings, CLIPTextTransformer, \
     CLIPTextModel
 
-from custom_nodes.ClipStuff.lib.actions.base import PromptSegment, Action
-from custom_nodes.ClipStuff.lib.tokenizer import TokenDict
+from custom_nodes.ClipStuff.lib.action.base import Action
+from custom_nodes.ClipStuff.lib.actions.types import SegOrAction
 
 def slerp(val, low, high):
     low = low.unsqueeze(0)
@@ -20,18 +19,20 @@ def slerp(val, low, high):
     res = (torch.sin((1.0-val)*omega)/so).unsqueeze(1)*low + (torch.sin(val*omega)/so).unsqueeze(1) * high
     return res
 
-class MyCLIPTextEmbeddings(CLIPTextEmbeddings):
+class PromptLangCLIPTextEmbeddings(CLIPTextEmbeddings):
     def __init__(self, config: CLIPTextConfig):
         super().__init__(config)
 
     def forward(
             self,
-            input_dicts: Optional[list[list[tuple[TokenDict]]]] = None,
+            input_dicts: Optional[list[list[SegOrAction]]] = None,
             input_ids: Optional[torch.LongTensor] = None,
             position_ids: Optional[torch.LongTensor] = None,
             inputs_embeds: Optional[torch.FloatTensor] = None,
     ) -> torch.Tensor:
 
+        if input_dicts is None:
+            raise ValueError("You have to specify input_dicts")
 
         batches = []
         for batch_idx, batch in enumerate(input_dicts):
@@ -48,29 +49,6 @@ class MyCLIPTextEmbeddings(CLIPTextEmbeddings):
         if position_ids is None:
             position_ids = self.position_ids[:, :seq_length]
 
-        # if inputs_embeds is None:
-        #     inputs_embeds = self.token_embedding(input_ids)
-
-        # for batch_idx, batch in enumerate(input_dicts):
-        #     for token_idx, token in enumerate(batch):
-        #         if token[0].nudge_id is not None:
-        #             nudged_embed = inputs_embeds[batch_idx, token_idx][:] + self.token_embedding(torch.LongTensor([token[0].nudge_id]).to(torch.device('cpu')))[0]
-        #             if token[0].nudge_index_start is not None and token[0].nudge_index_stop is not None:
-        #                 nudge_start = token[0].nudge_index_start
-        #                 nudge_end = token[0].nudge_index_stop
-        #             else:
-        #                 nudge_start = 0
-        #                 nudge_end = 768
-        #             inputs_embeds[batch_idx, token_idx][nudge_start:nudge_end] = (slerp(token[0].nudge_weight, inputs_embeds[batch_idx, token_idx][:], nudged_embed)[0][nudge_start:nudge_end])
-        #         elif token[0].arith_ops is not None:
-        #             for op, id_list in token[0].arith_ops.items():
-        #                 if op == '+':
-        #                     for this_id in id_list:
-        #                         inputs_embeds[batch_idx, token_idx] += self.token_embedding(torch.LongTensor([this_id]).to(torch.device('cpu')))[0]
-        #                 elif op == '-':
-        #                     for this_id in id_list:
-        #                         inputs_embeds[batch_idx, token_idx] -= self.token_embedding(torch.LongTensor([this_id]).to(torch.device('cpu')))[0]
-
         embeds = []
         for batch in batches:
             if len(batch) == 1:
@@ -84,14 +62,14 @@ class MyCLIPTextEmbeddings(CLIPTextEmbeddings):
         return embeddings
 
 
-class MyCLIPTextTransformer(CLIPTextTransformer):
+class PrompLangCLIPTextTransformer(CLIPTextTransformer):
     def __init__(self, config: CLIPTextConfig):
         super().__init__(config)
-        self.embeddings = MyCLIPTextEmbeddings(config)
+        self.embeddings = PromptLangCLIPTextEmbeddings(config)
 
     def forward(
             self,
-            input_ids: Optional[list[list[PromptSegment | Action]]] = None,
+            input_ids: Optional[list[list[SegOrAction]]] = None,
             attention_mask: Optional[torch.Tensor] = None,
             position_ids: Optional[torch.Tensor] = None,
             output_attentions: Optional[bool] = None,
@@ -174,37 +152,21 @@ class MyCLIPTextTransformer(CLIPTextTransformer):
         )
 
 
-class MyCLIPTextModel(CLIPTextModel):
+# This is necessary to pass the PromptLangCLIPTextTransformer
+class PromptLangTextModel(CLIPTextModel):
     def __init__(self, config: CLIPTextConfig):
         super().__init__(config)
-        self.text_model = MyCLIPTextTransformer(config)
+        self.text_model = PrompLangCLIPTextTransformer(config)
 
     def forward(
             self,
-            input_ids: Optional[list[list[tuple[TokenDict]]]] = None,
+            input_ids: Optional[list[list[SegOrAction]]] = None,
             attention_mask: Optional[torch.Tensor] = None,
             position_ids: Optional[torch.Tensor] = None,
             output_attentions: Optional[bool] = None,
             output_hidden_states: Optional[bool] = None,
             return_dict: Optional[bool] = None,
     ) -> Union[Tuple, BaseModelOutputWithPooling]:
-        r"""
-        Returns:
-
-        Examples:
-
-        ```python
-        >>> from transformers import AutoTokenizer, CLIPTextModel
-
-        >>> model = CLIPTextModel.from_pretrained("openai/clip-vit-base-patch32")
-        >>> tokenizer = AutoTokenizer.from_pretrained("openai/clip-vit-base-patch32")
-
-        >>> inputs = tokenizer(["a photo of a cat", "a photo of a dog"], padding=True, return_tensors="pt")
-
-        >>> outputs = model(**inputs)
-        >>> last_hidden_state = outputs.last_hidden_state
-        >>> pooled_output = outputs.pooler_output  # pooled (EOS token) states
-        ```"""
         return_dict = return_dict if return_dict is not None else self.config.use_return_dict
 
         return self.text_model(

lib/parser/__init__.py

@@ -0,0 +1,5 @@
+from lark import Lark
+
+from .grammar import grammar
+
+PromptParser = Lark(grammar, start="start", parser="earley")

lib/parser/grammar.py

@@ -0,0 +1,28 @@
+grammar = """
+?start: item+
+
+item: embedding
+    | WORD
+    | function
+    | QUOTED_STRING
+
+function: sum_function
+    | neg_function
+    | norm_function
+    | diff_function
+
+sum_function: "sum(" arg ("|" arg)* ")"
+neg_function: "neg(" arg ")"
+norm_function: "norm(" arg ")"
+diff_function: "diff(" arg ("|" arg)* ")"
+
+arg: item+
+
+embedding: "embedding:" WORD
+
+WORD: /[A-Za-z0-9,_-]+/
+QUOTED_STRING: /"([^"\\\]*(\\\.[^"\\\]*)*)"|'([^'\\\]*(\\\.[^'\\\]*)*)'/
+
+%import common.WS
+%ignore WS
+"""

lib/parser/prompt_segment.py

@@ -0,0 +1,31 @@
+from typing import Union
+
+import torch
+from torch import Tensor
+from torch.nn import Embedding
+
+
+class PromptSegment:
+    def __init__(self, text: str, tokens: list[Union[int, Tensor]]):
+        self.text = text
+        self.tokens = tokens
+
+    def __repr__(self):
+        return f'"{self.text}"{self.tokens}'
+
+    def token_length(self):
+        return len(self.tokens)
+
+    def get_embeddings(self, embedding_module: Embedding) -> Tensor:
+        tensors = torch.LongTensor(self.tokens).to(torch.device('cpu'))
+        unsqueezed_tensors = tensors.unsqueeze(0)
+        return embedding_module(unsqueezed_tensors)
+
+    def depth_repr(self, depth=1):
+        out = f'"{self.text}"('
+
+        cleaned_tokens = list(map(lambda x: str(x) if isinstance(x, int) else "EMBD", self.tokens))
+        out += ", ".join(cleaned_tokens)
+
+        out += ")"
+        return out

lib/parser/transformer.py

@@ -0,0 +1,61 @@
+from lark import Transformer, Token
+
+from comfy.sd1_clip import SD1Tokenizer
+from custom_nodes.ClipStuff.lib.action.base import Action
+from custom_nodes.ClipStuff.lib.actions.diff import DiffAction
+from custom_nodes.ClipStuff.lib.parser.utils import build_prompt_segment
+from custom_nodes.ClipStuff.lib.actions.neg import NegAction
+from custom_nodes.ClipStuff.lib.actions.norm import NormAction
+from custom_nodes.ClipStuff.lib.actions.sum import SumAction
+from custom_nodes.ClipStuff.lib.parser.prompt_segment import PromptSegment
+
+
+class PromptTransformer(Transformer):
+    # def WORD(self, items):
+    #     return items
+
+    def __init__(self, tokenizer: SD1Tokenizer):
+        super().__init__()
+        self.tokenizer = tokenizer
+
+    def item(self, items: list[Token]):
+        for item in items:
+            if isinstance(item, Action):
+                return item
+
+            if isinstance(item, PromptSegment):
+                return item
+
+            if item.type == "WORD":
+                return build_prompt_segment(str(item), self.tokenizer)
+            elif item.type == "QUOTED_STRING":
+                # Remove the quotes
+                unquoted = item[1:-1]
+                # Replace escaped quotes with quotes
+                unescaped = unquoted.replace("\\\"", "\"").replace("\\\'", "\'")
+                return build_prompt_segment(unescaped, self.tokenizer)
+            elif item.type == "embedding":
+                return build_prompt_segment(item, self.tokenizer)
+            elif item.type == "function":
+                return item
+            else:
+                raise Exception("Unknown item type: " + str(item.type))
+
+    def arg(self, items):
+        return items
+
+    def embedding(self, items):
+        return build_prompt_segment(f'{self.tokenizer.embedding_identifier}{items[0]}', self.tokenizer)
+
+    def function(self, items):
+        for item in items:
+            if item.data == 'sum_function':
+                return SumAction(item.children[0][:], item.children[1:][:])
+            elif item.data == 'neg_function':
+                return NegAction(item.children[0])
+            elif item.data == 'norm_function':
+                return NormAction(item.children[0])
+            elif item.data == 'diff_function':
+                return DiffAction(item.children[0][:], item.children[1:][:])
+            else:
+                raise Exception("Unknown function type: " + str(item.data))

lib/parser/utils.py

@@ -0,0 +1,38 @@
+from lark import Token
+
+from comfy.sd1_clip import SD1Tokenizer
+from custom_nodes.ClipStuff.lib.parser.prompt_segment import PromptSegment
+
+
+def flatten_tree(tree):
+    if isinstance(tree, Token):
+        return [str(tree)]
+    else:
+        return [str(tree.data)] + sum([flatten_tree(child) for child in tree.children], [])
+
+
+def build_prompt_segment(text: str, tokenizer: SD1Tokenizer) -> PromptSegment:
+    split_text = text.split(" ")
+    tokens = []
+    for word in split_text:
+        if word.startswith(tokenizer.embedding_identifier) and tokenizer.embedding_directory is not None:
+            embedding_name = word[len(tokenizer.embedding_identifier):].strip('\n')
+
+            get_embed_ret = tokenizer._try_get_embedding(embedding_name)
+            embedding = get_embed_ret[0]
+            leftover = get_embed_ret[1]
+            if embedding is None:
+                print(f"warning, embedding:{embedding_name} does not exist, ignoring")
+            else:
+                if len(embedding.shape) == 1:
+                    tokens.append(embedding)
+                else:
+                    tokens.extend(embedding)
+
+            if leftover != "":
+                word = leftover
+            else:
+                continue
+        tokens.extend(tokenizer.tokenizer(word)["input_ids"][1:-1])
+
+    return PromptSegment(text, tokens)

lib/tokenizer.py

@@ -1,116 +1,13 @@
-import re
-from typing import Union
+from lark import Tree
 
 from comfy.sd1_clip import SD1Tokenizer
-from custom_nodes.ClipStuff.lib.actions import (
-    NudgeAction,
-    ArithAction,
-    ALL_START_CHARS,
-    ALL_END_CHARS,
-    ALL_ACTIONS,
-)
-from custom_nodes.ClipStuff.lib.actions.base import (
-    Action,
-    PromptSegment,
-    build_prompt_segment,
-)
-from custom_nodes.ClipStuff.lib.actions.lib import (
-    is_any_action_segment,
-    is_action_segment,
-)
-from custom_nodes.ClipStuff.lib.actions.utils import batch_size_info
+from custom_nodes.ClipStuff.lib.actions.types import SegOrAction
 
-arith_action = r'(<[a-zA-Z0-9\-_]+:[a-zA-Z0-9\-_]+>)'
+from custom_nodes.ClipStuff.lib.parser import PromptParser
+from custom_nodes.ClipStuff.lib.parser.transformer import PromptTransformer
+from custom_nodes.ClipStuff.lib.parser.prompt_segment import PromptSegment
 
-# TODO: Get embedding identifier from tokenizer
-tokenizer_regex = re.compile(
-    fr"""
-    \d+\.\d+                  # Capture decimals
-    |
-    (?:(?!embedding:)[\w\s]|embedding:[a-zA-Z0-9_]+)+ # Capture sequences of characters, including "embedding:"
-    |
-    \d+                       # Capture whole numbers
-    |
-    [:+-{re.escape("".join(ALL_START_CHARS))}{re.escape("".join(ALL_END_CHARS))}] # Capture special characters including start and end characters
-    """,
-    re.VERBOSE
-)
-def tokenize(text: str) -> list[str]:
-    # Captures:
-    # 1. Words
-    # 2. Numbers(1.0, 1)
-    # 3. Special characters(ALL_START_CHARS, ALL_END_CHARS, :, +, -)
-    tokens = re.findall(tokenizer_regex, text)
-    print(tokens)
-    return [token.strip() for token in tokens]
-
-
-
-def parse_segment(tokens: list[str], tokenizer: SD1Tokenizer) -> PromptSegment | Action:
-    print("Parse segment: Checking token: " + tokens[0])
-    for action in ALL_ACTIONS:
-        if tokens[0] == action.START_CHAR:
-            return action.parse_segment(tokens, ALL_START_CHARS, ALL_END_CHARS, parse_segment, tokenizer)
-    # If we get here, it's a text segment
-    return build_prompt_segment(tokens.pop(0), tokenizer)
-
-def parse(tokens: list[str], tokenizer: SD1Tokenizer) -> list[PromptSegment | Action]:
-    parsed = []
-    while tokens:
-        if tokens[0] == '':
-            tokens.pop(0)
-            continue
-        print("Parse: Checking token: " + tokens[0])
-        if tokens[0] in ALL_START_CHARS:
-            parsed.append(parse_segment(tokens, tokenizer))
-        else:
-            parsed.append(build_prompt_segment(tokens.pop(0), tokenizer))
-    return parsed
-
-
-def parse_special_tokens(string) -> list[str]:
-    out = []
-    current = ""
-
-    for char in string:
-        if char in ALL_START_CHARS:
-            out += [current]
-            current = char
-        elif char in ALL_END_CHARS:
-            out += [current + char]
-            current = ""
-        else:
-            current += char
-    out += [current]
-    return out
-
-
-def parse_segment_actions(string, tokenizer: SD1Tokenizer) -> list[PromptSegment | NudgeAction | ArithAction]:
-    tokens = tokenize(string)
-    parsed = parse(tokens, tokenizer)
-    return parsed
-
-class TokenDict:
-    def __init__(self,
-                 token_id: int,
-                 weight: float = None,
-                 nudge_id=None, nudge_weight=None, nudge_start: int = None, nudge_end: int = None,
-                 arith_ops: dict[str, list[str]] = None):
-        if weight is None:
-            self.weight = 1.0
-        else:
-            self.weight = weight
-
-        self.token_id = token_id
-        self.nudge_id = nudge_id
-        self.nudge_weight = nudge_weight
-        self.nudge_index_start = nudge_start
-        self.nudge_index_stop = nudge_end
-
-        self.arith_ops = arith_ops
-
-
-class MyTokenizer(SD1Tokenizer):
+class PromptLangTokenizer(SD1Tokenizer):
     def __init__(self, tokenizer_path=None, max_length=77, pad_with_end=True, embedding_directory=None, embedding_size=768, embedding_key='clip_l', special_tokens=None):
         super().__init__(tokenizer_path, max_length, pad_with_end, embedding_directory, embedding_size, embedding_key)
 
@@ -119,34 +16,25 @@ class MyTokenizer(SD1Tokenizer):
     Returns batches of segments and actions
     :return: List of list(batches) of segments and actions
     """
-    def tokenize_with_weights(self, text:str, return_word_ids=False, **kwargs) -> list[list[PromptSegment | Action]]:
+    def tokenize_with_weights(self, text:str, return_word_ids=False, **kwargs) -> list[list[SegOrAction]]:
         if self.pad_with_end:
             pad_token = self.end_token
         else:
             pad_token = 0
 
-        parsed_actions = parse_segment_actions(text, self)
-
-        # nudge_start = kwargs.get("nudge_start")
-        # nudge_end = kwargs.get("nudge_end")
-        #
-        # if nudge_start is not None and nudge_end is not None:
-        #     nudge_start = int(nudge_start)
-        #     nudge_end = int(nudge_end)
-        #
-        # # tokenize words
-        for segment in parsed_actions:
-            if isinstance(segment, Action):
-                print(segment.depth_repr())
-            else:
-                print(segment.depth_repr())
+        parsed_prompt = PromptParser.parse(text)
+        parsed_actions = PromptTransformer(self).transform(parsed_prompt)
 
         # reshape token array to CLIP input size
         batched_segments = []
         batch = [PromptSegment(text="[SOT]", tokens=[self.start_token])]
         # batched_segments.append(batch)
         batch_size = 1
-        for segment in parsed_actions:
+        if isinstance(parsed_actions, Tree):
+            segments_to_process = parsed_actions.children
+        else:
+            segments_to_process = [parsed_actions]
+        for segment in segments_to_process:
             num_tokens = segment.token_length()
             # determine if we're going to try and keep the tokens in a single batch
             is_large = num_tokens >= self.max_word_length
@@ -163,7 +51,7 @@ class MyTokenizer(SD1Tokenizer):
                 batch_size = num_tokens + 1 # +1 for start token
                 continue
 
-            # If the segment is small enough to fit in the current batch, add it
+            # Since the segment fits in the current batch, add it
             batch.append(segment)
             batch_size += num_tokens
 
@@ -172,7 +60,7 @@ class MyTokenizer(SD1Tokenizer):
         batch.append(PromptSegment("__PAD__", [self.end_token] + [pad_token] * remaining_length))
         batched_segments.append(batch)
 
-        for batch in batched_segments:
-            batch_size_info(batch)
+        # for batch in batched_segments:
+        #     batch_size_info(batch)
 
         return batched_segments

nodes.py

@@ -6,9 +6,9 @@ from PIL import Image
 import folder_paths
 import comfy.sd
 import comfy.ops
-from custom_nodes.ClipStuff.lib.clip_model import SD1FunClipModel
+from custom_nodes.ClipStuff.lib.clip_model import PromptLangClipModel
 
-from custom_nodes.ClipStuff.lib.tokenizer import MyTokenizer
+from custom_nodes.ClipStuff.lib.tokenizer import PromptLangTokenizer
 
 
 class EmptyClass:
@@ -33,10 +33,9 @@ class SpecialClipLoader:
     def load_clip(source_clip):
         clip_target = EmptyClass()
         clip_target.params = {}
-        clip_target.clip = SD1FunClipModel
-        clip_target.tokenizer = MyTokenizer
+        clip_target.clip = PromptLangClipModel
+        clip_target.tokenizer = PromptLangTokenizer
 
-        # TODO: Extract embedding directory from source_clip
         clip = comfy.sd.CLIP(clip_target, embedding_directory=source_clip.tokenizer.embedding_directory)
         comfy.sd.load_clip_weights(
             clip.cond_stage_model, source_clip.cond_stage_model.state_dict()
@@ -44,49 +43,6 @@ class SpecialClipLoader:
         return (clip,)
 
 
-class KepAdvTextEncode:
-    @classmethod
-    def INPUT_TYPES(s):
-        return {
-            "required": {
-                "text": ("STRING", {"multiline": True}),
-                "clip": ("CLIP",),
-                "nudge_start": ("INT", {}),
-                "nudge_end": ("INT", {}),
-                "split_newlines": ("BOOL", {"default": True}),
-            }
-        }
-
-    RETURN_TYPES = ("CONDITIONING",)
-    FUNCTION = "encode"
-    OUTPUT_IS_LIST = (True,)
-    CATEGORY = "conditioning"
-
-    @staticmethod
-    def encode(clip, text, nudge_start, nudge_end, split_newlines):
-        ret = []
-        if split_newlines:
-            prompts = text.split("\n")
-        else:
-            prompts = [text]
-
-        for prompt in prompts:
-            if prompt.strip() == "":
-                continue
-            tokens = clip.tokenizer.tokenize_with_weights(
-                text,
-                return_word_ids=False,
-                nudge_start=nudge_start,
-                nudge_end=nudge_end,
-            )
-            cond, pooled = clip.encode_from_tokens(
-                tokens, return_pooled=True, position_ids=[0] * 77
-            )
-            cond = [[cond, {"pooled_output": pooled}]]
-            ret.append(cond)
-        return (ret,)
-
-
 def tensor2img(tensor_img):
     i = 255.0 * tensor_img.cpu().numpy()
     i_np_arr = np.clip(i, 0, 255, out=i).astype(np.uint8, copy=False)

requirements.txt

@@ -0,0 +1 @@
+lark