From 2f535a772fea5907e084dcb39edf1022832b5f6c Mon Sep 17 00:00:00 2001 From: James Pine Date: Fri, 13 Mar 2026 05:40:18 -0700 Subject: [PATCH] fix: load model into local var before patching to avoid half-initialised state 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. --- backend/backends/chatterbox_backend.py | 17 ++++++++++++----- backend/backends/chatterbox_turbo_backend.py | 13 +++++++++---- backend/main.py | 15 ++++++++++++--- 3 files changed, 33 insertions(+), 12 deletions(-) diff --git a/backend/backends/chatterbox_backend.py b/backend/backends/chatterbox_backend.py index a440f998..ed264e4d 100644 --- a/backend/backends/chatterbox_backend.py +++ b/backend/backends/chatterbox_backend.py @@ -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: diff --git a/backend/backends/chatterbox_turbo_backend.py b/backend/backends/chatterbox_turbo_backend.py index 3ce54f38..a87441a0 100644 --- a/backend/backends/chatterbox_turbo_backend.py +++ b/backend/backends/chatterbox_turbo_backend.py @@ -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: diff --git a/backend/main.py b/backend/main.py index 79a1d43d..27d200d5 100644 --- a/backend/main.py +++ b/backend/main.py @@ -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