diff --git a/lib/__init__.py b/lib/__init__.py index e66a35f..157815b 100644 --- a/lib/__init__.py +++ b/lib/__init__.py @@ -3,6 +3,7 @@ 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.rand import RandAction +from custom_nodes.KepPromptLang.lib.actions.slerp import SlerpAction from custom_nodes.KepPromptLang.lib.actions.sum import SumAction from custom_nodes.KepPromptLang.lib.parser.registration import register_action @@ -12,3 +13,4 @@ register_action(NegAction) register_action(NormAction) register_action(RandAction) register_action(SumAction) +register_action(SlerpAction) diff --git a/lib/actions/slerp.py b/lib/actions/slerp.py new file mode 100644 index 0000000..15acc43 --- /dev/null +++ b/lib/actions/slerp.py @@ -0,0 +1,100 @@ +from typing import List, Union + +import torch +from torch.nn import Embedding + +from custom_nodes.KepPromptLang.lib.action.base import MultiArgAction, Action +from custom_nodes.KepPromptLang.lib.actions.action_utils import get_embedding +from custom_nodes.KepPromptLang.lib.actions.types import SegOrAction +from custom_nodes.KepPromptLang.lib.actions.utils import slerp +from custom_nodes.KepPromptLang.lib.parser.prompt_segment import PromptSegment + + +class SlerpAction(MultiArgAction): + grammar = 'slerp(" arg "|" arg "|" arg ")"' + name = "slerp" + chars = ["+", "+"] + + def __init__(self, args: List[List[Union[PromptSegment, Action]]]) -> None: + super().__init__(args) + + if len(args) != 3: + raise ValueError("Slerp action should have exactly three arguments(2 vectors and a weight)") + + self.start_argument = args[0] + self.end_argument = args[1] + self._parse_weight(args[2]) + + self._validate_args() + + + def _parse_weight(self, arg: List[SegOrAction]) -> None: + if len(arg) != 1: + raise ValueError("Slerp weight should have exactly one segment") + + weight_seg_or_action = arg[0] + + if isinstance(weight_seg_or_action, Action): + raise ValueError("Slerp weight should not have an action as an argument") + + try: + self.parsed_weight = float(weight_seg_or_action.text) + except ValueError: + raise ValueError("Slerp should have an integer/float as the weight") + + def _validate_args(self) -> None: + start_arg_token_length = sum(seg_or_action.token_length() for seg_or_action in self.start_argument) + end_arg_token_length = sum(seg_or_action.token_length() for seg_or_action in self.end_argument) + if start_arg_token_length != end_arg_token_length: + raise ValueError(f"Slerp start and end arguments should have the same length. Got {start_arg_token_length} and {end_arg_token_length}") + + if self.parsed_weight < 0 or self.parsed_weight > 1: + print(f"WARNING: Slerp weight should be between 0 and 1. Got {self.parsed_weight}") + + def token_length(self) -> int: + # Slerp interpolates between the embeddings of the start and end segments, so the length is the length of the start segment + return sum(seg_or_action.token_length() for seg_or_action in self.start_argument) + + def get_result(self, embedding_module: Embedding) -> torch.Tensor: + # Calculate the embeddings for the start segment + all_start_embeddings = [ + get_embedding(seg_or_action, embedding_module) + for seg_or_action in self.start_argument + ] + + start_embedding = torch.cat(all_start_embeddings, dim=1) + + # Calculate the embeddings for the end segment + all_end_embeddings = [ + get_embedding(seg_or_action, embedding_module) + for seg_or_action in self.end_argument + ] + + end_embedding = torch.cat(all_end_embeddings, dim=1) + + # Perform the slerp + result = slerp(self.parsed_weight, start_embedding, end_embedding) + + 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.additional_args))})" + + def depth_repr(self, depth=1): + out = "NudgeAction(\n" + if isinstance(self.base_arg, Action): + base_segment_repr = self.base_arg.depth_repr(depth + 1) + out += "\t" * depth + f"base_segment={base_segment_repr}\n" + else: + out += "\t" * depth + f"base_segment={self.base_arg.depth_repr()},\n" + + if isinstance(self.additional_args, Action): + target_repr = self.additional_args.depth_repr(depth + 1) + out += "\t" * depth + f"target={target_repr},\n" + else: + out += "\t" * depth + f"target={self.additional_args.depth_repr()},\n" + out += "\t" * depth + f"weight={self.weight},\n" + out += "\t" * (depth - 1) + ")" + return out diff --git a/lib/actions/utils.py b/lib/actions/utils.py index e974086..2ff7e2d 100644 --- a/lib/actions/utils.py +++ b/lib/actions/utils.py @@ -1,8 +1,38 @@ from typing import List +import torch + from custom_nodes.KepPromptLang.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()) + +def slerp(val: float, low: torch.Tensor, high: torch.Tensor, epsilon=1e-5): + # Convert val to tensor and clamp between 0 and 1 + val = torch.tensor(val, dtype=torch.float32).clamp(0, 1) + + # Normalize the vectors + low_norm = low / torch.norm(low, dim=-1, keepdim=True) + high_norm = high / torch.norm(high, dim=-1, keepdim=True) + + # Calculate the cosine of the angle between the vectors + dot = (low_norm * high_norm).sum(-1, keepdim=True) + + # Clamp to prevent numerical errors + dot = torch.clamp(dot, -1, 1) + + omega = torch.acos(dot) + + # Slerp formula + sin_omega = torch.sin(omega) + scale_0 = torch.sin((1.0 - val) * omega) / (sin_omega + epsilon) + scale_1 = torch.sin(val * omega) / (sin_omega + epsilon) + + # Handle the case where omega is small (the vectors are close) + close_condition = sin_omega < epsilon + scale_0 = torch.where(close_condition, 1.0 - val, scale_0) + scale_1 = torch.where(close_condition, val, scale_1) + + return scale_0 * low + scale_1 * high