update

inference/tts/ps_flow.py

@@ -11,7 +11,8 @@ class PortaSpeechFlowInfer(BaseTTSInfer):
         word_dict_size = len(self.word_encoder)
         model = PortaSpeechFlow(ph_dict_size, word_dict_size, self.hparams)
         load_ckpt(model, hparams['work_dir'], 'model')
-        model.post_flow.store_inverse()
+        with torch.no_grad():
+            model.store_inverse_all()
         model.eval()
         return model
 

modules/tts/portaspeech/portaspeech.py

@@ -212,4 +212,15 @@ class PortaSpeech(FastSpeech):
         x_pos = build_word_mask(word2word, x2word).float()  # [B, T_word, T_ph]
         x_pos = (x_pos.cumsum(-1) / x_pos.sum(-1).clamp(min=1)[..., None] * x_pos).sum(1)
         x_pos = self.sin_pos(x_pos.float())  # [B, T_ph, H]
-        return x_pos
+        return x_pos
+
+    def store_inverse_all(self):
+        def remove_weight_norm(m):
+            try:
+                if hasattr(m, 'store_inverse'):
+                    m.store_inverse()
+                nn.utils.remove_weight_norm(m)
+            except ValueError:  # this module didn't have weight norm
+                return
+
+        self.apply(remove_weight_norm)

tasks/tts/ps.py

@@ -156,14 +156,7 @@ class PortaSpeechTask(FastSpeechTask):
         super().test_start()
         if hparams.get('save_attn', False):
             os.makedirs(f'{self.gen_dir}/attn', exist_ok=True)
-
-        def remove_weight_norm(m):
-            try:
-                nn.utils.remove_weight_norm(m)
-            except ValueError:
-                return
-
-        self.apply(remove_weight_norm)
+        self.model.store_inverse_all()
 
     def test_step(self, sample, batch_idx):
         assert sample['txt_tokens'].shape[0] == 1, 'only support batch_size=1 in inference'

tasks/tts/ps_flow.py

@@ -131,12 +131,4 @@ class PortaSpeechFlowTask(PortaSpeechTask):
             return [self.optimizer]
 
     def build_scheduler(self, optimizer):
-        return FastSpeechTask.build_scheduler(self, optimizer[0])
-
-    ############
-    # infer
-    ############
-    def test_start(self):
-        super().test_start()
-        if hparams['use_post_flow']:
-            self.model.post_flow.store_inverse()
+        return FastSpeechTask.build_scheduler(self, optimizer[0])