from suno_utils.tasks.dac_vae_fixed_25hz import preload_models as preload_vae_models_
from suno_utils.tasks.mert_25 import preload_models as preload_semantic_models_
from suno_utils.tasks.musicfm_v3 import preload_models as preload_musicfm_models_

_vae_model_loaded = False
_mert_model_loaded = False
_musicfm_model_loaded = False


def load_vae_model():
    global _vae_model_loaded
    if not _vae_model_loaded:
        preload_vae_models_(checkpoint_filepath="s3://suno-data/minz/models/dac_vae_tuned_25hz.pth")
        _vae_model_loaded = True


def preload_mert_models():
    global _mert_model_loaded
    if not _mert_model_loaded:
        preload_semantic_models_(
            checkpoint_filepath="s3://suno-data/georg/models/semantic/mert_25.pt",
            centroids_filepath="s3://suno-data/georg/models/semantic/mert_25_2x4k.npy",
        )
        _mert_model_loaded = True


def preload_musicfm_models():
    global _musicfm_model_loaded
    if not _musicfm_model_loaded:
        preload_musicfm_models_(
            checkpoint_filepath="/app2/suno/checkpoints/2025-08-18_20-07-29/last_ckpt_infer.pt",
            centroids_filepath="/home/minz/temp/musicfm3_5s_1024d_centroids_4000_50.npy",
            token_type="emb_pre",
        )
        _musicfm_model_loaded = True