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

Shrink text embeds to max token length for LTX-2. Drastically reduces cached text embedding sizes

Browse Source
This commit is contained in:
Jaret Burkett
2026-01-28 12:54:49 -07:00
parent ea912d2d7b
commit 1ce2428722
7 changed files with 130 additions and 27 deletions
Download Patch File Download Diff File Expand all files Collapse all files

20
toolkit/data_loader.py
View File

@@ -487,6 +487,23 @@ class AiToolkitDataset(LatentCachingMixin, ControlCachingMixin, CLIPCachingMixin
self.size_database["__version__"] = dataloader_version
# set latent space version
latent_space_version = "sd1"
if self.sd is not None and self.sd.model_config.latent_space_version is not None:
latent_space_version = self.sd.model_config.latent_space_version
elif self.sd.is_xl:
latent_space_version = 'sdxl'
elif self.sd.is_v3:
latent_space_version = 'sd3'
elif self.sd.is_auraflow:
latent_space_version = 'sdxl'
elif self.sd.is_flux:
latent_space_version = 'flux1'
elif self.sd.model_config.is_pixart_sigma:
latent_space_version = 'sdxl'
else:
latent_space_version = self.sd.model_config.arch if self.sd is not None else "sd1"
bad_count = 0
for file in tqdm(file_list):
try:
@@ -498,6 +515,9 @@ class AiToolkitDataset(LatentCachingMixin, ControlCachingMixin, CLIPCachingMixin
size_database=self.size_database,
dataset_root=dataset_folder,
encode_control_in_text_embeddings=self.sd.encode_control_in_text_embeddings if self.sd else False,
text_embedding_space_version=self.sd.model_config.arch if self.sd else "sd1",
te_padding_side=self.sd.te_padding_side if self.sd else "right",
latent_space_version=latent_space_version,
)
self.file_list.append(file_item)
except Exception as e: