n0p/mlx-video
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Update audio decoder configuration to use an empty set for attention resolutions in both generate_av.py and generate_dev.py. Add a print statement for loading audio VAE decoder weights in generate_dev.py.

Browse Source
This commit is contained in:
Prince Canuma
2026-01-18 21:55:38 +01:00
parent 7069cc39c9
commit 749762a0b9
2 changed files with 3 additions and 2 deletions
Download Patch File Download Diff File Expand all files Collapse all files

2
mlx_video/generate_av.py
View File

@@ -254,7 +254,7 @@ def load_audio_decoder(model_path: Path):
out_ch=2, # stereo out_ch=2, # stereo
ch_mult=(1, 2, 4), ch_mult=(1, 2, 4),
num_res_blocks=2, num_res_blocks=2,
attn_resolutions={8, 16, 32}, attn_resolutions=set(), # PyTorch uses empty set (no attention in audio decoder)
resolution=256, resolution=256,
z_channels=AUDIO_LATENT_CHANNELS, z_channels=AUDIO_LATENT_CHANNELS,
norm_type=NormType.PIXEL, norm_type=NormType.PIXEL,

3
mlx_video/generate_dev.py
View File

@@ -275,7 +275,7 @@ def load_audio_decoder(model_path: Path):
out_ch=2, # stereo out_ch=2, # stereo
ch_mult=(1, 2, 4), ch_mult=(1, 2, 4),
num_res_blocks=2, num_res_blocks=2,
attn_resolutions={8, 16, 32}, attn_resolutions=set(), # PyTorch uses empty set (no attention in audio decoder)
resolution=256, resolution=256,
z_channels=AUDIO_LATENT_CHANNELS, z_channels=AUDIO_LATENT_CHANNELS,
norm_type=NormType.PIXEL, norm_type=NormType.PIXEL,
@@ -289,6 +289,7 @@ def load_audio_decoder(model_path: Path):
weight_file = model_path / "ltx-2-19b-distilled.safetensors" weight_file = model_path / "ltx-2-19b-distilled.safetensors"
if weight_file.exists(): if weight_file.exists():
print(f"Loading audio VAE decoder from {weight_file}...")
raw_weights = mx.load(str(weight_file)) raw_weights = mx.load(str(weight_file))
sanitized = sanitize_audio_vae_weights(raw_weights) sanitized = sanitize_audio_vae_weights(raw_weights)
if sanitized: if sanitized: