refactor: both UIs use shared tts_segments() from domain/conversion_pipeline.py

- domain/conversion_pipeline.py: add tts_segments() for pre-normalized text;
  emit_text_segments() now delegates to tts_segments() internally
- pyqt/conversion.py: inner TTS loop replaced with tts_segments() iterator;
  removed FakeToken import (handled by domain)
- webui/conversion_runner.py: emit_text() inner loop replaced with
  tts_segments() iterator; removed FakeToken import
- Both UIs now share the same TTS emission logic — normalization + backend
  invocation + segment iteration + token extraction
- +3 tests for tts_segments (no-normalization, chunk_start, basic)
- 1188 tests pass
This commit is contained in:
Artem Akymenko
2026-07-19 11:47:14 +00:00
parent 1193185833
commit aec3462f1f
4 changed files with 176 additions and 158 deletions
+72 -38
View File
@@ -30,6 +30,75 @@ class SegmentResult:
tokens: List[Dict[str, Any]] = field(default_factory=list)
def tts_segments(
text: str,
*,
backend: Any,
voice: Any,
speed: float,
split_pattern: str,
current_time: float = 0.0,
) -> Iterator[SegmentResult]:
"""Invoke TTS backend on (already normalized) text and yield SegmentResults.
Use this when you've already normalized the text yourself (e.g. after
spaCy sentence segmentation). For raw text, use emit_text_segments() instead.
Args:
text: Already-normalized text to synthesize.
backend: TTS pipeline callable.
voice: Resolved voice.
speed: TTS speed multiplier.
split_pattern: Regex pattern for sentence splitting.
current_time: Current position in the audio timeline (seconds).
Yields:
SegmentResult for each non-empty TTS segment.
"""
segment_iter = backend(
text,
voice=voice,
speed=speed,
split_pattern=split_pattern,
)
chunk_start = current_time
for segment in segment_iter:
graphemes_raw = getattr(segment, "graphemes", "") or ""
graphemes = graphemes_raw.strip()
audio = to_float32(getattr(segment, "audio", None))
if audio.size == 0:
continue
duration = len(audio) / SAMPLE_RATE
tokens_list = getattr(segment, "tokens", [])
if not tokens_list and graphemes:
tokens_list = [FakeToken(graphemes, 0, duration)]
tokens = [
{
"start": chunk_start + (tok.start_ts or 0),
"end": chunk_start + (tok.end_ts or 0),
"text": tok.text,
"whitespace": tok.whitespace,
}
for tok in tokens_list
]
yield SegmentResult(
graphemes=graphemes,
audio=audio,
duration=duration,
chunk_start=chunk_start,
tokens=tokens,
)
chunk_start += duration
def emit_text_segments(
text: str,
*,
@@ -81,50 +150,15 @@ def emit_text_segments(
usage_counter=usage_counter,
)
segment_iter = backend(
yield from tts_segments(
normalized,
backend=backend,
voice=voice,
speed=speed,
split_pattern=split_pattern,
current_time=current_time,
)
chunk_start = current_time
for segment in segment_iter:
graphemes_raw = getattr(segment, "graphemes", "") or ""
graphemes = graphemes_raw.strip()
audio = to_float32(getattr(segment, "audio", None))
if audio.size == 0:
continue
duration = len(audio) / SAMPLE_RATE
# Extract tokens with timestamps
tokens_list = getattr(segment, "tokens", [])
if not tokens_list and graphemes:
tokens_list = [FakeToken(graphemes, 0, duration)]
tokens = [
{
"start": chunk_start + (tok.start_ts or 0),
"end": chunk_start + (tok.end_ts or 0),
"text": tok.text,
"whitespace": tok.whitespace,
}
for tok in tokens_list
]
yield SegmentResult(
graphemes=graphemes,
audio=audio,
duration=duration,
chunk_start=chunk_start,
tokens=tokens,
)
chunk_start += duration
def emit_text_to_sinks(
text: str,
+32 -70
View File
@@ -36,7 +36,7 @@ from abogen.domain.output_paths import (
)
from abogen.domain.audio_helpers import build_ffmpeg_command, to_float32
from abogen.domain.audio_sink import AudioSink, open_audio_sink
from abogen.domain.tokens import FakeToken
from abogen.domain.conversion_pipeline import tts_segments
from abogen.domain.audio_buffer import (
create_silence,
mix_audio,
@@ -936,8 +936,7 @@ class ConversionThread(QThread):
print("Using split pattern: (unprintable)")
for text_segment in text_segments:
# Normalize text through the shared pipeline
# (heteronym + pronunciation + apostrophe normalization)
# Normalize text before TTS
try:
text_segment = prepare_text_for_tts(
text_segment,
@@ -951,118 +950,81 @@ class ConversionThread(QThread):
(f"Warning: Text normalization failed: {exc}", "orange")
)
for result in self.backend(
for seg in tts_segments(
text_segment,
backend=self.backend,
voice=loaded_voice,
speed=self.speed,
split_pattern=active_split_pattern,
current_time=current_time,
):
# Print the result for debugging
# print(f"Result: {result}")
if self.cancel_requested:
sink_stack.close()
self.conversion_finished.emit("Cancelled", None)
return
current_segment += 1
grapheme_len = len(result.graphemes)
self.processed_char_count += grapheme_len
# Log progress with both character counts and the graphemes content
self.processed_char_count += len(seg.graphemes)
self.log_updated.emit(
f"\n{self.processed_char_count:,}/{self.total_char_count:,}: {result.graphemes}"
f"\n{self.processed_char_count:,}/{self.total_char_count:,}: {seg.graphemes}"
)
chunk_dur = len(result.audio) / rate
chunk_start = current_time
audio_np = to_float32(result.audio)
# Write audio directly to merged file ONLY if merging
# Write audio
if merge_chapters_at_end and merged_sink:
merged_sink.write(audio_np)
merged_sink.write(seg.audio)
if chapter_sink:
chapter_sink.write(audio_np)
chapter_sink.write(seg.audio)
# Subtitle logic
if self.subtitle_mode != "Disabled":
tokens_list = getattr(result, "tokens", [])
# Fallback for languages without token support (non-English)
# Create a single token representing the entire segment duration
if not tokens_list and result.graphemes:
tokens_list = [
FakeToken(result.graphemes, 0, chunk_dur)
]
tokens_with_timestamps = []
chapter_tokens_with_timestamps = []
# Process every token, regardless of text or timestamps
for tok in tokens_list:
tokens_with_timestamps.append(
{
"start": chunk_start + (tok.start_ts or 0),
"end": chunk_start + (tok.end_ts or 0),
"text": tok.text,
"whitespace": tok.whitespace,
}
)
if chapter_sink:
chapter_tokens_with_timestamps.append(
{
"start": chapter_current_time
+ (tok.start_ts or 0),
"end": chapter_current_time
+ (tok.end_ts or 0),
"text": tok.text,
"whitespace": tok.whitespace,
}
)
# Process tokens according to subtitle mode
# Global subtitle processing ONLY if merging
if self.subtitle_mode != "Disabled" and seg.tokens:
tokens_with_timestamps = list(seg.tokens)
chapter_tokens_with_timestamps = [
{
"start": chapter_current_time + (t["start"] - seg.chunk_start),
"end": chapter_current_time + (t["end"] - seg.chunk_start),
"text": t["text"],
"whitespace": t["whitespace"],
}
for t in seg.tokens
]
if merge_chapters_at_end:
# Incremental subtitle writing for merged output
new_entries = []
self._process_subtitle_tokens(
tokens_with_timestamps,
new_entries,
self.max_subtitle_words,
fallback_end_time=chunk_start + chunk_dur,
fallback_end_time=seg.chunk_start + seg.duration,
)
if merged_subtitle_writer:
for start, end, text in new_entries:
merged_subtitle_writer.write_entry(start, end, text)
# Per-chapter subtitle processing for both file and sink
if chapter_sink:
new_chapter_entries = []
self._process_subtitle_tokens(
chapter_tokens_with_timestamps,
new_chapter_entries,
self.max_subtitle_words,
fallback_end_time=chapter_current_time + chunk_dur,
fallback_end_time=chapter_current_time + seg.duration,
)
if chapter_subtitle_writer:
for start, end, text in new_chapter_entries:
chapter_subtitle_writer.write_entry(start, end, text)
# Update timing
if merge_chapters_at_end:
current_time += chunk_dur
if chapter_sink:
chapter_current_time += chunk_dur
else:
if chapter_sink:
chapter_current_time += chunk_dur
# Calculate percentage based on characters processed
current_time += seg.duration
if chapter_sink:
chapter_current_time += seg.duration
# Progress
percent = min(
int(
self.processed_char_count / self.total_char_count * 100
),
int(self.processed_char_count / self.total_char_count * 100),
99,
)
# Calculate ETR based on characters processed
etr_str = calc_etr_str(
time.time() - self.etr_start_time,
self.processed_char_count,
self.total_char_count,
)
# Update progress more frequently (after each result)
self.progress_updated.emit(percent, etr_str)
# Add silence between chapters for merged output (except after the last chapter)
+24 -48
View File
@@ -117,8 +117,8 @@ from abogen.domain.audio_buffer import (
SAMPLE_RATE,
)
from abogen.domain.audio_sink import AudioSink, open_audio_sink
from abogen.domain.tokens import FakeToken
from abogen.domain.pipeline_factory import PipelinePool
from abogen.domain.conversion_pipeline import tts_segments
from abogen.domain.voice_utils import resolve_voice_target as _resolve_voice_target
@@ -424,44 +424,29 @@ def run_conversion_job(job: Job) -> None:
if provider == "supertonic":
supertonic_pipeline = pipeline_pool.get("supertonic", job.language, job.use_gpu, job=job)
voice_name = _supertonic_voice_from_spec(voice_choice, getattr(job, "voice", "M1"))
segment_iter = supertonic_pipeline(
normalized,
voice=voice_name,
speed=float(speed_override if speed_override is not None else job.speed),
split_pattern=split_pattern,
total_steps=int(supertonic_steps_override if supertonic_steps_override is not None else getattr(job, "supertonic_total_steps", 5)),
)
backend = supertonic_pipeline
resolved_voice = voice_name
effective_speed = float(speed_override if speed_override is not None else job.speed)
else:
kokoro_backend = pipeline_pool.get("kokoro", job.language, job.use_gpu, job=job)
segment_iter = kokoro_backend(
normalized,
voice=voice_choice,
speed=float(speed_override if speed_override is not None else job.speed),
split_pattern=split_pattern,
)
backend = kokoro_backend
resolved_voice = voice_choice
effective_speed = float(speed_override if speed_override is not None else job.speed)
try:
# Accumulate tokens for subtitle processing (token-level grouping)
accumulated_tokens: List[dict] = []
for segment in segment_iter:
for seg in tts_segments(
normalized,
backend=backend,
voice=resolved_voice,
speed=effective_speed,
split_pattern=split_pattern,
current_time=current_time,
):
canceller()
graphemes_raw = getattr(segment, "graphemes", "") or ""
graphemes = graphemes_raw.strip()
audio = _to_float32(getattr(segment, "audio", None))
if audio.size == 0:
continue
local_segments += 1
if chapter_sink:
chapter_sink.write(audio)
if audio_sink:
audio_sink.write(audio)
duration = len(audio) / SAMPLE_RATE
chunk_start = current_time
processed_chars += len(graphemes)
processed_chars += len(seg.graphemes)
job.processed_characters = processed_chars
if job.total_characters:
job.progress = min(processed_chars / job.total_characters, 0.999)
@@ -473,30 +458,21 @@ def run_conversion_job(job: Job) -> None:
else:
job.progress = 0.0 if processed_chars == 0 else 0.999
preview_text = graphemes or (graphemes_raw[:80] if graphemes_raw else "[silence]")
preview_text = seg.graphemes or "[silence]"
prefix = f"{preview_prefix} · " if preview_prefix else ""
job.add_log(f"{prefix}{processed_chars:,}/{job.total_characters or ''}: {preview_text[:80]}")
# Accumulate tokens from this segment for subtitle processing
if chapter_sink:
chapter_sink.write(seg.audio)
if audio_sink:
audio_sink.write(seg.audio)
if subtitle_writer and audio_sink:
tokens_list = getattr(segment, "tokens", [])
# Fallback for languages without token support: create a single token
if not tokens_list and graphemes:
tokens_list = [FakeToken(graphemes, 0, duration)]
for tok in tokens_list:
accumulated_tokens.append({
"start": chunk_start + (tok.start_ts or 0),
"end": chunk_start + (tok.end_ts or 0),
"text": tok.text,
"whitespace": tok.whitespace,
})
accumulated_tokens.extend(seg.tokens)
if audio_sink:
current_time += duration
current_time += seg.duration
# Flush accumulated tokens through process_subtitle_tokens
if subtitle_writer and audio_sink and accumulated_tokens:
_use_spacy = job.subtitle_mode not in ("Disabled", "Line")
new_entries: List[tuple] = []
+48 -2
View File
@@ -1,4 +1,4 @@
"""Tests for domain/conversion_pipeline.py — emit_text_segments."""
"""Tests for domain/conversion_pipeline.py — tts_segments, emit_text_segments."""
from dataclasses import dataclass, field
from typing import Any, List, Optional
@@ -7,7 +7,7 @@ from unittest.mock import MagicMock
import numpy as np
import pytest
from abogen.domain.conversion_pipeline import emit_text_segments, SegmentResult
from abogen.domain.conversion_pipeline import tts_segments, emit_text_segments, SegmentResult
@dataclass
@@ -135,3 +135,49 @@ class TestEmitTextSegments:
assert seg.graphemes == "Test"
assert seg.duration == 2.0
assert len(seg.audio) == 48000
class TestTtsSegments:
def test_yields_segments_from_normalized_text(self):
audio = np.ones(24000, dtype="float32")
segments = [FakeSegment("Hello", audio)]
results = list(tts_segments(
"Already normalized text",
backend=make_backend(segments),
voice="A",
speed=1.0,
split_pattern=r"\s+",
))
assert len(results) == 1
assert results[0].graphemes == "Hello"
def test_no_normalization_performed(self):
"""tts_segments should NOT normalize — it passes text directly to backend."""
received_texts = []
def capture_backend(text, voice=None, speed=1.0, split_pattern=None):
received_texts.append(text)
yield FakeSegment("ok", np.ones(24000, dtype="float32"))
list(tts_segments(
"Raw unnormalized text",
backend=capture_backend,
voice="A",
speed=1.0,
split_pattern=r"\s+",
))
assert received_texts[0] == "Raw unnormalized text"
def test_chunk_start_increments(self):
audio = np.ones(24000, dtype="float32")
segments = [FakeSegment("A", audio), FakeSegment("B", audio)]
results = list(tts_segments(
"test",
backend=make_backend(segments),
voice="A",
speed=1.0,
split_pattern=r"\s+",
current_time=10.0,
))
assert results[0].chunk_start == 10.0
assert results[1].chunk_start == 11.0