Merge branch 'openvpi:master' into master

Author: oxygen-dioxide

.gitignore

@@ -4,6 +4,7 @@ __pycache__/
 *.sh
 local_tools/
 *.ckpt
+*.pth
 *.wav
 infer_out/
 config.yaml

README.md

@@ -2,6 +2,7 @@
 [![arXiv](https://img.shields.io/badge/arXiv-Paper-<COLOR>.svg)](https://arxiv.org/abs/2105.02446)
 [![GitHub Stars](https://img.shields.io/github/stars/MoonInTheRiver/DiffSinger?style=social)](https://github.com/MoonInTheRiver/DiffSinger)
 [![downloads](https://img.shields.io/github/downloads/MoonInTheRiver/DiffSinger/total.svg)](https://github.com/MoonInTheRiver/DiffSinger/releases)
+ | [![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/drive/1kfmZ6Y018c5trSwQAbhdQtZ7Il8W_4BU)
  | [Interactive🤗 TTS](https://huggingface.co/spaces/NATSpeech/DiffSpeech)
  | [Interactive🤗 SVS](https://huggingface.co/spaces/Silentlin/DiffSinger)
 

inference/svs/base_svs_infer.py

@@ -47,15 +47,19 @@ def forward_model(self, inp):
     def build_vocoder(self):
         base_dir = hparams['vocoder_ckpt']
         config_path = f'{base_dir}/config.yaml'
-        ckpt = sorted(glob.glob(f'{base_dir}/model_ckpt_steps_*.ckpt'), key=
-        lambda x: int(re.findall(f'{base_dir}/model_ckpt_steps_(\d+).ckpt', x.replace('\\','/'))[0]))[-1]
-        print('| load HifiGAN: ', ckpt)
-        ckpt_dict = torch.load(ckpt, map_location="cpu")
-        config = set_hparams(config_path, global_hparams=False)
-        state = ckpt_dict["state_dict"]["model_gen"]
-        vocoder = HifiGanGenerator(config)
-        vocoder.load_state_dict(state, strict=True)
-        vocoder.remove_weight_norm()
+        file_path = sorted(glob.glob(f'{base_dir}/model_ckpt_steps_*.*'), key=
+        lambda x: int(re.findall(f'{base_dir}/model_ckpt_steps_(\d+).*', x.replace('\\','/'))[0]))[-1]
+        print('| load HifiGAN: ', file_path)
+        ext = os.path.splitext(file_path)[-1]
+        if ext == '.pth':
+            vocoder = torch.load(file_path, map_location="cpu")
+        elif ext == '.ckpt':
+            ckpt_dict = torch.load(file_path, map_location="cpu")
+            config = set_hparams(config_path, global_hparams=False)
+            state = ckpt_dict["state_dict"]["model_gen"]
+            vocoder = HifiGanGenerator(config)
+            vocoder.load_state_dict(state, strict=True)
+            vocoder.remove_weight_norm()
         vocoder = vocoder.eval().to(self.device)
         return vocoder
 

tasks/base_task.py

@@ -251,7 +251,7 @@ def start(cls):
             t = datetime.now().strftime('%Y%m%d%H%M%S')
             code_dir = f'{work_dir}/codes/{t}'
             # TODO: test filesystem calls
-            os.mkdir(code_dir)
+            os.makedirs(code_dir, exist_ok=True)
             # subprocess.check_call(f'mkdir "{code_dir}"', shell=True)
             for c in hparams['save_codes']:
                 shutil.copytree(c, code_dir, dirs_exist_ok=True)

vocoders/hifigan.py

@@ -14,21 +14,28 @@
 from vocoders.vocoder_utils import denoise
 
 
-def load_model(config_path, checkpoint_path):
+def load_model(config_path, file_path):
     device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
-    ckpt_dict = torch.load(checkpoint_path, map_location="cpu")
-    if '.yaml' in config_path:
-        config = set_hparams(config_path, global_hparams=False)
-        state = ckpt_dict["state_dict"]["model_gen"]
-    elif '.json' in config_path:
-        config = json.load(open(config_path, 'r', encoding='utf-8'))
-        state = ckpt_dict["generator"]
-
-    model = HifiGanGenerator(config)
-    model.load_state_dict(state, strict=True)
-    model.remove_weight_norm()
+    ext = os.path.splitext(file_path)[-1]
+    if ext == '.pth':
+        if '.yaml' in config_path:
+            config = set_hparams(config_path, global_hparams=False)
+        elif '.json' in config_path:
+            config = json.load(open(config_path, 'r', encoding='utf-8'))
+        model = torch.load(file_path, map_location="cpu")
+    elif ext == '.ckpt':
+        ckpt_dict = torch.load(file_path, map_location="cpu")
+        if '.yaml' in config_path:
+            config = set_hparams(config_path, global_hparams=False)
+            state = ckpt_dict["state_dict"]["model_gen"]
+        elif '.json' in config_path:
+            config = json.load(open(config_path, 'r', encoding='utf-8'))
+            state = ckpt_dict["generator"]
+        model = HifiGanGenerator(config)
+        model.load_state_dict(state, strict=True)
+        model.remove_weight_norm()
     model = model.eval().to(device)
-    print(f"| Loaded model parameters from {checkpoint_path}.")
+    print(f"| Loaded model parameters from {file_path}.")
     print(f"| HifiGAN device: {device}.")
     return model, config, device
 
@@ -42,15 +49,15 @@ def __init__(self):
         base_dir = hparams['vocoder_ckpt']
         config_path = f'{base_dir}/config.yaml'
         if os.path.exists(config_path):
-            ckpt = sorted(glob.glob(f'{base_dir}/model_ckpt_steps_*.ckpt'), key=
-            lambda x: int(re.findall(f'{base_dir}/model_ckpt_steps_(\d+).ckpt', x.replace('\\','/'))[0]))[-1]
-            print('| load HifiGAN: ', ckpt)
-            self.model, self.config, self.device = load_model(config_path=config_path, checkpoint_path=ckpt)
+            file_path = sorted(glob.glob(f'{base_dir}/model_ckpt_steps_*.*'), key=
+            lambda x: int(re.findall(f'{base_dir}/model_ckpt_steps_(\d+).*', x.replace('\\','/'))[0]))[-1]
+            print('| load HifiGAN: ', file_path)
+            self.model, self.config, self.device = load_model(config_path=config_path, file_path=file_path)
         else:
             config_path = f'{base_dir}/config.json'
             ckpt = f'{base_dir}/generator_v1'
             if os.path.exists(config_path):
-                self.model, self.config, self.device = load_model(config_path=config_path, checkpoint_path=ckpt)
+                self.model, self.config, self.device = load_model(config_path=config_path, file_path=file_path)
 
     def spec2wav(self, mel, **kwargs):
         device = self.device