refactor(datasets): register formatters in dict

Author: eginhard

TTS/tts/datasets/__init__.py

@@ -9,7 +9,7 @@
 import numpy as np
 
 from TTS.tts.datasets.dataset import *
-from TTS.tts.datasets.formatters import *
+from TTS.tts.datasets.formatters import _FORMATTER_REGISTRY, Formatter, register_formatter
 
 logger = logging.getLogger(__name__)
 
@@ -162,32 +162,12 @@ def load_attention_mask_meta_data(metafile_path):
     return meta_data
 
 
-def add_formatter(name: str, formatter: Callable[[str, str, list[str] | None], list[dict]]):
-    """Add a formatter to the datasets module. If the formatter already exists, raise an error.
-    Args:
-        name (str): The name of the formatter.
-        formatter (Callable): The formatter function.
-    Raises:
-        ValueError: If the formatter already exists.
-    Returns:
-        None
-    """
-    thismodule = sys.modules[__name__]
-    if not hasattr(thismodule, name.lower()):
-        setattr(thismodule, name.lower(), formatter)
-    else:
-        raise ValueError(f"Formatter {name} already exists.")
-
-
-def _get_formatter_by_name(name):
+def _get_formatter_by_name(name: str) -> Formatter:
     """Returns the respective preprocessing function."""
-    thismodule = sys.modules[__name__]
-    if not hasattr(thismodule, name.lower()):
-        msg = (
-            f"{name} formatter not found. If it is a custom formatter, pass the function to load_tts_samples() instead."
-        )
+    if name.lower() not in _FORMATTER_REGISTRY:
+        msg = f"{name} formatter not found. If it is a custom formatter, make sure to call register_formatter() first."
         raise ValueError(msg)
-    return getattr(thismodule, name.lower())
+    return _FORMATTER_REGISTRY[name.lower()]
 
 
 def find_unique_chars(data_samples):

TTS/tts/datasets/formatters.py

@@ -5,11 +5,39 @@
 import xml.etree.ElementTree as ET
 from glob import glob
 from pathlib import Path
+from typing import Any, Protocol
 
 from tqdm import tqdm
 
 logger = logging.getLogger(__name__)
 
+
+class Formatter(Protocol):
+    def __call__(
+        self,
+        root_path: str | os.PathLike[Any],
+        meta_file: str | os.PathLike[Any],
+        ignored_speakers: list[str] | None,
+        **kwargs,
+    ) -> list[dict[str, Any]]: ...
+
+
+_FORMATTER_REGISTRY: dict[str, Formatter] = {}
+
+
+def register_formatter(name: str, formatter: Formatter) -> None:
+    """Add a formatter function to the registry.
+
+    Args:
+        name: Name of the formatter.
+        formatter: Formatter function.
+    """
+    if name.lower() in _FORMATTER_REGISTRY:
+        msg = f"Formatter {name} already exists."
+        raise ValueError(msg)
+    _FORMATTER_REGISTRY[name.lower()] = formatter
+
+
 ########################
 # DATASETS
 ########################
@@ -659,3 +687,35 @@ def bel_tts_formatter(root_path, meta_file, **kwargs):  # pylint: disable=unused
             text = cols[1]
             items.append({"text": text, "audio_file": wav_file, "speaker_name": speaker_name, "root_path": root_path})
     return items
+
+
+### Registrations
+register_formatter("cml_tts", cml_tts)
+register_formatter("coqui", coqui)
+register_formatter("tweb", tweb)
+register_formatter("mozilla", mozilla)
+register_formatter("mozilla_de", mozilla_de)
+register_formatter("mailabs", mailabs)
+register_formatter("ljspeech", ljspeech)
+register_formatter("ljspeech_test", ljspeech_test)
+register_formatter("thorsten", thorsten)
+register_formatter("sam_accenture", sam_accenture)
+register_formatter("ruslan", ruslan)
+register_formatter("css10", css10)
+register_formatter("nancy", nancy)
+register_formatter("common_voice", common_voice)
+register_formatter("libri_tts", libri_tts)
+register_formatter("custom_turkish", custom_turkish)
+register_formatter("brspeech", brspeech)
+register_formatter("vctk", vctk)
+register_formatter("vctk_old", vctk_old)
+register_formatter("synpaflex", synpaflex)
+register_formatter("open_bible", open_bible)
+register_formatter("mls", mls)
+register_formatter("voxceleb2", voxceleb2)
+register_formatter("voxceleb1", voxceleb1)
+register_formatter("emotion", emotion)
+register_formatter("baker", baker)
+register_formatter("kokoro", kokoro)
+register_formatter("kss", kss)
+register_formatter("bel_tts_formatter", bel_tts_formatter)

docs/source/datasets/formatting_your_dataset.md

@@ -101,7 +101,7 @@ train_samples, eval_samples = load_tts_samples(dataset_config, eval_split=True)
 Load a custom dataset with a custom formatter.
 
 ```python
-from TTS.tts.datasets import load_tts_samples, add_formatter
+from TTS.tts.datasets import load_tts_samples, register_formatter
 
 
 # custom formatter implementation
@@ -119,7 +119,7 @@ def formatter(root_path, manifest_file, **kwargs):  # pylint: disable=unused-arg
             items.append({"text":text, "audio_file":wav_file, "speaker_name":speaker_name, "root_path": root_path})
     return items
 
-add_formatter("custom_formatter_name", formatter) # Use the custom formatter name in the dataset config
+register_formatter("custom_formatter_name", formatter) # Use the custom formatter name in the dataset config
 dataset_config = BaseDatasetConfig(
     formatter="custom_formatter_name", meta_file_train="", language="en-us", path="dataset-path")
 )

tests/data_tests/test_dataset_formatters.py

@@ -2,7 +2,7 @@
 import unittest
 
 from tests import get_tests_input_path
-from TTS.tts.datasets.formatters import common_voice
+from TTS.tts.datasets.formatters import common_voice, register_formatter
 
 
 class TestTTSFormatters(unittest.TestCase):
@@ -17,12 +17,10 @@ def test_common_voice_preprocessor(self):  # pylint: disable=no-self-use
         assert items[-1]["audio_file"] == os.path.join(get_tests_input_path(), "clips", "common_voice_en_19737074.wav")
 
     def test_custom_formatter_with_existing_name(self):
-        from TTS.tts.datasets import add_formatter
-
         def custom_formatter(root_path, meta_file, ignored_speakers=None):
             return []
 
-        add_formatter("custom_formatter", custom_formatter)
+        register_formatter("custom_formatter", custom_formatter)
 
         with self.assertRaises(ValueError):
-            add_formatter("custom_formatter", custom_formatter)
+            register_formatter("custom_formatter", custom_formatter)

tests/data_tests/test_loader.py

@@ -8,7 +8,7 @@
 
 from tests import get_tests_data_path
 from TTS.tts.configs.shared_configs import BaseDatasetConfig, BaseTTSConfig
-from TTS.tts.datasets import add_formatter, load_tts_samples
+from TTS.tts.datasets import load_tts_samples, register_formatter
 from TTS.tts.datasets.dataset import TTSDataset
 from TTS.tts.utils.text.tokenizer import TTSTokenizer
 from TTS.utils.audio import AudioProcessor
@@ -268,8 +268,8 @@ def custom_formatter2(x, *args, **kwargs):
         [item.update({"audio_file": f"{item['audio_file']}.wav"}) for item in items]
         return items
 
-    add_formatter("custom_formatter1", custom_formatter)
-    add_formatter("custom_formatter2", custom_formatter2)
+    register_formatter("custom_formatter1", custom_formatter)
+    register_formatter("custom_formatter2", custom_formatter2)
     dataset1 = BaseDatasetConfig(
         formatter="custom_formatter1",
         meta_file_train="metadata.csv",