mirror of
https://github.com/jamiepine/voicebox.git
synced 2026-10-04 01:25:18 -07:00
fix(mlx): release the model binding inside the sync bodies so a failing op's traceback cannot pin it
Review follow-up: the post-op drain ran in a finally, but on the error path the exception's traceback kept the _generate_sync/_transcribe_sync frame (and its model local) alive, so mx.clear_cache() had nothing to return. Each sync body now drops its model (and tokenizer) binding in its own finally before the exception propagates.
This commit is contained in:
committed by
capy-ai-staging[bot]
parent
fb67ad436c
commit
63899fd865
@@ -241,6 +241,7 @@ class MLXTTSBackend:
|
|||||||
# thread mid-generation cannot turn a later self.model read into
|
# thread mid-generation cannot turn a later self.model read into
|
||||||
# None (the fallback path below runs seconds into a request).
|
# None (the fallback path below runs seconds into a request).
|
||||||
model = self.model
|
model = self.model
|
||||||
|
try:
|
||||||
# MLX generate() returns a generator yielding GenerationResult objects
|
# MLX generate() returns a generator yielding GenerationResult objects
|
||||||
audio_chunks = []
|
audio_chunks = []
|
||||||
sample_rate = 24000
|
sample_rate = 24000
|
||||||
@@ -305,6 +306,11 @@ class MLXTTSBackend:
|
|||||||
audio = np.array([], dtype=np.float32)
|
audio = np.array([], dtype=np.float32)
|
||||||
|
|
||||||
return audio, sample_rate
|
return audio, sample_rate
|
||||||
|
finally:
|
||||||
|
# Drop the local binding here, inside the frame a propagating
|
||||||
|
# traceback would keep alive, so the caller's drain can
|
||||||
|
# actually return the model's buffers to MLX.
|
||||||
|
del model
|
||||||
|
|
||||||
def _reload_and_generate_sync():
|
def _reload_and_generate_sync():
|
||||||
"""Ensure the configured model is loaded, then generate — as ONE
|
"""Ensure the configured model is loaded, then generate — as ONE
|
||||||
@@ -324,9 +330,10 @@ class MLXTTSBackend:
|
|||||||
finally:
|
finally:
|
||||||
# An unload_model() that landed while we were generating only
|
# An unload_model() that landed while we were generating only
|
||||||
# dropped the backend's reference; the model's buffers were
|
# dropped the backend's reference; the model's buffers were
|
||||||
# kept alive by _generate_sync's local and have just been
|
# kept alive by _generate_sync's local, which that function
|
||||||
# returned to MLX's pool (on success or failure). Drain it now,
|
# drops in its own finally (so a propagating traceback cannot
|
||||||
# or they stay resident until the next load/unload cycle.
|
# pin it), and have just been returned to MLX's pool. Drain it
|
||||||
|
# now, or they stay resident until the next load/unload cycle.
|
||||||
if self.model is None:
|
if self.model is None:
|
||||||
empty_mlx_cache()
|
empty_mlx_cache()
|
||||||
|
|
||||||
@@ -423,6 +430,7 @@ class MLXSTTBackend:
|
|||||||
# MLX Whisper transcription using generate method
|
# MLX Whisper transcription using generate method
|
||||||
# The generate method accepts audio path directly
|
# The generate method accepts audio path directly
|
||||||
model = self.model
|
model = self.model
|
||||||
|
try:
|
||||||
decode_options = {}
|
decode_options = {}
|
||||||
if language:
|
if language:
|
||||||
decode_options["language"] = language
|
decode_options["language"] = language
|
||||||
@@ -441,6 +449,11 @@ class MLXSTTBackend:
|
|||||||
return result.text.strip()
|
return result.text.strip()
|
||||||
else:
|
else:
|
||||||
return str(result).strip()
|
return str(result).strip()
|
||||||
|
finally:
|
||||||
|
# Drop the local binding here, inside the frame a propagating
|
||||||
|
# traceback would keep alive, so the caller's drain can
|
||||||
|
# actually return the model's buffers to MLX.
|
||||||
|
del model
|
||||||
|
|
||||||
def _reload_and_transcribe_sync():
|
def _reload_and_transcribe_sync():
|
||||||
"""Ensure the requested model is loaded, then transcribe — as ONE
|
"""Ensure the requested model is loaded, then transcribe — as ONE
|
||||||
|
|||||||
@@ -305,6 +305,7 @@ class MLXQwenLLMBackend:
|
|||||||
from mlx_lm.sample_utils import make_sampler
|
from mlx_lm.sample_utils import make_sampler
|
||||||
|
|
||||||
model, tokenizer = self.model, self.tokenizer
|
model, tokenizer = self.model, self.tokenizer
|
||||||
|
try:
|
||||||
messages = _build_messages(prompt, system, examples)
|
messages = _build_messages(prompt, system, examples)
|
||||||
chat_prompt = tokenizer.apply_chat_template(
|
chat_prompt = tokenizer.apply_chat_template(
|
||||||
messages,
|
messages,
|
||||||
@@ -323,3 +324,8 @@ class MLXQwenLLMBackend:
|
|||||||
verbose=False,
|
verbose=False,
|
||||||
)
|
)
|
||||||
return text.strip()
|
return text.strip()
|
||||||
|
finally:
|
||||||
|
# Drop the local binding here, inside the frame a propagating
|
||||||
|
# traceback would keep alive, so the caller's drain can
|
||||||
|
# actually return the model's buffers to MLX.
|
||||||
|
del model, tokenizer
|
||||||
|
|||||||
Reference in New Issue
Block a user