Skip to content

Commit 6c62ca0

Browse files
authored
fix: error when embedding is loaded with models using llama_template (#14744)
1 parent 3fe9f5f commit 6c62ca0

1 file changed

Lines changed: 11 additions & 5 deletions

File tree

comfy/sd1_clip.py

Lines changed: 11 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -543,18 +543,24 @@ def __init__(self, tokenizer_path=None, max_length=77, pad_with_end=True, embedd
543543
def _try_get_embedding(self, embedding_name:str):
544544
'''
545545
Takes a potential embedding name and tries to retrieve it.
546-
Returns a Tuple consisting of the embedding and any leftover string, embedding can be None.
546+
Returns a Tuple consisting of the embedding, the cleaned embedding name, and any leftover string, embedding can be None.
547547
'''
548548
split_embed = embedding_name.split()
549549
embedding_name = split_embed[0]
550550
leftover = ' '.join(split_embed[1:])
551+
552+
match = re.search(r'[<\[]', embedding_name)
553+
if match is not None:
554+
leftover = embedding_name[match.start():] + (" " + leftover if leftover else "")
555+
embedding_name = embedding_name[:match.start()]
556+
551557
embed = load_embed(embedding_name, self.embedding_directory, self.embedding_size, self.embedding_key)
552558
if embed is None:
553559
stripped = embedding_name.strip(',')
554560
if len(stripped) < len(embedding_name):
555561
embed = load_embed(stripped, self.embedding_directory, self.embedding_size, self.embedding_key)
556-
return (embed, "{} {}".format(embedding_name[len(stripped):], leftover))
557-
return (embed, leftover)
562+
return (embed, embedding_name, "{} {}".format(embedding_name[len(stripped):], leftover))
563+
return (embed, embedding_name, leftover)
558564

559565
def pad_tokens(self, tokens, amount):
560566
if self.pad_left:
@@ -585,7 +591,7 @@ def tokenize_with_weights(self, text:str, return_word_ids=False, tokenizer_optio
585591
tokens = []
586592
for weighted_segment, weight in parsed_weights:
587593
to_tokenize = unescape_important(weighted_segment)
588-
split = re.split(' {0}|\n{0}'.format(self.embedding_identifier), to_tokenize)
594+
split = re.split(r'(?<=\s){}'.format(re.escape(self.embedding_identifier)), to_tokenize)
589595
to_tokenize = [split[0]]
590596
for i in range(1, len(split)):
591597
to_tokenize.append("{}{}".format(self.embedding_identifier, split[i]))
@@ -595,7 +601,7 @@ def tokenize_with_weights(self, text:str, return_word_ids=False, tokenizer_optio
595601
# if we find an embedding, deal with the embedding
596602
if word.startswith(self.embedding_identifier) and self.embedding_directory is not None:
597603
embedding_name = word[len(self.embedding_identifier):].strip('\n')
598-
embed, leftover = self._try_get_embedding(embedding_name)
604+
embed, embedding_name, leftover = self._try_get_embedding(embedding_name)
599605
if embed is None:
600606
logging.warning(f"warning, embedding:{embedding_name} does not exist, ignoring")
601607
else:

0 commit comments

Comments
 (0)