fix(llama): release a vision chat handler's mtmd context on close (#2342)

[Anai-Guo]

Summary

Fixes #2342.

Llama.close() unwinds self._stack, which frees the model, context and batch — but the chat handler's own _exit_stack is never touched. That stack holds the mtmd_free callback for the mtmd/clip context, so a vision handler's context is simply leaked on close.

The crash in #2342 follows from that leak. A handler routinely outlives the Llama that initialized it (callers build one handler and reuse it across loads, as ComfyUI's keep_model_alive: false path does), and _init_mtmd_context() returns early while mtmd_ctx is set:

def _init_mtmd_context(self, llama_model: llama.Llama):
    if self.mtmd_ctx is not None:
        return  # Already initialized

The context was built from the first model (mtmd_init_from_file(..., llama_model.model, ...)). After close() frees that model, the handler still believes it is initialized and hands the stale context to the second Llama — which is where mtmd_ctx shows up null on the C++ side.

Fix

Register the handler's exit stack on the Llama's _stack in __init__:

handler_stack = getattr(chat_handler, "_exit_stack", None)
if handler_stack is not None:
    self._stack.callback(handler_stack.close)

• Ordering is correct by construction. ExitStack unwinds LIFO and the model/context/batch are entered earlier in __init__, so the mtmd context is released before the model it was built from.
• Re-initialization now works. mtmd_free sets self.mtmd_ctx = None, so the next load re-initializes against the new model instead of reusing a dangling one. ExitStack is reusable after close(), so the handler stays usable.
• Covers both handler families. Duck-typed on _exit_stack, which Llava15ChatHandler and MTMDChatHandler — and therefore every subclass (Gemma4ChatHandler, Qwen25VLChatHandler, MiniCPMv26ChatHandler, …) — expose. getattr keeps plain non-vision handlers and chat_handler=None unaffected.

This is the upstream equivalent of the workaround the reporter applied downstream in _drop_local_llama.

Verification

I don't have the reporter's vision model to hand, so I verified the lifecycle contract directly by reproducing the object graph (handler with _exit_stack + mtmd_ctx early-return, Llama with an _exit_stack-based teardown) and driving two load/close cycles against it.

Without the fix, the second cycle skips re-initialization and reuses the context bound to the freed model — the reported failure. With it, the trace is:

init ctx@m1 / free ctx@m1 / free model m1
init ctx@m2 / free ctx@m2 / free model m2

mtmd freed ahead of the model on both cycles, and a fresh context on the second load.

Scope

One file, +10/-0, no behavior change for non-vision handlers.

🤖 Generated with Claude Code