Refactor weight loading and sanitization processes for audio models

2026-01-23 17:31:25 +01:00
parent 2681f75d2f
commit 02bfa228d9
18 changed files with 510 additions and 498 deletions
--- a/mlx_video/models/ltx/audio_vae/downsample.py
+++ b/mlx_video/models/ltx/audio_vae/downsample.py
@@ -6,7 +6,7 @@ import mlx.core as mx
 import mlx.nn as nn

 from .attention import AttentionType, make_attn
-from .causality_axis import CausalityAxis
+from ..config import CausalityAxis
 from .normalization import NormType
 from .resnet import ResnetBlock