def load_waveglow(filename, waveglow_config):
class RenamingUnpickler(pickle.Unpickler):
def find_class(self, module, name):
if module == 'glow':
module = 'waveglow.model'
return super().find_class(module, name)
class RenamingPickleModule:
def load(self, f, *args, **kw_args):
return self.Unpickler(f, *args, **kw_args).load()
def Unpickler(self, f, **pickle_load_args):
return RenamingUnpickler(f, **pickle_load_args)
pickle_module = RenamingPickleModule()
blob = torch.load(filename, pickle_module=pickle_module)
if 'state_dict' in blob:
waveglow = WaveGlow(**waveglow_config).cuda()
state_dict = {}
for key, value in blob["state_dict"].items():
newKey = key
if key.startswith("module."):
newKey = key[len("module."):]
state_dict[newKey] = value
waveglow.load_state_dict(state_dict)
else:
waveglow = blob['model']
waveglow = split_cond_layers(waveglow)
waveglow = waveglow.remove_weightnorm(waveglow)
waveglow.cuda().eval()
return waveglow
def get_model(model_name, model_config, to_cuda):
""" Code chooses a model based on name"""
model = None
if model_name == 'Tacotron2':
model = Tacotron2(**model_config)
elif model_name == 'WaveGlow':
model = WaveGlow(**model_config)
else:
raise NotImplementedError(model_name)
if to_cuda:
model = model.cuda()
return model
def get_model(model_name, model_config, to_cuda,
uniform_initialize_bn_weight=False):
""" Code chooses a model based on name"""
model = None
if model_name == 'Tacotron2':
model = Tacotron2(**model_config)
elif model_name == 'WaveGlow':
model = WaveGlow(**model_config)
else:
raise NotImplementedError(model_name)
if uniform_initialize_bn_weight:
init_bn(model)
if to_cuda:
model = model.cuda()
return model
def get_model(model_name, model_config, to_fp16, to_cuda, training=True):
""" Code chooses a model based on name"""
model = None
if model_name == 'Tacotron2':
model = Tacotron2(**model_config)
elif model_name == 'WaveGlow':
model = WaveGlow(**model_config)
else:
raise NotImplementedError(model_name)
if to_fp16:
model = batchnorm_to_float(model.half())
model = lstmcell_to_float(model)
if model_name == "WaveGlow":
for k in model.convinv:
k.float()
if to_cuda:
model = model.cuda()
return model
def get_model(model_name, model_config, to_cuda,
uniform_initialize_bn_weight=False, rename=False):
""" Code chooses a model based on name"""
model = None
if model_name == 'Tacotron2':
if rename:
class Tacotron2_extra(Tacotron2):
def forward(self, inputs, input_lengths):
return self.infer(inputs, input_lengths)
model = Tacotron2_extra(**model_config)
else:
model = Tacotron2(**model_config)
elif model_name == 'WaveGlow':
model = WaveGlow(**model_config)
else:
raise NotImplementedError(model_name)
if uniform_initialize_bn_weight:
init_bn(model)
if to_cuda:
model = model.cuda()
return model