feat: Enhance voice formula parsing and validation, implement voice asset caching, and add tests for new functionality

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

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)]

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)
 

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,11 +671,33 @@ 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
+                if not has_formula:
                     payload["resolved_voice"] = voice_value
             else:
                 payload.pop("voice", None)
+                if not has_formula:
                     payload.pop("resolved_voice", None)
 
             lang_key = f"speaker-{speaker_id}-languages"
@@ -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")

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():

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"

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)