mirror of
https://github.com/storytold/storyteller-ml.git
synced 2026-10-17 00:23:14 +00:00
Brandon Thomas
96 lines
2.6 KiB
Python
96 lines
2.6 KiB
Python
from __future__ import absolute_import, division, print_function, unicode_literals
|
|
|
|
import glob
|
|
import os
|
|
import argparse
|
|
import json
|
|
import torch
|
|
from scipy.io.wavfile import write
|
|
from env import AttrDict
|
|
from meldataset import mel_spectrogram, MAX_WAV_VALUE, load_wav
|
|
from models import Generator
|
|
|
|
h = None
|
|
device = None
|
|
|
|
|
|
def load_checkpoint(filepath, device):
|
|
assert os.path.isfile(filepath)
|
|
print("Loading '{}'".format(filepath))
|
|
checkpoint_dict = torch.load(filepath, map_location=device)
|
|
print("Complete.")
|
|
return checkpoint_dict
|
|
|
|
|
|
def get_mel(x):
|
|
return mel_spectrogram(x, h.n_fft, h.num_mels, h.sampling_rate, h.hop_size, h.win_size, h.fmin, h.fmax)
|
|
|
|
|
|
def scan_checkpoint(cp_dir, prefix):
|
|
pattern = os.path.join(cp_dir, prefix + '*')
|
|
cp_list = glob.glob(pattern)
|
|
if len(cp_list) == 0:
|
|
return ''
|
|
return sorted(cp_list)[-1]
|
|
|
|
|
|
def inference(a):
|
|
generator = Generator(h).to(device)
|
|
|
|
state_dict_g = load_checkpoint(a.checkpoint_file, device)
|
|
generator.load_state_dict(state_dict_g['generator'])
|
|
|
|
filelist = os.listdir(a.input_wavs_dir)
|
|
|
|
os.makedirs(a.output_dir, exist_ok=True)
|
|
|
|
generator.eval()
|
|
generator.remove_weight_norm()
|
|
with torch.no_grad():
|
|
for i, filname in enumerate(filelist):
|
|
wav, sr = load_wav(os.path.join(a.input_wavs_dir, filname))
|
|
wav = wav / MAX_WAV_VALUE
|
|
wav = torch.FloatTensor(wav).to(device)
|
|
x = get_mel(wav.unsqueeze(0))
|
|
y_g_hat = generator(x)
|
|
audio = y_g_hat.squeeze()
|
|
audio = audio * MAX_WAV_VALUE
|
|
audio = audio.cpu().numpy().astype('int16')
|
|
|
|
output_file = os.path.join(a.output_dir, os.path.splitext(filname)[0] + '_generated.wav')
|
|
write(output_file, h.sampling_rate, audio)
|
|
print(output_file)
|
|
|
|
|
|
def main():
|
|
print('Initializing Inference Process..')
|
|
|
|
parser = argparse.ArgumentParser()
|
|
parser.add_argument('--input_wavs_dir', default='test_files')
|
|
parser.add_argument('--output_dir', default='generated_files')
|
|
parser.add_argument('--checkpoint_file', required=True)
|
|
a = parser.parse_args()
|
|
|
|
config_file = os.path.join(os.path.split(a.checkpoint_file)[0], 'config.json')
|
|
with open(config_file) as f:
|
|
data = f.read()
|
|
|
|
global h
|
|
json_config = json.loads(data)
|
|
h = AttrDict(json_config)
|
|
|
|
torch.manual_seed(h.seed)
|
|
global device
|
|
if torch.cuda.is_available():
|
|
torch.cuda.manual_seed(h.seed)
|
|
device = torch.device('cuda')
|
|
else:
|
|
device = torch.device('cpu')
|
|
|
|
inference(a)
|
|
|
|
|
|
if __name__ == '__main__':
|
|
main()
|
|
|