| import torch |
| import torch.nn as nn |
| from dataclasses import asdict |
|
|
| from .utils.audio import LogMelSpectrogram |
| from .config import ModelConfig, MelConfig |
| from .models.model import StableTTS |
|
|
| from .text import symbols |
| from .text import cleaned_text_to_sequence |
| from .text.burmese import burmese_to_ipa2 |
| from .datas.dataset import intersperse |
| from .utils.audio import load_and_resample_audio |
|
|
|
|
| def get_vocoder(model_path, model_name='vocos'): |
| if model_name == 'vocos': |
| from .vocoders.vocos.models.model import Vocos |
| from .config import VocosConfig, MelConfig |
| vocoder = Vocos(VocosConfig(), MelConfig()) |
| vocoder.load_state_dict(torch.load(model_path, weights_only=True, map_location='cpu')) |
| vocoder.eval() |
| else: |
| raise NotImplementedError(f"Unsupported vocoder: {model_name}") |
| return vocoder |
|
|
|
|
| class StableTTSAPI(nn.Module): |
| def __init__(self, tts_model_path, vocoder_model_path, vocoder_name='vocos'): |
| super().__init__() |
| self.mel_config = MelConfig() |
| self.tts_model_config = ModelConfig() |
|
|
| self.mel_extractor = LogMelSpectrogram(**asdict(self.mel_config)) |
|
|
| self.tts_model = StableTTS(len(symbols), self.mel_config.n_mels, **asdict(self.tts_model_config)) |
|
|
| if tts_model_path.endswith(".safetensors"): |
| from safetensors.torch import load_file |
| state = load_file(tts_model_path) |
| else: |
| state = torch.load(tts_model_path, map_location='cpu', weights_only=True) |
| self.tts_model.load_state_dict(state) |
| self.tts_model.eval() |
|
|
| self.vocoder_model = get_vocoder(vocoder_model_path, vocoder_name) |
| self.vocoder_model.eval() |
|
|
| self.g2p_mapping = { |
| 'burmese': burmese_to_ipa2, |
| } |
| self.supported_languages = self.g2p_mapping.keys() |
|
|
| @torch.inference_mode() |
| def inference(self, text, ref_audio, language, step, temperature=1.0, |
| length_scale=1.0, solver=None, cfg=3.0): |
| device = next(self.parameters()).device |
| phonemizer = self.g2p_mapping.get(language) |
| if phonemizer is None: |
| raise ValueError(f"Unsupported language: {language}") |
|
|
| text = phonemizer(text) |
| text = torch.tensor( |
| intersperse(cleaned_text_to_sequence(text), item=0), |
| dtype=torch.long, device=device |
| ).unsqueeze(0) |
| text_length = torch.tensor([text.size(-1)], dtype=torch.long, device=device) |
|
|
| ref_audio = load_and_resample_audio(ref_audio, self.mel_config.sample_rate).to(device) |
| ref_audio = self.mel_extractor(ref_audio) |
|
|
| mel_output = self.tts_model.synthesise( |
| text, text_length, step, temperature, ref_audio, |
| length_scale, solver, cfg |
| )['decoder_outputs'] |
| audio_output = self.vocoder_model(mel_output) |
| return audio_output.cpu(), mel_output.cpu() |
|
|
| def get_params(self): |
| tts_param = sum(p.numel() for p in self.tts_model.parameters()) / 1e6 |
| vocoder_param = sum(p.numel() for p in self.vocoder_model.parameters()) / 1e6 |
| return tts_param, vocoder_param |
|
|
