132 lines
4.8 KiB
Python
132 lines
4.8 KiB
Python
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)
|