from suno_utils.tasks.dac_vae_100hz_peaq import (
    encode_overlap as encode_audio_to_vae_v1,
    decode_stream_to_full_audio as decode_vae_to_audio_v1,
    preload_models as preload_vae_models_v1,
)
from suno_utils.tasks.dac_vae_fixed_25hz import (
    encode_overlap as encode_audio_to_vae_v2,
    decode_stream_to_full_audio as decode_vae_to_audio_v2,
    preload_models as preload_vae_models_v2,
)
from suno_utils.worker.modal_model_configs import VAEVersion

vae_version_to_preload_fn = {
    VAEVersion.V_VAE_25_PEAQ_1: preload_vae_models_v1,
    VAEVersion.V_VAE_25_TUNED_2: preload_vae_models_v2,
    # Also support string keys for backward compatibility
    VAEVersion.V_VAE_25_PEAQ_1.value: preload_vae_models_v1,
    VAEVersion.V_VAE_25_TUNED_2.value: preload_vae_models_v2,
}

vae_version_to_encode_fn = {
    VAEVersion.V_VAE_25_PEAQ_1: encode_audio_to_vae_v1,
    VAEVersion.V_VAE_25_TUNED_2: encode_audio_to_vae_v2,
    # Also support string keys for backward compatibility
    VAEVersion.V_VAE_25_PEAQ_1.value: encode_audio_to_vae_v1,
    VAEVersion.V_VAE_25_TUNED_2.value: encode_audio_to_vae_v2,
}

vae_version_to_decode_fn = {
    VAEVersion.V_VAE_25_PEAQ_1: decode_vae_to_audio_v1,
    VAEVersion.V_VAE_25_TUNED_2: decode_vae_to_audio_v2,
    # Also support string keys for backward compatibility
    VAEVersion.V_VAE_25_PEAQ_1.value: decode_vae_to_audio_v1,
    VAEVersion.V_VAE_25_TUNED_2.value: decode_vae_to_audio_v2,
}