Spaces:
Running
Running
| import os | |
| import numpy as np | |
| import torch | |
| import torch.nn as nn | |
| import torch.nn.functional as F | |
| import yaml | |
| from diffusion_onnx import GaussianDiffusion | |
| class DotDict(dict): | |
| def __getattr__(*args): | |
| val = dict.get(*args) | |
| return DotDict(val) if type(val) is dict else val | |
| __setattr__ = dict.__setitem__ | |
| __delattr__ = dict.__delitem__ | |
| def load_model_vocoder( | |
| model_path, | |
| device='cpu'): | |
| config_file = os.path.join(os.path.split(model_path)[0], 'config.yaml') | |
| with open(config_file, "r") as config: | |
| args = yaml.safe_load(config) | |
| args = DotDict(args) | |
| # load model | |
| model = Unit2Mel( | |
| args.data.encoder_out_channels, | |
| args.model.n_spk, | |
| args.model.use_pitch_aug, | |
| 128, | |
| args.model.n_layers, | |
| args.model.n_chans, | |
| args.model.n_hidden, | |
| args.model.timesteps, | |
| args.model.k_step_max) | |
| print(' [Loading] ' + model_path) | |
| ckpt = torch.load(model_path, map_location=torch.device(device)) | |
| model.to(device) | |
| model.load_state_dict(ckpt['model']) | |
| model.eval() | |
| return model, args | |
| class Unit2Mel(nn.Module): | |
| def __init__( | |
| self, | |
| input_channel, | |
| n_spk, | |
| use_pitch_aug=False, | |
| out_dims=128, | |
| n_layers=20, | |
| n_chans=384, | |
| n_hidden=256, | |
| timesteps=1000, | |
| k_step_max=1000): | |
| super().__init__() | |
| self.unit_embed = nn.Linear(input_channel, n_hidden) | |
| self.f0_embed = nn.Linear(1, n_hidden) | |
| self.volume_embed = nn.Linear(1, n_hidden) | |
| if use_pitch_aug: | |
| self.aug_shift_embed = nn.Linear(1, n_hidden, bias=False) | |
| else: | |
| self.aug_shift_embed = None | |
| self.n_spk = n_spk | |
| if n_spk is not None and n_spk > 1: | |
| self.spk_embed = nn.Embedding(n_spk, n_hidden) | |
| self.timesteps = timesteps if timesteps is not None else 1000 | |
| self.k_step_max = k_step_max if k_step_max is not None and k_step_max>0 and k_step_max<self.timesteps else self.timesteps | |
| # diffusion | |
| self.decoder = GaussianDiffusion(out_dims, n_layers, n_chans, n_hidden,self.timesteps,self.k_step_max) | |
| self.hidden_size = n_hidden | |
| self.speaker_map = torch.zeros((self.n_spk,1,1,n_hidden)) | |
| def forward(self, units, mel2ph, f0, volume, g = None): | |
| ''' | |
| input: | |
| B x n_frames x n_unit | |
| return: | |
| dict of B x n_frames x feat | |
| ''' | |
| decoder_inp = F.pad(units, [0, 0, 1, 0]) | |
| mel2ph_ = mel2ph.unsqueeze(2).repeat([1, 1, units.shape[-1]]) | |
| units = torch.gather(decoder_inp, 1, mel2ph_) # [B, T, H] | |
| x = self.unit_embed(units) + self.f0_embed((1 + f0.unsqueeze(-1) / 700).log()) + self.volume_embed(volume.unsqueeze(-1)) | |
| if self.n_spk is not None and self.n_spk > 1: # [N, S] * [S, B, 1, H] | |
| g = g.reshape((g.shape[0], g.shape[1], 1, 1, 1)) # [N, S, B, 1, 1] | |
| g = g * self.speaker_map # [N, S, B, 1, H] | |
| g = torch.sum(g, dim=1) # [N, 1, B, 1, H] | |
| g = g.transpose(0, -1).transpose(0, -2).squeeze(0) # [B, H, N] | |
| x = x.transpose(1, 2) + g | |
| return x | |
| else: | |
| return x.transpose(1, 2) | |
| def init_spkembed(self, units, f0, volume, spk_id = None, spk_mix_dict = None, aug_shift = None, | |
| gt_spec=None, infer=True, infer_speedup=10, method='dpm-solver', k_step=300, use_tqdm=True): | |
| ''' | |
| input: | |
| B x n_frames x n_unit | |
| return: | |
| dict of B x n_frames x feat | |
| ''' | |
| x = self.unit_embed(units) + self.f0_embed((1+ f0 / 700).log()) + self.volume_embed(volume) | |
| if self.n_spk is not None and self.n_spk > 1: | |
| if spk_mix_dict is not None: | |
| spk_embed_mix = torch.zeros((1,1,self.hidden_size)) | |
| for k, v in spk_mix_dict.items(): | |
| spk_id_torch = torch.LongTensor(np.array([[k]])).to(units.device) | |
| spk_embeddd = self.spk_embed(spk_id_torch) | |
| self.speaker_map[k] = spk_embeddd | |
| spk_embed_mix = spk_embed_mix + v * spk_embeddd | |
| x = x + spk_embed_mix | |
| else: | |
| x = x + self.spk_embed(spk_id - 1) | |
| self.speaker_map = self.speaker_map.unsqueeze(0) | |
| self.speaker_map = self.speaker_map.detach() | |
| return x.transpose(1, 2) | |
| def OnnxExport(self, project_name=None, init_noise=None, export_encoder=True, export_denoise=True, export_pred=True, export_after=True): | |
| hubert_hidden_size = 768 | |
| n_frames = 100 | |
| hubert = torch.randn((1, n_frames, hubert_hidden_size)) | |
| mel2ph = torch.arange(end=n_frames).unsqueeze(0).long() | |
| f0 = torch.randn((1, n_frames)) | |
| volume = torch.randn((1, n_frames)) | |
| spk_mix = [] | |
| spks = {} | |
| if self.n_spk is not None and self.n_spk > 1: | |
| for i in range(self.n_spk): | |
| spk_mix.append(1.0/float(self.n_spk)) | |
| spks.update({i:1.0/float(self.n_spk)}) | |
| spk_mix = torch.tensor(spk_mix) | |
| spk_mix = spk_mix.repeat(n_frames, 1) | |
| self.init_spkembed(hubert, f0.unsqueeze(-1), volume.unsqueeze(-1), spk_mix_dict=spks) | |
| self.forward(hubert, mel2ph, f0, volume, spk_mix) | |
| if export_encoder: | |
| torch.onnx.export( | |
| self, | |
| (hubert, mel2ph, f0, volume, spk_mix), | |
| f"{project_name}_encoder.onnx", | |
| input_names=["hubert", "mel2ph", "f0", "volume", "spk_mix"], | |
| output_names=["mel_pred"], | |
| dynamic_axes={ | |
| "hubert": [1], | |
| "f0": [1], | |
| "volume": [1], | |
| "mel2ph": [1], | |
| "spk_mix": [0], | |
| }, | |
| opset_version=16 | |
| ) | |
| self.decoder.OnnxExport(project_name, init_noise=init_noise, export_denoise=export_denoise, export_pred=export_pred, export_after=export_after) | |
| def ExportOnnx(self, project_name=None): | |
| hubert_hidden_size = 768 | |
| n_frames = 100 | |
| hubert = torch.randn((1, n_frames, hubert_hidden_size)) | |
| mel2ph = torch.arange(end=n_frames).unsqueeze(0).long() | |
| f0 = torch.randn((1, n_frames)) | |
| volume = torch.randn((1, n_frames)) | |
| spk_mix = [] | |
| spks = {} | |
| if self.n_spk is not None and self.n_spk > 1: | |
| for i in range(self.n_spk): | |
| spk_mix.append(1.0/float(self.n_spk)) | |
| spks.update({i:1.0/float(self.n_spk)}) | |
| spk_mix = torch.tensor(spk_mix) | |
| self.orgforward(hubert, f0.unsqueeze(-1), volume.unsqueeze(-1), spk_mix_dict=spks) | |
| self.forward(hubert, mel2ph, f0, volume, spk_mix) | |
| torch.onnx.export( | |
| self, | |
| (hubert, mel2ph, f0, volume, spk_mix), | |
| f"{project_name}_encoder.onnx", | |
| input_names=["hubert", "mel2ph", "f0", "volume", "spk_mix"], | |
| output_names=["mel_pred"], | |
| dynamic_axes={ | |
| "hubert": [1], | |
| "f0": [1], | |
| "volume": [1], | |
| "mel2ph": [1] | |
| }, | |
| opset_version=16 | |
| ) | |
| condition = torch.randn(1,self.decoder.n_hidden,n_frames) | |
| noise = torch.randn((1, 1, self.decoder.mel_bins, condition.shape[2]), dtype=torch.float32) | |
| pndm_speedup = torch.LongTensor([100]) | |
| K_steps = torch.LongTensor([1000]) | |
| self.decoder = torch.jit.script(self.decoder) | |
| self.decoder(condition, noise, pndm_speedup, K_steps) | |
| torch.onnx.export( | |
| self.decoder, | |
| (condition, noise, pndm_speedup, K_steps), | |
| f"{project_name}_diffusion.onnx", | |
| input_names=["condition", "noise", "pndm_speedup", "K_steps"], | |
| output_names=["mel"], | |
| dynamic_axes={ | |
| "condition": [2], | |
| "noise": [3], | |
| }, | |
| opset_version=16 | |
| ) | |
| if __name__ == "__main__": | |
| project_name = "dddsp" | |
| model_path = f'{project_name}/model_500000.pt' | |
| model, _ = load_model_vocoder(model_path) | |
| # 分开Diffusion导出（需要使用MoeSS/MoeVoiceStudio或者自己编写Pndm/Dpm采样） | |
| model.OnnxExport(project_name, export_encoder=True, export_denoise=True, export_pred=True, export_after=True) | |
| # 合并Diffusion导出（Encoder和Diffusion分开，直接将Encoder的结果和初始噪声输入Diffusion即可） | |
| # model.ExportOnnx(project_name) | |