def parse_model_args(model_name, parser, add_help=False):
if model_name == 'Tacotron2':
return parse_tacotron2_args(parser, add_help)
if model_name == 'WaveGlow':
return parse_waveglow_args(parser, add_help)
else:
raise NotImplementedError(model_name)
def parse_model_args(model_name, parser, add_help=False):
if model_name == 'WaveGlow':
from waveglow.arg_parser import parse_waveglow_args
return parse_waveglow_args(parser, add_help)
elif model_name == 'FastPitch':
from fastpitch.arg_parser import parse_fastpitch_args
return parse_fastpitch_args(parser, add_help)
else:
raise NotImplementedError(model_name)
def parse_model_args(model_name, parser, add_help=False):
if model_name == 'Tacotron2' or model_name == 'Tacotron2mmi':
from tacotron2.arg_parser import parse_tacotron2_args
print('+++++++++++++++++++++++++')
return parse_tacotron2_args(parser, add_help)
if model_name == 'WaveGlow':
from waveglow.arg_parser import parse_waveglow_args
return parse_waveglow_args(parser, add_help)
elif model_name == 'FastPitch':
from fastpitch.arg_parser import parse_fastpitch_args
return parse_fastpitch_args(parser, add_help)
else:
raise NotImplementedError(model_name)
def __init__(self, ckpt_file, device='cuda', use_fp16=False, use_denoiser=False):
self.ckpt_file = ckpt_file
self.device = device
self.use_fp16 = use_fp16
self.use_denoiser = use_denoiser
# model
# sys.path.append('waveglow')
from waveglow.arg_parser import parse_waveglow_args
parser = parser = argparse.ArgumentParser()
model_parser= parse_waveglow_args(parser)
args, _ = model_parser.parse_known_args()
model_config = dict(
n_mel_channels=args.n_mel_channels,
n_flows=args.flows,
n_group=args.groups,
n_early_every=args.early_every,
n_early_size=args.early_size,
WN_config=dict(
n_layers=args.wn_layers,
kernel_size=args.wn_kernel_size,
n_channels=args.wn_channels
)
)
self.model = WaveGlow(**model_config)
state_dict = torch.load(self.ckpt_file, map_location=self.device)['state_dict']
state_dict = unwrap_distributed(state_dict)
self.model.load_state_dict(state_dict)
self.model = to_device_async(self.model, self.device)
self.model = self.model.remove_weightnorm(self.model)
self.model.eval()
if self.use_fp16:
self.model = self.model.half()
self.model = self.model
if self.use_denoiser:
self.denoiser = Denoiser(self.model, device=device)
self.denoiser = to_device_async(self.denoiser, self.device)
tprint('Using WaveGlow denoiser.')