mirror of
https://github.com/denizsafak/abogen.git
synced 2026-07-18 13:40:27 +02:00
feat: add KokoroBackend implementing TTSBackend protocol
- Create KokoroBackend class implementing TTSBackend protocol - Move all KPipeline interaction inside KokoroBackend - Update LoadPipelineThread to create backend via create_backend() - Update ConversionThread and VoicePreviewThread to accept backend - Replace np_module/kpipeline_class parameters with single backend - Add 24 unit tests for KokoroBackend - KPipeline is now an internal implementation detail of KokoroBackend
This commit is contained in:
+28
-56
@@ -5,6 +5,7 @@ import hashlib # For generating unique cache filenames
|
|||||||
from platformdirs import user_desktop_dir
|
from platformdirs import user_desktop_dir
|
||||||
from PyQt6.QtCore import QThread, pyqtSignal, Qt, QTimer
|
from PyQt6.QtCore import QThread, pyqtSignal, Qt, QTimer
|
||||||
from PyQt6.QtWidgets import QCheckBox, QVBoxLayout, QDialog, QLabel, QDialogButtonBox
|
from PyQt6.QtWidgets import QCheckBox, QVBoxLayout, QDialog, QLabel, QDialogButtonBox
|
||||||
|
import numpy as np
|
||||||
import soundfile as sf
|
import soundfile as sf
|
||||||
from abogen.utils import (
|
from abogen.utils import (
|
||||||
create_process,
|
create_process,
|
||||||
@@ -259,8 +260,7 @@ class ConversionThread(QThread):
|
|||||||
output_folder,
|
output_folder,
|
||||||
subtitle_mode,
|
subtitle_mode,
|
||||||
output_format,
|
output_format,
|
||||||
np_module,
|
backend,
|
||||||
kpipeline_class,
|
|
||||||
start_time,
|
start_time,
|
||||||
total_char_count,
|
total_char_count,
|
||||||
use_gpu=True,
|
use_gpu=True,
|
||||||
@@ -270,8 +270,7 @@ class ConversionThread(QThread):
|
|||||||
super().__init__()
|
super().__init__()
|
||||||
self._chapter_options_event = threading.Event()
|
self._chapter_options_event = threading.Event()
|
||||||
self._timestamp_response_event = threading.Event()
|
self._timestamp_response_event = threading.Event()
|
||||||
self.np = np_module
|
self.backend = backend
|
||||||
self.KPipeline = kpipeline_class
|
|
||||||
self.file_name = file_name
|
self.file_name = file_name
|
||||||
self.lang_code = lang_code
|
self.lang_code = lang_code
|
||||||
self.speed = speed
|
self.speed = speed
|
||||||
@@ -490,19 +489,6 @@ class ConversionThread(QThread):
|
|||||||
|
|
||||||
self.log_updated.emit(("\nInitializing TTS pipeline...", "grey"))
|
self.log_updated.emit(("\nInitializing TTS pipeline...", "grey"))
|
||||||
|
|
||||||
# Set device based on use_gpu setting and platform
|
|
||||||
if self.use_gpu:
|
|
||||||
if platform.system() == "Darwin" and platform.processor() == "arm":
|
|
||||||
device = "mps" # Use MPS for Apple Silicon
|
|
||||||
else:
|
|
||||||
device = "cuda" # Use CUDA for other platforms
|
|
||||||
else:
|
|
||||||
device = "cpu"
|
|
||||||
|
|
||||||
tts = self.KPipeline(
|
|
||||||
lang_code=self.lang_code, repo_id="hexgrad/Kokoro-82M", device=device
|
|
||||||
)
|
|
||||||
|
|
||||||
# Check if the input is a subtitle file or timestamp text file
|
# Check if the input is a subtitle file or timestamp text file
|
||||||
is_subtitle_file = False
|
is_subtitle_file = False
|
||||||
is_timestamp_text = False
|
is_timestamp_text = False
|
||||||
@@ -538,7 +524,7 @@ class ConversionThread(QThread):
|
|||||||
|
|
||||||
# Process subtitle files separately
|
# Process subtitle files separately
|
||||||
if is_subtitle_file or is_timestamp_text:
|
if is_subtitle_file or is_timestamp_text:
|
||||||
self._process_subtitle_file(tts, base_path, is_timestamp_text)
|
self._process_subtitle_file(self.backend, base_path, is_timestamp_text)
|
||||||
return
|
return
|
||||||
|
|
||||||
if self.is_direct_text:
|
if self.is_direct_text:
|
||||||
@@ -1071,7 +1057,7 @@ class ConversionThread(QThread):
|
|||||||
for segment_idx, (voice_name, segment_text) in enumerate(voice_segments):
|
for segment_idx, (voice_name, segment_text) in enumerate(voice_segments):
|
||||||
# Load voice for this segment (with caching)
|
# Load voice for this segment (with caching)
|
||||||
try:
|
try:
|
||||||
loaded_voice = self.load_voice_cached(voice_name, tts)
|
loaded_voice = self.load_voice_cached(voice_name, self.backend)
|
||||||
if segment_idx > 0:
|
if segment_idx > 0:
|
||||||
voice_display = voice_name if len(voice_name) < 50 else voice_name[:47] + "..."
|
voice_display = voice_name if len(voice_name) < 50 else voice_name[:47] + "..."
|
||||||
self.log_updated.emit((f" → Voice: {voice_display}", "grey"))
|
self.log_updated.emit((f" → Voice: {voice_display}", "grey"))
|
||||||
@@ -1080,7 +1066,7 @@ class ConversionThread(QThread):
|
|||||||
(f"⚠ Voice loading error for '{voice_name}', continuing with previous", "orange")
|
(f"⚠ Voice loading error for '{voice_name}', continuing with previous", "orange")
|
||||||
)
|
)
|
||||||
if segment_idx == 0:
|
if segment_idx == 0:
|
||||||
loaded_voice = self.load_voice_cached(self.voice, tts)
|
loaded_voice = self.load_voice_cached(self.voice, self.backend)
|
||||||
|
|
||||||
# Determine if spaCy segmentation should be used for PRE-TTS segmentation
|
# Determine if spaCy segmentation should be used for PRE-TTS segmentation
|
||||||
# Only non-English languages use spaCy for pre-segmentation
|
# Only non-English languages use spaCy for pre-segmentation
|
||||||
@@ -1166,7 +1152,7 @@ class ConversionThread(QThread):
|
|||||||
print("Using split pattern: (unprintable)")
|
print("Using split pattern: (unprintable)")
|
||||||
|
|
||||||
for text_segment in text_segments:
|
for text_segment in text_segments:
|
||||||
for result in tts(
|
for result in self.backend(
|
||||||
text_segment,
|
text_segment,
|
||||||
voice=loaded_voice,
|
voice=loaded_voice,
|
||||||
speed=self.speed,
|
speed=self.speed,
|
||||||
@@ -1368,7 +1354,7 @@ class ConversionThread(QThread):
|
|||||||
silence_samples = int(
|
silence_samples = int(
|
||||||
self.silence_duration * 24000
|
self.silence_duration * 24000
|
||||||
) # Silence duration at 24,000 Hz
|
) # Silence duration at 24,000 Hz
|
||||||
silence_audio = self.np.zeros(silence_samples, dtype="float32")
|
silence_audio = np.zeros(silence_samples, dtype="float32")
|
||||||
silence_bytes = silence_audio.tobytes()
|
silence_bytes = silence_audio.tobytes()
|
||||||
|
|
||||||
if merged_out_file:
|
if merged_out_file:
|
||||||
@@ -1707,7 +1693,7 @@ class ConversionThread(QThread):
|
|||||||
max_end_time = max(
|
max_end_time = max(
|
||||||
(end for _, end, _ in subtitles if end is not None), default=0
|
(end for _, end, _ in subtitles if end is not None), default=0
|
||||||
)
|
)
|
||||||
audio_buffer = self.np.zeros(
|
audio_buffer = np.zeros(
|
||||||
int(max_end_time * rate) + rate, dtype="float32"
|
int(max_end_time * rate) + rate, dtype="float32"
|
||||||
)
|
)
|
||||||
|
|
||||||
@@ -1771,7 +1757,7 @@ class ConversionThread(QThread):
|
|||||||
# Generate TTS audio
|
# Generate TTS audio
|
||||||
tts_results = [
|
tts_results = [
|
||||||
r
|
r
|
||||||
for r in tts(
|
for r in self.backend(
|
||||||
processed_text,
|
processed_text,
|
||||||
voice=loaded_voice,
|
voice=loaded_voice,
|
||||||
speed=self.speed,
|
speed=self.speed,
|
||||||
@@ -1789,11 +1775,11 @@ class ConversionThread(QThread):
|
|||||||
|
|
||||||
# Concatenate audio and determine duration
|
# Concatenate audio and determine duration
|
||||||
full_audio = (
|
full_audio = (
|
||||||
self.np.concatenate(
|
np.concatenate(
|
||||||
[a.numpy() if hasattr(a, "numpy") else a for a in audio_chunks]
|
[a.numpy() if hasattr(a, "numpy") else a for a in audio_chunks]
|
||||||
)
|
)
|
||||||
if audio_chunks
|
if audio_chunks
|
||||||
else self.np.zeros(
|
else np.zeros(
|
||||||
int((subtitle_duration or 0) * rate), dtype="float32"
|
int((subtitle_duration or 0) * rate), dtype="float32"
|
||||||
)
|
)
|
||||||
)
|
)
|
||||||
@@ -1827,8 +1813,8 @@ class ConversionThread(QThread):
|
|||||||
num_stages = max(
|
num_stages = max(
|
||||||
1,
|
1,
|
||||||
int(
|
int(
|
||||||
self.np.ceil(
|
np.ceil(
|
||||||
self.np.log(speed_factor) / self.np.log(2.0)
|
np.log(speed_factor) / np.log(2.0)
|
||||||
)
|
)
|
||||||
),
|
),
|
||||||
)
|
)
|
||||||
@@ -1861,7 +1847,7 @@ class ConversionThread(QThread):
|
|||||||
stdout=subprocess.PIPE,
|
stdout=subprocess.PIPE,
|
||||||
stderr=subprocess.PIPE,
|
stderr=subprocess.PIPE,
|
||||||
)
|
)
|
||||||
full_audio = self.np.frombuffer(
|
full_audio = np.frombuffer(
|
||||||
speed_proc.communicate(input=full_audio.tobytes())[0],
|
speed_proc.communicate(input=full_audio.tobytes())[0],
|
||||||
dtype="float32",
|
dtype="float32",
|
||||||
)
|
)
|
||||||
@@ -1875,7 +1861,7 @@ class ConversionThread(QThread):
|
|||||||
|
|
||||||
tts_results = [
|
tts_results = [
|
||||||
r
|
r
|
||||||
for r in tts(
|
for r in self.backend(
|
||||||
processed_text,
|
processed_text,
|
||||||
voice=loaded_voice,
|
voice=loaded_voice,
|
||||||
speed=new_speed,
|
speed=new_speed,
|
||||||
@@ -1886,14 +1872,14 @@ class ConversionThread(QThread):
|
|||||||
audio_chunks = [r.audio for r in tts_results]
|
audio_chunks = [r.audio for r in tts_results]
|
||||||
|
|
||||||
full_audio = (
|
full_audio = (
|
||||||
self.np.concatenate(
|
np.concatenate(
|
||||||
[
|
[
|
||||||
a.numpy() if hasattr(a, "numpy") else a
|
a.numpy() if hasattr(a, "numpy") else a
|
||||||
for a in audio_chunks
|
for a in audio_chunks
|
||||||
]
|
]
|
||||||
)
|
)
|
||||||
if audio_chunks
|
if audio_chunks
|
||||||
else self.np.zeros(
|
else np.zeros(
|
||||||
int(subtitle_duration * rate), dtype="float32"
|
int(subtitle_duration * rate), dtype="float32"
|
||||||
)
|
)
|
||||||
)
|
)
|
||||||
@@ -1910,10 +1896,10 @@ class ConversionThread(QThread):
|
|||||||
# Pad or trim to subtitle duration
|
# Pad or trim to subtitle duration
|
||||||
target_samples = int(subtitle_duration * rate)
|
target_samples = int(subtitle_duration * rate)
|
||||||
if len(full_audio) < target_samples:
|
if len(full_audio) < target_samples:
|
||||||
full_audio = self.np.concatenate(
|
full_audio = np.concatenate(
|
||||||
[
|
[
|
||||||
full_audio,
|
full_audio,
|
||||||
self.np.zeros(
|
np.zeros(
|
||||||
target_samples - len(full_audio), dtype="float32"
|
target_samples - len(full_audio), dtype="float32"
|
||||||
),
|
),
|
||||||
]
|
]
|
||||||
@@ -1926,10 +1912,10 @@ class ConversionThread(QThread):
|
|||||||
end_sample = start_sample + len(full_audio)
|
end_sample = start_sample + len(full_audio)
|
||||||
if end_sample > len(audio_buffer):
|
if end_sample > len(audio_buffer):
|
||||||
# Extend buffer if needed
|
# Extend buffer if needed
|
||||||
audio_buffer = self.np.concatenate(
|
audio_buffer = np.concatenate(
|
||||||
[
|
[
|
||||||
audio_buffer,
|
audio_buffer,
|
||||||
self.np.zeros(
|
np.zeros(
|
||||||
end_sample - len(audio_buffer), dtype="float32"
|
end_sample - len(audio_buffer), dtype="float32"
|
||||||
),
|
),
|
||||||
]
|
]
|
||||||
@@ -1971,7 +1957,7 @@ class ConversionThread(QThread):
|
|||||||
self.progress_updated.emit(percent, etr_str)
|
self.progress_updated.emit(percent, etr_str)
|
||||||
|
|
||||||
# Normalize audio buffer to prevent clipping from mixed overlaps
|
# Normalize audio buffer to prevent clipping from mixed overlaps
|
||||||
max_amplitude = self.np.abs(audio_buffer).max()
|
max_amplitude = np.abs(audio_buffer).max()
|
||||||
if max_amplitude > 1.0:
|
if max_amplitude > 1.0:
|
||||||
self.log_updated.emit(
|
self.log_updated.emit(
|
||||||
f"\n -> Normalizing audio (peak: {max_amplitude:.2f})"
|
f"\n -> Normalizing audio (peak: {max_amplitude:.2f})"
|
||||||
@@ -2440,8 +2426,7 @@ class VoicePreviewThread(QThread):
|
|||||||
|
|
||||||
def __init__(
|
def __init__(
|
||||||
self,
|
self,
|
||||||
np_module,
|
backend,
|
||||||
kpipeline_class,
|
|
||||||
lang_code,
|
lang_code,
|
||||||
voice,
|
voice,
|
||||||
speed,
|
speed,
|
||||||
@@ -2449,8 +2434,7 @@ class VoicePreviewThread(QThread):
|
|||||||
parent=None,
|
parent=None,
|
||||||
):
|
):
|
||||||
super().__init__(parent)
|
super().__init__(parent)
|
||||||
self.np_module = np_module
|
self.backend = backend
|
||||||
self.kpipeline_class = kpipeline_class
|
|
||||||
self.lang_code = lang_code
|
self.lang_code = lang_code
|
||||||
self.voice = voice
|
self.voice = voice
|
||||||
self.speed = speed
|
self.speed = speed
|
||||||
@@ -2484,31 +2468,19 @@ class VoicePreviewThread(QThread):
|
|||||||
# Generate the preview and save to cache
|
# Generate the preview and save to cache
|
||||||
try:
|
try:
|
||||||
|
|
||||||
# Set device based on use_gpu setting and platform
|
|
||||||
if self.use_gpu:
|
|
||||||
if platform.system() == "Darwin" and platform.processor() == "arm":
|
|
||||||
device = "mps" # Use MPS for Apple Silicon
|
|
||||||
else:
|
|
||||||
device = "cuda" # Use CUDA for other platforms
|
|
||||||
else:
|
|
||||||
device = "cpu"
|
|
||||||
|
|
||||||
tts = self.kpipeline_class(
|
|
||||||
lang_code=self.lang_code, repo_id="hexgrad/Kokoro-82M", device=device
|
|
||||||
)
|
|
||||||
# Enable voice formula support for preview
|
# Enable voice formula support for preview
|
||||||
if "*" in self.voice:
|
if "*" in self.voice:
|
||||||
loaded_voice = get_new_voice(tts, self.voice, self.use_gpu)
|
loaded_voice = get_new_voice(self.backend, self.voice, self.use_gpu)
|
||||||
else:
|
else:
|
||||||
loaded_voice = self.voice
|
loaded_voice = self.voice
|
||||||
sample_text = get_sample_voice_text(self.lang_code)
|
sample_text = get_sample_voice_text(self.lang_code)
|
||||||
audio_segments = []
|
audio_segments = []
|
||||||
for result in tts(
|
for result in self.backend(
|
||||||
sample_text, voice=loaded_voice, speed=self.speed, split_pattern=None
|
sample_text, voice=loaded_voice, speed=self.speed, split_pattern=None
|
||||||
):
|
):
|
||||||
audio_segments.append(result.audio)
|
audio_segments.append(result.audio)
|
||||||
if audio_segments:
|
if audio_segments:
|
||||||
audio = self.np_module.concatenate(audio_segments)
|
audio = np.concatenate(audio_segments)
|
||||||
# Save directly to the cache path
|
# Save directly to the cache path
|
||||||
sf.write(self.cache_path, audio, 24000)
|
sf.write(self.cache_path, audio, 24000)
|
||||||
self.temp_wav = self.cache_path
|
self.temp_wav = self.cache_path
|
||||||
|
|||||||
+32
-11
@@ -2316,9 +2316,9 @@ class abogen(QWidget):
|
|||||||
file_size_str = "Unknown"
|
file_size_str = "Unknown"
|
||||||
|
|
||||||
# pipeline_loaded_callback remains unchanged
|
# pipeline_loaded_callback remains unchanged
|
||||||
def pipeline_loaded_callback(np_module, kpipeline_class, error):
|
def pipeline_loaded_callback(backend, error):
|
||||||
if error:
|
if error:
|
||||||
self.update_log((f"Error loading numpy or KPipeline: {error}", "red"))
|
self.update_log((f"Error loading TTS backend: {error}", "red"))
|
||||||
prevent_sleep_end()
|
prevent_sleep_end()
|
||||||
return
|
return
|
||||||
|
|
||||||
@@ -2341,8 +2341,7 @@ class abogen(QWidget):
|
|||||||
self.selected_output_folder,
|
self.selected_output_folder,
|
||||||
subtitle_mode=actual_subtitle_mode,
|
subtitle_mode=actual_subtitle_mode,
|
||||||
output_format=self.selected_format,
|
output_format=self.selected_format,
|
||||||
np_module=np_module,
|
backend=backend,
|
||||||
kpipeline_class=kpipeline_class,
|
|
||||||
start_time=self.start_time,
|
start_time=self.start_time,
|
||||||
total_char_count=self.char_count,
|
total_char_count=self.char_count,
|
||||||
use_gpu=self.gpu_ok,
|
use_gpu=self.gpu_ok,
|
||||||
@@ -2426,7 +2425,20 @@ class abogen(QWidget):
|
|||||||
self.gpu_ok = gpu_ok
|
self.gpu_ok = gpu_ok
|
||||||
self.update_log((gpu_msg, gpu_ok))
|
self.update_log((gpu_msg, gpu_ok))
|
||||||
self.update_log("Loading modules...")
|
self.update_log("Loading modules...")
|
||||||
load_thread = LoadPipelineThread(pipeline_loaded_callback)
|
|
||||||
|
# Determine device based on GPU availability
|
||||||
|
if gpu_ok:
|
||||||
|
if platform.system() == "Darwin" and platform.processor() == "arm":
|
||||||
|
device = "mps"
|
||||||
|
else:
|
||||||
|
device = "cuda"
|
||||||
|
else:
|
||||||
|
device = "cpu"
|
||||||
|
|
||||||
|
lang_code = self.selected_lang or "a"
|
||||||
|
load_thread = LoadPipelineThread(
|
||||||
|
pipeline_loaded_callback, lang_code=lang_code, device=device
|
||||||
|
)
|
||||||
load_thread.start()
|
load_thread.start()
|
||||||
|
|
||||||
threading.Thread(target=gpu_and_load, daemon=True).start()
|
threading.Thread(target=gpu_and_load, daemon=True).start()
|
||||||
@@ -2863,18 +2875,27 @@ class abogen(QWidget):
|
|||||||
)
|
)
|
||||||
self.loading_movie.start()
|
self.loading_movie.start()
|
||||||
|
|
||||||
def pipeline_loaded_callback(np_module, kpipeline_class, error):
|
# Determine device based on GPU availability
|
||||||
self._on_pipeline_loaded_for_preview(np_module, kpipeline_class, error)
|
if self.gpu_ok:
|
||||||
|
if platform.system() == "Darwin" and platform.processor() == "arm":
|
||||||
|
device = "mps"
|
||||||
|
else:
|
||||||
|
device = "cuda"
|
||||||
|
else:
|
||||||
|
device = "cpu"
|
||||||
|
|
||||||
load_thread = LoadPipelineThread(pipeline_loaded_callback)
|
lang = self.selected_lang or "a"
|
||||||
|
load_thread = LoadPipelineThread(
|
||||||
|
self._on_pipeline_loaded_for_preview, lang_code=lang, device=device
|
||||||
|
)
|
||||||
load_thread.start()
|
load_thread.start()
|
||||||
|
|
||||||
def _on_pipeline_loaded_for_preview(self, np_module, kpipeline_class, error):
|
def _on_pipeline_loaded_for_preview(self, backend, error):
|
||||||
# stop loading animation and restore icon on error
|
# stop loading animation and restore icon on error
|
||||||
if error:
|
if error:
|
||||||
self.loading_movie.stop()
|
self.loading_movie.stop()
|
||||||
self._show_error_message_box(
|
self._show_error_message_box(
|
||||||
"Loading Error", f"Error loading numpy or KPipeline: {error}"
|
"Loading Error", f"Error loading TTS backend: {error}"
|
||||||
)
|
)
|
||||||
self.btn_preview.setIcon(self.play_icon)
|
self.btn_preview.setIcon(self.play_icon)
|
||||||
self.btn_preview.setEnabled(True)
|
self.btn_preview.setEnabled(True)
|
||||||
@@ -2912,7 +2933,7 @@ class abogen(QWidget):
|
|||||||
gpu_msg, gpu_ok = get_gpu_acceleration(self.use_gpu)
|
gpu_msg, gpu_ok = get_gpu_acceleration(self.use_gpu)
|
||||||
|
|
||||||
self.preview_thread = VoicePreviewThread(
|
self.preview_thread = VoicePreviewThread(
|
||||||
np_module, kpipeline_class, lang, voice, speed, gpu_ok
|
backend, lang, voice, speed, gpu_ok
|
||||||
)
|
)
|
||||||
self.preview_thread.finished.connect(self._play_preview_audio)
|
self.preview_thread.finished.connect(self._play_preview_audio)
|
||||||
self.preview_thread.error.connect(self._preview_error)
|
self.preview_thread.error.connect(self._preview_error)
|
||||||
|
|||||||
+105
-18
@@ -1,31 +1,118 @@
|
|||||||
def load_numpy_kpipeline():
|
"""
|
||||||
|
Kokoro TTS Backend
|
||||||
|
|
||||||
|
Encapsulates the Kokoro KPipeline as a TTSBackend implementation.
|
||||||
|
"""
|
||||||
|
|
||||||
|
from __future__ import annotations
|
||||||
|
|
||||||
|
from typing import Any, Dict, Iterator, List, Optional
|
||||||
|
|
||||||
import numpy as np
|
import numpy as np
|
||||||
|
|
||||||
|
|
||||||
|
def _load_kpipeline():
|
||||||
|
"""Lazy-load Kokoro dependencies."""
|
||||||
from kokoro import KPipeline # type: ignore[import-not-found]
|
from kokoro import KPipeline # type: ignore[import-not-found]
|
||||||
|
|
||||||
return np, KPipeline
|
return KPipeline
|
||||||
|
|
||||||
|
|
||||||
def create_kokoro_backend(**kwargs):
|
class KokoroBackend:
|
||||||
"""Create a Kokoro TTS backend instance.
|
"""TTSBackend implementation wrapping the Kokoro KPipeline.
|
||||||
|
|
||||||
Args:
|
All interaction with KPipeline is encapsulated here.
|
||||||
lang_code: Language code (e.g. "a" for American English).
|
The rest of the project depends only on this class.
|
||||||
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(
|
def __init__(self, **kwargs: Any) -> None:
|
||||||
lang_code=kwargs["lang_code"],
|
lang_code = kwargs["lang_code"]
|
||||||
repo_id=kwargs.get("repo_id", "hexgrad/Kokoro-82M"),
|
repo_id = kwargs.get("repo_id", "hexgrad/Kokoro-82M")
|
||||||
device=kwargs.get("device", "cpu"),
|
device = kwargs.get("device", "cpu")
|
||||||
|
|
||||||
|
KPipeline = _load_kpipeline()
|
||||||
|
self._pipeline = KPipeline(
|
||||||
|
lang_code=lang_code,
|
||||||
|
repo_id=repo_id,
|
||||||
|
device=device,
|
||||||
|
)
|
||||||
|
self._lang_code = lang_code
|
||||||
|
|
||||||
|
@property
|
||||||
|
def metadata(self):
|
||||||
|
from abogen.tts_backend import TTSBackendMetadata
|
||||||
|
|
||||||
|
return TTSBackendMetadata(
|
||||||
|
id="kokoro",
|
||||||
|
name="Kokoro",
|
||||||
|
description="Kokoro TTS engine",
|
||||||
)
|
)
|
||||||
|
|
||||||
|
def __call__(
|
||||||
|
self,
|
||||||
|
text: str,
|
||||||
|
*,
|
||||||
|
voice: Any,
|
||||||
|
speed: float = 1.0,
|
||||||
|
split_pattern: Optional[str] = None,
|
||||||
|
) -> Iterator[Any]:
|
||||||
|
"""Delegate to KPipeline's __call__."""
|
||||||
|
return self._pipeline(
|
||||||
|
text,
|
||||||
|
voice=voice,
|
||||||
|
speed=speed,
|
||||||
|
split_pattern=split_pattern,
|
||||||
|
)
|
||||||
|
|
||||||
from abogen.tts_backend import TTSBackendMetadata
|
def load_single_voice(self, voice_name: str) -> Any:
|
||||||
from abogen.tts_backend_registry import register_backend
|
"""Load a single voice tensor. Used by voice formula system."""
|
||||||
|
return self._pipeline.load_single_voice(voice_name)
|
||||||
|
|
||||||
|
def synthesize(self, text: str, **kwargs: Any) -> bytes:
|
||||||
|
"""Synthesize speech from text. Returns raw audio bytes."""
|
||||||
|
voice = kwargs.get("voice", "")
|
||||||
|
speed = kwargs.get("speed", 1.0)
|
||||||
|
split_pattern = kwargs.get("split_pattern", None)
|
||||||
|
|
||||||
|
audio_parts: list[np.ndarray] = []
|
||||||
|
for segment in self(text, voice=voice, speed=speed, split_pattern=split_pattern):
|
||||||
|
audio = segment.audio
|
||||||
|
if hasattr(audio, "numpy"):
|
||||||
|
audio = audio.numpy()
|
||||||
|
audio_parts.append(np.asarray(audio, dtype="float32"))
|
||||||
|
|
||||||
|
if not audio_parts:
|
||||||
|
return b""
|
||||||
|
|
||||||
|
combined = np.concatenate(audio_parts).astype("float32", copy=False)
|
||||||
|
return combined.tobytes()
|
||||||
|
|
||||||
|
def get_available_voices(self) -> List[str]:
|
||||||
|
"""Return known Kokoro voice identifiers."""
|
||||||
|
from abogen.constants import VOICES_INTERNAL
|
||||||
|
|
||||||
|
return list(VOICES_INTERNAL)
|
||||||
|
|
||||||
|
def get_supported_formats(self) -> List[str]:
|
||||||
|
"""Kokoro outputs raw PCM float32 audio."""
|
||||||
|
return ["pcm_float32"]
|
||||||
|
|
||||||
|
def get_info(self) -> Dict[str, Any]:
|
||||||
|
return {
|
||||||
|
"id": "kokoro",
|
||||||
|
"name": "Kokoro",
|
||||||
|
"lang_code": self._lang_code,
|
||||||
|
}
|
||||||
|
|
||||||
|
|
||||||
|
def create_kokoro_backend(**kwargs: Any) -> KokoroBackend:
|
||||||
|
"""Factory callable registered with TTSBackendRegistry."""
|
||||||
|
return KokoroBackend(**kwargs)
|
||||||
|
|
||||||
|
|
||||||
|
# --- Registration ---
|
||||||
|
from abogen.tts_backend import TTSBackendMetadata # noqa: E402
|
||||||
|
from abogen.tts_backend_registry import register_backend # noqa: E402
|
||||||
|
|
||||||
register_backend(
|
register_backend(
|
||||||
metadata=TTSBackendMetadata(
|
metadata=TTSBackendMetadata(
|
||||||
|
|||||||
+9
-5
@@ -530,15 +530,19 @@ def prevent_sleep_end():
|
|||||||
|
|
||||||
|
|
||||||
class LoadPipelineThread(Thread):
|
class LoadPipelineThread(Thread):
|
||||||
def __init__(self, callback):
|
def __init__(self, callback, lang_code="a", device="cpu"):
|
||||||
super().__init__()
|
super().__init__()
|
||||||
self.callback = callback
|
self.callback = callback
|
||||||
|
self.lang_code = lang_code
|
||||||
|
self.device = device
|
||||||
|
|
||||||
def run(self):
|
def run(self):
|
||||||
try:
|
try:
|
||||||
from abogen.tts_backends.kokoro import load_numpy_kpipeline
|
from abogen.tts_backend_registry import create_backend
|
||||||
|
|
||||||
np_module, kpipeline_class = load_numpy_kpipeline()
|
backend = create_backend(
|
||||||
self.callback(np_module, kpipeline_class, None)
|
"kokoro", lang_code=self.lang_code, device=self.device
|
||||||
|
)
|
||||||
|
self.callback(backend, None)
|
||||||
except Exception as e:
|
except Exception as e:
|
||||||
self.callback(None, None, str(e))
|
self.callback(None, str(e))
|
||||||
|
|||||||
@@ -0,0 +1,216 @@
|
|||||||
|
"""Tests for KokoroBackend class."""
|
||||||
|
|
||||||
|
from __future__ import annotations
|
||||||
|
|
||||||
|
from dataclasses import dataclass
|
||||||
|
from typing import Any, Iterator, List
|
||||||
|
from unittest.mock import MagicMock, patch
|
||||||
|
|
||||||
|
import numpy as np
|
||||||
|
import pytest
|
||||||
|
|
||||||
|
from abogen.tts_backend import TTSBackendMetadata
|
||||||
|
|
||||||
|
|
||||||
|
# ---------------------------------------------------------------------------
|
||||||
|
# Helpers
|
||||||
|
# ---------------------------------------------------------------------------
|
||||||
|
|
||||||
|
@dataclass
|
||||||
|
class _FakeSegment:
|
||||||
|
graphemes: str
|
||||||
|
audio: Any # np.ndarray or torch-like tensor
|
||||||
|
|
||||||
|
|
||||||
|
class _FakePipeline:
|
||||||
|
"""Minimal mock for kokoro.KPipeline."""
|
||||||
|
|
||||||
|
def __init__(self, *, lang_code: str, repo_id: str, device: str):
|
||||||
|
self.lang_code = lang_code
|
||||||
|
self.repo_id = repo_id
|
||||||
|
self.device = device
|
||||||
|
self._voices: dict[str, np.ndarray] = {}
|
||||||
|
|
||||||
|
def __call__(
|
||||||
|
self,
|
||||||
|
text: str,
|
||||||
|
*,
|
||||||
|
voice: Any = "",
|
||||||
|
speed: float = 1.0,
|
||||||
|
split_pattern: str | None = None,
|
||||||
|
) -> Iterator[_FakeSegment]:
|
||||||
|
yield _FakeSegment(graphemes=text, audio=np.zeros(100, dtype="float32"))
|
||||||
|
|
||||||
|
def load_single_voice(self, name: str) -> np.ndarray:
|
||||||
|
if name not in self._voices:
|
||||||
|
self._voices[name] = np.ones((1, 256), dtype="float32") * 0.5
|
||||||
|
return self._voices[name]
|
||||||
|
|
||||||
|
|
||||||
|
def _make_backend(**kwargs: Any):
|
||||||
|
"""Create KokoroBackend with mocked KPipeline."""
|
||||||
|
with patch("abogen.tts_backends.kokoro._load_kpipeline") as load:
|
||||||
|
load.return_value = _FakePipeline
|
||||||
|
from abogen.tts_backends.kokoro import KokoroBackend
|
||||||
|
|
||||||
|
return KokoroBackend(**kwargs)
|
||||||
|
|
||||||
|
|
||||||
|
# ---------------------------------------------------------------------------
|
||||||
|
# Tests
|
||||||
|
# ---------------------------------------------------------------------------
|
||||||
|
|
||||||
|
class TestKokoroBackendMetadata:
|
||||||
|
def test_metadata_returns_tts_backend_metadata(self):
|
||||||
|
backend = _make_backend(lang_code="a")
|
||||||
|
meta = backend.metadata
|
||||||
|
assert isinstance(meta, TTSBackendMetadata)
|
||||||
|
|
||||||
|
def test_metadata_fields(self):
|
||||||
|
backend = _make_backend(lang_code="a")
|
||||||
|
meta = backend.metadata
|
||||||
|
assert meta.id == "kokoro"
|
||||||
|
assert meta.name == "Kokoro"
|
||||||
|
assert "Kokoro" in meta.description
|
||||||
|
|
||||||
|
|
||||||
|
class TestKokoroBackendInit:
|
||||||
|
def test_stores_lang_code(self):
|
||||||
|
backend = _make_backend(lang_code="b")
|
||||||
|
assert backend._lang_code == "b"
|
||||||
|
|
||||||
|
def test_default_repo_id(self):
|
||||||
|
with patch("abogen.tts_backends.kokoro._load_kpipeline") as load:
|
||||||
|
load.return_value = _FakePipeline
|
||||||
|
from abogen.tts_backends.kokoro import KokoroBackend
|
||||||
|
|
||||||
|
b = KokoroBackend(lang_code="a")
|
||||||
|
assert b._pipeline.repo_id == "hexgrad/Kokoro-82M"
|
||||||
|
|
||||||
|
def test_custom_repo_id(self):
|
||||||
|
backend = _make_backend(lang_code="a", repo_id="custom/repo")
|
||||||
|
assert backend._pipeline.repo_id == "custom/repo"
|
||||||
|
|
||||||
|
def test_default_device(self):
|
||||||
|
backend = _make_backend(lang_code="a")
|
||||||
|
assert backend._pipeline.device == "cpu"
|
||||||
|
|
||||||
|
def test_custom_device(self):
|
||||||
|
backend = _make_backend(lang_code="a", device="cuda")
|
||||||
|
assert backend._pipeline.device == "cuda"
|
||||||
|
|
||||||
|
|
||||||
|
class TestKokoroBackendCall:
|
||||||
|
def test_call_delegates_to_pipeline(self):
|
||||||
|
backend = _make_backend(lang_code="a")
|
||||||
|
results = list(backend("hello", voice="af_heart", speed=1.2, split_pattern=r"\n"))
|
||||||
|
assert len(results) == 1
|
||||||
|
assert results[0].graphemes == "hello"
|
||||||
|
|
||||||
|
def test_call_returns_iterator(self):
|
||||||
|
backend = _make_backend(lang_code="a")
|
||||||
|
result = backend("test", voice="af_heart")
|
||||||
|
assert hasattr(result, "__iter__")
|
||||||
|
|
||||||
|
def test_call_with_voice_tensor(self):
|
||||||
|
backend = _make_backend(lang_code="a")
|
||||||
|
voice_tensor = np.ones((1, 256), dtype="float32")
|
||||||
|
results = list(backend("test", voice=voice_tensor))
|
||||||
|
assert len(results) == 1
|
||||||
|
|
||||||
|
def test_call_default_speed(self):
|
||||||
|
backend = _make_backend(lang_code="a")
|
||||||
|
# Should not raise with default speed
|
||||||
|
list(backend("text", voice="af_heart"))
|
||||||
|
|
||||||
|
def test_call_default_split_pattern_is_none(self):
|
||||||
|
backend = _make_backend(lang_code="a")
|
||||||
|
# split_pattern defaults to None
|
||||||
|
list(backend("text", voice="af_heart"))
|
||||||
|
|
||||||
|
|
||||||
|
class TestLoadSingleVoice:
|
||||||
|
def test_load_single_voice_delegates(self):
|
||||||
|
backend = _make_backend(lang_code="a")
|
||||||
|
tensor = backend.load_single_voice("af_heart")
|
||||||
|
assert isinstance(tensor, np.ndarray)
|
||||||
|
assert tensor.shape == (1, 256)
|
||||||
|
|
||||||
|
def test_load_single_voice_caches(self):
|
||||||
|
backend = _make_backend(lang_code="a")
|
||||||
|
t1 = backend.load_single_voice("af_heart")
|
||||||
|
t2 = backend.load_single_voice("af_heart")
|
||||||
|
assert t1 is t2 # same object
|
||||||
|
|
||||||
|
|
||||||
|
class TestSynthesize:
|
||||||
|
def test_synthesize_returns_bytes(self):
|
||||||
|
backend = _make_backend(lang_code="a")
|
||||||
|
result = backend.synthesize("hello", voice="af_heart")
|
||||||
|
assert isinstance(result, bytes)
|
||||||
|
|
||||||
|
def test_synthesize_nonempty(self):
|
||||||
|
backend = _make_backend(lang_code="a")
|
||||||
|
result = backend.synthesize("hello", voice="af_heart")
|
||||||
|
assert len(result) > 0
|
||||||
|
|
||||||
|
def test_synthesize_with_speed(self):
|
||||||
|
backend = _make_backend(lang_code="a")
|
||||||
|
result = backend.synthesize("hello", voice="af_heart", speed=1.5)
|
||||||
|
assert isinstance(result, bytes)
|
||||||
|
|
||||||
|
def test_synthesize_empty_text(self):
|
||||||
|
backend = _make_backend(lang_code="a")
|
||||||
|
# Empty text produces no segments
|
||||||
|
result = backend.synthesize("", voice="af_heart")
|
||||||
|
assert isinstance(result, bytes)
|
||||||
|
|
||||||
|
|
||||||
|
class TestProtocolMethods:
|
||||||
|
def test_get_available_voices(self):
|
||||||
|
backend = _make_backend(lang_code="a")
|
||||||
|
voices = backend.get_available_voices()
|
||||||
|
assert isinstance(voices, list)
|
||||||
|
assert len(voices) > 0
|
||||||
|
assert all(isinstance(v, str) for v in voices)
|
||||||
|
|
||||||
|
def test_get_supported_formats(self):
|
||||||
|
backend = _make_backend(lang_code="a")
|
||||||
|
formats = backend.get_supported_formats()
|
||||||
|
assert "pcm_float32" in formats
|
||||||
|
|
||||||
|
def test_get_info(self):
|
||||||
|
backend = _make_backend(lang_code="a")
|
||||||
|
info = backend.get_info()
|
||||||
|
assert info["id"] == "kokoro"
|
||||||
|
assert info["name"] == "Kokoro"
|
||||||
|
assert info["lang_code"] == "a"
|
||||||
|
|
||||||
|
|
||||||
|
class TestRegistration:
|
||||||
|
def test_factory_creates_kokoro_backend(self):
|
||||||
|
from abogen.tts_backends.kokoro import create_kokoro_backend, KokoroBackend
|
||||||
|
|
||||||
|
with patch("abogen.tts_backends.kokoro._load_kpipeline") as load:
|
||||||
|
load.return_value = _FakePipeline
|
||||||
|
backend = create_kokoro_backend(lang_code="a")
|
||||||
|
assert isinstance(backend, KokoroBackend)
|
||||||
|
|
||||||
|
def test_registry_has_kokoro(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"
|
||||||
|
|
||||||
|
def test_registry_factory_returns_kokoro_backend(self):
|
||||||
|
import abogen.tts_backends # noqa: F401
|
||||||
|
from abogen.tts_backend_registry import _registry
|
||||||
|
from abogen.tts_backends.kokoro import KokoroBackend
|
||||||
|
|
||||||
|
factory = _registry._factories["kokoro"]
|
||||||
|
with patch("abogen.tts_backends.kokoro._load_kpipeline") as load:
|
||||||
|
load.return_value = _FakePipeline
|
||||||
|
backend = factory(lang_code="a")
|
||||||
|
assert isinstance(backend, KokoroBackend)
|
||||||
Reference in New Issue
Block a user