From aec3462f1f6d6fa067c869992e3cba2b4301066d Mon Sep 17 00:00:00 2001 From: Artem Akymenko Date: Sun, 19 Jul 2026 11:47:14 +0000 Subject: [PATCH] refactor: both UIs use shared tts_segments() from domain/conversion_pipeline.py MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - 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 --- abogen/domain/conversion_pipeline.py | 110 +++++++++++++++-------- abogen/pyqt/conversion.py | 102 +++++++-------------- abogen/webui/conversion_runner.py | 72 +++++---------- tests/test_domain_conversion_pipeline.py | 50 ++++++++++- 4 files changed, 176 insertions(+), 158 deletions(-) diff --git a/abogen/domain/conversion_pipeline.py b/abogen/domain/conversion_pipeline.py index 03fa93e..2fbea66 100644 --- a/abogen/domain/conversion_pipeline.py +++ b/abogen/domain/conversion_pipeline.py @@ -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, diff --git a/abogen/pyqt/conversion.py b/abogen/pyqt/conversion.py index 790bd53..85d1dcd 100644 --- a/abogen/pyqt/conversion.py +++ b/abogen/pyqt/conversion.py @@ -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) diff --git a/abogen/webui/conversion_runner.py b/abogen/webui/conversion_runner.py index 2aa2b9e..e6fc1d5 100644 --- a/abogen/webui/conversion_runner.py +++ b/abogen/webui/conversion_runner.py @@ -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] = [] diff --git a/tests/test_domain_conversion_pipeline.py b/tests/test_domain_conversion_pipeline.py index 66cd364..c455a6c 100644 --- a/tests/test_domain_conversion_pipeline.py +++ b/tests/test_domain_conversion_pipeline.py @@ -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