github/ai-toolkit
1
0
Fork 0
Code Issues Pull Requests Actions 1 Packages Projects Releases Wiki Activity

Add support for training Z-Image Turbo with a de-distill training adapter

Browse Source
This commit is contained in:
Jaret Burkett
2025-11-28 08:08:53 -07:00
parent 21bb8a2bf4
commit 4e62c38df5
11 changed files with 459 additions and 7 deletions
Download Patch File Download Diff File Expand all files Collapse all files

6
toolkit/models/base_model.py
View File

@@ -805,7 +805,11 @@ class BaseModel:
# check if batch size of embeddings matches batch size of latents
if isinstance(text_embeddings.text_embeds, list):
te_batch_size = text_embeddings.text_embeds[0].shape[0]
if len(text_embeddings.text_embeds) == latents.shape[0]:
# handle list of embeddings
te_batch_size = len(text_embeddings.text_embeds)
else:
te_batch_size = text_embeddings.text_embeds[0].shape[0]
else:
te_batch_size = text_embeddings.text_embeds.shape[0]
if latents.shape[0] == te_batch_size: