diff --git a/modules/prompt_parser.py b/modules/prompt_parser.py index 870218db8..696653725 100644 --- a/modules/prompt_parser.py +++ b/modules/prompt_parser.py @@ -274,6 +274,7 @@ re_attention = re.compile(r""" : """, re.X) +re_break = re.compile(r"\s*\bBREAK\b\s*", re.S) def parse_prompt_attention(text): """ @@ -339,7 +340,11 @@ def parse_prompt_attention(text): elif text == ']' and len(square_brackets) > 0: multiply_range(square_brackets.pop(), square_bracket_multiplier) else: - res.append([text, 1.0]) + parts = re.split(re_break, text) + for i, part in enumerate(parts): + if i > 0: + res.append(["BREAK", -1]) + res.append([part, 1.0]) for pos in round_brackets: multiply_range(pos, round_bracket_multiplier) diff --git a/modules/sd_hijack_clip.py b/modules/sd_hijack_clip.py index 852afc665..9fa5c5c50 100644 --- a/modules/sd_hijack_clip.py +++ b/modules/sd_hijack_clip.py @@ -96,13 +96,18 @@ class FrozenCLIPEmbedderWithCustomWordsBase(torch.nn.Module): token_count = 0 last_comma = -1 - def next_chunk(): - """puts current chunk into the list of results and produces the next one - empty""" + def next_chunk(is_last=False): + """puts current chunk into the list of results and produces the next one - empty; + if is_last is true, tokens tokens at the end won't add to token_count""" nonlocal token_count nonlocal last_comma nonlocal chunk - token_count += len(chunk.tokens) + if is_last: + token_count += len(chunk.tokens) + else: + token_count += self.chunk_length + to_add = self.chunk_length - len(chunk.tokens) if to_add > 0: chunk.tokens += [self.id_end] * to_add @@ -116,6 +121,10 @@ class FrozenCLIPEmbedderWithCustomWordsBase(torch.nn.Module): chunk = PromptChunk() for tokens, (text, weight) in zip(tokenized, parsed): + if text == 'BREAK' and weight == -1: + next_chunk() + continue + position = 0 while position < len(tokens): token = tokens[position] @@ -159,7 +168,7 @@ class FrozenCLIPEmbedderWithCustomWordsBase(torch.nn.Module): position += embedding_length_in_tokens if len(chunk.tokens) > 0 or len(chunks) == 0: - next_chunk() + next_chunk(is_last=True) return chunks, token_count