🌹 Add and for all config class.

Author: dathudeptrai

tensorflow_tts/configs/__init__.py

@@ -1,3 +1,4 @@
+from tensorflow_tts.configs.base_config import BaseConfig
 from tensorflow_tts.configs.fastspeech import FastSpeechConfig
 from tensorflow_tts.configs.fastspeech2 import FastSpeech2Config
 from tensorflow_tts.configs.melgan import (

tensorflow_tts/configs/base_config.py

@@ -0,0 +1,31 @@
+# -*- coding: utf-8 -*-
+# Copyright 2020 TensorFlowTTS Team.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""Base Config for all config."""
+
+import abc
+import yaml
+import os
+
+from tensorflow_tts.utils.utils import CONFIG_FILE_NAME
+
+
+class BaseConfig(abc.ABC):
+    def set_pretrained_config(self, config):
+        self.config = config
+
+    def save_pretrained(self, saved_path):
+        """Save config to file"""
+        with open(os.path.join(saved_path, CONFIG_FILE_NAME), "w") as file:
+            yaml.dump(self.config, file)

tensorflow_tts/configs/fastspeech.py

@@ -16,6 +16,7 @@
 
 import collections
 
+from tensorflow_tts.configs import BaseConfig
 from tensorflow_tts.processor.ljspeech import LJSPEECH_SYMBOLS as lj_symbols
 from tensorflow_tts.processor.kss import KSS_SYMBOLS as kss_symbols
 from tensorflow_tts.processor.baker import BAKER_SYMBOLS as bk_symbols
@@ -44,12 +45,12 @@
 )
 
 
-class FastSpeechConfig(object):
+class FastSpeechConfig(BaseConfig):
     """Initialize FastSpeech Config."""
 
     def __init__(
         self,
-        dataset='ljspeech',
+        dataset="ljspeech",
         vocab_size=len(lj_symbols),
         n_speakers=1,
         encoder_hidden_size=384,

tensorflow_tts/configs/fastspeech2.py

@@ -14,6 +14,7 @@
 # limitations under the License.
 """FastSpeech2 Config object."""
 
+
 from tensorflow_tts.configs import FastSpeechConfig
 
 

tensorflow_tts/configs/hifigan.py

@@ -15,7 +15,10 @@
 """HifiGAN Config object."""
 
 
-class HifiGANGeneratorConfig(object):
+from tensorflow_tts.configs import BaseConfig
+
+
+class HifiGANGeneratorConfig(BaseConfig):
     """Initialize HifiGAN Generator Config."""
 
     def __init__(

tensorflow_tts/configs/melgan.py

@@ -15,7 +15,10 @@
 """MelGAN Config object."""
 
 
-class MelGANGeneratorConfig(object):
+from tensorflow_tts.configs import BaseConfig
+
+
+class MelGANGeneratorConfig(BaseConfig):
     """Initialize MelGAN Generator Config."""
 
     def __init__(

tensorflow_tts/configs/parallel_wavegan.py

@@ -15,7 +15,10 @@
 """ParallelWaveGAN Config object."""
 
 
-class ParallelWaveGANGeneratorConfig(object):
+from tensorflow_tts.configs import BaseConfig
+
+
+class ParallelWaveGANGeneratorConfig(BaseConfig):
     """Initialize ParallelWaveGAN Generator Config."""
 
     def __init__(

tensorflow_tts/configs/tacotron2.py

@@ -14,18 +14,20 @@
 # limitations under the License.
 """Tacotron-2 Config object."""
 
+
+from tensorflow_tts.configs import BaseConfig
 from tensorflow_tts.processor.ljspeech import LJSPEECH_SYMBOLS as lj_symbols
 from tensorflow_tts.processor.kss import KSS_SYMBOLS as kss_symbols
 from tensorflow_tts.processor.baker import BAKER_SYMBOLS as bk_symbols
 from tensorflow_tts.processor.libritts import LIBRITTS_SYMBOLS as lbri_symbols
 
 
-class Tacotron2Config(object):
+class Tacotron2Config(BaseConfig):
     """Initialize Tacotron-2 Config."""
 
     def __init__(
         self,
-        dataset='ljspeech',
+        dataset="ljspeech",
         vocab_size=len(lj_symbols),
         embedding_hidden_size=512,
         initializer_range=0.02,
@@ -60,7 +62,7 @@ def __init__(
             self.vocab_size = vocab_size
         elif dataset == "kss":
             self.vocab_size = len(kss_symbols)
-        elif dataset == 'baker':
+        elif dataset == "baker":
             self.vocab_size = len(bk_symbols)
         elif dataset == "libritts":
             self.vocab_size = len(lbri_symbols)

tensorflow_tts/inference/auto_config.py

@@ -41,7 +41,7 @@
         ("melgan_generator", MelGANGeneratorConfig),
         ("hifigan_generator", HifiGANGeneratorConfig),
         ("tacotron2", Tacotron2Config),
-        ("parallel_wavegan_generator", ParallelWaveGANGeneratorConfig)
+        ("parallel_wavegan_generator", ParallelWaveGANGeneratorConfig),
     ]
 )
 
@@ -58,7 +58,9 @@ def from_pretrained(cls, pretrained_path, **kwargs):
         # load weights from hf hub
         if not os.path.isfile(pretrained_path):
             # retrieve correct hub url
-            download_url = hf_hub_url(repo_id=pretrained_path, filename=CONFIG_FILE_NAME)
+            download_url = hf_hub_url(
+                repo_id=pretrained_path, filename=CONFIG_FILE_NAME
+            )
 
             pretrained_path = str(
                 cached_download(
@@ -76,6 +78,7 @@ def from_pretrained(cls, pretrained_path, **kwargs):
             model_type = config["model_type"]
             config_class = CONFIG_MAPPING[model_type]
             config_class = config_class(**config[model_type + "_params"], **kwargs)
+            config_class.set_pretrained_config(config)
             return config_class
         except Exception:
             raise ValueError(