Merge pull request #179 from k0sm0naft/refactor/backend-registry

Add TTS backend registry and automatic backend registration

abogen/tts_backend_registry.py

@@ -0,0 +1,71 @@
+"""
+TTS Backend Registry
+
+Provides a global registry for TTS backend factories.
+Backends register themselves with metadata and a factory callable.
+The registry is universal and does not know about backend constructors.
+"""
+
+from typing import Callable, Any
+
+from abogen.tts_backend import TTSBackend, TTSBackendMetadata
+
+
+class TTSBackendRegistry:
+    """Registry of TTS backend factories.
+
+    Stores metadata and factory callables for registered backends.
+    """
+
+    def __init__(self) -> None:
+        self._backends: dict[str, TTSBackendMetadata] = {}
+        self._factories: dict[str, Callable[..., TTSBackend]] = {}
+
+    def register(
+        self,
+        metadata: TTSBackendMetadata,
+        factory: Callable[..., TTSBackend],
+    ) -> None:
+        """Register a backend with its metadata and factory callable."""
+        self._backends[metadata.id] = metadata
+        self._factories[metadata.id] = factory
+
+    def list_backends(self) -> list[TTSBackendMetadata]:
+        """Return metadata for all registered backends."""
+        return list(self._backends.values())
+
+    def get_metadata(self, backend_id: str) -> TTSBackendMetadata:
+        """Get metadata for a specific backend.
+
+        Raises:
+            KeyError: If backend with given id is not registered.
+        """
+        if backend_id not in self._backends:
+            raise KeyError(f"Unknown backend: {backend_id}")
+        return self._backends[backend_id]
+
+    def create_backend(self, backend_id: str, **kwargs: Any) -> TTSBackend:
+        """Create a backend instance by id.
+
+        Raises:
+            KeyError: If backend with given id is not registered.
+        """
+        if backend_id not in self._factories:
+            raise KeyError(f"Unknown backend: {backend_id}")
+        return self._factories[backend_id](**kwargs)
+
+
+_registry = TTSBackendRegistry()
+
+
+def register_backend(
+    metadata: TTSBackendMetadata,
+    factory: Callable[..., TTSBackend],
+) -> None:
+    """Register a TTS backend in the global registry."""
+    _registry.register(metadata, factory)
+
+
+def create_backend(backend_id: str, **kwargs: Any) -> TTSBackend:
+    """Create a TTS backend instance by provider id."""
+    return _registry.create_backend(backend_id, **kwargs)

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()
+

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,
+)

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,
+)

tests/test_tts_backend.py

@@ -1,29 +1,149 @@
-from dataclasses import dataclass
-
-from abogen.tts_backend import TTSBackendMetadata
-
-
-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"
+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)