From 5642b656b926edcb231f306f656f11eb8398a73d Mon Sep 17 00:00:00 2001 From: Jaret Burkett Date: Mon, 23 Mar 2026 20:08:33 +0000 Subject: [PATCH] Fix audio issues with ltx2 models. Silent codec fails now raised. Auto convert surround sound audio to stereo. Invalidate old caches just to be safe so they recache now. --- dgx_requirements.txt | 2 +- .../diffusion_models/ltx2/ltx2.py | 7 ++++-- requirements.txt | 2 +- toolkit/data_loader.py | 2 ++ toolkit/dataloader_mixins.py | 25 +++++++++++++++---- toolkit/models/base_model.py | 3 +++ toolkit/stable_diffusion_model.py | 3 +++ version.py | 2 +- 8 files changed, 36 insertions(+), 10 deletions(-) diff --git a/dgx_requirements.txt b/dgx_requirements.txt index 1906121..91a69a7 100644 --- a/dgx_requirements.txt +++ b/dgx_requirements.txt @@ -50,4 +50,4 @@ pytorch-wavelets==1.3.0 matplotlib==3.10.1 setuptools==69.5.1 av==16.0.1 -torchcodec \ No newline at end of file +torchcodec==0.9.1 \ No newline at end of file diff --git a/extensions_built_in/diffusion_models/ltx2/ltx2.py b/extensions_built_in/diffusion_models/ltx2/ltx2.py index 4359116..49692f1 100644 --- a/extensions_built_in/diffusion_models/ltx2/ltx2.py +++ b/extensions_built_in/diffusion_models/ltx2/ltx2.py @@ -222,6 +222,9 @@ class LTX2Model(BaseModel): # gemma needs left side padding self.te_padding_side = "left" + + # invalidate older caches + self.latent_space_version = f"{self.arch}_v2" # static method to get the noise scheduler @staticmethod @@ -890,8 +893,8 @@ class LTX2Model(BaseModel): # set video timestep video_timestep = timestep.unsqueeze(-1) * (1 - packed_conditioning_mask) - # todo get this somehow - frame_rate = 24 + + 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]. packed_latents = self.pipeline._pack_latents( diff --git a/requirements.txt b/requirements.txt index 54f05f2..5550011 100644 --- a/requirements.txt +++ b/requirements.txt @@ -38,4 +38,4 @@ matplotlib==3.10.1 setuptools==69.5.1 scipy==1.12.0 av==16.0.1 -torchcodec \ No newline at end of file +torchcodec==0.9.1 \ No newline at end of file diff --git a/toolkit/data_loader.py b/toolkit/data_loader.py index 51605c9..970fae1 100644 --- a/toolkit/data_loader.py +++ b/toolkit/data_loader.py @@ -491,6 +491,8 @@ class AiToolkitDataset(LatentCachingMixin, ControlCachingMixin, CLIPCachingMixin 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 not None and self.sd.latent_space_version is not None: + latent_space_version = self.sd.latent_space_version elif self.sd.is_xl: latent_space_version = 'sdxl' elif self.sd.is_v3: diff --git a/toolkit/dataloader_mixins.py b/toolkit/dataloader_mixins.py index 20140eb..1ab9cc0 100644 --- a/toolkit/dataloader_mixins.py +++ b/toolkit/dataloader_mixins.py @@ -106,6 +106,22 @@ def clean_caption(caption): # caption = ', '.join(caption_split) return caption +def waveform_to_stereo(waveform): + c = waveform.shape[0] + if c == 2: + return waveform + if c == 1: + return waveform.expand(2, -1) + if c == 6: # 5.1: FL, FR, FC, LFE, BL, BR + fl, fr, fc, _, bl, br = waveform + k = 0.7071 + return torch.stack([fl + k * fc + k * bl, fr + k * fc + k * br]) + if c == 8: # 7.1: FL, FR, FC, LFE, BL, BR, SL, SR + fl, fr, fc, _, bl, br, sl, sr = waveform + k = 0.7071 + return torch.stack([fl + k * fc + k * (bl + sl), fr + k * fc + k * (br + sr)]) + return waveform.mean(0, keepdim=True).expand(2, -1) + class CaptionMixin: def get_caption_item(self: 'AiToolkitDataset', index): @@ -632,6 +648,8 @@ class ImageProcessingDTOMixin: waveform, sample_rate = torchaudio.load(self.path) # [channels, samples] + waveform = waveform_to_stereo(waveform) # Convert to stereo if not already + if self.dataset_config.audio_normalize: peak = waveform.abs().amax() # global peak across channels eps = 1e-9 @@ -670,11 +688,8 @@ class ImageProcessingDTOMixin: self.audio_data = {"waveform": waveform, "sample_rate": int(sample_rate)} except Exception as e: - # Keep behavior identical for non-audio datasets; for audio datasets, just skip if missing/broken. - if hasattr(self.dataset_config, 'debug') and self.dataset_config.debug: - print_acc(f"Could not extract/stretch audio for {self.path}: {e}") - self.audio_data = None - self.audio_tensor = None + # if issue with libtorchcodec "Could not load libtorchcodec" + raise Exception(f"** WARNING ** - Error Processing audio for {self.path}. Error: {e}") # Only log success in debug mode if hasattr(self.dataset_config, 'debug') and self.dataset_config.debug: diff --git a/toolkit/models/base_model.py b/toolkit/models/base_model.py index ae117c8..938649b 100644 --- a/toolkit/models/base_model.py +++ b/toolkit/models/base_model.py @@ -194,6 +194,9 @@ class BaseModel: # 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" + + # can be used on models to invalidate cache if things change. + self.latent_space_version = None # properties for old arch for backwards compatibility @property diff --git a/toolkit/stable_diffusion_model.py b/toolkit/stable_diffusion_model.py index e235da9..d0b5488 100644 --- a/toolkit/stable_diffusion_model.py +++ b/toolkit/stable_diffusion_model.py @@ -233,6 +233,9 @@ class StableDiffusion: # some llms need left side padding, others need right side self.te_padding_side = "right" + # can be used on models to invalidate cache if things change. + self.latent_space_version = None + # properties for old arch for backwards compatibility @property def is_xl(self): diff --git a/version.py b/version.py index 1c11438..482e124 100644 --- a/version.py +++ b/version.py @@ -1 +1 @@ -VERSION = "0.7.28" +VERSION = "0.7.29"