Refactor weight loading and sanitization processes for audio models

mlx_video/__init__.py

@@ -1,11 +1,59 @@
 from mlx_video.models.ltx import LTXModel, LTXModelConfig
-from mlx_video.convert import load_transformer_weights, load_vae_weights
-import os
+from mlx_video.convert import (
+    load_transformer_weights,
+    load_vae_weights,
+    load_audio_vae_weights,
+    load_vocoder_weights,
+    sanitize_audio_vae_weights,
+    sanitize_vocoder_weights,
+)
+
+# Audio VAE components
+from mlx_video.models.ltx.audio_vae import (
+    AudioEncoder,
+    AudioDecoder,
+    Vocoder,
+    AudioProcessor,
+    decode_audio,
+)
+
+# Patchifiers
+from mlx_video.components.patchifiers import (
+    VideoLatentPatchifier,
+    AudioPatchifier,
+    VideoLatentShape,
+    AudioLatentShape,
+)
+
+# Conditioning
+from mlx_video.conditioning import (
+    VideoConditionByKeyframeIndex,
+    VideoConditionByLatentIndex,
+)
+
 __all__ = [
+    # Models
     "LTXModel",
     "LTXModelConfig",
+    # Weight loading
     "load_transformer_weights",
     "load_vae_weights",
-]
-
-os.environ["TRANSFORMERS_NO_ADVISORY_WARNINGS"] = "1"
+    "load_audio_vae_weights",
+    "load_vocoder_weights",
+    "sanitize_audio_vae_weights",
+    "sanitize_vocoder_weights",
+    # Audio VAE
+    "AudioEncoder",
+    "AudioDecoder",
+    "Vocoder",
+    "AudioProcessor",
+    "decode_audio",
+    # Patchifiers
+    "VideoLatentPatchifier",
+    "AudioPatchifier",
+    "VideoLatentShape",
+    "AudioLatentShape",
+    # Conditioning
+    "VideoConditionByKeyframeIndex",
+    "VideoConditionByLatentIndex",
+]