Cpu cache for unload

CHANGELOG.md

@@ -5,6 +5,10 @@ Notable changes to this project will be documented in this file.
 Per-PR attribution and contributor credits are published automatically on the corresponding GitHub release page; this file is the curated, human-readable summary.
 
 ## [Unreleased]
+### Added
+- Optional model auto-unload after an idle timeout (`MODEL_AUTO_UNLOAD_TIMEOUT_SECONDS`, default off) to release VRAM. Reloads on the next request. `/dev/model` reports load/idle state and `POST /dev/reload` pre-warms the model, both behind `ALLOW_DEV_UNLOAD`.
+- `MODEL_UNLOAD_STRATEGY=move_to_cpu` moves model weights from GPU to system RAM on unload, so reloads are faster than the default `destroy` strategy, which reloads from disk.
+
 ### Fixed
 - Native Windows installs (`start-cpu.ps1` etc) no longer need a C++ toolchain: `pyopenjtalk-plus` (a drop-in fork with prebuilt Windows wheels) replaces `pyopenjtalk` on win32 only (#508, proposed by @siliconfps). Needs a recent `uv`. Linux, macOS, and Docker are unchanged.
 

README.md

@@ -712,6 +712,11 @@ Key Performance Metrics:
 
 Floor is host + CUDA context. Reproduce with `uv run --extra benchmarks assorted_checks/benchmarks/benchmark_model_unload.py` from `examples/`.
 
+`POST /dev/reload` reloads the model and `GET /dev/model` reports model load state and auto-unload settings. Set `ALLOW_DEV_UNLOAD=true` to expose these controls.
+
+To automatically unload the model after an idle timeout, set `MODEL_AUTO_UNLOAD_TIMEOUT_SECONDS` to a positive number of seconds. The default `0` disables auto-unload.
+Set `MODEL_UNLOAD_STRATEGY=move_to_cpu` to move model weights from GPU to system RAM on unload for faster reload while still clearing GPU memory. The default `destroy` strategy releases model objects completely and will load from disk.
+
 ### Transcription roundtrip (WER/CER)
 
 End-to-end roundtrip: synthesize with Kokoro, transcribe the result back with [`faster-whisper`](https://github.com/SYSTRAN/faster-whisper), compare to the source text. Scripts and data live under `examples/assorted_checks/test_transcription/`.
@@ -759,7 +764,9 @@ System state and resource usage, for debugging exhaustion or performance issues.
 - `/debug/threads` - Get thread information and stack traces
 - `/debug/storage` - Disk usage per mounted partition
 - `/debug/system` - Get system information (CPU, memory, GPU)
+- `/dev/model` - Get model load state and auto-unload timing. Off by default; set `ALLOW_DEV_UNLOAD=true` to enable
 - `POST /dev/unload` - Release model from VRAM; reloads lazily on next request. Off by default; set `ALLOW_DEV_UNLOAD=true` to enable
+- `POST /dev/reload` - Load the model into VRAM. Off by default; set `ALLOW_DEV_UNLOAD=true` to enable
 
 Stability: the `/v1/*` OpenAI-compatible routes are the stable API. `/dev/*` and `/debug/*` are operational helpers, and may change or move behind flags between minor releases.
 </details>

api/src/core/config.py

@@ -39,7 +39,11 @@ class Settings(BaseSettings):
     allow_local_voice_saving: bool = (
         False  # Whether to allow saving combined voices locally
     )
-    allow_dev_unload: bool = False  # Whether to expose the POST /dev/unload endpoint
+    allow_dev_unload: bool = False  # Whether to expose /dev/model, POST /dev/unload, and POST /dev/reload
+    model_auto_unload_timeout_seconds: float = (
+        0.0  # Idle seconds before unloading; 0 disables auto-unload
+    )
+    model_unload_strategy: str = "destroy"  # "destroy" or "move_to_cpu"
     enable_debug_endpoints: bool = (
         False  # Whether to expose /debug/* host and process introspection routes
     )

api/src/inference/kokoro_v1.py

@@ -89,6 +89,7 @@ def __init__(self):
         # Strictly respect settings.use_gpu
         self._device = settings.get_device()
         self._model: Optional[KModel] = None
+        self._model_cpu_cached = False
         self._pipelines: Dict[str, KPipeline] = {}  # Store pipelines by lang_code
         self._voice_cache: Dict[str, torch.Tensor] = {}  # Cache voice tensors by path
 
@@ -144,12 +145,53 @@ async def load_model(self, path: str) -> None:
                 self._model = self._model.cuda()
             else:
                 self._model = self._model.cpu()
+            self._model_cpu_cached = False
 
         except FileNotFoundError:
             raise
         except Exception as e:
             raise RuntimeError(f"Failed to load Kokoro model: {e}")
 
+    def _move_model_to_device(self) -> None:
+        """Move a CPU-cached model back to the configured inference device."""
+        if self._model is None or not self._model_cpu_cached:
+            return
+
+        logger.info(f"Moving CPU-cached Kokoro model back to {self._device}")
+        if self._device == "mps":
+            self._model = self._model.to(torch.device("mps"))
+        elif self._device == "cuda":
+            self._model = self._model.cuda()
+            torch.cuda.synchronize()
+        else:
+            self._model = self._model.cpu()
+        self._model_cpu_cached = False
+        logger.info(f"CPU-cached Kokoro model restored to {self._device}")
+
+    def restore_to_device(self) -> None:
+        """Restore a CPU-cached model to the configured inference device."""
+        self._move_model_to_device()
+
+    def _clear_runtime_caches(self) -> None:
+        """Release cached objects that can hold device tensors."""
+        for pipeline in self._pipelines.values():
+            del pipeline
+        self._pipelines.clear()
+        self._voice_cache.clear()
+
+    def _offload_model_to_cpu(self) -> bool:
+        """Move the model out of VRAM while retaining weights in system RAM."""
+        if self._model is None or self._device not in {"cuda", "mps"}:
+            return False
+
+        logger.info("Moving Kokoro model to CPU cache")
+        self._model = self._model.cpu()
+        self._model_cpu_cached = True
+        self._clear_runtime_caches()
+        self._clear_memory()
+        logger.info("Kokoro model offloaded to CPU cache and device caches cleared")
+        return True
+
     def _get_pipeline(self, lang_code: str) -> KPipeline:
         """Get or create pipeline for language code.
 
@@ -197,6 +239,7 @@ async def generate_from_tokens(
             raise RuntimeError("Model not loaded")
 
         try:
+            self._move_model_to_device()
             # Memory management for GPU
             if self._device == "cuda":
                 if self._check_memory():
@@ -295,6 +338,7 @@ async def generate(
         if not self.is_loaded:
             raise RuntimeError("Model not loaded")
         try:
+            self._move_model_to_device()
             # Memory management for GPU
             if self._device == "cuda":
                 if self._check_memory():
@@ -455,18 +499,26 @@ def _clear_memory(self) -> None:
             if hasattr(torch.mps, "empty_cache"):
                 torch.mps.empty_cache()
 
-    def unload(self) -> None:
+    def unload(self, strategy: str = "destroy") -> None:
         """Unload model and free resources."""
+        if strategy == "move_to_cpu" and self._offload_model_to_cpu():
+            return
+
+        logger.info("Destroying Kokoro model backend state")
         if self._model is not None:
             del self._model
             self._model = None
-        for pipeline in self._pipelines.values():
-            del pipeline
-        self._pipelines.clear()
-        self._voice_cache.clear()
+        self._model_cpu_cached = False
+        self._clear_runtime_caches()
         if torch.cuda.is_available():
             torch.cuda.empty_cache()
             torch.cuda.synchronize()
+        logger.info("Kokoro model backend state destroyed")
+
+    @property
+    def is_cpu_cached(self) -> bool:
+        """Check if model weights are retained in CPU RAM after device unload."""
+        return self._model_cpu_cached
 
     @property
     def is_loaded(self) -> bool: