From bdb4d00910d5ee244ec910d81c73fb1ed5268be4 Mon Sep 17 00:00:00 2001 From: Michael Poutre Date: Fri, 25 Aug 2023 18:59:28 -0700 Subject: [PATCH] fix(tokenizer): Fix issue when prompt is just an action, and has no tree --- lib/tokenizer.py | 8 +++++++- 1 file changed, 7 insertions(+), 1 deletion(-) diff --git a/lib/tokenizer.py b/lib/tokenizer.py index 02c0f5c..1a5720d 100644 --- a/lib/tokenizer.py +++ b/lib/tokenizer.py @@ -1,3 +1,5 @@ +from lark import Tree + from comfy.sd1_clip import SD1Tokenizer from custom_nodes.ClipStuff.lib.actions.types import SegOrAction @@ -28,7 +30,11 @@ class PromptLangTokenizer(SD1Tokenizer): batch = [PromptSegment(text="[SOT]", tokens=[self.start_token])] # batched_segments.append(batch) batch_size = 1 - for segment in parsed_actions.children: + 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