Refactor weight loading and utility functions for LTX-2 model; remove deprecated weight loading file and update imports accordingly

This commit is contained in:
Prince Canuma
2026-03-16 22:25:22 +01:00
parent dd573d53d2
commit 7a576bfbf4
5 changed files with 182 additions and 789 deletions

View File

@@ -1,16 +1,9 @@
from mlx_video.models.ltx_2 import LTXModel, LTXModelConfig
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_2.audio_vae import (
AudioDecoder,
AudioEncoder,
Vocoder,
decode_audio,
AudioPatchifier,
@@ -23,19 +16,22 @@ from mlx_video.models.ltx_2.conditioning import (
VideoConditionByLatentIndex,
)
# Utilities
from mlx_video.models.ltx_2.utils import (
convert_audio_encoder,
get_model_path,
load_safetensors,
load_config,
save_weights,
)
__all__ = [
# Models
"LTXModel",
"LTXModelConfig",
# Weight loading
"load_transformer_weights",
"load_vae_weights",
"load_audio_vae_weights",
"load_vocoder_weights",
"sanitize_audio_vae_weights",
"sanitize_vocoder_weights",
# Audio VAE
"AudioDecoder",
"AudioEncoder",
"Vocoder",
"decode_audio",
"AudioPatchifier",
@@ -43,4 +39,10 @@ __all__ = [
"PerChannelStatistics",
# Conditioning
"VideoConditionByLatentIndex",
]
# Utilities
"convert_audio_encoder",
"get_model_path",
"load_safetensors",
"load_config",
"save_weights",
]