""" Pre-download dialog and worker for Abogen This module consolidates pre-download logic for Kokoro voices and model and spaCy language models. The code favors clarity, avoids duplication, and handles optional dependencies gracefully. """ from typing import List, Optional, Tuple import importlib import importlib.util from PyQt6.QtWidgets import ( QDialog, QVBoxLayout, QHBoxLayout, QLabel, QPushButton, QSpacerItem, QSizePolicy, ) from PyQt6.QtCore import QThread, pyqtSignal from abogen.constants import COLORS from abogen.tts_backend_registry import get_metadata from abogen.spacy_utils import SPACY_MODELS import abogen.hf_tracker # Helpers def _unique_sorted_models() -> List[str]: """Return a sorted list of unique spaCy model package names.""" return sorted(set(SPACY_MODELS.values())) def _is_package_installed(pkg_name: str) -> bool: """Return True if a package with the given name can be imported (site-packages).""" try: return importlib.util.find_spec(pkg_name) is not None except Exception: return False # NOTE: explicit HF cache helper removed; we use try_to_load_from_cache in-scope where needed class PreDownloadWorker(QThread): """Worker thread to download required models/voices. Emits human-readable messages via `progress`. Uses `category_done` to indicate a category (voices/model/spacy) finished successfully. Emits `error` on exception and `finished` after all work completes. """ # Emit (category, status, message) progress = pyqtSignal(str, str, str) category_done = pyqtSignal(str) finished = pyqtSignal() error = pyqtSignal(str) def __init__(self, parent=None): super().__init__(parent) self._cancelled = False # repo and filenames used for Kokoro model self._repo_id = "hexgrad/Kokoro-82M" self._model_files = ["kokoro-v1_0.pth", "config.json"] # Track download success per category self._voices_success = False self._model_success = False self._spacy_success = False # Suppress HF tracker warnings during downloads self._original_emitter = abogen.hf_tracker.show_warning_signal_emitter def cancel(self) -> None: self._cancelled = True def run(self) -> None: # Suppress HF tracker warnings during downloads abogen.hf_tracker.show_warning_signal_emitter = None try: self._download_kokoro_voices() if self._cancelled: return if self._voices_success: self.category_done.emit("voices") self._download_kokoro_model() if self._cancelled: return if self._model_success: self.category_done.emit("model") self._download_spacy_models() if self._cancelled: return if self._spacy_success: self.category_done.emit("spacy") self.finished.emit() except Exception as exc: # pragma: no cover - best-effort reporting self.error.emit(str(exc)) finally: # Restore original emitter abogen.hf_tracker.show_warning_signal_emitter = self._original_emitter # Kokoro voices def _download_kokoro_voices(self) -> None: self._voices_success = True try: from huggingface_hub import hf_hub_download, try_to_load_from_cache except Exception: self.progress.emit( "voice", "warning", "huggingface_hub not installed, skipping voices..." ) self._voices_success = False return voice_list = get_metadata("kokoro").voices for idx, voice in enumerate(voice_list, start=1): if self._cancelled: self._voices_success = False return filename = f"voices/{voice}.pt" if try_to_load_from_cache(repo_id=self._repo_id, filename=filename): self.progress.emit( "voice", "installed", f"{idx}/{len(voice_list)}: {voice} already present", ) continue self.progress.emit( "voice", "downloading", f"{idx}/{len(voice_list)}: {voice}..." ) try: hf_hub_download(repo_id=self._repo_id, filename=filename) self.progress.emit("voice", "downloaded", f"{voice} downloaded") except Exception as exc: self.progress.emit( "voice", "warning", f"could not download {voice}: {exc}" ) self._voices_success = False # Kokoro model def _download_kokoro_model(self) -> None: self._model_success = True try: from huggingface_hub import hf_hub_download, try_to_load_from_cache except Exception: self.progress.emit( "model", "warning", "huggingface_hub not installed, skipping model..." ) self._model_success = False return for fname in self._model_files: if self._cancelled: self._model_success = False return category = "config" if fname == "config.json" else "model" if try_to_load_from_cache(repo_id=self._repo_id, filename=fname): self.progress.emit( category, "installed", f"file {fname} already present" ) continue self.progress.emit(category, "downloading", f"file {fname}...") try: hf_hub_download(repo_id=self._repo_id, filename=fname) self.progress.emit(category, "downloaded", f"file {fname} downloaded") except Exception as exc: self.progress.emit( category, "warning", f"could not download file {fname}: {exc}" ) self._model_success = False # spaCy models def _download_spacy_models(self) -> None: """Download spaCy models. Prefer missing models provided by parent. Parent dialog will populate _spacy_models_missing during checking. """ self._spacy_success = True # Determine which models to process: prefer parent-provided missing list to avoid # re-checking everything; otherwise use the full unique list. parent = self.parent() models_to_process: List[str] = _unique_sorted_models() try: if ( parent is not None and hasattr(parent, "_spacy_models_missing") and parent._spacy_models_missing ): models_to_process = list(dict.fromkeys(parent._spacy_models_missing)) except Exception: pass # If spaCy is not available to run the CLI, skip gracefully try: import spacy.cli as _spacy_cli except Exception: self.progress.emit( "spacy", "warning", "spaCy not available, skipping spaCy models..." ) self._spacy_success = False return for idx, model_name in enumerate(models_to_process, start=1): if self._cancelled: self._spacy_success = False return if _is_package_installed(model_name): self.progress.emit( "spacy", "installed", f"{idx}/{len(models_to_process)}: {model_name} already installed", ) continue self.progress.emit( "spacy", "downloading", f"{idx}/{len(models_to_process)}: {model_name}...", ) try: _spacy_cli.download(model_name) self.progress.emit("spacy", "downloaded", f"{model_name} downloaded") except Exception as exc: self.progress.emit( "spacy", "warning", f"could not download {model_name}: {exc}" ) self._spacy_success = False class PreDownloadDialog(QDialog): """Dialog to show and control pre-download process.""" VOICE_PREFIX = "Kokoro voices: " MODEL_PREFIX = "Kokoro model: " CONFIG_PREFIX = "Kokoro config: " SPACY_PREFIX = "spaCy models: " def __init__(self, parent=None): super().__init__(parent) self.setWindowTitle("Pre-download Models and Voices") self.setMinimumWidth(500) self.worker: Optional[PreDownloadWorker] = None self.has_missing = False self._spacy_models_checked: List[tuple] = [] self._spacy_models_missing: List[str] = [] self._status_worker = None # Map keywords to (label, prefix) - labels filled after UI creation self.status_map = { "voice": (None, self.VOICE_PREFIX), "spacy": (None, self.SPACY_PREFIX), "model": (None, self.MODEL_PREFIX), "config": (None, self.CONFIG_PREFIX), } self.category_map = { "voices": ["voice"], "model": ["model", "config"], "spacy": ["spacy"], } self._setup_ui() self._start_status_check() def _setup_ui(self) -> None: layout = QVBoxLayout(self) layout.setSpacing(0) layout.setContentsMargins(15, 0, 15, 15) desc = QLabel( "You can pre-download all required models and voices for offline use.\n" "This includes Kokoro voices, Kokoro model (and config), and spaCy models." ) desc.setWordWrap(True) layout.addWidget(desc) # Status rows status_layout = QVBoxLayout() status_title = QLabel("Current Status:") status_layout.addWidget(status_title) self.voices_status = QLabel(self.VOICE_PREFIX + "⏳ Checking...") row = QHBoxLayout() row.addWidget(self.voices_status) row.addStretch() status_layout.addLayout(row) self.model_status = QLabel(self.MODEL_PREFIX + "⏳ Checking...") row = QHBoxLayout() row.addWidget(self.model_status) row.addStretch() status_layout.addLayout(row) self.config_status = QLabel(self.CONFIG_PREFIX + "⏳ Checking...") row = QHBoxLayout() row.addWidget(self.config_status) row.addStretch() status_layout.addLayout(row) self.spacy_status = QLabel(self.SPACY_PREFIX + "⏳ Checking...") row = QHBoxLayout() row.addWidget(self.spacy_status) row.addStretch() status_layout.addLayout(row) # register labels self.status_map["voice"] = (self.voices_status, self.VOICE_PREFIX) self.status_map["model"] = (self.model_status, self.MODEL_PREFIX) self.status_map["config"] = (self.config_status, self.CONFIG_PREFIX) self.status_map["spacy"] = (self.spacy_status, self.SPACY_PREFIX) layout.addLayout(status_layout) layout.addItem( QSpacerItem(0, 20, QSizePolicy.Policy.Minimum, QSizePolicy.Policy.Fixed) ) # Buttons button_row = QHBoxLayout() button_row.setSpacing(10) self.download_btn = QPushButton("Download all") self.download_btn.setMinimumWidth(100) self.download_btn.setMinimumHeight(35) self.download_btn.setEnabled(False) self.download_btn.clicked.connect(self._start_download) button_row.addWidget(self.download_btn) self.close_btn = QPushButton("Close") self.close_btn.setMinimumWidth(100) self.close_btn.setMinimumHeight(35) self.close_btn.clicked.connect(self._handle_close) button_row.addWidget(self.close_btn) layout.addLayout(button_row) self.adjustSize() # Status checking worker class StatusCheckWorker(QThread): voices_checked = pyqtSignal(bool, list) model_checked = pyqtSignal(bool) config_checked = pyqtSignal(bool) spacy_model_checking = pyqtSignal(str) spacy_model_result = pyqtSignal(str, bool) spacy_checked = pyqtSignal(bool, list) def run(self): parent = self.parent() if parent is None: return voices_ok, missing_voices = parent._check_kokoro_voices() self.voices_checked.emit(voices_ok, missing_voices) model_ok = parent._check_kokoro_model() self.model_checked.emit(model_ok) config_ok = parent._check_kokoro_config() self.config_checked.emit(config_ok) # Check spaCy models by package name to detect site-package installs unique = _unique_sorted_models() missing: List[str] = [] for name in unique: self.spacy_model_checking.emit(name) ok = _is_package_installed(name) self.spacy_model_result.emit(name, ok) if not ok: missing.append(name) parent._spacy_models_missing = missing self.spacy_checked.emit(len(missing) == 0, missing) def _start_status_check(self) -> None: self._status_worker = self.StatusCheckWorker(self) self._status_worker.voices_checked.connect(self._update_voices_status) self._status_worker.model_checked.connect(self._update_model_status) self._status_worker.config_checked.connect(self._update_config_status) self._status_worker.spacy_model_checking.connect(self._spacy_model_checking) self._status_worker.spacy_model_result.connect(self._spacy_model_result) self._status_worker.spacy_checked.connect(self._update_spacy_status) # These are initialized in __init__ to keep consistent object state # Set checking visual state for lbl in ( self.voices_status, self.model_status, self.config_status, self.spacy_status, ): lbl.setStyleSheet(f"color: {COLORS['ORANGE']};") self.spacy_status.setText(self.SPACY_PREFIX + "⏳ Checking...") self._status_worker.start() # UI update callbacks def _spacy_model_checking(self, name: str) -> None: self.spacy_status.setText(f"{self.SPACY_PREFIX}Checking {name}...") def _spacy_model_result(self, name: str, ok: bool) -> None: self._spacy_models_checked.append((name, ok)) if not ok and name not in self._spacy_models_missing: self._spacy_models_missing.append(name) checked = len(self._spacy_models_checked) missing_count = len(self._spacy_models_missing) if missing_count: self.spacy_status.setText( f"{self.SPACY_PREFIX}{checked} checked, {missing_count} missing..." ) else: self.spacy_status.setText(f"{self.SPACY_PREFIX}{checked} checked...") def _update_voices_status(self, ok: bool, missing: List[str]) -> None: if ok: self._set_status("voice", "✓ Downloaded", COLORS["GREEN"]) else: self.has_missing = True if missing: self._set_status( "voice", f"✗ Missing {len(missing)} voices", COLORS["RED"] ) else: self._set_status("voice", "✗ Not downloaded", COLORS["RED"]) def _update_model_status(self, ok: bool) -> None: if ok: self._set_status("model", "✓ Downloaded", COLORS["GREEN"]) else: self.has_missing = True self._set_status("model", "✗ Not downloaded", COLORS["RED"]) def _update_config_status(self, ok: bool) -> None: if ok: self._set_status("config", "✓ Downloaded", COLORS["GREEN"]) else: self.has_missing = True self._set_status("config", "✗ Not downloaded", COLORS["RED"]) def _update_spacy_status(self, ok: bool, missing: List[str]) -> None: if ok: self._set_status("spacy", "✓ Downloaded", COLORS["GREEN"]) else: self.has_missing = True if missing: self._set_status( "spacy", f"✗ Missing {len(missing)} model(s)", COLORS["RED"] ) else: self._set_status("spacy", "✗ Not downloaded", COLORS["RED"]) self.download_btn.setEnabled(self.has_missing) def _set_status(self, key: str, text: str, color: str) -> None: lbl, prefix = self.status_map.get(key, (None, "")) if not lbl: return lbl.setText(prefix + text) lbl.setStyleSheet(f"color: {color};") # Helper checks def _check_kokoro_voices(self) -> Tuple[bool, List[str]]: """Return (ok, missing_list) for Kokoro voices check.""" missing = [] try: from huggingface_hub import try_to_load_from_cache for voice in get_metadata("kokoro").voices: if not try_to_load_from_cache( repo_id="hexgrad/Kokoro-82M", filename=f"voices/{voice}.pt" ): missing.append(voice) except Exception: # If HF missing, report all as missing return False, list(get_metadata("kokoro").voices) return (len(missing) == 0), missing def _check_kokoro_model(self) -> bool: try: from huggingface_hub import try_to_load_from_cache return ( try_to_load_from_cache( repo_id="hexgrad/Kokoro-82M", filename="kokoro-v1_0.pth" ) is not None ) except Exception: return False def _check_kokoro_config(self) -> bool: try: from huggingface_hub import try_to_load_from_cache return ( try_to_load_from_cache( repo_id="hexgrad/Kokoro-82M", filename="config.json" ) is not None ) except Exception: return False def _check_spacy_models(self) -> bool: unique = _unique_sorted_models() missing = [m for m in unique if not _is_package_installed(m)] self._spacy_models_missing = missing return len(missing) == 0 # Download control def _start_download(self) -> None: self.download_btn.setEnabled(False) self.download_btn.setText("Downloading...") # mark the start of downloads; this triggers the labels self._on_progress("system", "starting", "Processing, please wait...") self.worker = PreDownloadWorker(self) self.worker.progress.connect(self._on_progress) self.worker.category_done.connect(self._on_category_done) self.worker.finished.connect(self._on_download_finished) self.worker.error.connect(self._on_download_error) self.worker.start() def _on_progress(self, category: str, status: str, message: str) -> None: """Map worker (category, status, message) to UI label updates. Status is one of: 'downloading', 'installed', 'downloaded', 'warning', 'starting'. Category is one of: 'voice', 'model', 'spacy', 'config', or 'system'. """ try: # If the category targets a specific label, update directly if category in self.status_map: lbl, prefix = self.status_map[category] if not lbl: return # Compose message and set color based on status token full_text = prefix + message if len(full_text) > 60: display_text = full_text[:57] + "..." lbl.setText(display_text) lbl.setToolTip(full_text) else: lbl.setText(full_text) lbl.setToolTip("") # Clear tooltip if not needed if status == "downloading": lbl.setStyleSheet(f"color: {COLORS['ORANGE']};") elif status in ("installed", "downloaded"): lbl.setStyleSheet(f"color: {COLORS['GREEN']};") elif status == "warning": lbl.setStyleSheet(f"color: {COLORS['RED']};") elif status == "error": lbl.setStyleSheet(f"color: {COLORS['RED']};") return # System-level messages if category == "system": if status == "starting": for k in self.status_map: lbl, prefix = self.status_map[k] if lbl: lbl.setText(prefix + "Processing, please wait...") lbl.setStyleSheet(f"color: {COLORS['ORANGE']};") # other system statuses don't require action return except Exception: # Do not let UI thread crash on unexpected worker message pass def _on_category_done(self, category: str) -> None: for key in self.category_map.get(category, []): self._set_status(key, "✓ Downloaded", COLORS["GREEN"]) def _on_download_finished(self) -> None: self.has_missing = False self.download_btn.setText("Download all") self.download_btn.setEnabled(False) def _on_download_error(self, error_msg: str) -> None: self.download_btn.setText("Download all") self.download_btn.setEnabled(True) for key in self.status_map: self._set_status(key, f"✗ Error - {error_msg}", COLORS["RED"]) def _handle_close(self) -> None: if self.worker and self.worker.isRunning(): self.worker.cancel() self.worker.wait(2000) self.accept() def closeEvent(self, event) -> None: if self.worker and self.worker.isRunning(): self.worker.cancel() self.worker.wait(2000) super().closeEvent(event)