diff --git a/comfy/sd1_clip.py b/comfy/sd1_clip.py index f0fdf1aa5..044fd513c 100644 --- a/comfy/sd1_clip.py +++ b/comfy/sd1_clip.py @@ -540,6 +540,14 @@ class SDTokenizer: self.disable_weights = disable_weights + @property + def embedding_identifiers(self): + identifiers = [self.embedding_identifier] + for item in ["EMB:", "TE:"]: + if item not in identifiers: + identifiers.append(item) + return identifiers + def _try_get_embedding(self, embedding_name:str): ''' Takes a potential embedding name and tries to retrieve it. @@ -591,16 +599,24 @@ class SDTokenizer: tokens = [] for weighted_segment, weight in parsed_weights: to_tokenize = unescape_important(weighted_segment) - split = re.split(r'(?<=\s){}'.format(re.escape(self.embedding_identifier)), to_tokenize) + pattern = r'(?<=\s)({})'.format('|'.join(map(re.escape, self.embedding_identifiers))) + split = re.split(pattern, to_tokenize) to_tokenize = [split[0]] - for i in range(1, len(split)): - to_tokenize.append("{}{}".format(self.embedding_identifier, split[i])) + for i in range(1, len(split), 2): + if i + 1 < len(split): + to_tokenize.append("{}{}".format(split[i], split[i+1])) to_tokenize = [x for x in to_tokenize if x != ""] for word in to_tokenize: # if we find an embedding, deal with the embedding - if word.startswith(self.embedding_identifier) and self.embedding_directory is not None: - embedding_name = word[len(self.embedding_identifier):].strip('\n') + matched_id = None + for identifier in self.embedding_identifiers: + if word.startswith(identifier): + matched_id = identifier + break + + if matched_id is not None and self.embedding_directory is not None: + embedding_name = word[len(matched_id):].strip('\n') embed, embedding_name, leftover = self._try_get_embedding(embedding_name) if embed is None: logging.warning(f"warning, embedding:{embedding_name} does not exist, ignoring")