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

Add support for audo frame count so datasets can have varrying length videos. Varous ltx 2.3 VAE optimizations such as removing tiling articacts, and doing frame split encoding to reduce vram on encoding/decoding.

Browse Source
This commit is contained in:
Jaret Burkett
2026-03-24 12:20:09 -06:00
parent 4ad14d211a
commit 7f3309b291
10 changed files with 132 additions and 36 deletions
Download Patch File Download Diff File Expand all files Collapse all files

50
extensions_built_in/diffusion_models/ltx2/ltx2.py
View File

@@ -77,7 +77,7 @@ dit_prefix = "model.diffusion_model."
vae_prefix = "vae."
audio_vae_prefix = "audio_vae."
vocoder_prefix = "vocoder."
base_te_path = "google/gemma-3-12b-it-qat-q4_0-unquantized"
base_te_path = "Lightricks/gemma-3-12b-it-qat-q4_0-unquantized"
HF_TOKEN = os.getenv("HF_TOKEN", None)
@@ -122,6 +122,10 @@ class ComboVae(torch.nn.Module):
def dtype(self):
return self.vae.dtype
@property
def config(self):
return self.vae.config
def encode(
self,
*args,
@@ -222,7 +226,7 @@ class LTX2Model(BaseModel):
# gemma needs left side padding
self.te_padding_side = "left"
# invalidate older caches
self.latent_space_version = f"{self.arch}_v2"
@@ -554,6 +558,13 @@ class LTX2Model(BaseModel):
self.pipeline.vae.eval()
self.pipeline.vae.requires_grad_(False)
if self.model_config.low_vram:
self.pipeline.vae.tile_sample_min_num_frames = 64
self.pipeline.vae.tile_sample_stride_num_frames = 16
# they check the wrong flat on encode currently so set both to future proof
self.pipeline.vae.use_framewise_decoding = True
self.pipeline.vae.use_framewise_encoding = True
image_list = [image.to(device, dtype=dtype) for image in image_list]
# Normalize shapes
@@ -571,7 +582,7 @@ class LTX2Model(BaseModel):
# Stack to (B, C, T, H, W)
images = torch.stack(norm_images)
latents = self.pipeline.vae.encode(images).latent_dist.mode()
latents = self.pipeline.vae.encode(images).latent_dist.sample()
# Normalize latents across the channel dimension [B, C, F, H, W]
scaling_factor = 1.0
@@ -582,6 +593,10 @@ class LTX2Model(BaseModel):
latents.device, latents.dtype
)
latents = (latents - latents_mean) * scaling_factor / latents_std
if self.model_config.low_vram:
self.pipeline.vae.use_framewise_decoding = False
self.pipeline.vae.use_framewise_encoding = False
return latents.to(device, dtype=dtype)
@@ -666,14 +681,17 @@ class LTX2Model(BaseModel):
if self.low_vram:
# set vae to tile decode
pipeline.vae.enable_tiling(
tile_sample_min_height=256,
tile_sample_min_width=256,
tile_sample_min_num_frames=8,
tile_sample_stride_height=224,
tile_sample_stride_width=224,
tile_sample_stride_num_frames=4,
)
# pipeline.vae.enable_tiling(
# tile_sample_min_height=256,
# tile_sample_min_width=256,
# tile_sample_min_num_frames=8,
# tile_sample_stride_height=224,
# tile_sample_stride_width=224,
# tile_sample_stride_num_frames=4,
# )
self.pipeline.vae.tile_sample_min_num_frames = 16
self.pipeline.vae.tile_sample_stride_num_frames = 8
self.pipeline.vae.use_framewise_decoding = True
# We only encode and store the minimum prompt tokens, but need them padded to 1024 for LTX2
conditional_embeds = self.pad_embeds(conditional_embeds)
@@ -718,7 +736,8 @@ class LTX2Model(BaseModel):
)
if self.low_vram:
# Restore no tiling
pipeline.vae.use_tiling = False
# pipeline.vae.use_tiling = False
self.pipeline.vae.use_framewise_decoding = False
if is_video:
# redurn as a dict, we will handle it with an override function
@@ -775,7 +794,7 @@ class LTX2Model(BaseModel):
# Encode mel spectrogram to latents
latents = self.pipeline.audio_vae.encode(
mel_spectrogram.to(self.device_torch, dtype=self.torch_dtype)
).latent_dist.mode()
).latent_dist.sample()
if audio_num_frames is None:
audio_num_frames = latents.shape[2] # (latents is [B, C, T, F])
@@ -841,7 +860,7 @@ class LTX2Model(BaseModel):
video_timestep = timestep.clone()
# i2v from first frame
if batch.dataset_config.do_i2v and batch.dataset_config.num_frames > 1:
if batch.dataset_config.do_i2v and batch.num_frames > 1:
# check to see if we had it cached
if batch.first_frame_latents is not None:
init_latents = batch.first_frame_latents.to(
@@ -893,7 +912,6 @@ class LTX2Model(BaseModel):
# set video timestep
video_timestep = timestep.unsqueeze(-1) * (1 - packed_conditioning_mask)
frame_rate = batch.dataset_config.fps
# check frame dimension
# Unpacked latents of shape are [B, C, F, H, W] are patched into tokens of shape [B, C, F // p_t, p_t, H // p, p, W // p, p].
@@ -930,7 +948,7 @@ class LTX2Model(BaseModel):
num_channels_latents_audio = (
self.pipeline.audio_vae.config.latent_channels
)
duration_s = batch.dataset_config.num_frames / frame_rate
duration_s = batch.num_frames / frame_rate
audio_latents_per_second = (
self.pipeline.audio_sampling_rate
/ self.pipeline.audio_hop_length

8
toolkit/config_modules.py
View File

@@ -977,6 +977,11 @@ class DatasetConfig:
# I recommend trimming your videos to the desired length and using shrink_video_to_frames(default)
self.fps: int = kwargs.get('fps', 24)
# auto_frame_count pull as many frames as in the video at given fps
# Important, make sure fps for dataset is set correctly.
# this wont work with bucketing for now until I can handle this before bucketing.
self.auto_frame_count: bool = kwargs.get('auto_frame_count', False)
# debug the frame count and frame selection. You dont need this. It is for debugging.
self.debug: bool = kwargs.get('debug', False)
@@ -1360,5 +1365,8 @@ def validate_configs(
if train_config.diff_output_preservation and train_config.blank_prompt_preservation:
raise ValueError("Cannot use both differential output preservation and blank prompt preservation at the same time. Please set one of them to False.")
if train_config.batch_size > 1 and any(dataset_config.auto_frame_count for dataset_config in dataset_configs):
raise ValueError("Cannot use batch size greater than 1 with auto_frame_count. Please set batch_size to 1 or auto_frame_count to False.")

10
toolkit/data_loader.py
View File

@@ -389,7 +389,7 @@ class AiToolkitDataset(LatentCachingMixin, ControlCachingMixin, CLIPCachingMixin
self.dataset_config = dataset_config
# update bucket divisibility
self.dataset_config.bucket_tolerance = sd.get_bucket_divisibility()
self.is_video = dataset_config.num_frames > 1
self.is_video = dataset_config.num_frames > 1 or dataset_config.auto_frame_count
super().__init__()
folder_path = dataset_config.folder_path
self.dataset_path = dataset_config.dataset_path
@@ -505,6 +505,13 @@ class AiToolkitDataset(LatentCachingMixin, ControlCachingMixin, CLIPCachingMixin
latent_space_version = 'sdxl'
else:
latent_space_version = self.sd.model_config.arch if self.sd is not None else "sd1"
temporal_compression = 8
if self.sd is not None:
if hasattr(self.sd.vae.config, 'scale_factor_temporal'):
temporal_compression = self.sd.vae.config.scale_factor_temporal
if hasattr(self.sd.unet.config, 'temporal_compression_ratio'):
temporal_compression = self.sd.unet.config.temporal_compression_ratio
bad_count = 0
for file in tqdm(file_list):
@@ -520,6 +527,7 @@ class AiToolkitDataset(LatentCachingMixin, ControlCachingMixin, CLIPCachingMixin
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,
temporal_compression=temporal_compression,
)
self.file_list.append(file_item)
except Exception as e:

6
toolkit/data_transfer_object/data_loader.py
View File

@@ -54,7 +54,9 @@ class FileItemDTO(
def __init__(self, *args, **kwargs):
self.path = kwargs.get("path", "")
self.dataset_config: "DatasetConfig" = kwargs.get("dataset_config", None)
self.is_video = self.dataset_config.num_frames > 1
self.is_video = self.dataset_config.num_frames > 1 or self.dataset_config.auto_frame_count
self.num_frames = self.dataset_config.num_frames
self.temporal_compression = kwargs.get("temporal_compression", 8)
size_database = kwargs.get("size_database", {})
dataset_root = kwargs.get("dataset_root", None)
self.encode_control_in_text_embeddings = kwargs.get(
@@ -197,6 +199,8 @@ class DataLoaderBatchDTO:
# just for holding noise and preds during training
self.audio_target: Union[torch.Tensor, None] = None
self.audio_pred: Union[torch.Tensor, None] = None
self.num_frames: int = self.file_items[0].num_frames
if not is_latents_cached:
# only return a tensor if latents are not cached

63
toolkit/dataloader_mixins.py
View File

@@ -507,11 +507,26 @@ class ImageProcessingDTOMixin:
frames_to_extract = []
if self.dataset_config.auto_frame_count:
# allow for any length video here but make sure it is temporally compressable.
vid_length_seconds = total_frames / video_fps
desired_num_frames = int(vid_length_seconds * self.dataset_config.fps)
# make sure it is divisible by temporal_compression
desired_num_frames = desired_num_frames // self.temporal_compression * self.temporal_compression
# TODO, all models currently add a key frame, but future models may not, update here if this changes.
desired_num_frames += 1 # add one for the key frame that is always added
self.num_frames = desired_num_frames
# Always stretch/shrink to the requested number of frames if needed
if self.dataset_config.shrink_video_to_frames or total_frames < self.dataset_config.num_frames:
if self.dataset_config.shrink_video_to_frames or total_frames < self.num_frames:
# Distribute frames evenly across the entire video
interval = max_frame_index / (self.dataset_config.num_frames - 1) if self.dataset_config.num_frames > 1 else 0
frames_to_extract = [min(int(round(i * interval)), max_frame_index) for i in range(self.dataset_config.num_frames)]
interval = max_frame_index / (self.num_frames - 1) if self.num_frames > 1 else 0
frames_to_extract = [min(int(round(i * interval)), max_frame_index) for i in range(self.num_frames)]
else:
# Calculate frame interval based on FPS ratio
fps_ratio = video_fps / self.dataset_config.fps
@@ -520,17 +535,17 @@ class ImageProcessingDTOMixin:
# Calculate max consecutive frames we can extract at desired FPS
max_consecutive_frames = (total_frames // frame_interval)
if max_consecutive_frames < self.dataset_config.num_frames:
if max_consecutive_frames < self.num_frames:
# Not enough frames at desired FPS, so stretch instead
interval = max_frame_index / (self.dataset_config.num_frames - 1) if self.dataset_config.num_frames > 1 else 0
frames_to_extract = [min(int(round(i * interval)), max_frame_index) for i in range(self.dataset_config.num_frames)]
interval = max_frame_index / (self.num_frames - 1) if self.num_frames > 1 else 0
frames_to_extract = [min(int(round(i * interval)), max_frame_index) for i in range(self.num_frames)]
else:
# Calculate max start frame to ensure we can get all num_frames
max_start_frame = max_frame_index - ((self.dataset_config.num_frames - 1) * frame_interval)
max_start_frame = max_frame_index - ((self.num_frames - 1) * frame_interval)
start_frame = random.randint(0, max(0, max_start_frame))
# Generate list of frames to extract
frames_to_extract = [start_frame + (i * frame_interval) for i in range(self.dataset_config.num_frames)]
frames_to_extract = [start_frame + (i * frame_interval) for i in range(self.num_frames)]
# Final safety check - ensure no frame exceeds max valid index
frames_to_extract = [min(frame_idx, max_frame_index) for frame_idx in frames_to_extract]
@@ -642,7 +657,7 @@ class ImageProcessingDTOMixin:
# Target duration is how this sampled/stretched clip is interpreted for training
# (i.e. num_frames at the configured dataset FPS).
if hasattr(self.dataset_config, "fps") and self.dataset_config.fps and self.dataset_config.fps > 0:
target_duration = float(self.dataset_config.num_frames) / float(self.dataset_config.fps)
target_duration = float(self.num_frames) / float(self.dataset_config.fps)
else:
target_duration = source_duration
@@ -757,7 +772,7 @@ class ImageProcessingDTOMixin:
if self.has_unconditional:
self.load_unconditional_image()
return
if self.dataset_config.num_frames > 1:
if self.dataset_config.num_frames > 1 or self.dataset_config.auto_frame_count:
self.load_and_process_video(transform, only_load_latents)
return
try:
@@ -1751,16 +1766,25 @@ class LatentCachingFileItemDTOMixin:
("latent_space_version", self.latent_space_version),
("latent_version", self.latent_version),
])
is_video = False
# when adding items, do it after so we dont change old latents
if self.flip_x:
item["flip_x"] = True
if self.flip_y:
item["flip_y"] = True
if self.dataset_config.num_frames > 1:
if self.dataset_config.auto_frame_count:
# don't store num frames here as it is calculated dynamically
item["auto_frame_count"] = True
is_video = True
elif self.dataset_config.num_frames > 1:
item["num_frames"] = self.dataset_config.num_frames
if self.dataset_config.do_i2v:
is_video = True
if is_video and self.dataset_config.fps != 24:
# only add fps if it deviates from the default
item["fps"] = self.dataset_config.fps
if is_video and self.dataset_config.do_i2v:
item["do_i2v"] = True
if self.dataset_config.do_audio:
if is_video and self.dataset_config.do_audio:
item["do_audio"] = True
if self.dataset_config.audio_normalize:
item["audio_normalize"] = True
@@ -1815,6 +1839,8 @@ class LatentCachingFileItemDTOMixin:
self._cached_first_frame_latent = state_dict['first_frame_latent']
if 'audio_latent' in state_dict:
self._cached_audio_latent = state_dict['audio_latent']
if 'num_frames' in state_dict:
self.num_frames = int(state_dict['num_frames'].item())
return self._encoded_latent
@@ -1866,6 +1892,7 @@ class LatentCachingMixin:
state_dict = OrderedDict()
first_frame_latent = None
audio_latent = None
frames = None
# add batch dimension
try:
imgs = file_item.tensor.unsqueeze(0).to(device, dtype=dtype)
@@ -1877,7 +1904,8 @@ class LatentCachingMixin:
print_acc(f"Error: {str(e)}")
raise e
# do first frame
if self.dataset_config.num_frames > 1 and self.dataset_config.do_i2v:
is_video = self.dataset_config.auto_frame_count or self.dataset_config.num_frames > 1
if is_video and self.dataset_config.do_i2v:
frames = file_item.tensor.unsqueeze(0).to(device, dtype=dtype)
if len(frames.shape) == 4:
first_frames = frames
@@ -1895,6 +1923,9 @@ class LatentCachingMixin:
if to_disk:
state_dict['audio_latent'] = audio_latent.clone().detach().cpu()
if is_video:
state_dict['num_frames'] = torch.tensor(file_item.num_frames, dtype=torch.int32)
# save_latent
if to_disk:
# metadata
@@ -1912,7 +1943,11 @@ class LatentCachingMixin:
del imgs
del latent
del frames
del file_item.tensor
del state_dict
del first_frame_latent
del audio_latent
file_item.cleanup()
file_item.is_latent_cached = True

10
ui/src/app/jobs/new/SimpleJob.tsx
View File

@@ -846,7 +846,7 @@ export default function SimpleJob({
min={0}
required
/>
{modelArch?.additionalSections?.includes('datasets.num_frames') && (
{modelArch?.additionalSections?.includes('datasets.num_frames') && !dataset.auto_frame_count && (
<NumberInput
label="Num Frames"
className="pt-2"
@@ -873,6 +873,14 @@ export default function SimpleJob({
checked={dataset.is_reg || false}
onChange={value => setJobConfig(value, `config.process[0].datasets[${i}].is_reg`)}
/>
{modelArch?.additionalSections?.includes('datasets.auto_frame_count') && (
<Checkbox
label="Auto Frame Count"
checked={dataset.auto_frame_count || false}
onChange={value => setJobConfig(value, `config.process[0].datasets[${i}].auto_frame_count`)}
docKey="datasets.auto_frame_count"
/>
)}
{modelArch?.additionalSections?.includes('datasets.do_i2v') && (
<Checkbox
label="Do I2V"

7
ui/src/app/jobs/new/options.ts
View File

@@ -20,6 +20,7 @@ type AdditionalSections =
| 'datasets.do_audio'
| 'datasets.audio_normalize'
| 'datasets.audio_preserve_pitch'
| 'datasets.auto_frame_count'
| 'sample.ctrl_img'
| 'sample.multi_ctrl_imgs'
| 'train.audio_loss_multiplier'
@@ -664,9 +665,10 @@ export const modelArchs: ModelArch[] = [
'config.process[0].datasets[x].do_i2v': [false, undefined],
'config.process[0].datasets[x].do_audio': [true, undefined],
'config.process[0].datasets[x].fps': [24, undefined],
'config.process[0].datasets[x].auto_frame_count': [false, undefined],
},
disableSections: ['network.conv'],
additionalSections: ['sample.ctrl_img', 'datasets.num_frames', 'model.layer_offloading', 'model.low_vram', 'datasets.do_audio', 'datasets.audio_normalize', 'datasets.audio_preserve_pitch', 'datasets.do_i2v', 'train.audio_loss_multiplier'],
additionalSections: ['sample.ctrl_img', 'datasets.num_frames', 'model.layer_offloading', 'model.low_vram', 'datasets.do_audio', 'datasets.audio_normalize', 'datasets.audio_preserve_pitch', 'datasets.do_i2v', 'train.audio_loss_multiplier', 'datasets.auto_frame_count'],
},
{
name: 'ltx2.3',
@@ -691,9 +693,10 @@ export const modelArchs: ModelArch[] = [
'config.process[0].datasets[x].do_i2v': [false, undefined],
'config.process[0].datasets[x].do_audio': [true, undefined],
'config.process[0].datasets[x].fps': [24, undefined],
'config.process[0].datasets[x].auto_frame_count': [false, undefined],
},
disableSections: ['network.conv'],
additionalSections: ['sample.ctrl_img', 'datasets.num_frames', 'model.layer_offloading', 'model.low_vram', 'datasets.do_audio', 'datasets.audio_normalize', 'datasets.audio_preserve_pitch', 'datasets.do_i2v', 'train.audio_loss_multiplier'],
additionalSections: ['sample.ctrl_img', 'datasets.num_frames', 'model.layer_offloading', 'model.low_vram', 'datasets.do_audio', 'datasets.audio_normalize', 'datasets.audio_preserve_pitch', 'datasets.do_i2v', 'train.audio_loss_multiplier', 'datasets.auto_frame_count'],
},
{
name: 'flux2_klein_4b',

11
ui/src/docs.tsx
View File

@@ -329,6 +329,17 @@ const docs: { [key: string]: ConfigDoc } = {
</>
),
},
'datasets.auto_frame_count': {
title: 'Auto Frame Count',
description: (
<>
This will automatically determine the number of frames to use for each video in your dataset instead of relying
on a fixed num_frames. This allows you to include videos of different lengths in the dataset, and each video
will be processed without speeding up or slowing down. Be careful about adding long videos into your dataset, as
they use up more VRAM. This currently will not work with a batch size greater than 1.
</>
),
},
};
export const getDoc = (key: string | null | undefined): ConfigDoc | null => {

1
ui/src/types.ts
View File

@@ -107,6 +107,7 @@ export interface DatasetConfig {
control_path_1?: string | null;
control_path_2?: string | null;
control_path_3?: string | null;
auto_frame_count?: boolean;
}
export interface EMAConfig {

2
version.py
View File

@@ -1 +1 @@
VERSION = "0.7.29"
VERSION = "0.7.30"