mirror of
https://github.com/denizsafak/abogen.git
synced 2026-07-18 05:40:26 +02:00
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
This commit is contained in:
Artem Akymenko
@@ -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()
|
||||
|
||||
|
||||
@@ -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,
|
||||
)
|
||||
|
||||
@@ -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,
|
||||
)
|
||||
|
||||
+149
-99
@@ -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)
|
||||
|
||||
Reference in New Issue
Block a user