From 7f3309b29114a2f47c65d25cfeeb697a0617199f Mon Sep 17 00:00:00 2001 From: Jaret Burkett Date: Tue, 24 Mar 2026 12:20:09 -0600 Subject: [PATCH] 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. --- .../diffusion_models/ltx2/ltx2.py | 50 ++++++++++----- toolkit/config_modules.py | 8 +++ toolkit/data_loader.py | 10 ++- toolkit/data_transfer_object/data_loader.py | 6 +- toolkit/dataloader_mixins.py | 63 ++++++++++++++----- ui/src/app/jobs/new/SimpleJob.tsx | 10 ++- ui/src/app/jobs/new/options.ts | 7 ++- ui/src/docs.tsx | 11 ++++ ui/src/types.ts | 1 + version.py | 2 +- 10 files changed, 132 insertions(+), 36 deletions(-) diff --git a/extensions_built_in/diffusion_models/ltx2/ltx2.py b/extensions_built_in/diffusion_models/ltx2/ltx2.py index 49692f1..f407a7c 100644 --- a/extensions_built_in/diffusion_models/ltx2/ltx2.py +++ b/extensions_built_in/diffusion_models/ltx2/ltx2.py @@ -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 diff --git a/toolkit/config_modules.py b/toolkit/config_modules.py index 228ad3f..a581417 100644 --- a/toolkit/config_modules.py +++ b/toolkit/config_modules.py @@ -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.") diff --git a/toolkit/data_loader.py b/toolkit/data_loader.py index 970fae1..dbfde71 100644 --- a/toolkit/data_loader.py +++ b/toolkit/data_loader.py @@ -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: diff --git a/toolkit/data_transfer_object/data_loader.py b/toolkit/data_transfer_object/data_loader.py index ebd53a7..5768ba1 100644 --- a/toolkit/data_transfer_object/data_loader.py +++ b/toolkit/data_transfer_object/data_loader.py @@ -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 diff --git a/toolkit/dataloader_mixins.py b/toolkit/dataloader_mixins.py index 1ab9cc0..55bbff8 100644 --- a/toolkit/dataloader_mixins.py +++ b/toolkit/dataloader_mixins.py @@ -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 diff --git a/ui/src/app/jobs/new/SimpleJob.tsx b/ui/src/app/jobs/new/SimpleJob.tsx index 0bbb2a2..36e3039 100644 --- a/ui/src/app/jobs/new/SimpleJob.tsx +++ b/ui/src/app/jobs/new/SimpleJob.tsx @@ -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 && ( setJobConfig(value, `config.process[0].datasets[${i}].is_reg`)} /> + {modelArch?.additionalSections?.includes('datasets.auto_frame_count') && ( + setJobConfig(value, `config.process[0].datasets[${i}].auto_frame_count`)} + docKey="datasets.auto_frame_count" + /> + )} {modelArch?.additionalSections?.includes('datasets.do_i2v') && ( ), }, + '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 => { diff --git a/ui/src/types.ts b/ui/src/types.ts index 5cc0c44..a22fe91 100644 --- a/ui/src/types.ts +++ b/ui/src/types.ts @@ -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 { diff --git a/version.py b/version.py index 482e124..1347f40 100644 --- a/version.py +++ b/version.py @@ -1 +1 @@ -VERSION = "0.7.29" +VERSION = "0.7.30"