feat(PosScale: Add node

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/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}