From f0b6976d1235a1d5cebc779a18a1856679bd5d74 Mon Sep 17 00:00:00 2001 From: JB Date: Thu, 9 Oct 2025 13:37:36 -0700 Subject: [PATCH] feat: Enhance voice formula parsing and validation, implement voice asset caching, and add tests for new functionality --- abogen/voice_cache.py | 126 ++++++++++++++++++++++++++++++++ abogen/voice_formulas.py | 68 +++++++++++------ abogen/web/conversion_runner.py | 70 +++++++++++++++++- abogen/web/routes.py | 48 +++++++----- abogen/web/service.py | 19 +++++ tests/test_prepare_form.py | 61 ++++++++++++++++ tests/test_voice_cache.py | 56 ++++++++++++++ 7 files changed, 404 insertions(+), 44 deletions(-) create mode 100644 abogen/voice_cache.py create mode 100644 tests/test_prepare_form.py create mode 100644 tests/test_voice_cache.py diff --git a/abogen/voice_cache.py b/abogen/voice_cache.py new file mode 100644 index 0000000..9642a9f --- /dev/null +++ b/abogen/voice_cache.py @@ -0,0 +1,126 @@ +from __future__ import annotations + +import threading +from typing import Callable, Dict, Iterable, Optional, Set, Tuple + +try: # pragma: no cover - optional dependency guard + from huggingface_hub import hf_hub_download # type: ignore + from huggingface_hub.utils import LocalEntryNotFoundError # type: ignore +except Exception: # pragma: no cover - import fallback + hf_hub_download = None # type: ignore[assignment] + LocalEntryNotFoundError = None # type: ignore[assignment] + +from abogen.constants import VOICES_INTERNAL + +_CACHE_LOCK = threading.Lock() +_CACHED_VOICES: Set[str] = set() +_BOOTSTRAP_LOCK = threading.Lock() +_BOOTSTRAPPED = False + + +def _normalize_targets(voices: Optional[Iterable[str]]) -> Set[str]: + if not voices: + return set(VOICES_INTERNAL) + normalized: Set[str] = set() + for voice in voices: + if not voice: + continue + voice_id = str(voice).strip() + if not voice_id: + continue + if voice_id in VOICES_INTERNAL: + normalized.add(voice_id) + return normalized + + +def ensure_voice_assets( + voices: Optional[Iterable[str]] = None, + *, + repo_id: str = "hexgrad/Kokoro-82M", + cache_dir: Optional[str] = None, + on_progress: Optional[Callable[[str], None]] = None, +) -> Tuple[Set[str], Dict[str, str]]: + """Ensure Kokoro voice weight files are present locally. + + Returns a tuple of (downloaded voices, errors) where errors maps the + voice id to the underlying exception message. + """ + + if hf_hub_download is None: + raise RuntimeError("huggingface_hub is required to cache voices") + + targets = _normalize_targets(voices) + if not targets: + return set(), {} + + with _CACHE_LOCK: + missing = [voice for voice in targets if voice not in _CACHED_VOICES] + + downloaded: Set[str] = set() + errors: Dict[str, str] = {} + + for voice_id in missing: + if on_progress: + on_progress(f"Fetching voice asset '{voice_id}'") + try: + downloaded_flag = _ensure_single_voice_asset( + voice_id, + repo_id=repo_id, + cache_dir=cache_dir, + ) + except Exception as exc: # pragma: no cover - network variance + errors[voice_id] = str(exc) + continue + + if downloaded_flag: + downloaded.add(voice_id) + with _CACHE_LOCK: + _CACHED_VOICES.add(voice_id) + + return downloaded, errors + + +def bootstrap_voice_cache( + voices: Optional[Iterable[str]] = None, + *, + repo_id: str = "hexgrad/Kokoro-82M", + cache_dir: Optional[str] = None, + on_progress: Optional[Callable[[str], None]] = None, +) -> Tuple[Set[str], Dict[str, str]]: + """Ensure voices are cached once per process. + + Subsequent calls are no-ops and return empty structures. + """ + + global _BOOTSTRAPPED + with _BOOTSTRAP_LOCK: + if _BOOTSTRAPPED: + return set(), {} + downloaded, errors = ensure_voice_assets( + voices, + repo_id=repo_id, + cache_dir=cache_dir, + on_progress=on_progress, + ) + _BOOTSTRAPPED = True + return downloaded, errors + + +def _ensure_single_voice_asset( + voice_id: str, + *, + repo_id: str, + cache_dir: Optional[str], +) -> bool: + if hf_hub_download is None: + raise RuntimeError("huggingface_hub is required to cache voices") + + filename = f"voices/{voice_id}.pt" + + hf_hub_download( + repo_id=repo_id, + filename=filename, + cache_dir=cache_dir, + resume_download=True, + ) + return True \ No newline at end of file diff --git a/abogen/voice_formulas.py b/abogen/voice_formulas.py index 813d27c..b91cffe 100644 --- a/abogen/voice_formulas.py +++ b/abogen/voice_formulas.py @@ -1,4 +1,6 @@ import re +from typing import List, Tuple + from abogen.constants import VOICES_INTERNAL @@ -15,38 +17,56 @@ def get_new_voice(pipeline, formula, use_gpu): raise ValueError(f"Failed to create voice: {str(e)}") -# Parse the formula and get the combined voice tensor -def parse_voice_formula(pipeline, formula): - if not formula.strip(): +def parse_formula_terms(formula: str) -> List[Tuple[str, float]]: + if not formula or not formula.strip(): raise ValueError("Empty voice formula") - # Initialize the weighted sum - weighted_sum = None - - total_weight = calculate_sum_from_formula(formula) - - # Split the formula into terms - voices = formula.split("+") - - for term in voices: - # Parse each term (format: "voice_name*0.333") - voice_name, weight = term.strip().split("*") - weight = float(weight.strip()) - # normalize the weight - weight /= total_weight if total_weight > 0 else 1.0 + terms: List[Tuple[str, float]] = [] + for segment in formula.split("+"): + part = segment.strip() + if not part: + continue + if "*" not in part: + raise ValueError("Each component must be in the form voice*weight") + voice_name, raw_weight = part.split("*", 1) voice_name = voice_name.strip() - - # Get the voice tensor if voice_name not in VOICES_INTERNAL: raise ValueError(f"Unknown voice: {voice_name}") + try: + weight = float(raw_weight.strip()) + except ValueError as exc: + raise ValueError(f"Invalid weight for {voice_name}") from exc + if weight <= 0: + raise ValueError(f"Weight for {voice_name} must be positive") + terms.append((voice_name, weight)) + + if not terms: + raise ValueError("Voice weights must sum to a positive value") + + return terms + + +def parse_voice_formula(pipeline, formula): + terms = parse_formula_terms(formula) + + total_weight = sum(weight for _, weight in terms) + if total_weight <= 0: + raise ValueError("Voice weights must sum to a positive value") + + weighted_sum = None + + for voice_name, weight in terms: + normalized_weight = weight / total_weight if total_weight > 0 else weight voice_tensor = pipeline.load_single_voice(voice_name) - # Add to weighted sum if weighted_sum is None: - weighted_sum = weight * voice_tensor + weighted_sum = normalized_weight * voice_tensor else: - weighted_sum += weight * voice_tensor + weighted_sum += normalized_weight * voice_tensor + + if weighted_sum is None: + raise ValueError("Voice formula produced no components") return weighted_sum @@ -55,3 +75,7 @@ def calculate_sum_from_formula(formula): weights = re.findall(r"\* *([\d.]+)", formula) total_sum = sum(float(weight) for weight in weights) return total_sum + + +def extract_voice_ids(formula: str) -> List[str]: + return [voice for voice, _ in parse_formula_terms(formula)] diff --git a/abogen/web/conversion_runner.py b/abogen/web/conversion_runner.py index 06d85e0..5701d16 100644 --- a/abogen/web/conversion_runner.py +++ b/abogen/web/conversion_runner.py @@ -12,12 +12,13 @@ from collections import defaultdict from contextlib import ExitStack from dataclasses import dataclass from pathlib import Path -from typing import Any, Callable, Dict, Iterable, List, Optional, cast +from typing import Any, Callable, Dict, Iterable, List, Optional, Set, cast import numpy as np import soundfile as sf import static_ffmpeg +from abogen.constants import VOICES_INTERNAL from abogen.epub3.exporter import build_epub3_package from abogen.kokoro_text_normalization import ( ApostropheConfig, @@ -36,7 +37,8 @@ from abogen.utils import ( load_config, load_numpy_kpipeline, ) -from abogen.voice_formulas import get_new_voice +from abogen.voice_cache import ensure_voice_assets +from abogen.voice_formulas import extract_voice_ids, get_new_voice from .service import Job, JobStatus @@ -69,6 +71,69 @@ def _coerce_truthy(value: Any, default: bool = True) -> bool: return bool(value) +def _spec_to_voice_ids(spec: Any) -> Set[str]: + text = str(spec or "").strip() + if not text: + return set() + if "*" in text: + try: + return set(extract_voice_ids(text)) + except ValueError: + return set() + if text in VOICES_INTERNAL: + return {text} + return set() + + +def _collect_required_voice_ids(job: Job) -> Set[str]: + voices: Set[str] = set() + voices.update(_spec_to_voice_ids(job.voice)) + + for chapter in getattr(job, "chapters", []) or []: + if not isinstance(chapter, dict): + continue + for key in ("resolved_voice", "voice_formula", "voice"): + voices.update(_spec_to_voice_ids(chapter.get(key))) + + for chunk in getattr(job, "chunks", []) or []: + if not isinstance(chunk, dict): + continue + for key in ("resolved_voice", "voice_formula", "voice"): + voices.update(_spec_to_voice_ids(chunk.get(key))) + + speakers = getattr(job, "speakers", {}) + if isinstance(speakers, dict): + for payload in speakers.values() or []: + if not isinstance(payload, dict): + continue + for key in ("resolved_voice", "voice_formula", "voice"): + voices.update(_spec_to_voice_ids(payload.get(key))) + + voices.update(VOICES_INTERNAL) + return voices + + +def _initialize_voice_cache(job: Job) -> None: + try: + targets = _collect_required_voice_ids(job) + downloaded, errors = ensure_voice_assets( + targets, + on_progress=lambda message: job.add_log(message, level="debug"), + ) + except RuntimeError as exc: + job.add_log(f"Voice cache unavailable: {exc}", level="warning") + return + + if downloaded: + job.add_log( + f"Cached {len(downloaded)} voice asset{'s' if len(downloaded) != 1 else ''} locally.", + level="info", + ) + + for voice_id, error in errors.items(): + job.add_log(f"Failed to cache voice '{voice_id}': {error}", level="warning") + + _SIGNIFICANT_LENGTH_THRESHOLDS: Dict[str, int] = {"epub": 1000, "markdown": 500} _MIN_SHORT_CONTENT: Dict[str, int] = {"epub": 240, "markdown": 160} _STRUCTURAL_KEYWORDS = ( @@ -631,6 +696,7 @@ def run_conversion_job(job: Job) -> None: active_chapter_configs: List[Dict[str, Any]] = [] try: pipeline = _load_pipeline(job) + _initialize_voice_cache(job) extraction = extract_from_path(job.stored_path) file_type = _infer_file_type(job.stored_path) diff --git a/abogen/web/routes.py b/abogen/web/routes.py index 9f761bc..3fdb234 100644 --- a/abogen/web/routes.py +++ b/abogen/web/routes.py @@ -58,7 +58,7 @@ from abogen.voice_profiles import ( serialize_profiles, ) -from abogen.voice_formulas import get_new_voice +from abogen.voice_formulas import get_new_voice, parse_formula_terms from abogen.speaker_analysis import analyze_speakers from abogen.speaker_configs import ( delete_config, @@ -656,6 +656,8 @@ def _apply_prepare_form( pending.applied_speaker_config = selected_config or None + errors: List[str] = [] + if isinstance(pending.speakers, dict): for speaker_id, payload in list(pending.speakers.items()): if not isinstance(payload, dict): @@ -669,12 +671,34 @@ def _apply_prepare_form( payload.pop("pronunciation", None) voice_value = (form.get(f"speaker-{speaker_id}-voice") or "").strip() + formula_key = f"speaker-{speaker_id}-formula" + formula_value = (form.get(formula_key) or "").strip() + has_formula = False + if formula_value: + try: + _parse_voice_formula(formula_value) + except ValueError as exc: + label = payload.get("label") or speaker_id.replace("_", " ").title() + errors.append(f"Invalid custom mix for {label}: {exc}") + else: + payload["voice_formula"] = formula_value + payload["resolved_voice"] = formula_value + payload.pop("voice_profile", None) + has_formula = True + else: + payload.pop("voice_formula", None) + + if voice_value == "__custom_mix": + voice_value = "" + if voice_value: payload["voice"] = voice_value - payload["resolved_voice"] = voice_value + if not has_formula: + payload["resolved_voice"] = voice_value else: payload.pop("voice", None) - payload.pop("resolved_voice", None) + if not has_formula: + payload.pop("resolved_voice", None) lang_key = f"speaker-{speaker_id}-languages" languages: List[str] = [] @@ -689,7 +713,6 @@ def _apply_prepare_form( payload["config_languages"] = languages profiles = serialize_profiles() - errors: List[str] = [] raw_delay = form.get("chapter_intro_delay") if raw_delay is not None: raw_normalized = raw_delay.strip() @@ -1135,22 +1158,7 @@ def _persist_cover_image(extraction_result: Any, stored_path: Path) -> tuple[Opt def _parse_voice_formula(formula: str) -> List[tuple[str, float]]: - parts = [segment.strip() for segment in formula.split("+") if segment.strip()] - voices: List[tuple[str, float]] = [] - for part in parts: - if "*" not in part: - raise ValueError("Each component must be in the form voice*weight") - name, weight_str = part.split("*", 1) - name = name.strip() - if name not in VOICES_INTERNAL: - raise ValueError(f"Unknown voice '{name}'") - try: - weight = float(weight_str.strip()) - except ValueError as exc: # pragma: no cover - validated via form - raise ValueError(f"Invalid weight for {name}") from exc - if weight <= 0: - raise ValueError(f"Weight for {name} must be positive") - voices.append((name, weight)) + voices = parse_formula_terms(formula) total = sum(weight for _, weight in voices) if total <= 0: raise ValueError("Voice weights must sum to a positive value") diff --git a/abogen/web/service.py b/abogen/web/service.py index 9542e14..6366eec 100644 --- a/abogen/web/service.py +++ b/abogen/web/service.py @@ -15,6 +15,7 @@ from pathlib import Path from typing import Any, Callable, Dict, Iterable, List, Optional, Mapping from abogen.utils import get_internal_cache_path, get_user_settings_dir +from abogen.voice_cache import bootstrap_voice_cache def _create_set_event() -> threading.Event: @@ -262,6 +263,7 @@ class ConversionService: self._pending_jobs: Dict[str, PendingJob] = {} self._state_path = self._determine_state_path() self._ensure_directories() + self._bootstrap_voice_cache() self._load_state() # Public API --------------------------------------------------------- @@ -562,6 +564,23 @@ class ConversionService: self._uploads_root.mkdir(parents=True, exist_ok=True) self._state_path.parent.mkdir(parents=True, exist_ok=True) + def _bootstrap_voice_cache(self) -> None: + try: + downloaded, errors = bootstrap_voice_cache( + on_progress=lambda msg: _JOB_LOGGER.debug("[voice cache] %s", msg) + ) + except RuntimeError as exc: + _JOB_LOGGER.warning("Voice cache bootstrap skipped: %s", exc) + return + + if downloaded: + count = len(downloaded) + suffix = "s" if count != 1 else "" + _JOB_LOGGER.info("Voice cache ready: downloaded %d new asset%s.", count, suffix) + if errors: + for voice_id, message in errors.items(): + _JOB_LOGGER.warning("Voice cache failed for %s: %s", voice_id, message) + def _ensure_worker(self) -> None: with self._lock: if self._worker_thread and self._worker_thread.is_alive(): diff --git a/tests/test_prepare_form.py b/tests/test_prepare_form.py new file mode 100644 index 0000000..ae3798c --- /dev/null +++ b/tests/test_prepare_form.py @@ -0,0 +1,61 @@ +from pathlib import Path + +from werkzeug.datastructures import MultiDict + +from abogen.web.routes import _apply_prepare_form +from abogen.web.service import PendingJob + + +def _make_pending_job() -> PendingJob: + return PendingJob( + id="pending", + original_filename="example.epub", + stored_path=Path("example.epub"), + language="a", + voice="af_nova", + speed=1.0, + use_gpu=False, + subtitle_mode="none", + output_format="mp3", + save_mode="save_next_to_input", + output_folder=None, + replace_single_newlines=False, + subtitle_format="srt", + total_characters=0, + save_chapters_separately=False, + merge_chapters_at_end=True, + separate_chapters_format="wav", + silence_between_chapters=2.0, + save_as_project=False, + voice_profile=None, + max_subtitle_words=50, + metadata_tags={}, + chapters=[], + created_at=0.0, + ) + + +def test_apply_prepare_form_handles_custom_mix_for_speakers(): + pending = _make_pending_job() + pending.speakers = { + "hero": { + "id": "hero", + "label": "Hero", + } + } + + form = MultiDict( + { + "chapter_intro_delay": "0.5", + "speaker-hero-voice": "__custom_mix", + "speaker-hero-formula": "af_nova*0.6+am_liam*0.4", + } + ) + + _, _, _, errors, *_ = _apply_prepare_form(pending, form) + + assert not errors + hero = pending.speakers["hero"] + assert hero["voice_formula"] == "af_nova*0.6+am_liam*0.4" + assert hero["resolved_voice"] == "af_nova*0.6+am_liam*0.4" + assert "voice" not in hero or hero["voice"] != "__custom_mix" diff --git a/tests/test_voice_cache.py b/tests/test_voice_cache.py new file mode 100644 index 0000000..40909c4 --- /dev/null +++ b/tests/test_voice_cache.py @@ -0,0 +1,56 @@ +from types import SimpleNamespace +from typing import cast + +import pytest + +from abogen.constants import VOICES_INTERNAL +from abogen.voice_cache import _CACHED_VOICES, ensure_voice_assets +from abogen.web.conversion_runner import _collect_required_voice_ids +from abogen.web.service import Job + + +@pytest.fixture(autouse=True) +def clear_voice_cache(): + _CACHED_VOICES.clear() + yield + _CACHED_VOICES.clear() + + +def test_ensure_voice_assets_downloads_missing(monkeypatch): + recorded = [] + + def fake_download(**kwargs): + recorded.append(kwargs["filename"]) + return "/tmp/fake" + + monkeypatch.setattr("abogen.voice_cache.hf_hub_download", fake_download) + + downloaded, errors = ensure_voice_assets(["af_nova", "am_liam"]) + + assert downloaded == {"af_nova", "am_liam"} + assert errors == {} + assert recorded == ["voices/af_nova.pt", "voices/am_liam.pt"] + + recorded.clear() + downloaded_again, errors_again = ensure_voice_assets(["af_nova"]) + + assert downloaded_again == set() + assert errors_again == {} + assert recorded == [] + + +def test_collect_required_voice_ids_includes_all(): + job = SimpleNamespace( + voice="af_nova", + chapters=[{"voice_formula": "af_nova*0.7+am_liam*0.3"}], + chunks=[{"voice": "am_michael"}], + speakers={ + "hero": {"voice_formula": "af_nova*0.6+am_liam*0.4"}, + "narrator": {"voice": "af_nova"}, + }, + ) + + voices = _collect_required_voice_ids(cast(Job, job)) + + assert {"af_nova", "am_liam", "am_michael"}.issubset(voices) + assert voices.issuperset(VOICES_INTERNAL)