File size: 3,372 Bytes
f2cd0c4 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 |
import torch
from vdecoder.nsf_hifigan.nvSTFT import STFT
from vdecoder.nsf_hifigan.models import load_model,load_config
from torchaudio.transforms import Resample
class Vocoder:
def __init__(self, vocoder_type, vocoder_ckpt, device = None):
if device is None:
device = 'cuda' if torch.cuda.is_available() else 'cpu'
self.device = device
if vocoder_type == 'nsf-hifigan':
self.vocoder = NsfHifiGAN(vocoder_ckpt, device = device)
elif vocoder_type == 'nsf-hifigan-log10':
self.vocoder = NsfHifiGANLog10(vocoder_ckpt, device = device)
else:
raise ValueError(f" [x] Unknown vocoder: {vocoder_type}")
self.resample_kernel = {}
self.vocoder_sample_rate = self.vocoder.sample_rate()
self.vocoder_hop_size = self.vocoder.hop_size()
self.dimension = self.vocoder.dimension()
def extract(self, audio, sample_rate, keyshift=0):
# resample
if sample_rate == self.vocoder_sample_rate:
audio_res = audio
else:
key_str = str(sample_rate)
if key_str not in self.resample_kernel:
self.resample_kernel[key_str] = Resample(sample_rate, self.vocoder_sample_rate, lowpass_filter_width = 128).to(self.device)
audio_res = self.resample_kernel[key_str](audio)
# extract
mel = self.vocoder.extract(audio_res, keyshift=keyshift) # B, n_frames, bins
return mel
def infer(self, mel, f0):
f0 = f0[:,:mel.size(1),0] # B, n_frames
audio = self.vocoder(mel, f0)
return audio
class NsfHifiGAN(torch.nn.Module):
def __init__(self, model_path, device=None):
super().__init__()
if device is None:
device = 'cuda' if torch.cuda.is_available() else 'cpu'
self.device = device
self.model_path = model_path
self.model = None
self.h = load_config(model_path)
self.stft = STFT(
self.h.sampling_rate,
self.h.num_mels,
self.h.n_fft,
self.h.win_size,
self.h.hop_size,
self.h.fmin,
self.h.fmax)
def sample_rate(self):
return self.h.sampling_rate
def hop_size(self):
return self.h.hop_size
def dimension(self):
return self.h.num_mels
def extract(self, audio, keyshift=0):
mel = self.stft.get_mel(audio, keyshift=keyshift).transpose(1, 2) # B, n_frames, bins
return mel
def forward(self, mel, f0):
if self.model is None:
print('| Load HifiGAN: ', self.model_path)
self.model, self.h = load_model(self.model_path, device=self.device)
with torch.no_grad():
c = mel.transpose(1, 2)
audio = self.model(c, f0)
return audio
class NsfHifiGANLog10(NsfHifiGAN):
def forward(self, mel, f0):
if self.model is None:
print('| Load HifiGAN: ', self.model_path)
self.model, self.h = load_model(self.model_path, device=self.device)
with torch.no_grad():
c = 0.434294 * mel.transpose(1, 2)
audio = self.model(c, f0)
return audio |