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

@@ -50,4 +50,4 @@ pytorch-wavelets==1.3.0
 matplotlib==3.10.1
 setuptools==69.5.1
 av==16.0.1
-torchcodec
+torchcodec==0.9.1

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(

requirements.txt

@@ -38,4 +38,4 @@ matplotlib==3.10.1
 setuptools==69.5.1
 scipy==1.12.0
 av==16.0.1
-torchcodec
+torchcodec==0.9.1

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:

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:

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

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):

version.py

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