Automated PR - 2026-03-04

packages/ltx-trainer/scripts/decode_latents.py

@@ -287,7 +287,7 @@ class LatentsDecoder:
 
         # Save as WAV
         output_path = output_dir / f"{latent_file.stem}.wav"
-        sample_rate = self.vocoder.output_sample_rate
+        sample_rate = self.vocoder.output_sampling_rate
         torchaudio.save(str(output_path), waveform[0].cpu(), sample_rate)
 
 

packages/ltx-trainer/scripts/inference.py

@@ -410,7 +410,7 @@ def main() -> None:  # noqa: PLR0912, PLR0915
     # Get audio sample rate from vocoder if audio was generated
     audio_sample_rate = None
     if audio is not None and components.vocoder is not None:
-        audio_sample_rate = components.vocoder.output_sample_rate
+        audio_sample_rate = components.vocoder.output_sampling_rate
 
     save_video(
         video_tensor=video,

packages/ltx-trainer/scripts/process_captions.py

@@ -303,14 +303,13 @@ def compute_captions_embeddings(  # noqa: PLR0913
     ) as progress:
         task = progress.add_task("Processing captions", total=len(dataloader))
         for batch in dataloader:
-            # Encode prompts using _preprocess_text (returns embeddings before connector)
-            # This is what we want to save - the connector is applied during training
+            # Encode prompts using precompute() (returns video/audio features before connector)
+            # The connector is applied during training via embeddings_processor
             with torch.inference_mode():
-                # TODO(batch-tokenization): When tokenizer supports batching, encode all prompts at once:
-                #   prompt_embeds, prompt_attention_mask = text_encoder._preprocess_text(batch["prompt"]) # noqa: ERA001
+                # TODO(batch-tokenization): When tokenizer supports batching, encode all prompts at once.
                 # For now, process one at a time:
                 for i in range(len(batch["prompt"])):
-                    prompt_embeds, prompt_attention_mask = text_encoder._preprocess_text(
+                    video_prompt_embeds, audio_prompt_embeds, prompt_attention_mask = text_encoder.precompute(
                         batch["prompt"][i], padding_side="left"
                     )
 
@@ -321,9 +320,11 @@ def compute_captions_embeddings(  # noqa: PLR0913
                     output_dir_path.mkdir(parents=True, exist_ok=True)
 
                     embedding_data = {
-                        "prompt_embeds": prompt_embeds[0].cpu().contiguous(),
+                        "video_prompt_embeds": video_prompt_embeds[0].cpu().contiguous(),
                         "prompt_attention_mask": prompt_attention_mask[0].cpu().contiguous(),
                     }
+                    if audio_prompt_embeds is not None:
+                        embedding_data["audio_prompt_embeds"] = audio_prompt_embeds[0].cpu().contiguous()
 
                     output_file = output_path / output_rel_path
                     torch.save(embedding_data, output_file)

packages/ltx-trainer/scripts/process_videos.py

@@ -41,6 +41,7 @@ from torchvision.transforms.functional import crop, resize, to_tensor
 from transformers.utils.logging import disable_progress_bar
 
 from ltx_core.model.audio_vae import AudioProcessor
+from ltx_core.types import Audio
 from ltx_trainer import logger
 from ltx_trainer.model_loader import load_audio_vae_encoder, load_video_vae_encoder
 from ltx_trainer.utils import open_image_as_srgb
@@ -503,7 +504,7 @@ def compute_latents(  # noqa: PLR0913, PLR0915
             )
             # Create audio processor for waveform-to-spectrogram conversion
             audio_processor = AudioProcessor(
-                sample_rate=audio_vae_encoder.sample_rate,
+                target_sample_rate=audio_vae_encoder.sample_rate,
                 mel_bins=audio_vae_encoder.mel_bins,
                 mel_hop_length=audio_vae_encoder.mel_hop_length,
                 n_fft=audio_vae_encoder.n_fft,
@@ -567,10 +568,10 @@ def compute_latents(  # noqa: PLR0913, PLR0915
                     if audio_batch is not None:
                         # Extract the i-th item from batched audio data
                         # DataLoader collates [channels, samples] -> [batch, channels, samples]
-                        audio_data = {
-                            "waveform": audio_batch["waveform"][i],
-                            "sample_rate": audio_batch["sample_rate"][i].item(),
-                        }
+                        audio_data = Audio(
+                            waveform=audio_batch["waveform"][i],
+                            sampling_rate=audio_batch["sample_rate"][i].item(),
+                        )
 
                         # Encode audio
                         with torch.inference_mode():
@@ -822,13 +823,13 @@ def tiled_encode_video(  # noqa: PLR0912, PLR0915
 def encode_audio(
     audio_vae_encoder: torch.nn.Module,
     audio_processor: torch.nn.Module,
-    audio_data: dict[str, torch.Tensor | int],
+    audio: Audio,
 ) -> dict[str, torch.Tensor | int | float]:
     """Encode audio waveform into latent representation.
     Args:
         audio_vae_encoder: Audio VAE encoder model from ltx-core
         audio_processor: AudioProcessor for waveform-to-spectrogram conversion
-        audio_data: Dict with {"waveform": Tensor[channels, samples], "sample_rate": int}
+        audio: Audio container with waveform tensor and sampling rate.
     Returns:
         Dict containing audio latents and shape information:
         {
@@ -841,18 +842,17 @@ def encode_audio(
     device = next(audio_vae_encoder.parameters()).device
     dtype = next(audio_vae_encoder.parameters()).dtype
 
-    waveform = audio_data["waveform"].to(device=device, dtype=dtype)
-    sample_rate = audio_data["sample_rate"]
+    waveform = audio.waveform.to(device=device, dtype=dtype)
 
     # Add batch dimension if needed: [channels, samples] -> [batch, channels, samples]
     if waveform.dim() == 2:
         waveform = waveform.unsqueeze(0)
 
     # Calculate duration
-    duration = waveform.shape[-1] / sample_rate
+    duration = waveform.shape[-1] / audio.sampling_rate
 
     # Convert waveform to mel spectrogram using AudioProcessor
-    mel_spectrogram = audio_processor.waveform_to_mel(waveform, waveform_sample_rate=sample_rate)
+    mel_spectrogram = audio_processor.waveform_to_mel(Audio(waveform=waveform, sampling_rate=audio.sampling_rate))
     mel_spectrogram = mel_spectrogram.to(dtype=dtype)
 
     # Encode mel spectrogram to latents