github/ai-toolkit
1
0
Fork 0
Code Issues Pull Requests Actions 1 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 = [] text_embeds = []
pooled_embeds = None pooled_embeds = None
attention_mask = [] attention_mask = []
is_list = False
for key in sorted(state_dict.keys()): for key in sorted(state_dict.keys()):
if key.startswith("text_embed_"): if key.startswith("text_embed_"):
is_list = True
text_embeds.append(state_dict[key]) text_embeds.append(state_dict[key])
elif key == "text_embed": elif key == "text_embed":
text_embeds.append(state_dict[key]) text_embeds.append(state_dict[key])
@@ -155,7 +157,7 @@ class PromptEmbeds:
attention_mask.append(state_dict[key]) attention_mask.append(state_dict[key])
pe = cls(None) pe = cls(None)
pe.text_embeds = text_embeds 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] pe.text_embeds = text_embeds[0]
if pooled_embeds is not None: if pooled_embeds is not None:
pe.pooled_embeds = pooled_embeds pe.pooled_embeds = pooled_embeds