fix: Update to work with transformers changes to attention masks

Fix end token calculation for long prompts

lib/__init__.py

@@ -3,6 +3,7 @@ from custom_nodes.KepPromptLang.lib.actions.diff import DiffAction
 from custom_nodes.KepPromptLang.lib.actions.mult import MultiplyAction
 from custom_nodes.KepPromptLang.lib.actions.neg import NegAction
 from custom_nodes.KepPromptLang.lib.actions.norm import NormAction
+from custom_nodes.KepPromptLang.lib.actions.pos_scale import PosScaleAction
 from custom_nodes.KepPromptLang.lib.actions.rand import RandAction
 from custom_nodes.KepPromptLang.lib.actions.scale_dims import ScaleDims
 from custom_nodes.KepPromptLang.lib.actions.set_dims import SetDims
@@ -20,3 +21,4 @@ register_action(SlerpAction)
 register_action(AverageAction)
 register_action(ScaleDims)
 register_action(SetDims)
+register_action(PosScaleAction)

lib/action/base.py

@@ -1,6 +1,6 @@
 from abc import ABC, abstractmethod
 from enum import Enum
-from typing import Union, List
+from typing import Union, List, TypedDict, Tuple
 
 from torch import Tensor
 from torch.nn import Embedding
@@ -13,6 +13,13 @@ class ActionArity(Enum):
     SINGLE = 1
     MULTI = 2
 
+class PostModifiers(TypedDict):
+    """
+    A dictionary of post modifiers for an action result.
+    """
+    position_embed_scale: Union[float, None]
+
+
 class Action(ABC):
     @property
     @abstractmethod
@@ -67,7 +74,7 @@ class Action(ABC):
         pass
 
     @abstractmethod
-    def get_result(self, embedding_module: Embedding) -> Tensor:
+    def get_result(self, embedding_module: Embedding) -> Union[Tensor, Tuple[Tensor, PostModifiers]]:
         """
         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.

lib/actions/pos_scale.py

@@ -0,0 +1,65 @@
+from typing import Tuple, List
+
+import torch
+from torch.nn import Embedding
+
+from custom_nodes.KepPromptLang.lib.action.base import (
+    Action,
+    PostModifiers,
+    MultiArgAction,
+)
+from custom_nodes.KepPromptLang.lib.actions.types import SegOrAction
+
+
+class PosScaleAction(MultiArgAction):
+    grammar = 'posScale(" arg+ ")"'
+    name = "posScale"
+    chars = ["[", "]"]
+
+    def __init__(self, args: List[List[SegOrAction]]) -> None:
+        super().__init__(args)
+        if len(args) != 2:
+            raise ValueError("PosScale action should have exactly two arguments")
+
+        self.target_arg = args[0]
+        self._parse_multiplier(args[1])
+
+    def _parse_multiplier(self, arg: List[SegOrAction]) -> None:
+        if len(arg) != 1:
+            raise ValueError(
+                "PosScale actions multiplier should have exactly one segment"
+            )
+
+        multiplier_seg_or_action = arg[0]
+
+        if isinstance(multiplier_seg_or_action, Action):
+            raise ValueError("PosScale actions multiplier must be a number")
+
+        try:
+            self.parsed_multiplier = float(multiplier_seg_or_action.text)
+        except ValueError:
+            raise ValueError(
+                "PosScale action should have an integer/float as the multiplier"
+            )
+
+    def token_length(self) -> int:
+        """
+        PosScale modifies the posional 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.target_arg:
+            total_length += seg_or_action.token_length()
+
+        return total_length
+
+    def get_result(self, embedding_module: Embedding) -> Tuple[torch.Tensor, PostModifiers]:
+        all_embeddings = []
+        for seg_or_action in self.target_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, {"position_embed_scale": self.parsed_multiplier}

lib/clip_model.py

@@ -7,6 +7,7 @@ from transformers import CLIPTextConfig, modeling_utils
 
 from comfy import model_management
 import comfy.ops
+from comfy.sd1_clip import SD1ClipModel
 from comfy.sdxl_clip import SDXLClipModel
 from custom_nodes.KepPromptLang.lib.action.base import Action
 from custom_nodes.KepPromptLang.lib.actions.types import SegOrAction
@@ -15,7 +16,7 @@ from custom_nodes.KepPromptLang.lib.parser.prompt_segment import PromptSegment
 
 
 # Methods with no comment can be assumed to be the same as comfy.sd1_clip.SD1ClipModel
-class PromptLangClipModel(torch.nn.Module):
+class PromptLangSDClipModel(torch.nn.Module):
     """Uses the CLIP transformer encoder for text (from huggingface)"""
     LAYERS = [
         "last",
@@ -211,15 +212,23 @@ class PromptLangClipModel(torch.nn.Module):
             return z_empty.cpu(), first_pooled.cpu()
         return torch.cat(output, dim=-2).cpu(), first_pooled.cpu()
 
+class PromptLangSD1ClipModel(SD1ClipModel):
+    def __init__(self, device="cpu", dtype=None, clip_name="l", clip_model=PromptLangSDClipModel):
+        super().__init__()
+        self.clip_name = clip_name
+        self.clip = "clip_{}".format(self.clip_name)
+        setattr(self, self.clip, clip_model(device=device, dtype=dtype))
+
+
 class PromptLangSDXLClipModel(SDXLClipModel):
     def __init__(self, device="cpu", dtype=None) -> None:
         # Skip SDXLClipModel's init
         super(SDXLClipModel, self).__init__()
-        self.clip_l = PromptLangClipModel(layer="hidden", layer_idx=11, device=device, dtype=dtype)
+        self.clip_l = PromptLangSDClipModel(layer="hidden", layer_idx=11, device=device, dtype=dtype)
         self.clip_l.layer_norm_hidden_state = False
         self.clip_g = PromptLangSDXLClipG(device, dtype)
 
-class PromptLangSDXLClipG(PromptLangClipModel):
+class PromptLangSDXLClipG(PromptLangSDClipModel):
     def __init__(self, device="cpu", max_length=77, freeze=True, layer="penultimate", layer_idx=None, textmodel_path=None, dtype=None):
         if layer == "penultimate":
             layer="hidden"

lib/fun_clip_stuff.py

@@ -1,10 +1,11 @@
-from typing import Optional, Tuple, Union, List
+from typing import Optional, Tuple, Union, List, TypedDict
+from importlib_metadata import version as import_version
+from packaging import version
 
 import torch
 from transformers import CLIPTextConfig
 from transformers.modeling_outputs import BaseModelOutputWithPooling
 from transformers.models.clip.modeling_clip import (
-    _expand_mask,
     CLIPTextEmbeddings,
     CLIPTextTransformer,
     CLIPTextModel,
@@ -23,6 +24,17 @@ 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 PosModifier(TypedDict):
+    """
+    A dictionary of post modifiers for an action result.
+    """
+
+    position_embed_scale: Union[float]
+    start_idx: Union[int]
+    end_idx: Union[int]
+
+
 class PromptLangCLIPTextEmbeddings(CLIPTextEmbeddings):
     def __init__(self, config: CLIPTextConfig):
         super().__init__(config)
@@ -39,14 +51,30 @@ class PromptLangCLIPTextEmbeddings(CLIPTextEmbeddings):
             raise ValueError("You have to specify input_dicts")
 
         batches = []
+        pos_modifiers: List[List[PosModifier]] = []
         for batch_idx, batch in enumerate(input_dicts):
             results = []
+            batch_pos_modifiers = []
+            token_idx = 0
             for seg_or_action in batch:
                 if isinstance(seg_or_action, Action):
-                    results.append(seg_or_action.get_result(self.token_embedding))
+                    action_result = seg_or_action.get_result(self.token_embedding)
+                    if isinstance(action_result, tuple):
+                        result, post_modifiers = action_result
+                        if post_modifiers["position_embed_scale"] is not None:
+                            post_modifiers["start_idx"] = token_idx
+                            post_modifiers["end_idx"] = (
+                                token_idx + seg_or_action.token_length()
+                            )
+                            batch_pos_modifiers.append(post_modifiers)
+                    else:
+                        result = action_result
                 else:
-                    results.append(seg_or_action.get_embeddings(self.token_embedding))
+                    result = seg_or_action.get_embeddings(self.token_embedding)
+                results.append(result)
+                token_idx += seg_or_action.token_length()
             batches.append(results)
+            pos_modifiers.append(batch_pos_modifiers)
 
         seq_length = batches[0][0].shape[-2]
 
@@ -60,8 +88,16 @@ class PromptLangCLIPTextEmbeddings(CLIPTextEmbeddings):
             else:
                 embeds.append(torch.cat(batch, dim=-2))
 
-        position_embeddings = self.position_embedding(position_ids)
-        embeddings = torch.cat(embeds, dim=0) + position_embeddings
+        for idx, batch_pos_modifiers in enumerate(pos_modifiers):
+            position_embeddings = self.position_embedding(position_ids)
+            if len(batch_pos_modifiers) > 0:
+                print(f"Found {len(batch_pos_modifiers)} pos modifiers for batch {idx}")
+                for post_modifier in batch_pos_modifiers:
+                    position_embeddings[
+                        0, post_modifier["start_idx"] : post_modifier["end_idx"]
+                    ] *= post_modifier["position_embed_scale"]
+            embeds[idx] = embeds[idx] + position_embeddings
+        embeddings = torch.cat(embeds, dim=0)
 
         return embeddings
 
@@ -70,6 +106,40 @@ class PrompLangCLIPTextTransformer(CLIPTextTransformer):
     def __init__(self, config: CLIPTextConfig):
         super().__init__(config)
         self.embeddings = PromptLangCLIPTextEmbeddings(config)
+        self.transformers_version = version.parse(import_version('transformers'))
+
+    def process_attention_mask(self, hidden_states, attention_mask, bsz, seq_len):
+        # Parse the transformer version
+        input_shape = torch.Size([bsz, seq_len])
+
+        v4_30 = version.parse('4.30.0')
+        v4_35 = version.parse('4.35')
+        if self.transformers_version < v4_30:
+            print("Using transformers < 4.30.0")
+            causal_attention_mask = self._build_causal_attention_mask(bsz, seq_len, hidden_states.dtype).to(
+                hidden_states.device)
+        elif v4_30 <= self.transformers_version < v4_35:
+            print("Using transformers >= 4.30.0 and <= 4.34.*")
+            from transformers.models.clip.modeling_clip import _make_causal_mask
+            causal_attention_mask = _make_causal_mask(input_shape, hidden_states.dtype, device=hidden_states.device)
+        else:
+            print("Using transformers >= 4.35")
+            from transformers.modeling_attn_mask_utils import _create_4d_causal_attention_mask
+            causal_attention_mask = _create_4d_causal_attention_mask(
+                input_shape, hidden_states.dtype, device=hidden_states.device
+            )
+
+        # Expand attention_mask if it exists
+        if attention_mask is not None:
+            # Import _expand_mask or _prepare_4d_attention_mask based on version
+            if self.transformers_version < v4_35:
+                from transformers.models.clip.modeling_clip import _expand_mask
+                attention_mask = _expand_mask(attention_mask, hidden_states.dtype)
+            else:
+                from transformers.modeling_attn_mask_utils import _prepare_4d_attention_mask
+                attention_mask = _prepare_4d_attention_mask(attention_mask, hidden_states.dtype)
+
+        return causal_attention_mask, attention_mask
 
     def forward(
             self,
@@ -101,24 +171,8 @@ class PrompLangCLIPTextTransformer(CLIPTextTransformer):
         bsz = len(input_ids)
         # TODO: Properly gather this
         seq_len = 77
-        input_shape = torch.Size([bsz, seq_len])
-        # CLIP's text model uses causal mask, prepare it here.
-        # https://github.com/openai/CLIP/blob/cfcffb90e69f37bf2ff1e988237a0fbe41f33c04/clip/model.py#L324
-        ## VERSION DIFF ##
-        # transformers < 4.30.0
-        if hasattr(self, "_build_causal_attention_mask"):
-            print("Using transformers < 4.30.0")
-            causal_attention_mask = self._build_causal_attention_mask(bsz, seq_len, hidden_states.dtype).to(hidden_states.device)
-        else:
-            # transformers >= 4.30.0
-            print("Using transformers >= 4.30.0")
-            from transformers.models.clip.modeling_clip import _make_causal_mask
-            causal_attention_mask = _make_causal_mask(input_shape, hidden_states.dtype, device=hidden_states.device)
 
-        # expand attention_mask
-        if attention_mask is not None:
-            # [bsz, seq_len] -> [bsz, 1, tgt_seq_len, src_seq_len]
-            attention_mask = _expand_mask(attention_mask, hidden_states.dtype)
+        causal_attention_mask, attention_mask = self.process_attention_mask(hidden_states, attention_mask, bsz, seq_len)
 
         encoder_outputs = self.encoder(
             inputs_embeds=hidden_states,

lib/parser/transformer.py

@@ -2,15 +2,10 @@ from typing import List
 
 from lark import Transformer, Token
 
-from comfy.sd1_clip import SD1Tokenizer
+from comfy.sd1_clip import SDTokenizer
 from custom_nodes.KepPromptLang.lib.action.base import Action, ActionArity
-from custom_nodes.KepPromptLang.lib.actions.diff import DiffAction
-from custom_nodes.KepPromptLang.lib.actions.rand import RandAction
 from custom_nodes.KepPromptLang.lib.parser.registration import get_action_by_name
 from custom_nodes.KepPromptLang.lib.parser.utils import build_prompt_segment
-from custom_nodes.KepPromptLang.lib.actions.neg import NegAction
-from custom_nodes.KepPromptLang.lib.actions.norm import NormAction
-from custom_nodes.KepPromptLang.lib.actions.sum import SumAction
 from custom_nodes.KepPromptLang.lib.parser.prompt_segment import PromptSegment
 
 
@@ -18,7 +13,7 @@ class PromptTransformer(Transformer):
     # def WORD(self, items):
     #     return items
 
-    def __init__(self, tokenizer: SD1Tokenizer):
+    def __init__(self, tokenizer: SDTokenizer):
         super().__init__()
         self.tokenizer = tokenizer
 

lib/parser/utils.py

@@ -1,6 +1,6 @@
 from lark import Token
 
-from comfy.sd1_clip import SD1Tokenizer
+from comfy.sd1_clip import SDTokenizer
 from custom_nodes.KepPromptLang.lib.parser.prompt_segment import PromptSegment
 
 
@@ -11,7 +11,7 @@ def flatten_tree(tree):
         return [str(tree.data)] + sum([flatten_tree(child) for child in tree.children], [])
 
 
-def build_prompt_segment(text: str, tokenizer: SD1Tokenizer) -> PromptSegment:
+def build_prompt_segment(text: str, tokenizer: SDTokenizer) -> PromptSegment:
     split_text = text.split(" ")
     tokens = []
     for word in split_text:

lib/tokenizer.py

@@ -2,17 +2,17 @@ from typing import List, Dict
 
 from lark import Tree
 
-from comfy.sd1_clip import SD1Tokenizer
+from comfy.sd1_clip import SD1Tokenizer, SDTokenizer
 from custom_nodes.KepPromptLang.lib.actions.types import SegOrAction
 
 from custom_nodes.KepPromptLang.lib.parser import PromptParser
 from custom_nodes.KepPromptLang.lib.parser.transformer import PromptTransformer
 from custom_nodes.KepPromptLang.lib.parser.prompt_segment import PromptSegment
 
-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) -> None:
-        super().__init__(tokenizer_path, max_length, pad_with_end, embedding_directory, embedding_size, embedding_key)
 
+class PromptLangSDTokenizer(SDTokenizer):
+    def __init__(self, tokenizer_path=None, max_length=77, pad_with_end=True, embedding_directory=None, embedding_size=768, embedding_key='clip_l'):
+        super().__init__(tokenizer_path, max_length, pad_with_end, embedding_directory, embedding_size, embedding_key)
     """
     Doesn't actually tokenize...
     Returns batches of segments and actions
@@ -43,9 +43,9 @@ class PromptLangTokenizer(SD1Tokenizer):
 
             # If the segment is too large to fit in a single batch, pad the current batch and start a new one
             if num_tokens + batch_size > self.max_length - 1:
-                remaining_length = self.max_length - batch_size - 1 # -1 for end token
+                remaining_length = self.max_length - batch_size
                 # Pad batch
-                batch.append(PromptSegment("__PAD__", [self.end_token] + [pad_token] * remaining_length - 1))
+                batch.append(PromptSegment("__PAD__", [self.end_token] + [pad_token] * (remaining_length - 1))) # -1 for end token
                 batched_segments.append(batch)
 
                 # start new batch
@@ -67,15 +67,18 @@ class PromptLangTokenizer(SD1Tokenizer):
 
         return batched_segments
 
+class PromptLangSD1Tokenizer(SD1Tokenizer):
+    def __init__(self, embedding_directory=None, clip_name='l', tokenizer=PromptLangSDTokenizer) -> None:
+        super().__init__(embedding_directory, clip_name, tokenizer)
 
-class PromptLangSDXLClipGTokenizer(PromptLangTokenizer):
-    def __init__(self, tokenizer_path=None, embedding_directory=None) -> None:
+
+class PromptLangSDXLClipGTokenizer(PromptLangSDTokenizer):
+    def __init__(self, tokenizer_path=None, embedding_directory=None):
         super().__init__(tokenizer_path, pad_with_end=False, embedding_directory=embedding_directory, embedding_size=1280, embedding_key='clip_g')
 
-
 class PromptLangSDXLTokenizer(SD1Tokenizer):
     def __init__(self, embedding_directory=None) -> None:
-        self.clip_l = PromptLangTokenizer(embedding_directory=embedding_directory)
+        self.clip_l = PromptLangSDTokenizer(embedding_directory=embedding_directory)
         self.clip_g = PromptLangSDXLClipGTokenizer(embedding_directory=embedding_directory)
 
     def tokenize_with_weights(self, text:str, return_word_ids=False) -> Dict[str, List[List[SegOrAction]]]:

nodes.py

@@ -1,4 +1,3 @@
-import random
 import os
 from typing import List, Tuple, Any
 
@@ -12,13 +11,13 @@ from comfy.sd2_clip import SD2ClipModel
 from comfy.sdxl_clip import SDXLClipModel
 from comfy.supported_models_base import ClipTarget
 from custom_nodes.KepPromptLang.lib.clip_model import (
-    PromptLangClipModel,
     PromptLangSDXLClipModel,
+    PromptLangSD1ClipModel,
 )
 
 from custom_nodes.KepPromptLang.lib.tokenizer import (
-    PromptLangTokenizer,
     PromptLangSDXLTokenizer,
+    PromptLangSD1Tokenizer,
 )
 
 
@@ -46,14 +45,15 @@ class SpecialClipLoader:
         if isinstance(source_clip.cond_stage_model, SDXLClipModel):
             clip_target = ClipTarget(PromptLangSDXLTokenizer, PromptLangSDXLClipModel)
             clip = comfy.sd.CLIP(clip_target, embedding_directory=source_clip.tokenizer.clip_g.embedding_directory)
+            comfy.sd.load_clip_weights(clip.cond_stage_model.clip_g,source_clip.cond_stage_model.clip_g.state_dict())
             comfy.sd.load_clip_weights(
-                clip.cond_stage_model, source_clip.cond_stage_model.state_dict()
+                clip.cond_stage_model.clip_l, source_clip.cond_stage_model.clip_l.state_dict()
             )
         elif isinstance(source_clip, SD2ClipModel):
             raise ValueError("SD2 Clip model is not supported.")
         else:
-            clip_target = ClipTarget(PromptLangTokenizer, PromptLangClipModel)
-            clip = comfy.sd.CLIP(clip_target, embedding_directory=source_clip.tokenizer.embedding_directory)
+            clip_target = ClipTarget(PromptLangSD1Tokenizer, PromptLangSD1ClipModel)
+            clip = comfy.sd.CLIP(clip_target, embedding_directory=source_clip.tokenizer.clip_l.embedding_directory)
             comfy.sd.load_clip_weights(
                 clip.cond_stage_model, source_clip.cond_stage_model.state_dict()
             )

requirements.txt

@@ -1 +1,2 @@
 lark
+packaging