mirror of
https://github.com/w-okada/voice-changer.git
synced 2025-01-24 05:55:01 +03:00
78 lines
2.9 KiB
Python
78 lines
2.9 KiB
Python
import os
|
|
import yaml # type: ignore
|
|
import torch
|
|
import torch.nn as nn
|
|
import numpy as np
|
|
from .diffusion import GaussianDiffusion
|
|
from .wavenet import WaveNet
|
|
from .vocoder import Vocoder
|
|
|
|
|
|
class DotDict(dict):
|
|
def __getattr__(*args): # type: ignore
|
|
val = dict.get(*args)
|
|
return DotDict(val) if type(val) is dict else val
|
|
|
|
__setattr__ = dict.__setitem__ # type: ignore
|
|
__delattr__ = dict.__delitem__ # type: ignore
|
|
|
|
|
|
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 vocoder
|
|
vocoder = Vocoder(args.vocoder.type, args.vocoder.ckpt, device=device)
|
|
|
|
# load model
|
|
model = Unit2Mel(args.data.encoder_out_channels, args.model.n_spk, args.model.use_pitch_aug, vocoder.dimension, args.model.n_layers, args.model.n_chans, args.model.n_hidden)
|
|
|
|
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, vocoder, 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):
|
|
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)
|
|
|
|
# diffusion
|
|
self.decoder = GaussianDiffusion(WaveNet(out_dims, n_layers, n_chans, n_hidden), out_dims=out_dims)
|
|
|
|
def forward(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:
|
|
for k, v in spk_mix_dict.items():
|
|
spk_id_torch = torch.LongTensor(np.array([[k]])).to(units.device)
|
|
x = x + v * self.spk_embed(spk_id_torch - 1)
|
|
else:
|
|
x = x + self.spk_embed(spk_id - 1)
|
|
if self.aug_shift_embed is not None and aug_shift is not None:
|
|
x = x + self.aug_shift_embed(aug_shift / 5)
|
|
x = self.decoder(x, gt_spec=gt_spec, infer=infer, infer_speedup=infer_speedup, method=method, k_step=k_step, use_tqdm=use_tqdm)
|
|
|
|
return x
|