refactor(shutdown): move shutdown logic in one place

- Monkey-patch transformers.AlbertModel in plugins/kokoro/__init__.py before kokoro imports it
- Works with transformers 4.x (no-op) and 5.x (adds moved symbol)
- No pinning, forks, or extra files needed

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."""

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:

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

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"

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

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

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

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

plugins/kokoro/__init__.py

@@ -34,6 +34,13 @@ from .engine import KokoroEngine
 
 def _load_kpipeline() -> Any:
     """Lazy-load Kokoro dependencies."""
+    # Transformers 5.x moved AlbertModel out of top-level imports.
+    # Monkey-patch before kokoro imports it.
+    import transformers
+    if not hasattr(transformers, "AlbertModel"):
+        from transformers.models.albert import AlbertModel as _AlbertModel
+        transformers.AlbertModel = _AlbertModel
+
     from kokoro import KPipeline  # type: ignore[import-not-found]
     return KPipeline