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:

7
toolkit/data_transfer_object/data_loader.py
View File

@@ -60,6 +60,9 @@ class FileItemDTO(
self.encode_control_in_text_embeddings = kwargs.get(
"encode_control_in_text_embeddings", False
)
self.te_padding_side = kwargs.get("te_padding_side", "right")
self.latent_space_version = kwargs.get("latent_space_version", "sd1")
self.text_embedding_space_version = kwargs.get("text_embedding_space_version", "sd1")
if dataset_root is not None:
# remove dataset root from path
file_key = self.path.replace(dataset_root, "")
@@ -399,7 +402,9 @@ class DataLoaderBatchDTO:
if not isinstance(y.text_embeds, list):
y.text_embeds = [y.text_embeds]
prompt_embeds_list.append(y)
self.prompt_embeds = concat_prompt_embeds(prompt_embeds_list)
padding_side = self.file_items[0].te_padding_side
self.prompt_embeds = concat_prompt_embeds(prompt_embeds_list, padding_side=padding_side)
if any([x.audio_tensor is not None for x in self.file_items]):
# find one to use as a base

19
toolkit/dataloader_mixins.py
View File

@@ -1721,8 +1721,6 @@ class LatentCachingFileItemDTOMixin:
self.is_caching_to_disk = False
self.is_caching_to_memory = False
self.latent_load_device = 'cpu'
# sd1 or sdxl or others
self.latent_space_version = 'sd1'
# todo, increment this if we change the latent format to invalidate cache
self.latent_version = 1
@@ -1829,21 +1827,6 @@ class LatentCachingMixin:
# use tqdm to show progress
i = 0
for file_item in tqdm(self.file_list, desc=f'Caching latents{" to disk" if to_disk else ""}'):
# set latent space version
if self.sd.model_config.latent_space_version is not None:
file_item.latent_space_version = self.sd.model_config.latent_space_version
elif self.sd.is_xl:
file_item.latent_space_version = 'sdxl'
elif self.sd.is_v3:
file_item.latent_space_version = 'sd3'
elif self.sd.is_auraflow:
file_item.latent_space_version = 'sdxl'
elif self.sd.is_flux:
file_item.latent_space_version = 'flux1'
elif self.sd.model_config.is_pixart_sigma:
file_item.latent_space_version = 'sdxl'
else:
file_item.latent_space_version = self.sd.model_config.arch
file_item.is_caching_to_disk = to_disk
file_item.is_caching_to_memory = to_memory
file_item.latent_load_device = self.sd.device
@@ -1933,7 +1916,6 @@ class TextEmbeddingFileItemDTOMixin:
self._text_embedding_path: Union[str, None] = None
self.is_text_embedding_cached = False
self.text_embedding_load_device = 'cpu'
self.text_embedding_space_version = 'sd1'
self.text_embedding_version = 1
def get_text_embedding_info_dict(self: 'FileItemDTO'):
@@ -1997,7 +1979,6 @@ class TextEmbeddingCachingMixin:
# use tqdm to show progress
i = 0
for file_item in tqdm(self.file_list, desc='Caching text embeddings to disk'):
file_item.text_embedding_space_version = self.sd.model_config.arch
file_item.latent_load_device = self.sd.device
text_embedding_path = file_item.get_text_embedding_path(recalculate=True)

4
toolkit/models/base_model.py
View File

@@ -190,6 +190,10 @@ class BaseModel:
# use new lokr format (default false for old models for backwards compatibility)
self.use_old_lokr_format = True
# when padding to make batch size work, which side padding to use, right or left
# some llms need left side padding, others need right side
self.te_padding_side = "right"
# properties for old arch for backwards compatibility
@property

17
toolkit/prompt_utils.py
View File

@@ -244,7 +244,7 @@ class EncodedPromptPair:
return self
def concat_prompt_embeds(prompt_embeds: list["PromptEmbeds"]):
def concat_prompt_embeds(prompt_embeds: list["PromptEmbeds"], padding_side: str = "right") -> PromptEmbeds:
# --- pad text_embeds ---
if isinstance(prompt_embeds[0].text_embeds, (list, tuple)):
embed_list = []
@@ -259,7 +259,10 @@ def concat_prompt_embeds(prompt_embeds: list["PromptEmbeds"]):
dtype=t.dtype,
device=t.device,
)
t = torch.cat([t, pad], dim=1)
if padding_side == "right":
t = torch.cat([t, pad], dim=1)
else:
t = torch.cat([pad, t], dim=1)
padded.append(t)
embed_list.append(torch.cat(padded, dim=0))
text_embeds = embed_list
@@ -274,7 +277,10 @@ def concat_prompt_embeds(prompt_embeds: list["PromptEmbeds"]):
dtype=t.dtype,
device=t.device,
)
t = torch.cat([t, pad], dim=1)
if padding_side == "right":
t = torch.cat([t, pad], dim=1)
else:
t = torch.cat([pad, t], dim=1)
padded.append(t)
text_embeds = torch.cat(padded, dim=0)
@@ -296,7 +302,10 @@ def concat_prompt_embeds(prompt_embeds: list["PromptEmbeds"]):
dtype=m.dtype,
device=m.device,
)
m = torch.cat([m, pad], dim=1)
if padding_side == "right":
m = torch.cat([m, pad], dim=1)
else:
m = torch.cat([pad, m], dim=1)
padded.append(m)
attention_mask = torch.cat(padded, dim=0)

4
toolkit/stable_diffusion_model.py
View File

@@ -229,6 +229,10 @@ class StableDiffusion:
# use new lokr format (default false for old models for backwards compatibility)
self.use_old_lokr_format = True
# when padding to make batch size work, which side padding to use, right or left
# some llms need left side padding, others need right side
self.te_padding_side = "right"
# properties for old arch for backwards compatibility
@property
def is_xl(self):