wonkitty/apple_oh
0
1import torch, traceback, os, pdb, sys2 3now_dir = os.getcwd()4sys.path.append(now_dir)5from collections import OrderedDict6from i18n import I18nAuto7 8i18n = I18nAuto()9 10 11def savee(ckpt, sr, if_f0, name, epoch, version, hps):12 try:13 opt = OrderedDict()14 opt["weight"] = {}15 for key in ckpt.keys():16 if "enc_q" in key:17 continue18 opt["weight"][key] = ckpt[key].half()19 opt["config"] = [20 hps.data.filter_length // 2 + 1,21 32,22 hps.model.inter_channels,23 hps.model.hidden_channels,24 hps.model.filter_channels,25 hps.model.n_heads,26 hps.model.n_layers,27 hps.model.kernel_size,28 hps.model.p_dropout,29 hps.model.resblock,30 hps.model.resblock_kernel_sizes,31 hps.model.resblock_dilation_sizes,32 hps.model.upsample_rates,33 hps.model.upsample_initial_channel,34 hps.model.upsample_kernel_sizes,35 hps.model.spk_embed_dim,36 hps.model.gin_channels,37 hps.data.sampling_rate,38 ]39 opt["info"] = "%sepoch" % epoch40 opt["sr"] = sr41 opt["f0"] = if_f042 opt["version"] = version43 torch.save(opt, "weights/%s.pth" % name)44 return "Success."45 except:46 return traceback.format_exc()47 48 49def show_info(path):50 try:51 a = torch.load(path, map_location="cpu")52 return "Epochs: %s\nSample rate: %s\nPitch guidance: %s\nRVC Version: %s" % (53 a.get("info", "None"),54 a.get("sr", "None"),55 a.get("f0", "None"),56 a.get("version", "None"),57 )58 except:59 return traceback.format_exc()60 61 62def extract_small_model(path, name, sr, if_f0, info, version):63 try:64 ckpt = torch.load(path, map_location="cpu")65 if "model" in ckpt:66 ckpt = ckpt["model"]67 opt = OrderedDict()68 opt["weight"] = {}69 for key in ckpt.keys():70 if "enc_q" in key:71 continue72 opt["weight"][key] = ckpt[key].half()73 if sr == "40k":74 opt["config"] = [75 1025,76 32,77 192,78 192,79 768,80 2,81 6,82 3,83 0,84 "1",85 [3, 7, 11],86 [[1, 3, 5], [1, 3, 5], [1, 3, 5]],87 [10, 10, 2, 2],88 512,89 [16, 16, 4, 4],90 109,91 256,92 40000,93 ]94 elif sr == "48k":95 if version == "v1":96 opt["config"] = [97 1025,98 32,99 192,100 192,101 768,102 2,103 6,104 3,105 0,106 "1",107 [3, 7, 11],108 [[1, 3, 5], [1, 3, 5], [1, 3, 5]],109 [10, 6, 2, 2, 2],110 512,111 [16, 16, 4, 4, 4],112 109,113 256,114 48000,115 ]116 else:117 opt["config"] = [118 1025,119 32,120 192,121 192,122 768,123 2,124 6,125 3,126 0,127 "1",128 [3, 7, 11],129 [[1, 3, 5], [1, 3, 5], [1, 3, 5]],130 [12, 10, 2, 2],131 512,132 [24, 20, 4, 4],133 109,134 256,135 48000,136 ]137 elif sr == "32k":138 if version == "v1":139 opt["config"] = [140 513,141 32,142 192,143 192,144 768,145 2,146 6,147 3,148 0,149 "1",150 [3, 7, 11],151 [[1, 3, 5], [1, 3, 5], [1, 3, 5]],152 [10, 4, 2, 2, 2],153 512,154 [16, 16, 4, 4, 4],155 109,156 256,157 32000,158 ]159 else:160 opt["config"] = [161 513,162 32,163 192,164 192,165 768,166 2,167 6,168 3,169 0,170 "1",171 [3, 7, 11],172 [[1, 3, 5], [1, 3, 5], [1, 3, 5]],173 [10, 8, 2, 2],174 512,175 [20, 16, 4, 4],176 109,177 256,178 32000,179 ]180 if info == "":181 info = "Extracted model."182 opt["info"] = info183 opt["version"] = version184 opt["sr"] = sr185 opt["f0"] = int(if_f0)186 torch.save(opt, "weights/%s.pth" % name)187 return "Success."188 except:189 return traceback.format_exc()190 191 192def change_info(path, info, name):193 try:194 ckpt = torch.load(path, map_location="cpu")195 ckpt["info"] = info196 if name == "":197 name = os.path.basename(path)198 torch.save(ckpt, "weights/%s" % name)199 return "Success."200 except:201 return traceback.format_exc()202 203 204def merge(path1, path2, alpha1, sr, f0, info, name, version):205 try:206 207 def extract(ckpt):208 a = ckpt["model"]209 opt = OrderedDict()210 opt["weight"] = {}211 for key in a.keys():212 if "enc_q" in key:213 continue214 opt["weight"][key] = a[key]215 return opt216 217 ckpt1 = torch.load(path1, map_location="cpu")218 ckpt2 = torch.load(path2, map_location="cpu")219 cfg = ckpt1["config"]220 if "model" in ckpt1:221 ckpt1 = extract(ckpt1)222 else:223 ckpt1 = ckpt1["weight"]224 if "model" in ckpt2:225 ckpt2 = extract(ckpt2)226 else:227 ckpt2 = ckpt2["weight"]228 if sorted(list(ckpt1.keys())) != sorted(list(ckpt2.keys())):229 return "Fail to merge the models. The model architectures are not the same."230 opt = OrderedDict()231 opt["weight"] = {}232 for key in ckpt1.keys():233 # try:234 if key == "emb_g.weight" and ckpt1[key].shape != ckpt2[key].shape:235 min_shape0 = min(ckpt1[key].shape[0], ckpt2[key].shape[0])236 opt["weight"][key] = (237 alpha1 * (ckpt1[key][:min_shape0].float())238 + (1 - alpha1) * (ckpt2[key][:min_shape0].float())239 ).half()240 else:241 opt["weight"][key] = (242 alpha1 * (ckpt1[key].float()) + (1 - alpha1) * (ckpt2[key].float())243 ).half()244 # except:245 # pdb.set_trace()246 opt["config"] = cfg247 """248 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,4], 109, 256, 40000]249 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], 109, 256, 48000]250 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]251 """252 opt["sr"] = sr253 opt["f0"] = 1 if f0 else 0254 opt["version"] = version255 opt["info"] = info256 torch.save(opt, "weights/%s.pth" % name)257 return "Success."258 except:259 return traceback.format_exc()260