"""Frozen VoxCPM2 AudioVAE loader.

FreyaTTS generates 64-dim latents at 25 Hz and relies on the pretrained
VoxCPM2 AudioVAE (openbmb/VoxCPM2) to decode them to 48 kHz waveforms.
Requires the `voxcpm` pip package.
"""

import os

import torch
from huggingface_hub import hf_hub_download
from voxcpm.modules.audiovae import AudioVAEV2, AudioVAEConfigV2


def load_audio_vae(device="cuda", token=None):
    """Download and return the frozen VoxCPM2 AudioVAE in eval mode on `device`."""
    path = hf_hub_download("openbmb/VoxCPM2", "audiovae.pth", token=token or os.environ.get("HF_TOKEN"))
    vae = AudioVAEV2(AudioVAEConfigV2())

    ckpt = torch.load(path, map_location="cpu", weights_only=True)
    vae.load_state_dict(ckpt.get("state_dict", ckpt), strict=False)

    vae = vae.to(device).float().eval()
    for p in vae.parameters():
        p.requires_grad = False
    return vae