diff --git a/abogen/main.py b/abogen/main.py index 90a55e1..7a737c1 100644 --- a/abogen/main.py +++ b/abogen/main.py @@ -2,13 +2,14 @@ from __future__ import annotations -import atexit import os import platform -import signal -import sys -from abogen.utils import load_config, prevent_sleep_end +# Initialise global shutdown handling (atexit, signals, Qt) as early as possible. +from abogen import shutdown # noqa: F401 +shutdown.register_shutdown() + +from abogen.utils import load_config from abogen.webui.app import main as _run_web_ui # Configure Hugging Face Hub behaviour (mirrors legacy GUI defaults). @@ -27,17 +28,6 @@ os.environ.setdefault("MIOPEN_CONV_PRECISE_ROCM_TUNING", "0") if platform.system() == "Darwin" and platform.processor() == "arm": os.environ.setdefault("PYTORCH_ENABLE_MPS_FALLBACK", "1") -atexit.register(prevent_sleep_end) - - -def _cleanup_sleep(signum, _frame): - prevent_sleep_end() - sys.exit(0) - - -signal.signal(signal.SIGINT, _cleanup_sleep) -signal.signal(signal.SIGTERM, _cleanup_sleep) - def main() -> None: """Launch the Flask-based web UI.""" diff --git a/abogen/pyqt/conversion.py b/abogen/pyqt/conversion.py index d435bce..811d56a 100644 --- a/abogen/pyqt/conversion.py +++ b/abogen/pyqt/conversion.py @@ -2397,6 +2397,10 @@ class ConversionThread(QThread): self.cancel_requested = True self.should_cancel = True self.waiting_for_user_input = False + # Clear voice cache (instance and module-level) + self.voice_cache.clear() + from abogen.voice_cache import clear_voice_cache + clear_voice_cache() # Terminate subprocess if running if self.process: try: diff --git a/abogen/pyqt/gui.py b/abogen/pyqt/gui.py index dce75b5..b5d7ed7 100644 --- a/abogen/pyqt/gui.py +++ b/abogen/pyqt/gui.py @@ -3236,12 +3236,16 @@ class abogen(QWidget): ) box.setDefaultButton(QMessageBox.StandardButton.No) if box.exec() == QMessageBox.StandardButton.Yes: + from abogen import shutdown + shutdown.request_shutdown() self.cleanup_conversion_thread() self.cleanup_preview_threads() event.accept() else: event.ignore() else: + from abogen import shutdown + shutdown.request_shutdown() self.cleanup_conversion_thread() self.cleanup_preview_threads() event.accept() diff --git a/abogen/pyqt/main.py b/abogen/pyqt/main.py index f788e15..9e11c0c 100644 --- a/abogen/pyqt/main.py +++ b/abogen/pyqt/main.py @@ -1,10 +1,10 @@ import os import sys import platform -import atexit -import signal -from abogen.utils import get_resource_path, load_config, prevent_sleep_end +# Initialise global shutdown handling (atexit, signals, Qt) as early as possible. +from abogen import shutdown # noqa: F401 +shutdown.register_shutdown() # Fix PyTorch DLL loading issue ([WinError 1114]) on Windows before importing PyQt6 if platform.system() == "Windows": @@ -94,6 +94,7 @@ os.environ["HF_HUB_DISABLE_TELEMETRY"] = "1" # Disable Hugging Face telemetry os.environ["HF_HUB_ETAG_TIMEOUT"] = "10" # Metadata request timeout (seconds) os.environ["HF_HUB_DOWNLOAD_TIMEOUT"] = "10" # File download timeout (seconds) os.environ["HF_HUB_DISABLE_SYMLINKS_WARNING"] = "1" # Disable symlinks warning +from abogen.utils import load_config if load_config().get("disable_kokoro_internet", False): print("INFO: Kokoro's internet access is disabled.") os.environ["HF_HUB_OFFLINE"] = "1" # Disable Hugging Face Hub internet access @@ -105,25 +106,6 @@ from abogen.constants import PROGRAM_NAME, VERSION os.environ["MIOPEN_FIND_MODE"] = "FAST" os.environ["MIOPEN_CONV_PRECISE_ROCM_TUNING"] = "0" -# Reset sleep states -atexit.register(prevent_sleep_end) - - -# Also handle signals (Ctrl+C, kill, etc.) -def _cleanup_sleep(signum, frame): - prevent_sleep_end() - sys.exit(0) - - -signal.signal(signal.SIGINT, _cleanup_sleep) -signal.signal(signal.SIGTERM, _cleanup_sleep) - -# Ensure sys.stdout and sys.stderr are valid in GUI mode -if sys.stdout is None: - sys.stdout = open(os.devnull, "w") -if sys.stderr is None: - sys.stderr = open(os.devnull, "w") - # Enable MPS GPU acceleration on Mac Apple Silicon if platform.system() == "Darwin" and platform.processor() == "arm": os.environ["PYTORCH_ENABLE_MPS_FALLBACK"] = "1" @@ -184,4 +166,4 @@ def main(): if __name__ == "__main__": - main() + main() \ No newline at end of file diff --git a/abogen/shutdown.py b/abogen/shutdown.py new file mode 100644 index 0000000..8ab078b --- /dev/null +++ b/abogen/shutdown.py @@ -0,0 +1,160 @@ +"""Graceful shutdown - single module, no over-engineering.""" +from __future__ import annotations + +import atexit +import gc +import signal +import sys +from typing import Callable + +_CLEANUP_FUNCS: list[Callable[[], None]] = [] +_EXECUTED = False + + +def register_cleanup(fn: Callable[[], None]) -> None: + """Register a cleanup function to run on shutdown.""" + _CLEANUP_FUNCS.append(fn) + + +def _run_cleanups() -> None: + global _EXECUTED + if _EXECUTED: + return + _EXECUTED = True + for fn in _CLEANUP_FUNCS: + try: + fn() + except Exception: + pass + + +# ---- Register built-in cleanup functions ---- + +# 1. Restore sleep prevention +def _restore_sleep() -> None: + try: + from abogen.utils import prevent_sleep_end + prevent_sleep_end() + except Exception: + pass + +register_cleanup(_restore_sleep) + +# 2. Shutdown web UI ConversionService +def _shutdown_conversion_service() -> None: + try: + from abogen.webui.service import get_service + svc = get_service() + if svc is not None: + svc.shutdown() + except Exception: + pass + +register_cleanup(_shutdown_conversion_service) + +# 3. Clear TTS pipelines and GPU memory +def _cleanup_tts_pipelines() -> None: + # Clear web UI pipeline cache + try: + from abogen.webui.conversion_runner import _PIPELINES + _PIPELINES.clear() + except Exception: + pass + + # Clear PyQt conversion thread voice cache + try: + from abogen.pyqt.conversion import ConversionThread + if hasattr(ConversionThread, "voice_cache"): + ConversionThread.voice_cache.clear() + except Exception: + pass + + gc.collect() + + # Release CUDA cache + try: + import torch + if torch.cuda.is_available(): + torch.cuda.empty_cache() + torch.cuda.ipc_collect() + except Exception: + pass + +register_cleanup(_cleanup_tts_pipelines) + +# 4. Clear global voice cache +def _clear_voice_cache() -> None: + try: + from abogen.voice_cache import clear_voice_cache + clear_voice_cache() + except Exception: + pass + +register_cleanup(_clear_voice_cache) + +# 5. Terminate child processes (ffmpeg, etc.) +def _terminate_subprocesses() -> None: + try: + import psutil + except Exception: + return + + try: + current = psutil.Process() + for child in current.children(recursive=True): + try: + child.terminate() + except Exception: + pass + gone, alive = psutil.wait_procs(current.children(recursive=True), timeout=3) + for proc in alive: + try: + proc.kill() + except Exception: + pass + except Exception: + pass + +register_cleanup(_terminate_subprocesses) + + +def register_shutdown() -> None: + """Install process-wide shutdown hooks (atexit, signals, Qt).""" + if register_shutdown._registered: + return + register_shutdown._registered = True + + atexit.register(_run_cleanups) + + # POSIX signals + for sig in (signal.SIGINT, signal.SIGTERM): + try: + signal.signal(sig, _on_signal) + except Exception: + pass + + # Qt hook + try: + from PyQt6.QtWidgets import QApplication + + app = QApplication.instance() + if app is not None: + app.aboutToQuit.connect(_run_cleanups) + except Exception: + pass + + +register_shutdown._registered = False + + +def _on_signal(signum: int, _frame) -> None: + _run_cleanups() + sys.exit(0) + + +def request_shutdown() -> None: + """Programmatically trigger cleanup (e.g., from GUI closeEvent).""" + _run_cleanups() + + +__all__ = ["register_shutdown", "request_shutdown", "register_cleanup"] diff --git a/abogen/voice_cache.py b/abogen/voice_cache.py index 838eb65..769b300 100644 --- a/abogen/voice_cache.py +++ b/abogen/voice_cache.py @@ -144,3 +144,11 @@ def _ensure_single_voice_asset( hf_hub_download(resume_download=True, **common_kwargs) return True + + +def clear_voice_cache() -> None: + """Clear the in‑process voice cache (used during shutdown).""" + with _CACHE_LOCK: + _CACHED_VOICES.clear() + global _BOOTSTRAPPED + _BOOTSTRAPPED = False diff --git a/abogen/webui/app.py b/abogen/webui/app.py index 211bfcb..2a336b9 100644 --- a/abogen/webui/app.py +++ b/abogen/webui/app.py @@ -1,6 +1,5 @@ from __future__ import annotations -import atexit import logging import os from pathlib import Path @@ -8,6 +7,8 @@ from typing import Any, Optional from flask import Flask +from abogen import shutdown # noqa: F401 +shutdown.register_shutdown() from abogen.utils import get_user_cache_path, get_user_output_path, get_user_settings_dir from .conversion_runner import run_conversion_job @@ -113,8 +114,6 @@ def create_app(config: Optional[dict[str, Any]] = None) -> Flask: app.register_blueprint(books_bp, url_prefix="/find-books") app.register_blueprint(api_bp, url_prefix="/api") - atexit.register(service.shutdown) - global _access_log_filter_attached if not _access_log_filter_attached: logging.getLogger("werkzeug").addFilter(_SuppressSuccessfulAccessFilter()) @@ -132,4 +131,4 @@ def main() -> None: if __name__ == "__main__": # pragma: no cover - main() + main() \ No newline at end of file diff --git a/abogen/webui/service.py b/abogen/webui/service.py index a04e9ce..bdf90b4 100644 --- a/abogen/webui/service.py +++ b/abogen/webui/service.py @@ -1609,10 +1609,12 @@ def build_service( output_root: Optional[Path] = None, uploads_root: Optional[Path] = None, ) -> ConversionService: + global _service_instance output_root = output_root or default_storage_root() service = ConversionService( output_root=output_root, uploads_root=uploads_root, runner=runner, ) + _service_instance = service return service