From 17b789dfbef0af39df95f4b7686c465e75776b77 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?=E6=BA=90=E6=96=87=E9=9B=A8?=
 <41315874+fumiama@users.noreply.github.com>
Date: Sun, 2 Apr 2023 19:16:34 +0800
Subject: [PATCH] fix size mismatch for dec.conv_post.weight

---
 train/process_ckpt.py | 20 ++++++++++----------
 1 file changed, 10 insertions(+), 10 deletions(-)

diff --git a/train/process_ckpt.py b/train/process_ckpt.py
index 3e4fc7e..3840345 100644
--- a/train/process_ckpt.py
+++ b/train/process_ckpt.py
@@ -69,10 +69,16 @@ def merge(path1,path2,alpha1,sr,f0,info,name):
             return opt
         ckpt1 = torch.load(path1, map_location="cpu")
         ckpt2 = torch.load(path2, map_location="cpu")
-        if("model"in ckpt1):ckpt1=extract(ckpt1)
-        else:ckpt1=ckpt1["weight"]
-        if("model"in ckpt2):ckpt2=extract(ckpt2)
-        else:ckpt2=ckpt2["weight"]
+        opt["config"] = ckpt1["config"]
+        '''
+        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]
+        '''
+        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"] = {}
@@ -85,12 +91,6 @@ def merge(path1,path2,alpha1,sr,f0,info,name):
                     opt["weight"][key] = (alpha1*(ckpt1[key].float())+(1-alpha1)*(ckpt2[key].float())).half()
             # except:
             #     pdb.set_trace()
-        opt["config"] = ckpt1["config"]
-        '''
-        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