diff --git a/train/process_ckpt.py b/train/process_ckpt.py index caf9e05..876be51 100644 --- a/train/process_ckpt.py +++ b/train/process_ckpt.py @@ -8,7 +8,7 @@ from i18n import I18nAuto i18n = I18nAuto() -def savee(ckpt, sr, if_f0, name, epoch, version): +def savee(ckpt, sr, if_f0, name, epoch, version,hps): try: opt = OrderedDict() opt["weight"] = {} @@ -16,69 +16,26 @@ def savee(ckpt, sr, if_f0, name, epoch, version): if "enc_q" in key: continue opt["weight"][key] = ckpt[key].half() - if sr == "40k": - opt["config"] = [ - 1025, - 32, - 192, - 192, - 768, - 2, - 6, - 3, - 0, - "1", - [3, 7, 11], - [[1, 3, 5], [1, 3, 5], [1, 3, 5]], - [10, 10, 2, 2], - 512, - [16, 16, 4, 4], - 109, - 256, - 40000, - ] - elif sr == "48k": - opt["config"] = [ - 1025, - 32, - 192, - 192, - 768, - 2, - 6, - 3, - 0, - "1", - [3, 7, 11], - [[1, 3, 5], [1, 3, 5], [1, 3, 5]], - [10, 6, 2, 2, 2], - 512, - [16, 16, 4, 4, 4], - 109, - 256, - 48000, - ] - elif sr == "32k": - opt["config"] = [ - 513, - 32, - 192, - 192, - 768, - 2, - 6, - 3, - 0, - "1", - [3, 7, 11], - [[1, 3, 5], [1, 3, 5], [1, 3, 5]], - [10, 4, 2, 2, 2], - 512, - [16, 16, 4, 4, 4], - 109, - 256, - 32000, - ] + opt["config"]=[ + hpt.data.filter_length//2+1, + 32, + hpt.model.inter_channels, + hpt.model.hidden_channels, + hpt.model.filter_channels, + hpt.model.n_heads, + hpt.model.n_layers, + hpt.model.kernel_size, + hpt.model.p_dropout, + hpt.model.resblock, + hpt.model.resblock_kernel_sizes, + hpt.model.resblock_dilation_sizes, + hpt.model.upsample_rates, + hpt.model.upsample_initial_channel, + hpt.model.upsample_kernel_sizes, + hpt.model.spk_embed_dim, + hpt.model.gin_channels, + hpt.data.sampling_rate + ] opt["info"] = "%sepoch" % epoch opt["sr"] = sr opt["f0"] = if_f0