vits

vits/tts_inferencer.py

@@ -33,21 +33,23 @@ def get_text(text, hps):
 
 class TTSInferencer:
     def __init__(self, hps_path, device="cpu"):
-        print("init")
         self.device = torch.device(device)
         self.hps = utils.get_hparams_from_file(hps_path)
-        self.model_paths = {}
-        self.models = {}
-        for key, value in mode_dict.items():
-            self.model_paths[key] = self.get_latest_model_path_by_mode(key)
-        self.load_models()
+        self.select_mode(default_mode)
+        self.load_model(self.latest_model_path)
 
-    def get_latest_model_path_by_mode(self, mode):
-        model_dir_path = os.path.join(ROOT_PATH, "models", mode_dict[mode])
-        return utils.latest_checkpoint_path(model_dir_path, "G_*.pth")
+    def select_mode(self, mode):
+        self.mode = mode
+        self.model_dir_path = os.path.join(ROOT_PATH, "models", mode_dict[mode])
+        self.models = []
+        for f in os.listdir(self.model_dir_path):
+            if (f.startswith("D_")):
+                continue
+            if (f.endswith(".pth")):
+                self.models.append(f)
+        self.latest_model_path = utils.latest_checkpoint_path(self.model_dir_path, "G_*.pth")
 
-    def infer(self, text, mode, noise_scale=.667, noise_scale_w=0.8, length_scale=1):
-        print(self.pth_path)
+    def infer(self, text, noise_scale=.667, noise_scale_w=0.8, length_scale=1):
         stn_tst = get_text(text, self.hps)
         with torch.no_grad():
             x_tst = stn_tst.unsqueeze(0).to(self.device)
@@ -65,6 +67,7 @@ class TTSInferencer:
     
     def render(self):
         choice_mode = gr.Radio(choices=["普通声线", "营业声线"], label="声线选择", value=default_mode)
+        choice_model = gr.Dropdown(choices=self.models, label=f"模型迭代版本选择", value=os.path.basename(self.pth_path))
         # with gr.Row():
         #     advanced = gr.Checkbox(label="显示高级设置（效果不可控）") 
         #     default = gr.Button("恢复默认设置").style(full_width=False)
@@ -77,6 +80,7 @@ class TTSInferencer:
             value="这里是爱喝奶茶，穿得也像奶茶魅力点是普通话二乙的星弥吼西咪，晚上齁。")
         tts_submit = gr.Button("合成", variant="primary")
         tts_output = gr.Audio(label="Output")
+        tts_model = gr.Markdown(f"载入模型：{os.path.basename(self.latest_model_path)}（{self.mode})")
         gr.HTML('''
             <div style="text-align:right;font-size:12px;color:#4D4D4D">
                 <div class="font-medium">版权声明</div>
@@ -89,19 +93,17 @@ class TTSInferencer:
         # default.click(fn=lambda visible: gr.update(value=default_noise_scale), inputs=advanced, outputs=noise_scale)
         # default.click(fn=lambda visible: gr.update(value=default_noise_scale_w), inputs=advanced, outputs=noise_scale_w)
         # default.click(fn=lambda visible: gr.update(value=default_length_scale), inputs=advanced, outputs=length_scale)        
-        tts_submit.click(self.infer, [choice_mode, tts_input, noise_scale, noise_scale_w, length_scale], [tts_output], api_name=f"infer")
+        choice_mode.change(self.change_mode, inputs=choice_mode, outputs=choice_model)
+        choice_model.change(self.change_model, inputs=[choice_model], outputs=[tts_model])
+        tts_submit.click(self.infer, [tts_input, noise_scale, noise_scale_w, length_scale], [tts_output], api_name=f"infer")
+
     
-    def load_models(self):
-        for key, model_path in self.latest_model_paths.items():
-            self.models[key] = SynthesizerTrn(
+    def load_model(self, model_path):
+        self.pth_path = model_path
+        self.net_g = SynthesizerTrn(
             len(symbols),
             self.hps.data.filter_length // 2 + 1,
             self.hps.train.segment_size // self.hps.data.hop_length,
             **self.hps.model).to(self.device)
-            _ = self.models[key].eval()
-            _ = utils.load_checkpoint(model_path, self.models[key], None)
-    
-    def __del__(self):
-        print("del")
-        del self.net_g
-        self.net_g = None
+        _ = self.net_g.eval()
+        _ = utils.load_checkpoint(self.pth_path, self.net_g, None)