import torch import torchaudio import numpy as np import re from hyperpyyaml import load_hyperpyyaml import uuid from collections import defaultdict def fade_in_out(fade_in_mel, fade_out_mel, window): device = fade_in_mel.device fade_in_mel, fade_out_mel = fade_in_mel.cpu(), fade_out_mel.cpu() mel_overlap_len = int(window.shape[0] / 2) fade_in_mel[..., :mel_overlap_len] = fade_in_mel[..., :mel_overlap_len] * window[:mel_overlap_len] + \ fade_out_mel[..., -mel_overlap_len:] * window[mel_overlap_len:] return fade_in_mel.to(device) class AudioDecoder: def __init__(self, config_path, flow_ckpt_path, hift_ckpt_path, device="cuda"): self.device = device with open(config_path, 'r') as f: self.scratch_configs = load_hyperpyyaml(f) # Load models self.flow = self.scratch_configs['flow'] self.flow.load_state_dict(torch.load(flow_ckpt_path, map_location=self.device)) self.hift = self.scratch_configs['hift'] self.hift.load_state_dict(torch.load(hift_ckpt_path, map_location=self.device)) # Move models to the appropriate device self.flow.to(self.device) self.hift.to(self.device) self.mel_overlap_dict = defaultdict(lambda: None) self.hift_cache_dict = defaultdict(lambda: None) self.token_min_hop_len = 2 * self.flow.input_frame_rate self.token_max_hop_len = 4 * self.flow.input_frame_rate self.token_overlap_len = 5 self.mel_overlap_len = int(self.token_overlap_len / self.flow.input_frame_rate * 22050 / 256) self.mel_window = np.hamming(2 * self.mel_overlap_len) # hift cache self.mel_cache_len = 1 self.source_cache_len = int(self.mel_cache_len * 256) # speech fade in out self.speech_window = np.hamming(2 * self.source_cache_len) def token2wav(self, token, uuid, prompt_token=torch.zeros(1, 0, dtype=torch.int32), prompt_feat=torch.zeros(1, 0, 80), embedding=torch.zeros(1, 192), finalize=False): tts_mel = self.flow.inference(token=token.to(self.device), token_len=torch.tensor([token.shape[1]], dtype=torch.int32).to(self.device), prompt_token=prompt_token.to(self.device), prompt_token_len=torch.tensor([prompt_token.shape[1]], dtype=torch.int32).to( self.device), prompt_feat=prompt_feat.to(self.device), prompt_feat_len=torch.tensor([prompt_feat.shape[1]], dtype=torch.int32).to( self.device), embedding=embedding.to(self.device)) # mel overlap fade in out if self.mel_overlap_dict[uuid] is not None: tts_mel = fade_in_out(tts_mel, self.mel_overlap_dict[uuid], self.mel_window) # append hift cache if self.hift_cache_dict[uuid] is not None: hift_cache_mel, hift_cache_source = self.hift_cache_dict[uuid]['mel'], self.hift_cache_dict[uuid]['source'] tts_mel = torch.concat([hift_cache_mel, tts_mel], dim=2) else: hift_cache_source = torch.zeros(1, 1, 0) # _tts_mel=tts_mel.contiguous() # keep overlap mel and hift cache if finalize is False: self.mel_overlap_dict[uuid] = tts_mel[:, :, -self.mel_overlap_len:] tts_mel = tts_mel[:, :, :-self.mel_overlap_len] tts_speech, tts_source = self.hift.inference(mel=tts_mel, cache_source=hift_cache_source) self.hift_cache_dict[uuid] = {'mel': tts_mel[:, :, -self.mel_cache_len:], 'source': tts_source[:, :, -self.source_cache_len:], 'speech': tts_speech[:, -self.source_cache_len:]} # if self.hift_cache_dict[uuid] is not None: # tts_speech = fade_in_out(tts_speech, self.hift_cache_dict[uuid]['speech'], self.speech_window) tts_speech = tts_speech[:, :-self.source_cache_len] else: tts_speech, tts_source = self.hift.inference(mel=tts_mel, cache_source=hift_cache_source) del self.hift_cache_dict[uuid] del self.mel_overlap_dict[uuid] # if uuid in self.hift_cache_dict.keys() and self.hift_cache_dict[uuid] is not None: # tts_speech = fade_in_out(tts_speech, self.hift_cache_dict[uuid]['speech'], self.speech_window) return tts_speech, tts_mel def offline_inference(self, token): this_uuid = str(uuid.uuid1()) tts_speech, tts_mel = self.token2wav(token, uuid=this_uuid, finalize=True) return tts_speech.cpu() def stream_inference(self, token): token.to(self.device) this_uuid = str(uuid.uuid1()) # Prepare other necessary input tensors llm_embedding = torch.zeros(1, 192).to(self.device) prompt_speech_feat = torch.zeros(1, 0, 80).to(self.device) flow_prompt_speech_token = torch.zeros(1, 0, dtype=torch.int32).to(self.device) tts_speechs = [] tts_mels = [] block_size = self.flow.encoder.block_size prev_mel = None for idx in range(0, token.size(1), block_size): # if idx>block_size: break tts_token = token[:, idx:idx + block_size] print(tts_token.size()) if prev_mel is not None: prompt_speech_feat = torch.cat(tts_mels, dim=-1).transpose(1, 2) flow_prompt_speech_token = token[:, :idx] if idx + block_size >= token.size(-1): is_finalize = True else: is_finalize = False tts_speech, tts_mel = self.token2wav(tts_token, uuid=this_uuid, prompt_token=flow_prompt_speech_token.to(self.device), prompt_feat=prompt_speech_feat.to(self.device), finalize=is_finalize) prev_mel = tts_mel prev_speech = tts_speech print(tts_mel.size()) tts_speechs.append(tts_speech) tts_mels.append(tts_mel) # Convert Mel spectrogram to audio using HiFi-GAN tts_speech = torch.cat(tts_speechs, dim=-1).cpu() return tts_speech.cpu()