From 57fec453e2384de1b95683e08ee1ba077857833b Mon Sep 17 00:00:00 2001 From: Artem Akymenko Date: Mon, 6 Jul 2026 09:02:40 +0000 Subject: [PATCH] feat: auto-register existing TTS backends - Add create_kokoro_backend() factory in kokoro.py - Add create_supertonic_backend() factory in supertonic.py - Auto-discover backend modules in __init__.py via pkgutil - Both backends register themselves on import - Tests verify registration and factory callables --- abogen/tts_backends/__init__.py | 20 +++ abogen/tts_backends/kokoro.py | 42 ++++- abogen/tts_backends/supertonic.py | 31 ++++ tests/test_tts_backend.py | 248 ++++++++++++++++++------------ 4 files changed, 237 insertions(+), 104 deletions(-) diff --git a/abogen/tts_backends/__init__.py b/abogen/tts_backends/__init__.py index e69de29..71982f2 100644 --- a/abogen/tts_backends/__init__.py +++ b/abogen/tts_backends/__init__.py @@ -0,0 +1,20 @@ +"""TTS backends package. + +Backend modules are auto-discovered and imported here. +Each backend module registers itself with the global registry +when imported. +""" + +import importlib +import pkgutil + + +def _discover_backends(): + """Import all modules in this package to trigger their registration.""" + package = __name__ + for _importer, modname, _ispkg in pkgutil.iter_modules(path=__path__): + importlib.import_module(f"{package}.{modname}") + + +_discover_backends() + diff --git a/abogen/tts_backends/kokoro.py b/abogen/tts_backends/kokoro.py index 7e3ab5a..8edfc5c 100644 --- a/abogen/tts_backends/kokoro.py +++ b/abogen/tts_backends/kokoro.py @@ -1,5 +1,37 @@ -def load_numpy_kpipeline(): - import numpy as np - from kokoro import KPipeline # type: ignore[import-not-found] - - return np, KPipeline +def load_numpy_kpipeline(): + import numpy as np + from kokoro import KPipeline # type: ignore[import-not-found] + + return np, KPipeline + + +def create_kokoro_backend(**kwargs): + """Create a Kokoro TTS backend instance. + + Args: + lang_code: Language code (e.g. "a" for American English). + repo_id: HuggingFace repo id. Defaults to "hexgrad/Kokoro-82M". + device: Device to use ("cpu", "cuda", "mps"). Defaults to "cpu". + + Returns: + KPipeline instance. + """ + _np, KPipeline = load_numpy_kpipeline() + return KPipeline( + lang_code=kwargs["lang_code"], + repo_id=kwargs.get("repo_id", "hexgrad/Kokoro-82M"), + device=kwargs.get("device", "cpu"), + ) + + +from abogen.tts_backend import TTSBackendMetadata +from abogen.tts_backend_registry import register_backend + +register_backend( + metadata=TTSBackendMetadata( + id="kokoro", + name="Kokoro", + description="Kokoro TTS engine", + ), + factory=create_kokoro_backend, +) diff --git a/abogen/tts_backends/supertonic.py b/abogen/tts_backends/supertonic.py index ed37c81..f225cf6 100644 --- a/abogen/tts_backends/supertonic.py +++ b/abogen/tts_backends/supertonic.py @@ -273,3 +273,34 @@ class SupertonicPipeline: audio = _resample_linear(audio, src_rate, self.sample_rate) yield SupertonicSegment(graphemes=chunk_to_speak, audio=audio) + + +def create_supertonic_backend(**kwargs): + """Create a SuperTonic TTS backend instance. + + Args: + sample_rate: Audio sample rate. Defaults to 24000. + auto_download: Auto-download models. Defaults to True. + total_steps: Inference steps. Defaults to 5. + + Returns: + SupertonicPipeline instance. + """ + return SupertonicPipeline( + sample_rate=kwargs.get("sample_rate", 24000), + auto_download=kwargs.get("auto_download", True), + total_steps=kwargs.get("total_steps", 5), + ) + + +from abogen.tts_backend import TTSBackendMetadata +from abogen.tts_backend_registry import register_backend + +register_backend( + metadata=TTSBackendMetadata( + id="supertonic", + name="SuperTonic", + description="SuperTonic TTS engine", + ), + factory=create_supertonic_backend, +) diff --git a/tests/test_tts_backend.py b/tests/test_tts_backend.py index fc777a4..c6cf4c1 100644 --- a/tests/test_tts_backend.py +++ b/tests/test_tts_backend.py @@ -1,99 +1,149 @@ -from dataclasses import dataclass - -from abogen.tts_backend import TTSBackendMetadata -from abogen.tts_backend_registry import TTSBackendRegistry - - -class TestTTSBackendMetadata: - def test_is_frozen_dataclass(self): - assert dataclass(TTSBackendMetadata) - - def test_fields_are_present(self): - meta = TTSBackendMetadata( - id="test", - name="Test Backend", - description="A test backend", - ) - assert meta.id == "test" - assert meta.name == "Test Backend" - assert meta.description == "A test backend" - - def test_is_immutable(self): - import pytest - - meta = TTSBackendMetadata( - id="kokoro", - name="Kokoro", - description="Test", - ) - with pytest.raises(Exception): - meta.id = "changed" - - -class TestTTSBackendRegistry: - def test_register_and_list(self): - registry = TTSBackendRegistry() - meta = TTSBackendMetadata(id="a", name="A", description="Backend A") - registry.register(metadata=meta, factory=lambda: None) - - backends = registry.list_backends() - assert len(backends) == 1 - assert backends[0].id == "a" - - def test_list_multiple(self): - registry = TTSBackendRegistry() - meta_a = TTSBackendMetadata(id="a", name="A", description="A") - meta_b = TTSBackendMetadata(id="b", name="B", description="B") - registry.register(metadata=meta_a, factory=lambda: None) - registry.register(metadata=meta_b, factory=lambda: None) - - backends = registry.list_backends() - ids = [b.id for b in backends] - assert "a" in ids - assert "b" in ids - - def test_get_metadata(self): - registry = TTSBackendRegistry() - meta = TTSBackendMetadata(id="x", name="X", description="X backend") - registry.register(metadata=meta, factory=lambda: None) - - result = registry.get_metadata("x") - assert result.id == "x" - assert result.name == "X" - - def test_get_metadata_unknown_raises(self): - import pytest - - registry = TTSBackendRegistry() - with pytest.raises(KeyError, match="Unknown backend: nope"): - registry.get_metadata("nope") - - def test_create_backend(self): - registry = TTSBackendRegistry() - meta = TTSBackendMetadata(id="test", name="Test", description="Test backend") - - def factory(**kwargs): - return {"created": True, "kwargs": kwargs} - - registry.register(metadata=meta, factory=factory) - result = registry.create_backend("test", foo="bar") - - assert result == {"created": True, "kwargs": {"foo": "bar"}} - - def test_create_backend_unknown_raises(self): - import pytest - - registry = TTSBackendRegistry() - with pytest.raises(KeyError, match="Unknown backend: missing"): - registry.create_backend("missing") - - def test_register_overwrites(self): - registry = TTSBackendRegistry() - meta1 = TTSBackendMetadata(id="x", name="V1", description="First") - meta2 = TTSBackendMetadata(id="x", name="V2", description="Second") - registry.register(metadata=meta1, factory=lambda: "v1") - registry.register(metadata=meta2, factory=lambda: "v2") - - result = registry.get_metadata("x") - assert result.name == "V2" - assert registry.create_backend("x") == "v2" +from dataclasses import dataclass + +from abogen.tts_backend import TTSBackendMetadata +from abogen.tts_backend_registry import TTSBackendRegistry + + +class TestTTSBackendMetadata: + def test_is_frozen_dataclass(self): + assert dataclass(TTSBackendMetadata) + + def test_fields_are_present(self): + meta = TTSBackendMetadata( + id="test", + name="Test Backend", + description="A test backend", + ) + assert meta.id == "test" + assert meta.name == "Test Backend" + assert meta.description == "A test backend" + + def test_is_immutable(self): + import pytest + + meta = TTSBackendMetadata( + id="kokoro", + name="Kokoro", + description="Test", + ) + with pytest.raises(Exception): + meta.id = "changed" + + +class TestTTSBackendRegistry: + def test_register_and_list(self): + registry = TTSBackendRegistry() + meta = TTSBackendMetadata(id="a", name="A", description="Backend A") + registry.register(metadata=meta, factory=lambda: None) + + backends = registry.list_backends() + assert len(backends) == 1 + assert backends[0].id == "a" + + def test_list_multiple(self): + registry = TTSBackendRegistry() + meta_a = TTSBackendMetadata(id="a", name="A", description="A") + meta_b = TTSBackendMetadata(id="b", name="B", description="B") + registry.register(metadata=meta_a, factory=lambda: None) + registry.register(metadata=meta_b, factory=lambda: None) + + backends = registry.list_backends() + ids = [b.id for b in backends] + assert "a" in ids + assert "b" in ids + + def test_get_metadata(self): + registry = TTSBackendRegistry() + meta = TTSBackendMetadata(id="x", name="X", description="X backend") + registry.register(metadata=meta, factory=lambda: None) + + result = registry.get_metadata("x") + assert result.id == "x" + assert result.name == "X" + + def test_get_metadata_unknown_raises(self): + import pytest + + registry = TTSBackendRegistry() + with pytest.raises(KeyError, match="Unknown backend: nope"): + registry.get_metadata("nope") + + def test_create_backend(self): + registry = TTSBackendRegistry() + meta = TTSBackendMetadata(id="test", name="Test", description="Test backend") + + def factory(**kwargs): + return {"created": True, "kwargs": kwargs} + + registry.register(metadata=meta, factory=factory) + result = registry.create_backend("test", foo="bar") + + assert result == {"created": True, "kwargs": {"foo": "bar"}} + + def test_create_backend_unknown_raises(self): + import pytest + + registry = TTSBackendRegistry() + with pytest.raises(KeyError, match="Unknown backend: missing"): + registry.create_backend("missing") + + def test_register_overwrites(self): + registry = TTSBackendRegistry() + meta1 = TTSBackendMetadata(id="x", name="V1", description="First") + meta2 = TTSBackendMetadata(id="x", name="V2", description="Second") + registry.register(metadata=meta1, factory=lambda: "v1") + registry.register(metadata=meta2, factory=lambda: "v2") + + result = registry.get_metadata("x") + assert result.name == "V2" + assert registry.create_backend("x") == "v2" + + +class TestBackendRegistration: + """Tests that existing backends are auto-registered.""" + + def test_import_triggers_registration(self): + import abogen.tts_backends # noqa: F401 + + from abogen.tts_backend_registry import _registry + + backends = _registry.list_backends() + ids = [b.id for b in backends] + assert "kokoro" in ids + assert "supertonic" in ids + + def test_kokoro_metadata(self): + import abogen.tts_backends # noqa: F401 + + from abogen.tts_backend_registry import _registry + + meta = _registry.get_metadata("kokoro") + assert meta.id == "kokoro" + assert meta.name == "Kokoro" + assert "Kokoro" in meta.description + + def test_supertonic_metadata(self): + import abogen.tts_backends # noqa: F401 + + from abogen.tts_backend_registry import _registry + + meta = _registry.get_metadata("supertonic") + assert meta.id == "supertonic" + assert meta.name == "SuperTonic" + assert "SuperTonic" in meta.description + + def test_kokoro_factory_callable(self): + import abogen.tts_backends # noqa: F401 + + from abogen.tts_backend_registry import _registry + + factory = _registry._factories["kokoro"] + assert callable(factory) + + def test_supertonic_factory_callable(self): + import abogen.tts_backends # noqa: F401 + + from abogen.tts_backend_registry import _registry + + factory = _registry._factories["supertonic"] + assert callable(factory)