fix: release image router models after compression

Author: eggrollofchaos

CHANGELOG.md

@@ -36,6 +36,11 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
   backup is missing, strips only the Headroom-managed block and leaves
   surrounding user content intact). Safe no-op when run without a prior
   wrap. Reported by @raenaryl in Discord.
+- **Image compressors now release shared router models after use and proxy shutdown** —
+  the proxy/image compression path no longer keeps global `technique-router`
+  and `SigLIP` model instances pinned in memory after one-off image
+  optimization work. The `get_compressor()` helper now returns a fresh,
+  caller-owned compressor instead of a process-lifetime singleton.
 - **`headroom learn` no longer clobbers prior recommendations on re-run** —
   the marker block in `CLAUDE.md` / `MEMORY.md` is now merged with the
   prior block instead of wholesale-replaced. Sections re-surfaced by the

headroom/image/compressor.py

@@ -102,6 +102,14 @@ def _get_router(self) -> TrainedRouter:
             )
         return self._router
 
+    def close(self, unload_models: bool = True) -> None:
+        """Release any router-held model state."""
+        if self._router is not None:
+            # Only loaded routers hold heavyweight image models; plain has_images()
+            # checks remain cheap and have nothing to release.
+            self._router.release_models(unload_registry=unload_models)
+            self._router = None
+
     def has_images(self, messages: list[dict[str, Any]]) -> bool:
         """Check if messages contain images."""
         for message in messages:
@@ -563,16 +571,13 @@ def compress(
         return compressed_messages
 
 
-# Singleton for convenience
-_default_compressor: ImageCompressor | None = None
-
-
 def get_compressor() -> ImageCompressor:
-    """Get the default ImageCompressor instance."""
-    global _default_compressor
-    if _default_compressor is None:
-        _default_compressor = ImageCompressor()
-    return _default_compressor
+    """Create an ImageCompressor instance.
+
+    Kept for backwards-compatible imports; callers that use it directly own
+    closing the returned compressor.
+    """
+    return ImageCompressor()
 
 
 def compress_images(
@@ -588,4 +593,8 @@ def compress_images(
     Returns:
         Messages with compressed images
     """
-    return get_compressor().compress(messages, provider)
+    compressor = ImageCompressor()
+    try:
+        return compressor.compress(messages, provider)
+    finally:
+        compressor.close()

headroom/image/trained_router.py

@@ -10,6 +10,7 @@
 
 from __future__ import annotations
 
+import gc
 import io
 from dataclasses import dataclass
 from enum import Enum
@@ -151,6 +152,8 @@ def __init__(
         self._siglip_model: Any = None
         self._siglip_processor: Any = None
         self._text_embeddings: Any = None
+        self._classifier_key: str | None = None
+        self._siglip_key: str | None = None
 
     def is_available(self) -> bool:
         """Check if required models can be loaded."""
@@ -186,6 +189,7 @@ def _load_models(self) -> None:
                 model_path=model_id,
                 device=self.device,
             )
+            self._classifier_key = f"technique_router:{model_id}"
 
         if self.use_siglip and self._siglip_model is None:
             # Use centralized registry for shared model instances
@@ -195,10 +199,34 @@ def _load_models(self) -> None:
                 model_name=self.siglip_model,
                 device=self.device,
             )
+            self._siglip_key = f"siglip:{self.siglip_model}"
 
             # Pre-compute text embeddings for image analysis
             self._compute_text_embeddings()
 
+    def release_models(self, unload_registry: bool = True) -> None:
+        """Release router-held model references and optional shared cache entries."""
+        classifier_key = self._classifier_key
+        siglip_key = self._siglip_key
+
+        self._text_embeddings = None
+        self._siglip_processor = None
+        self._siglip_model = None
+        self._tokenizer = None
+        self._classifier = None
+        self._classifier_key = None
+        self._siglip_key = None
+
+        if unload_registry:
+            from headroom.models.ml_models import MLModelRegistry
+
+            keys = [key for key in (classifier_key, siglip_key) if key]
+            MLModelRegistry.unload_many(keys)
+        else:
+            gc.collect()
+
+    close = release_models
+
     def _compute_text_embeddings(self) -> None:
         """Pre-compute SigLIP text embeddings for image analysis."""
         assert self._siglip_processor is not None

headroom/models/ml_models.py

@@ -26,6 +26,8 @@
 
 from __future__ import annotations
 
+import contextlib
+import gc
 import logging
 from threading import RLock
 from typing import TYPE_CHECKING, Any
@@ -76,6 +78,63 @@ def reset(cls) -> None:
                 cls._instance._models.clear()
             cls._instance = None
 
+    @classmethod
+    def _release_runtime_memory(cls) -> None:
+        """Best-effort cleanup after unloading heavyweight models."""
+        gc.collect()
+        try:
+            import torch
+        except ImportError:
+            return
+
+        with contextlib.suppress(Exception):
+            if torch.cuda.is_available():
+                torch.cuda.empty_cache()
+
+            mps = getattr(torch, "mps", None)
+            if mps is not None and hasattr(mps, "empty_cache"):
+                mps.empty_cache()
+
+    @classmethod
+    def unload(cls, key: str) -> bool:
+        """Unload one cached model entry."""
+        return bool(cls.unload_many([key]))
+
+    @classmethod
+    def unload_many(cls, keys: list[str]) -> list[str]:
+        """Unload several cached model entries with one runtime cleanup pass."""
+        instance = cls.get()
+        removed_keys: list[str] = []
+
+        with instance._model_lock:
+            for key in keys:
+                if key not in instance._models:
+                    continue
+                value = instance._models.pop(key)
+                del value
+                removed_keys.append(key)
+
+        if removed_keys:
+            cls._release_runtime_memory()
+        return removed_keys
+
+    @classmethod
+    def unload_prefix(cls, prefix: str) -> list[str]:
+        """Unload every cached model entry matching a prefix."""
+        instance = cls.get()
+        removed_keys: list[str] = []
+
+        with instance._model_lock:
+            for key in list(instance._models):
+                if key.startswith(prefix):
+                    value = instance._models.pop(key)
+                    del value
+                    removed_keys.append(key)
+
+        if removed_keys:
+            cls._release_runtime_memory()
+        return removed_keys
+
     # =========================================================================
     # Sentence Transformers
     # =========================================================================

headroom/proxy/handlers/anthropic.py

@@ -709,16 +709,21 @@ async def _finalize_pre_upstream() -> None:
                 and not _bypass
                 and not is_cache_mode(self.config.mode)
             ):
-                compressor = _get_image_compressor()
-                if compressor and compressor.has_images(messages):
-                    messages = compressor.compress(messages, provider="anthropic")
-                    if compressor.last_result:
-                        logger.info(
-                            f"Image compression: {compressor.last_result.technique.value} "
-                            f"({compressor.last_result.savings_percent:.0f}% saved, "
-                            f"{compressor.last_result.original_tokens} -> "
-                            f"{compressor.last_result.compressed_tokens} tokens)"
-                        )
+                compressor = None
+                try:
+                    compressor = _get_image_compressor()
+                    if compressor and compressor.has_images(messages):
+                        messages = compressor.compress(messages, provider="anthropic")
+                        if compressor.last_result:
+                            logger.info(
+                                f"Image compression: {compressor.last_result.technique.value} "
+                                f"({compressor.last_result.savings_percent:.0f}% saved, "
+                                f"{compressor.last_result.original_tokens} -> "
+                                f"{compressor.last_result.compressed_tokens} tokens)"
+                            )
+                finally:
+                    if compressor and hasattr(compressor, "close"):
+                        compressor.close()
 
             _compression_failed = False
             original_messages = messages  # Preserve for 400-retry fallback

headroom/proxy/handlers/openai.py

@@ -239,16 +239,21 @@ async def handle_openai_chat(
         if self.config.image_optimize and messages and not _bypass:
             from headroom.proxy.helpers import _get_image_compressor
 
-            compressor = _get_image_compressor()
-            if compressor and compressor.has_images(messages):
-                messages = compressor.compress(messages, provider="openai")
-                if compressor.last_result:
-                    logger.info(
-                        f"[{request_id}] Image: {compressor.last_result.technique.value} "
-                        f"({compressor.last_result.savings_percent:.0f}% saved, "
-                        f"{compressor.last_result.original_tokens} → "
-                        f"{compressor.last_result.compressed_tokens} tokens)"
-                    )
+            compressor = None
+            try:
+                compressor = _get_image_compressor()
+                if compressor and compressor.has_images(messages):
+                    messages = compressor.compress(messages, provider="openai")
+                    if compressor.last_result:
+                        logger.info(
+                            f"[{request_id}] Image: {compressor.last_result.technique.value} "
+                            f"({compressor.last_result.savings_percent:.0f}% saved, "
+                            f"{compressor.last_result.original_tokens} → "
+                            f"{compressor.last_result.compressed_tokens} tokens)"
+                        )
+            finally:
+                if compressor and hasattr(compressor, "close"):
+                    compressor.close()
 
         headers = dict(request.headers.items())
         headers.pop("host", None)

headroom/proxy/helpers.py

@@ -59,23 +59,31 @@ def jitter_delay_ms(base_ms: int, max_ms: int, attempt: int) -> float:
     return capped * (0.5 + random.random())
 
 
-# Image compression (lazy-loaded to avoid heavy dependencies at startup)
-_image_compressor = None
+# Image compression availability (do not retain a global compressor instance)
+_image_compressor_available: bool | None = None
 
 
 def _get_image_compressor():
-    """Lazy load image compressor to avoid startup overhead."""
-    global _image_compressor
-    if _image_compressor is None:
-        try:
-            from headroom.image import ImageCompressor
+    """Create a short-lived image compressor on demand."""
+    global _image_compressor_available
+    if _image_compressor_available is False:
+        return None
 
-            _image_compressor = ImageCompressor()
+    try:
+        from headroom.image import ImageCompressor
+
+        # Callers own closing the compressor; this helper only memoizes whether
+        # the optional image stack is importable.
+        compressor = ImageCompressor()
+        if _image_compressor_available is None:
             logger.info("Image compression enabled (model: chopratejas/technique-router)")
-        except ImportError as e:
+        _image_compressor_available = True
+        return compressor
+    except ImportError as e:
+        if _image_compressor_available is not False:
             logger.warning(f"Image compression not available: {e}")
-            _image_compressor = False  # Mark as unavailable
-    return _image_compressor if _image_compressor else None
+        _image_compressor_available = False
+        return None
 
 
 # Always-on file logging to the workspace logs directory for `headroom perf` analysis.

headroom/proxy/server.py

@@ -25,6 +25,7 @@
 
 import argparse
 import asyncio
+import contextlib
 import json
 import logging
 import os
@@ -844,6 +845,15 @@ async def shutdown(self):
         if self.memory_handler and hasattr(self.memory_handler, "close"):
             await self.memory_handler.close()
 
+        with contextlib.suppress(Exception):
+            from headroom.models.ml_models import MLModelRegistry
+
+            released_models: list[str] = []
+            released_models.extend(MLModelRegistry.unload_prefix("technique_router:"))
+            released_models.extend(MLModelRegistry.unload_prefix("siglip:"))
+            if released_models:
+                logger.info("Released image optimizer models: %s", ", ".join(released_models))
+
         # Stop all quota trackers via the registry
         await get_quota_registry().stop_all()