Labyricorn/voicebox
1
0
Fork 0
mirror of https://github.com/jamiepine/voicebox.git synced 2026-09-16 13:20:39 -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:
backend/main.py
+12 -3
View File
@@ -1512,7 +1512,10 @@ async def unload_model_by_name(model_name: str):
try:
if model_type == "tts":
tts_model = tts.get_tts_model()
if tts_model.is_loaded() and tts_model.model_size == model_size:
loaded_size = getattr(
tts_model, "_current_model_size", None
) or getattr(tts_model, "model_size", None)
if tts_model.is_loaded() and loaded_size == model_size:
tts.unload_tts_model()
else:
return {"message": f"Model {model_name} is not loaded"}
@@ -1595,7 +1598,10 @@ async def get_model_status():
"""Check if TTS model is loaded with specific size."""
try:
tts_model = tts.get_tts_model()
return tts_model.is_loaded() and getattr(tts_model, 'model_size', None) == model_size
loaded_size = getattr(
tts_model, "_current_model_size", None
) or getattr(tts_model, "model_size", None)
return tts_model.is_loaded() and loaded_size == model_size
except Exception:
return False
@@ -2072,7 +2078,10 @@ async def delete_model(model_name: str):
# Check if model is loaded and unload it first
if config["model_type"] == "tts":
tts_model = tts.get_tts_model()
if tts_model.is_loaded() and tts_model.model_size == config["model_size"]:
loaded_size = getattr(
tts_model, "_current_model_size", None
) or getattr(tts_model, "model_size", None)
if tts_model.is_loaded() and loaded_size == config["model_size"]:
tts.unload_tts_model()
elif config["model_type"] == "luxtts":
from .backends import get_tts_backend_for_engine