ostris/ai-toolkit
1
0
Fork 0
mirror of https://github.com/ostris/ai-toolkit.git synced 2026-01-26 16:39:47 +00:00
Code Issues Packages Projects Releases Wiki Activity

Fixed issue that prevented caching text embeddings on z-image

Browse Source
This commit is contained in:
Jaret Burkett
2025-11-28 09:19:39 -07:00
parent 4e62c38df5
commit 08a39754a4
1 changed files with 3 additions and 1 deletions
Download Patch File Download Diff File Expand all files Collapse all files

4
toolkit/prompt_utils.py
View File

@@ -142,8 +142,10 @@ class PromptEmbeds:
text_embeds = []
pooled_embeds = None
attention_mask = []
is_list = False
for key in sorted(state_dict.keys()):
if key.startswith("text_embed_"):
is_list = True
text_embeds.append(state_dict[key])
elif key == "text_embed":
text_embeds.append(state_dict[key])
@@ -155,7 +157,7 @@ class PromptEmbeds:
attention_mask.append(state_dict[key])
pe = cls(None)
pe.text_embeds = text_embeds
if len(text_embeds) == 1:
if len(text_embeds) == 1 and not is_list:
pe.text_embeds = text_embeds[0]
if pooled_embeds is not None:
pe.pooled_embeds = pooled_embeds