| import os
|
| from typing import Literal, Optional
|
|
|
| import torch
|
| import torch.nn as nn
|
| from mmgp import offload
|
| from shared.utils import files_locator as fl
|
|
|
| from ..autoencoder.vae import VAE, get_my_vae
|
| from ..bigvgan import BigVGAN
|
| from ..bigvgan_v2.bigvgan import BigVGAN as BigVGANv2
|
| from ...model.utils.distributions import DiagonalGaussianDistribution
|
|
|
| _BIGVGAN_V2_FOLDER = "bigvgan_v2_44khz_128band_512x"
|
|
|
|
|
| def _resolve_bigvgan_v2_files():
|
| weights_path = fl.locate_file(
|
| os.path.join(_BIGVGAN_V2_FOLDER, "bigvgan_generator.pt"), error_if_none=False
|
| )
|
| config_path = fl.locate_file(
|
| os.path.join(_BIGVGAN_V2_FOLDER, "config.json"), error_if_none=False
|
| )
|
| if weights_path is None or config_path is None:
|
| raise FileNotFoundError(
|
| f"Missing BigVGANv2 files in '{_BIGVGAN_V2_FOLDER}'. "
|
| "Expected 'config.json' and 'bigvgan_generator.pt'."
|
| )
|
| return weights_path, config_path
|
|
|
|
|
| def _preprocess_bigvgan_v2_state_dict(state_dict, quantization_map=None, tied_weights_map=None):
|
| if isinstance(state_dict, dict) and isinstance(state_dict.get("generator"), dict):
|
| state_dict = state_dict["generator"]
|
| return state_dict, quantization_map, tied_weights_map
|
|
|
|
|
| class AutoEncoderModule(nn.Module):
|
|
|
| def __init__(self,
|
| *,
|
| vae_ckpt_path,
|
| vocoder_ckpt_path: Optional[str] = None,
|
| mode: Literal['16k', '44k'],
|
| need_vae_encoder: bool = True):
|
| super().__init__()
|
| self.vae: VAE = get_my_vae(mode).eval()
|
| vae_state_dict = torch.load(vae_ckpt_path, weights_only=True, map_location='cpu')
|
| self.vae.load_state_dict(vae_state_dict)
|
| self.vae.remove_weight_norm()
|
|
|
| if mode == '16k':
|
| assert vocoder_ckpt_path is not None
|
| self.vocoder = BigVGAN(vocoder_ckpt_path).eval()
|
| elif mode == '44k':
|
| vocoder_ckpt_path, vocoder_config_path = _resolve_bigvgan_v2_files()
|
| self.vocoder = offload.fast_load_transformers_model(
|
| vocoder_ckpt_path,
|
| modelClass=BigVGANv2,
|
| forcedConfigPath=vocoder_config_path,
|
| preprocess_sd=_preprocess_bigvgan_v2_state_dict,
|
| configKwargs={"use_cuda_kernel": False},
|
| writable_tensors=False,
|
| default_dtype=torch.float32,
|
| )
|
| self.vocoder.remove_weight_norm()
|
| self.vocoder.eval()
|
| else:
|
| raise ValueError(f'Unknown mode: {mode}')
|
|
|
| for param in self.parameters():
|
| param.requires_grad = False
|
|
|
| if not need_vae_encoder:
|
| del self.vae.encoder
|
|
|
| @torch.inference_mode()
|
| def encode(self, x: torch.Tensor) -> DiagonalGaussianDistribution:
|
| return self.vae.encode(x)
|
|
|
| @torch.inference_mode()
|
| def decode(self, z: torch.Tensor) -> torch.Tensor:
|
| return self.vae.decode(z)
|
|
|
| @torch.inference_mode()
|
| def vocode(self, spec: torch.Tensor) -> torch.Tensor:
|
| return self.vocoder(spec)
|
|
|
