import torch, traceback, os, pdb from collections import OrderedDict def savee(ckpt, sr, if_f0, name, epoch): try: opt = OrderedDict() opt["weight"] = {} for key in ckpt.keys(): 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["info"] = "%sepoch" % epoch opt["sr"] = sr opt["f0"] = if_f0 torch.save(opt, "weights/%s.pth" % name) return "Success." except: return traceback.format_exc() def show_info(path): try: a = torch.load(path, map_location="cpu") return "模型信息:%s\n采样率:%s\n模型是否输入音高引导:%s" % ( a.get("info", "None"), a.get("sr", "None"), a.get("f0", "None"), ) except: return traceback.format_exc() def extract_small_model(path, name, sr, if_f0, info): try: ckpt = torch.load(path, map_location="cpu") if "model" in ckpt: ckpt = ckpt["model"] opt = OrderedDict() opt["weight"] = {} for key in ckpt.keys(): 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, ] if info == "": info = "Extracted model." opt["info"] = info opt["sr"] = sr opt["f0"] = int(if_f0) torch.save(opt, "weights/%s.pth" % name) return "Success." except: return traceback.format_exc() def change_info(path, info, name): try: ckpt = torch.load(path, map_location="cpu") ckpt["info"] = info if name == "": name = os.path.basename(path) torch.save(ckpt, "weights/%s" % name) return "Success." except: return traceback.format_exc() def merge(path1, path2, alpha1, sr, f0, info, name): try: def extract(ckpt): a = ckpt["model"] opt = OrderedDict() opt["weight"] = {} for key in a.keys(): if "enc_q" in key: continue opt["weight"][key] = a[key] return opt ckpt1 = torch.load(path1, map_location="cpu") ckpt2 = torch.load(path2, map_location="cpu") cfg = ckpt1["config"] if "model" in ckpt1: ckpt1 = extract(ckpt1) else: ckpt1 = ckpt1["weight"] if "model" in ckpt2: ckpt2 = extract(ckpt2) else: ckpt2 = ckpt2["weight"] if sorted(list(ckpt1.keys())) != sorted(list(ckpt2.keys())): return "Fail to merge the models. The model architectures are not the same." opt = OrderedDict() opt["weight"] = {} for key in ckpt1.keys(): # try: if key == "emb_g.weight" and ckpt1[key].shape != ckpt2[key].shape: min_shape0 = min(ckpt1[key].shape[0], ckpt2[key].shape[0]) opt["weight"][key] = ( alpha1 * (ckpt1[key][:min_shape0].float()) + (1 - alpha1) * (ckpt2[key][:min_shape0].float()) ).half() else: opt["weight"][key] = ( alpha1 * (ckpt1[key].float()) + (1 - alpha1) * (ckpt2[key].float()) ).half() # except: # pdb.set_trace() opt["config"] = cfg """ 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] 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] 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["sr"] = sr opt["f0"] = 1 if f0 == "是" else 0 opt["info"] = info torch.save(opt, "weights/%s.pth" % name) return "Success." except: return traceback.format_exc()