fix: resolve VRAM leak in multimodal models by explicitly closing mtmd context

- Remove the `ExitStack` closure in `Llava15ChatHandler` to break circular references preventing garbage collection of the vision context.
- Implement explicit `close()` and `__del__()` methods in the chat handler to safely free `mtmd_ctx`.
- Integrate `chat_handler.close()` into the main `Llama.close()` lifecycle and nullify related attributes for immediate memory reclamation.

Author: JamePeng

llama_cpp/llama.py

@@ -634,8 +634,12 @@ def close(self) -> None:
             self._candidates.close()
             self._candidates = None
 
+        if hasattr(self, "chat_handler") and hasattr(self.chat_handler, "close"):
+            self.chat_handler.close()
+
         self.model_params =None
         self.context_params = None
+        self.chat_handler = None
         self.input_ids = None
         self.scores = None
         self.tokenizer_ = None

llama_cpp/llama_chat_format.py

@@ -2850,24 +2850,24 @@ def _init_mtmd_context(self, llama_model: llama.Llama):
             self._mtmd_cpp.mtmd_helper_log_set(llama_log_callback, ctypes.c_void_p(0))
 
             # Get default parameters
-            mctx_params = self._mtmd_cpp.mtmd_context_params_default()
-            mctx_params.use_gpu = self.use_gpu
-            mctx_params.print_timings = self.verbose
-            mctx_params.n_threads = llama_model.n_threads
-            mctx_params.flash_attn_type  = self._mtmd_cpp.clip_flash_attn_type.CLIP_FLASH_ATTN_TYPE_AUTO
-            mctx_params.warmup = True
+            self.mctx_params = self._mtmd_cpp.mtmd_context_params_default()
+            self.mctx_params.use_gpu = self.use_gpu
+            self.mctx_params.print_timings = self.verbose
+            self.mctx_params.n_threads = llama_model.n_threads
+            self.mctx_params.flash_attn_type  = self._mtmd_cpp.clip_flash_attn_type.CLIP_FLASH_ATTN_TYPE_AUTO
+            self.mctx_params.warmup = True
             if self.image_min_tokens > 0:
-                mctx_params.image_min_tokens = self.image_min_tokens
+                self.mctx_params.image_min_tokens = self.image_min_tokens
             if self.image_max_tokens > 0:
-                mctx_params.image_max_tokens = self.image_max_tokens
+                self.mctx_params.image_max_tokens = self.image_max_tokens
             if (self.image_max_tokens < self.image_min_tokens) and self.image_max_tokens > 0:
                 raise ValueError(f"image_max_pixels {self.image_max_tokens} is less than image_min_pixels {self.image_min_tokens}")
 
             # Initialize mtmd context
             self.mtmd_ctx = self._mtmd_cpp.mtmd_init_from_file(
                 self.clip_model_path.encode(),
                 llama_model.model,
-                mctx_params
+                self.mctx_params
             )
 
             if self.mtmd_ctx is None:
@@ -2877,13 +2877,21 @@ def _init_mtmd_context(self, llama_model: llama.Llama):
             if not self._mtmd_cpp.mtmd_support_vision(self.mtmd_ctx):
                 raise ValueError("Vision is not supported by this model")
 
-            def mtmd_free():
-                with suppress_stdout_stderr(disable=self.verbose):
-                    if self.mtmd_ctx is not None:
-                        self._mtmd_cpp.mtmd_free(self.mtmd_ctx)
-                        self.mtmd_ctx = None
+    def close(self) -> None:
+        """Explicitly free the mtmd context and vision model resources."""
+        if getattr(self, "mtmd_ctx", None) is not None:
+            try:
+                with suppress_stdout_stderr(disable=getattr(self, "verbose", True)):
+                    self._mtmd_cpp.mtmd_free(self.mtmd_ctx)
+            except Exception:
+                pass
+            self.mtmd_ctx = None
+            self.mctx_params = None
+
+        self._exit_stack.close()
 
-            self._exit_stack.callback(mtmd_free)
+    def __del__(self) -> None:
+        self.close()
 
     def load_image(self, image_url: str) -> bytes:
         return self._load_image(image_url)