Refactor weight loading and utility functions for LTX-2 model; remove deprecated weight loading file and update imports accordingly
This commit is contained in:
@@ -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",
|
||||
]
|
||||
|
||||
Reference in New Issue
Block a user