| import os | |
| import torch | |
| import torch.multiprocessing | |
| from speechbrain.pretrained import EncoderClassifier | |
| from torchaudio.transforms import Resample | |
| from Modules.EmbeddingModel.StyleEmbedding import StyleEmbedding | |
| from Preprocessing.HiFiCodecAudioPreprocessor import CodecAudioPreprocessor | |
| from Utility.storage_config import MODELS_DIR | |
| class ProsodicConditionExtractor: | |
| def __init__(self, device=torch.device("cpu"), path_to_model=os.path.join(MODELS_DIR, "Embedding", "embedding_function.pt")): | |
| self.ap = CodecAudioPreprocessor(input_sr=100, output_sr=2) | |
| self.embed = StyleEmbedding() | |
| check_dict = torch.load(path_to_model, map_location="cpu") | |
| self.embed.load_state_dict(check_dict["style_emb_func"]) | |
| self.speaker_embedding_func_ecapa = EncoderClassifier.from_hparams(source="speechbrain/spkrec-ecapa-voxceleb", | |
| run_opts={"device": str(device)}, | |
| savedir=os.path.join(MODELS_DIR, "Embedding", "speechbrain_speaker_embedding_ecapa")) | |
| self.embed.to(device) | |
| self.device = device | |
| def extract_condition_from_reference_wave(self, wave, sr): | |
| wave_24khz = Resample(orig_freq=sr, new_freq=24000).to(self.device)(torch.tensor(wave, device=self.device, dtype=torch.float32)) | |
| spec = self.ap.audio_to_codec_tensor(wave_24khz, current_sampling_rate=24000).transpose(0, 1) | |
| spec_len = torch.LongTensor([len(spec)]) | |
| style_embedding = self.embed(spec.unsqueeze(0).to(self.device), spec_len.unsqueeze(0).to(self.device)).squeeze() | |
| wave_16kHz = Resample(orig_freq=sr, new_freq=16000).to(self.device)(torch.tensor(wave, device=self.device, dtype=torch.float32)) | |
| speaker_embedding = self.speaker_embedding_func_ecapa.encode_batch(wavs=wave_16kHz.to(self.device).unsqueeze(0)).squeeze() | |
| return torch.cat([style_embedding, speaker_embedding], dim=-1) | |