Labyricorn/voicebox
1
0
Fork 0
mirror of https://github.com/jamiepine/voicebox.git synced 2026-09-19 14:50:38 -07:00
Code Issues Packages Projects Releases Wiki Activity

fix: load model into local var before patching to avoid half-initialised state

Browse Source 
Apply local-var-then-assign pattern to chatterbox_backend.py (multilingual)
to match the turbo backend. Also use _current_model_size fallback in
unload, delete, and status endpoints for consistent Qwen model size checks.
This commit is contained in:
James Pine
2026-03-13 05:40:18 -07:00
parent bfd7b815a5
commit 2f535a772f
3 changed files with 33 additions and 12 deletions
Download Patch File Download Diff File Expand all files Collapse all files
backend/backends/chatterbox_backend.py
+12 -5
View File
@@ -136,6 +136,10 @@ class ChatterboxTTSBackend:
import torch
from chatterbox.mtl_tts import ChatterboxMultilingualTTS
# Load into a local variable first, apply all patches, then
# assign to self.model. This avoids leaving a half-initialised
# model on self.model if any patch step raises an exception.
#
# Monkey-patch torch.load for CPU loading. The model's .pt files
# were saved on CUDA; from_pretrained() doesn't pass map_location
# so loading on CPU fails without this.
@@ -150,13 +154,13 @@ class ChatterboxTTSBackend:
with ChatterboxTTSBackend._load_lock:
torch.load = _patched_load
try:
self.model = ChatterboxMultilingualTTS.from_pretrained(
model = ChatterboxMultilingualTTS.from_pretrained(
device=device,
)
finally:
torch.load = _orig_torch_load
else:
self.model = ChatterboxMultilingualTTS.from_pretrained(
model = ChatterboxMultilingualTTS.from_pretrained(
device=device,
)
finally:
@@ -165,7 +169,7 @@ class ChatterboxTTSBackend:
# Fix: transformers >= 4.36 defaults LlamaModel to sdpa attention
# which doesn't support output_attentions=True (needed by
# Chatterbox's AlignmentStreamAnalyzer). Force eager attention.
t3_tfmr = self.model.t3.tfmr
t3_tfmr = model.t3.tfmr
if hasattr(t3_tfmr, "config") and hasattr(
t3_tfmr.config, "_attn_implementation"
):
@@ -185,7 +189,7 @@ class ChatterboxTTSBackend:
import types
# Patch S3Tokenizer (used by s3gen.tokenizer)
_tokzr = self.model.s3gen.tokenizer
_tokzr = model.s3gen.tokenizer
_orig_log_mel = _tokzr.log_mel_spectrogram.__func__
def _f32_log_mel(self_tokzr, audio, padding=0):
@@ -197,7 +201,7 @@ class ChatterboxTTSBackend:
_tokzr.log_mel_spectrogram = types.MethodType(_f32_log_mel, _tokzr)
# Patch VoiceEncoder
_ve = self.model.ve
_ve = model.ve
_orig_ve_forward = _ve.forward.__func__
def _f32_ve_forward(self_ve, mels):
@@ -205,6 +209,9 @@ class ChatterboxTTSBackend:
_ve.forward = types.MethodType(_f32_ve_forward, _ve)
# All patches applied successfully — publish the model
self.model = model
logger.info("Chatterbox Multilingual TTS loaded successfully")
except ImportError as e:
backend/backends/chatterbox_turbo_backend.py
+9 -4
View File
@@ -154,6 +154,8 @@ class ChatterboxTurboTTSBackend:
# Monkey-patch torch.load for CPU loading. The model's .pt files
# were saved on CUDA; from_local() doesn't pass map_location
# so loading on CPU fails without this.
# Load into a local var, apply patches, then publish to
# self.model so a failed patch doesn't leave us half-initialised.
if device == "cpu":
_orig_torch_load = torch.load
@@ -164,13 +166,13 @@ class ChatterboxTurboTTSBackend:
with ChatterboxTurboTTSBackend._load_lock:
torch.load = _patched_load
try:
self.model = ChatterboxTurboTTS.from_local(
model = ChatterboxTurboTTS.from_local(
local_path, device,
)
finally:
torch.load = _orig_torch_load
else:
self.model = ChatterboxTurboTTS.from_local(
model = ChatterboxTurboTTS.from_local(
local_path, device,
)
@@ -191,7 +193,7 @@ class ChatterboxTurboTTSBackend:
import types
# Patch S3Tokenizer (used by s3gen.tokenizer)
_tokzr = self.model.s3gen.tokenizer
_tokzr = model.s3gen.tokenizer
_orig_log_mel = _tokzr.log_mel_spectrogram.__func__
def _f32_log_mel(self_tokzr, audio, padding=0):
@@ -203,7 +205,7 @@ class ChatterboxTurboTTSBackend:
_tokzr.log_mel_spectrogram = types.MethodType(_f32_log_mel, _tokzr)
# Patch VoiceEncoder
_ve = self.model.ve
_ve = model.ve
_orig_ve_forward = _ve.forward.__func__
def _f32_ve_forward(self_ve, mels):
@@ -211,6 +213,9 @@ class ChatterboxTurboTTSBackend:
_ve.forward = types.MethodType(_f32_ve_forward, _ve)
# Only publish after all patches succeed
self.model = model
logger.info("Chatterbox Turbo TTS loaded successfully")
except ImportError as e: