indigo/LTX-2
1
0
Fork 0
Code Issues Pull Requests 1 Actions Packages Projects Releases Wiki Activity

Automated PR - 2026-01-15

Browse Source
This commit is contained in:
sync-bot
2026-01-15 19:19:42 +00:00
parent bd92a5f408
commit 310103a53e
13 changed files with 722 additions and 207 deletions
Download Patch File Download Diff File Expand all files Collapse all files
packages/ltx-trainer/scripts/decode_latents.py
+93 -65
View File
@@ -15,6 +15,7 @@ import torch
import torchaudio
import torchvision.utils
import typer
from einops import rearrange
from rich.console import Console
from rich.progress import (
BarColumn,
@@ -27,10 +28,16 @@ from rich.progress import (
)
from transformers.utils.logging import disable_progress_bar
from ltx_core.model.video_vae import SpatialTilingConfig, TemporalTilingConfig, TilingConfig
from ltx_trainer import logger
from ltx_trainer.model_loader import load_audio_vae_decoder, load_video_vae_decoder, load_vocoder
from ltx_trainer.video_utils import save_video
DEFAULT_TILE_SIZE_PIXELS = 512 # Spatial tile size in pixels (must be ≥64 and divisible by 32)
DEFAULT_TILE_OVERLAP_PIXELS = 128 # Spatial tile overlap in pixels (must be divisible by 32)
DEFAULT_TILE_SIZE_FRAMES = 128 # Temporal tile size in frames (must be ≥16 and divisible by 8)
DEFAULT_TILE_OVERLAP_FRAMES = 24 # Temporal tile overlap in frames (must be divisible by 8)
disable_progress_bar()
console = Console()
app = typer.Typer(
@@ -60,19 +67,15 @@ class LatentsDecoder:
self.vae = None
self.audio_vae = None
self.vocoder = None
self._load_model(model_path, vae_tiling, with_audio)
self.vae_tiling = vae_tiling
def _load_model(self, model_path: str, vae_tiling: bool, with_audio: bool = False) -> None:
self._load_model(model_path, with_audio)
def _load_model(self, model_path: str, with_audio: bool = False) -> None:
"""Initialize and load the VAE model(s)."""
with console.status(f"[bold]Loading video VAE decoder from {model_path}...", spinner="dots"):
self.vae = load_video_vae_decoder(model_path, device=self.device, dtype=torch.bfloat16)
if vae_tiling:
logger.warning(
"VAE tiling is not yet implemented in this script. "
"Continuing without tiling - this may cause OOM errors for large resolutions."
)
if with_audio:
with console.status(f"[bold]Loading audio VAE decoder from {model_path}...", spinner="dots"):
self.audio_vae = load_audio_vae_decoder(model_path, device=self.device, dtype=torch.bfloat16)
@@ -125,61 +128,6 @@ class LatentsDecoder:
logger.info(f"Decoding complete! Videos saved to {output_dir}")
def _process_file(self, latent_file: Path, output_dir: Path, seed: int | None) -> None:
"""Process a single latent file."""
# Load the latent data
data = torch.load(latent_file, map_location=self.device, weights_only=False)
# Get latents - handle both old patchified [seq_len, C] and new [C, F, H, W] formats
latents = data["latents"]
num_frames = data["num_frames"]
height = data["height"]
width = data["width"]
# Check if latents need reshaping (old patchified format)
if latents.dim() == 2:
# Old format: [seq_len, C] -> reshape to [C, F, H, W]
_seq_len, channels = latents.shape
latents = latents.reshape(num_frames, height, width, channels)
latents = latents.permute(3, 0, 1, 2) # [F, H, W, C] -> [C, F, H, W]
# Add batch dimension: [C, F, H, W] -> [1, C, F, H, W]
latents = latents.unsqueeze(0).to(device=self.device, dtype=torch.bfloat16)
# Create generator only if seed is provided
generator = None
if seed is not None:
generator = torch.Generator(device=self.device)
generator.manual_seed(seed)
# Decode the video (VAE decoder uses forward/call, not decode method)
video = self.vae(latents) # [B, C, F, H, W]
# Convert to [F, C, H, W] format and normalize to [0, 1]
video = video[0] # Remove batch dimension -> [C, F, H, W]
video = video.permute(1, 0, 2, 3) # [C, F, H, W] -> [F, C, H, W]
video = (video + 1) / 2 # Denormalize from [-1, 1] to [0, 1]
video = video.clamp(0, 1)
# Determine output format and save
is_image = video.shape[0] == 1
if is_image:
# Save as PNG for single frame
output_path = output_dir / f"{latent_file.stem}.png"
torchvision.utils.save_image(
video[0], # [C, H, W] in [0, 1]
str(output_path),
)
else:
# Save as MP4 for video using PyAV-based save_video
output_path = output_dir / f"{latent_file.stem}.mp4"
fps = data.get("fps", 24) # Use stored FPS or default to 24
save_video(
video_tensor=video, # [F, C, H, W] in [0, 1]
output_path=output_path,
fps=fps,
)
@torch.inference_mode()
def decode_audio(self, latents_dir: Path, output_dir: Path) -> None:
"""Decode all audio latent files in the directory recursively.
@@ -229,6 +177,87 @@ class LatentsDecoder:
logger.info(f"Audio decoding complete! Audio files saved to {output_dir}")
def _process_file(self, latent_file: Path, output_dir: Path, seed: int | None) -> None:
"""Process a single latent file."""
# Load the latent data
data = torch.load(latent_file, map_location=self.device, weights_only=False)
# Get latents - handle both old patchified [seq_len, C] and new [C, F, H, W] formats
latents = data["latents"]
num_frames = data["num_frames"]
height = data["height"]
width = data["width"]
# Check if latents need reshaping (old patchified format)
if latents.dim() == 2:
# Old format: [seq_len, C] -> reshape to [C, F, H, W]
latents = rearrange(latents, "(f h w) c -> c f h w", f=num_frames, h=height, w=width)
# Add batch dimension: [C, F, H, W] -> [1, C, F, H, W]
latents = latents.unsqueeze(0).to(device=self.device, dtype=torch.bfloat16)
# Create generator only if seed is provided
generator = None
if seed is not None:
generator = torch.Generator(device=self.device)
generator.manual_seed(seed)
# Decode the video
video = self._decode_video(latents, generator)
# Determine output format and save
is_image = video.shape[0] == 1
if is_image:
# Save as PNG for single frame
output_path = output_dir / f"{latent_file.stem}.png"
torchvision.utils.save_image(
video[0], # [C, H, W] in [0, 1]
str(output_path),
)
else:
# Save as MP4 for video using PyAV-based save_video
output_path = output_dir / f"{latent_file.stem}.mp4"
fps = data.get("fps", 24) # Use stored FPS or default to 24
save_video(
video_tensor=video, # [F, C, H, W] in [0, 1]
output_path=output_path,
fps=fps,
)
def _decode_video(self, latents: torch.Tensor, generator: torch.Generator | None = None) -> torch.Tensor:
"""Decode latents to video frames."""
if self.vae_tiling:
# Use tiled decoding for reduced VRAM
tiling_config = TilingConfig(
spatial_config=SpatialTilingConfig(
tile_size_in_pixels=DEFAULT_TILE_SIZE_PIXELS,
tile_overlap_in_pixels=DEFAULT_TILE_OVERLAP_PIXELS,
),
temporal_config=TemporalTilingConfig(
tile_size_in_frames=DEFAULT_TILE_SIZE_FRAMES,
tile_overlap_in_frames=DEFAULT_TILE_OVERLAP_FRAMES,
),
)
chunks = list(
self.vae.tiled_decode(
latents,
tiling_config=tiling_config,
generator=generator,
)
)
# Concatenate along temporal dimension
video = torch.cat(chunks, dim=2) # [B, C, F, H, W]
else:
# Standard full decoding
video = self.vae(latents, generator=generator) # [B, C, F, H, W]
# Convert to [F, C, H, W] format and normalize to [0, 1]
video = rearrange(video, "1 c f h w -> f c h w")
video = (video + 1) / 2 # Denormalize from [-1, 1] to [0, 1]
video = video.clamp(0, 1)
return video
def _process_audio_file(self, latent_file: Path, output_dir: Path) -> None:
"""Process a single audio latent file."""
# Load the latent data
@@ -242,8 +271,7 @@ class LatentsDecoder:
if latents.dim() == 2:
# Old format: [seq_len, channels] where seq_len = time * freq
# Reshape to [C, T, F]
latents = latents.reshape(num_time_steps, freq_bins, -1) # [T, F, C]
latents = latents.permute(2, 0, 1) # [T, F, C] -> [C, T, F]
latents = rearrange(latents, "(t f) c -> c t f", t=num_time_steps, f=freq_bins)
# Add batch dimension: [C, T, F] -> [1, C, T, F]
latents = latents.unsqueeze(0)
packages/ltx-trainer/scripts/process_captions.py
+16 -3
View File
@@ -220,7 +220,7 @@ class CaptionsDataset(Dataset):
break
def compute_captions_embeddings(
def compute_captions_embeddings( # noqa: PLR0913
dataset_file: str | Path,
output_dir: str,
model_path: str,
@@ -231,6 +231,7 @@ def compute_captions_embeddings(
remove_llm_prefixes: bool = False,
batch_size: int = 8,
device: str = "cuda",
load_in_8bit: bool = False,
) -> None:
"""
Process captions and save text embeddings.
@@ -245,6 +246,7 @@ def compute_captions_embeddings(
remove_llm_prefixes: Whether to remove common LLM-generated prefixes
batch_size: Batch size for processing
device: Device to use for computation
load_in_8bit: Whether to load the Gemma text encoder in 8-bit precision
"""
console = Console()
@@ -264,7 +266,13 @@ def compute_captions_embeddings(
# Load text encoder
with console.status("[bold]Loading Gemma text encoder...", spinner="dots"):
text_encoder = load_text_encoder(model_path, text_encoder_path, device=device, dtype=torch.bfloat16)
text_encoder = load_text_encoder(
model_path,
text_encoder_path,
device=device,
dtype=torch.bfloat16,
load_in_8bit=load_in_8bit,
)
logger.info("Text encoder loaded successfully")
@@ -326,7 +334,7 @@ def compute_captions_embeddings(
@app.command()
def main(
def main( # noqa: PLR0913
dataset_file: str = typer.Argument(
...,
help="Path to metadata file (CSV/JSON/JSONL) containing captions and media paths",
@@ -368,6 +376,10 @@ def main(
default=False,
help="Remove common LLM-generated prefixes from captions",
),
load_text_encoder_in_8bit: bool = typer.Option(
default=False,
help="Load the Gemma text encoder in 8-bit precision to save GPU memory (requires bitsandbytes)",
),
) -> None:
"""Process text captions and save embeddings for video generation training.
This script processes captions from metadata files and saves text embeddings
@@ -408,6 +420,7 @@ def main(
remove_llm_prefixes=remove_llm_prefixes,
batch_size=batch_size,
device=device,
load_in_8bit=load_text_encoder_in_8bit,
)
packages/ltx-trainer/scripts/process_dataset.py
+44 -41
View File
@@ -20,6 +20,7 @@ from process_videos import compute_latents, parse_resolution_buckets
from rich.console import Console
from ltx_trainer import logger
from ltx_trainer.gpu_utils import free_gpu_memory_context
console = Console()
app = typer.Typer(
@@ -46,15 +47,9 @@ def preprocess_dataset( # noqa: PLR0913
remove_llm_prefixes: bool = False,
reference_column: str | None = None,
with_audio: bool = False,
load_text_encoder_in_8bit: bool = False,
) -> None:
"""Run the preprocessing pipeline with the given arguments."""
# VAE tiling is not yet implemented
if vae_tiling:
logger.warning(
"VAE tiling is not yet implemented in this script. "
"Continuing without tiling - this may cause OOM errors for large resolutions."
)
# Validate dataset file
_validate_dataset_file(dataset_file)
@@ -66,19 +61,21 @@ def preprocess_dataset( # noqa: PLR0913
if lora_trigger:
logger.info(f'LoRA trigger word "{lora_trigger}" will be prepended to all captions')
# Process captions using the dedicated function
compute_captions_embeddings(
dataset_file=dataset_file,
output_dir=str(conditions_dir),
model_path=model_path,
text_encoder_path=text_encoder_path,
caption_column=caption_column,
media_column=video_column,
lora_trigger=lora_trigger,
remove_llm_prefixes=remove_llm_prefixes,
batch_size=batch_size,
device=device,
)
with free_gpu_memory_context():
# Process captions using the dedicated function
compute_captions_embeddings(
dataset_file=dataset_file,
output_dir=str(conditions_dir),
model_path=model_path,
text_encoder_path=text_encoder_path,
caption_column=caption_column,
media_column=video_column,
lora_trigger=lora_trigger,
remove_llm_prefixes=remove_llm_prefixes,
batch_size=batch_size,
device=device,
load_in_8bit=load_text_encoder_in_8bit,
)
# Process videos using the dedicated function
audio_latents_dir = None
@@ -86,36 +83,37 @@ def preprocess_dataset( # noqa: PLR0913
logger.info("Audio preprocessing enabled - will extract and encode audio from videos")
audio_latents_dir = output_base / "audio_latents"
compute_latents(
dataset_file=dataset_file,
video_column=video_column,
resolution_buckets=resolution_buckets,
output_dir=str(latents_dir),
model_path=model_path,
batch_size=batch_size,
device=device,
vae_tiling=vae_tiling,
with_audio=with_audio,
audio_output_dir=str(audio_latents_dir) if audio_latents_dir else None,
)
# Process reference videos if reference_column is provided
if reference_column:
logger.info("Processing reference videos for IC-LoRA training...")
reference_latents_dir = output_base / "reference_latents"
with free_gpu_memory_context():
compute_latents(
dataset_file=dataset_file,
main_media_column=video_column,
video_column=reference_column,
video_column=video_column,
resolution_buckets=resolution_buckets,
output_dir=str(reference_latents_dir),
output_dir=str(latents_dir),
model_path=model_path,
batch_size=batch_size,
device=device,
vae_tiling=vae_tiling,
with_audio=with_audio,
audio_output_dir=str(audio_latents_dir) if audio_latents_dir else None,
)
# Process reference videos if reference_column is provided
if reference_column:
logger.info("Processing reference videos for IC-LoRA training...")
reference_latents_dir = output_base / "reference_latents"
compute_latents(
dataset_file=dataset_file,
main_media_column=video_column,
video_column=reference_column,
resolution_buckets=resolution_buckets,
output_dir=str(reference_latents_dir),
model_path=model_path,
batch_size=batch_size,
device=device,
vae_tiling=vae_tiling,
)
# Handle decoding if requested (for verification)
if decode:
logger.info("Decoding latents for verification...")
@@ -224,6 +222,10 @@ def main( # noqa: PLR0913
default=False,
help="Extract and encode audio from video files",
),
load_text_encoder_in_8bit: bool = typer.Option(
default=False,
help="Load the Gemma text encoder in 8-bit precision to save GPU memory (requires bitsandbytes)",
),
) -> None:
"""Preprocess a video dataset by computing and saving latents and text embeddings.
The dataset must be a CSV, JSON, or JSONL file with columns for captions and video paths.
@@ -269,6 +271,7 @@ def main( # noqa: PLR0913
remove_llm_prefixes=remove_llm_prefixes,
reference_column=reference_column,
with_audio=with_audio,
load_text_encoder_in_8bit=load_text_encoder_in_8bit,
)
packages/ltx-trainer/scripts/process_videos.py
+176 -9
View File
@@ -59,6 +59,9 @@ VAE_TEMPORAL_FACTOR = 8
AUDIO_LATENT_CHANNELS = 8
AUDIO_FREQUENCY_BINS = 16
DEFAULT_TILE_SIZE = 512 # Spatial tile size in pixels (must be ≥64 and divisible by 32)
DEFAULT_TILE_OVERLAP = 128 # Spatial tile overlap in pixels (must be divisible by 32)
app = typer.Typer(
pretty_exceptions_enable=False,
no_args_is_help=True,
@@ -488,12 +491,6 @@ def compute_latents( # noqa: PLR0913, PLR0915
with console.status(f"[bold]Loading video VAE encoder from [cyan]{model_path}[/]...", spinner="dots"):
vae = load_video_vae_encoder(model_path, device=torch_device, dtype=torch.bfloat16)
if vae_tiling:
logger.warning(
"VAE tiling is not yet implemented in this script. "
"Continuing without tiling - this may cause OOM errors for large resolutions."
)
# Load audio VAE encoder and audio processor if needed
audio_vae_encoder = None
audio_processor = None
@@ -543,7 +540,7 @@ def compute_latents( # noqa: PLR0913, PLR0915
# Encode video
with torch.inference_mode():
video_latent_data = encode_video(vae=vae, video=video)
video_latent_data = encode_video(vae=vae, video=video, use_tiling=vae_tiling)
# Save latents for each item in batch
for i in range(len(batch["relative_path"])):
@@ -611,6 +608,9 @@ def encode_video(
vae: torch.nn.Module,
video: torch.Tensor,
dtype: torch.dtype | None = None,
use_tiling: bool = False,
tile_size: int = DEFAULT_TILE_SIZE,
tile_overlap: int = DEFAULT_TILE_OVERLAP,
) -> dict[str, torch.Tensor | int]:
"""Encode video into non-patchified latent representation.
Args:
@@ -618,6 +618,9 @@ def encode_video(
video: Input tensor of shape [B, C, F, H, W] (batch, channels, frames, height, width)
This is the format expected by the VAE encoder.
dtype: Target dtype for output latents
use_tiling: Whether to use spatial tiling for memory efficiency
tile_size: Tile size in pixels (must be divisible by 32)
tile_overlap: Overlap between tiles in pixels (must be divisible by 32)
Returns:
Dict containing non-patchified latents and shape information:
{
@@ -636,8 +639,17 @@ def encode_video(
video = video.to(device=device, dtype=vae_dtype)
# Encode video - VAE expects [B, C, F, H, W], returns [B, C, F', H', W']
latents = vae(video)
# Choose encoding method based on tiling flag
if use_tiling:
latents = tiled_encode_video(
vae=vae,
video=video,
tile_size=tile_size,
tile_overlap=tile_overlap,
)
else:
# Encode video - VAE expects [B, C, F, H, W], returns [B, C, F', H', W']
latents = vae(video)
if dtype is not None:
latents = latents.to(dtype=dtype)
@@ -652,6 +664,161 @@ def encode_video(
}
def tiled_encode_video( # noqa: PLR0912, PLR0915
vae: torch.nn.Module,
video: torch.Tensor,
tile_size: int = DEFAULT_TILE_SIZE,
tile_overlap: int = DEFAULT_TILE_OVERLAP,
) -> torch.Tensor:
"""Encode video using spatial tiling for memory efficiency.
Splits the video into overlapping spatial tiles, encodes each tile separately,
and blends the results using linear feathering in the overlap regions.
Args:
vae: Video VAE encoder model
video: Input tensor of shape [B, C, F, H, W]
tile_size: Tile size in pixels (must be divisible by 32)
tile_overlap: Overlap between tiles in pixels (must be divisible by 32)
Returns:
Encoded latent tensor [B, C_latent, F_latent, H_latent, W_latent]
"""
batch, _channels, frames, height, width = video.shape
device = video.device
dtype = video.dtype
# Validate tile parameters
if tile_size % VAE_SPATIAL_FACTOR != 0:
raise ValueError(f"tile_size must be divisible by {VAE_SPATIAL_FACTOR}, got {tile_size}")
if tile_overlap % VAE_SPATIAL_FACTOR != 0:
raise ValueError(f"tile_overlap must be divisible by {VAE_SPATIAL_FACTOR}, got {tile_overlap}")
if tile_overlap >= tile_size:
raise ValueError(f"tile_overlap ({tile_overlap}) must be less than tile_size ({tile_size})")
# If video fits in a single tile, use regular encoding
if height <= tile_size and width <= tile_size:
return vae(video)
# Calculate output dimensions
# VAE compresses: H -> H/32, W -> W/32, F -> 1 + (F-1)/8
output_height = height // VAE_SPATIAL_FACTOR
output_width = width // VAE_SPATIAL_FACTOR
output_frames = 1 + (frames - 1) // VAE_TEMPORAL_FACTOR
# Latent channels (128 for LTX-2)
# Get from a small test encode or assume 128
latent_channels = 128
# Initialize output and weight tensors
output = torch.zeros(
(batch, latent_channels, output_frames, output_height, output_width),
device=device,
dtype=dtype,
)
weights = torch.zeros(
(batch, 1, output_frames, output_height, output_width),
device=device,
dtype=dtype,
)
# Calculate tile positions with overlap
# Step size is tile_size - tile_overlap
step_h = tile_size - tile_overlap
step_w = tile_size - tile_overlap
h_positions = list(range(0, max(1, height - tile_overlap), step_h))
w_positions = list(range(0, max(1, width - tile_overlap), step_w))
# Ensure last tile covers the edge
if h_positions[-1] + tile_size < height:
h_positions.append(height - tile_size)
if w_positions[-1] + tile_size < width:
w_positions.append(width - tile_size)
# Remove duplicates and sort
h_positions = sorted(set(h_positions))
w_positions = sorted(set(w_positions))
# Overlap in latent space
overlap_out_h = tile_overlap // VAE_SPATIAL_FACTOR
overlap_out_w = tile_overlap // VAE_SPATIAL_FACTOR
# Process each tile
for h_pos in h_positions:
for w_pos in w_positions:
# Calculate tile boundaries in input space
h_start = max(0, h_pos)
w_start = max(0, w_pos)
h_end = min(h_start + tile_size, height)
w_end = min(w_start + tile_size, width)
# Ensure tile dimensions are divisible by VAE_SPATIAL_FACTOR
tile_h = ((h_end - h_start) // VAE_SPATIAL_FACTOR) * VAE_SPATIAL_FACTOR
tile_w = ((w_end - w_start) // VAE_SPATIAL_FACTOR) * VAE_SPATIAL_FACTOR
if tile_h < VAE_SPATIAL_FACTOR or tile_w < VAE_SPATIAL_FACTOR:
continue
# Adjust end positions
h_end = h_start + tile_h
w_end = w_start + tile_w
# Extract tile
tile = video[:, :, :, h_start:h_end, w_start:w_end]
# Encode tile
encoded_tile = vae(tile)
# Get actual encoded dimensions
_, _, tile_out_frames, tile_out_height, tile_out_width = encoded_tile.shape
# Calculate output positions
out_h_start = h_start // VAE_SPATIAL_FACTOR
out_w_start = w_start // VAE_SPATIAL_FACTOR
out_h_end = min(out_h_start + tile_out_height, output_height)
out_w_end = min(out_w_start + tile_out_width, output_width)
# Trim encoded tile if necessary
actual_tile_h = out_h_end - out_h_start
actual_tile_w = out_w_end - out_w_start
encoded_tile = encoded_tile[:, :, :, :actual_tile_h, :actual_tile_w]
# Create blending mask with linear feathering at edges
mask = torch.ones(
(1, 1, tile_out_frames, actual_tile_h, actual_tile_w),
device=device,
dtype=dtype,
)
# Apply feathering at edges (linear blend in overlap regions)
# Left edge
if h_pos > 0 and overlap_out_h > 0 and overlap_out_h < actual_tile_h:
fade_in = torch.linspace(0.0, 1.0, overlap_out_h + 2, device=device, dtype=dtype)[1:-1]
mask[:, :, :, :overlap_out_h, :] *= fade_in.view(1, 1, 1, -1, 1)
# Right edge (bottom in height dimension)
if h_end < height and overlap_out_h > 0 and overlap_out_h < actual_tile_h:
fade_out = torch.linspace(1.0, 0.0, overlap_out_h + 2, device=device, dtype=dtype)[1:-1]
mask[:, :, :, -overlap_out_h:, :] *= fade_out.view(1, 1, 1, -1, 1)
# Top edge (left in width dimension)
if w_pos > 0 and overlap_out_w > 0 and overlap_out_w < actual_tile_w:
fade_in = torch.linspace(0.0, 1.0, overlap_out_w + 2, device=device, dtype=dtype)[1:-1]
mask[:, :, :, :, :overlap_out_w] *= fade_in.view(1, 1, 1, 1, -1)
# Bottom edge (right in width dimension)
if w_end < width and overlap_out_w > 0 and overlap_out_w < actual_tile_w:
fade_out = torch.linspace(1.0, 0.0, overlap_out_w + 2, device=device, dtype=dtype)[1:-1]
mask[:, :, :, :, -overlap_out_w:] *= fade_out.view(1, 1, 1, 1, -1)
# Accumulate weighted results
output[:, :, :, out_h_start:out_h_end, out_w_start:out_w_end] += encoded_tile * mask
weights[:, :, :, out_h_start:out_h_end, out_w_start:out_w_end] += mask
# Normalize by weights (avoid division by zero)
output = output / (weights + 1e-8)
return output
def encode_audio(
audio_vae_encoder: torch.nn.Module,
audio_processor: torch.nn.Module,