Merge pull request #264 from jamiepine/fix/chatterbox-float64-dtype

fix: Chatterbox float64 dtype mismatch + model unload button

app/src/components/ServerSettings/ModelManagement.tsx

@@ -13,6 +13,7 @@ import {
   RotateCcw,
   Scale,
   Trash2,
+  Unplug,
   X,
 } from 'lucide-react';
 import { useCallback, useMemo, useState } from 'react';
@@ -300,6 +301,27 @@ export function ModelManagement() {
     },
   });
 
+  const unloadMutation = useMutation({
+    mutationFn: async (modelName: string) => {
+      return await apiClient.unloadModel(modelName);
+    },
+    onSuccess: async (_data, modelName) => {
+      toast({
+        title: 'Model unloaded',
+        description: `${modelName} has been unloaded from memory.`,
+      });
+      await queryClient.invalidateQueries({ queryKey: ['modelStatus'], refetchType: 'all' });
+      await queryClient.refetchQueries({ queryKey: ['modelStatus'] });
+    },
+    onError: (error: Error) => {
+      toast({
+        title: 'Unload failed',
+        description: error.message,
+        variant: 'destructive',
+      });
+    },
+  });
+
   const formatSize = (sizeMb?: number): string => {
     if (!sizeMb) return 'Unknown size';
     if (sizeMb < 1024) return `${sizeMb.toFixed(1)} MB`;
@@ -697,26 +719,46 @@ export function ModelManagement() {
                       </Button>
                     </>
                   ) : freshSelectedModel.downloaded ? (
-                    <Button
-                      size="sm"
-                      onClick={() => {
-                        setModelToDelete({
-                          name: freshSelectedModel.model_name,
-                          displayName: freshSelectedModel.display_name,
-                          sizeMb: freshSelectedModel.size_mb,
-                        });
-                        setDeleteDialogOpen(true);
-                      }}
-                      variant="outline"
-                      disabled={freshSelectedModel.loaded}
-                      title={
-                        freshSelectedModel.loaded ? 'Unload model before deleting' : 'Delete model'
-                      }
-                      className="flex-1"
-                    >
-                      <Trash2 className="h-4 w-4 mr-2" />
-                      {freshSelectedModel.loaded ? 'Unload to Delete' : 'Delete Model'}
-                    </Button>
+                    <div className="flex gap-2 flex-1">
+                      {freshSelectedModel.loaded && (
+                        <Button
+                          size="sm"
+                          onClick={() => unloadMutation.mutate(freshSelectedModel.model_name)}
+                          variant="outline"
+                          disabled={unloadMutation.isPending}
+                          className="flex-1"
+                        >
+                          {unloadMutation.isPending ? (
+                            <Loader2 className="h-4 w-4 mr-2 animate-spin" />
+                          ) : (
+                            <Unplug className="h-4 w-4 mr-2" />
+                          )}
+                          {unloadMutation.isPending ? 'Unloading...' : 'Unload'}
+                        </Button>
+                      )}
+                      <Button
+                        size="sm"
+                        onClick={() => {
+                          setModelToDelete({
+                            name: freshSelectedModel.model_name,
+                            displayName: freshSelectedModel.display_name,
+                            sizeMb: freshSelectedModel.size_mb,
+                          });
+                          setDeleteDialogOpen(true);
+                        }}
+                        variant="outline"
+                        disabled={freshSelectedModel.loaded}
+                        title={
+                          freshSelectedModel.loaded
+                            ? 'Unload model before deleting'
+                            : 'Delete model'
+                        }
+                        className="flex-1"
+                      >
+                        <Trash2 className="h-4 w-4 mr-2" />
+                        Delete Model
+                      </Button>
+                    </div>
                   ) : (
                     <Button
                       size="sm"

app/src/lib/api/client.ts

@@ -337,6 +337,12 @@ class ApiClient {
     });
   }
 
+  async unloadModel(modelName: string): Promise<{ message: string }> {
+    return this.request<{ message: string }>(`/models/${modelName}/unload`, {
+      method: 'POST',
+    });
+  }
+
   async cancelDownload(modelName: string): Promise<{ message: string }> {
     return this.request<{ message: string }>('/models/download/cancel', {
       method: 'POST',

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"
             ):
@@ -178,6 +182,36 @@ class ChatterboxTTSBackend:
                 progress_manager.mark_complete(model_name)
                 task_manager.complete_download(model_name)
 
+            # Patch float64 → float32 dtype mismatches in upstream chatterbox.
+            # librosa.load returns float64 numpy; multiple upstream code paths
+            # convert it to a torch tensor via torch.from_numpy() without
+            # casting, then matmul it against float32 model weights.
+            import types
+
+            # Patch S3Tokenizer (used by s3gen.tokenizer)
+            _tokzr = model.s3gen.tokenizer
+            _orig_log_mel = _tokzr.log_mel_spectrogram.__func__
+
+            def _f32_log_mel(self_tokzr, audio, padding=0):
+                import torch as _torch
+                if _torch.is_tensor(audio):
+                    audio = audio.float()
+                return _orig_log_mel(self_tokzr, audio, padding)
+
+            _tokzr.log_mel_spectrogram = types.MethodType(_f32_log_mel, _tokzr)
+
+            # Patch VoiceEncoder
+            _ve = model.ve
+            _orig_ve_forward = _ve.forward.__func__
+
+            def _f32_ve_forward(self_ve, mels):
+                return _orig_ve_forward(self_ve, mels.float())
+
+            _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

@@ -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,
                 )
 
@@ -178,6 +180,42 @@ class ChatterboxTurboTTSBackend:
                 progress_manager.mark_complete(model_name)
                 task_manager.complete_download(model_name)
 
+            # Patch float64 → float32 dtype mismatches in upstream chatterbox.
+            # librosa.load returns float64 numpy; multiple upstream code paths
+            # convert it to a torch tensor via torch.from_numpy() without
+            # casting, then matmul it against float32 model weights.
+            # We patch the two known entry points:
+            #
+            # 1. S3Tokenizer.log_mel_spectrogram — the audio tensor from
+            #    librosa hits _mel_filters (float32) in a matmul.
+            # 2. VoiceEncoder.forward — float64 mel spectrograms hit the
+            #    float32 LSTM weights.
+            import types
+
+            # Patch S3Tokenizer (used by s3gen.tokenizer)
+            _tokzr = model.s3gen.tokenizer
+            _orig_log_mel = _tokzr.log_mel_spectrogram.__func__
+
+            def _f32_log_mel(self_tokzr, audio, padding=0):
+                import torch as _torch
+                if _torch.is_tensor(audio):
+                    audio = audio.float()
+                return _orig_log_mel(self_tokzr, audio, padding)
+
+            _tokzr.log_mel_spectrogram = types.MethodType(_f32_log_mel, _tokzr)
+
+            # Patch VoiceEncoder
+            _ve = model.ve
+            _orig_ve_forward = _ve.forward.__func__
+
+            def _f32_ve_forward(self_ve, mels):
+                return _orig_ve_forward(self_ve, mels.float())
+
+            _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

@@ -1479,7 +1479,7 @@ async def load_model(model_size: str = "1.7B"):
 
 @app.post("/models/unload")
 async def unload_model():
-    """Unload TTS model to free memory."""
+    """Unload the default Qwen TTS model to free memory."""
     try:
         tts.unload_tts_model()
         return {"message": "Model unloaded successfully"}
@@ -1487,6 +1487,71 @@ async def unload_model():
         raise HTTPException(status_code=500, detail=str(e))
 
 
+@app.post("/models/{model_name}/unload")
+async def unload_model_by_name(model_name: str):
+    """Unload a specific model from memory without deleting it from disk."""
+    # Map of model_name -> (model_type, model_size)
+    model_types = {
+        "qwen-tts-1.7B": ("tts", "1.7B"),
+        "qwen-tts-0.6B": ("tts", "0.6B"),
+        "luxtts": ("luxtts", "default"),
+        "chatterbox-tts": ("chatterbox", "default"),
+        "chatterbox-turbo": ("chatterbox_turbo", "default"),
+        "whisper-base": ("whisper", "base"),
+        "whisper-small": ("whisper", "small"),
+        "whisper-medium": ("whisper", "medium"),
+        "whisper-large": ("whisper", "large"),
+        "whisper-turbo": ("whisper", "turbo"),
+    }
+
+    if model_name not in model_types:
+        raise HTTPException(status_code=400, detail=f"Unknown model: {model_name}")
+
+    model_type, model_size = model_types[model_name]
+
+    try:
+        if model_type == "tts":
+            tts_model = tts.get_tts_model()
+            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"}
+        elif model_type == "luxtts":
+            from .backends import get_tts_backend_for_engine
+            backend = get_tts_backend_for_engine("luxtts")
+            if backend.is_loaded():
+                backend.unload_model()
+            else:
+                return {"message": f"Model {model_name} is not loaded"}
+        elif model_type == "chatterbox":
+            from .backends import get_tts_backend_for_engine
+            backend = get_tts_backend_for_engine("chatterbox")
+            if backend.is_loaded():
+                backend.unload_model()
+            else:
+                return {"message": f"Model {model_name} is not loaded"}
+        elif model_type == "chatterbox_turbo":
+            from .backends import get_tts_backend_for_engine
+            backend = get_tts_backend_for_engine("chatterbox_turbo")
+            if backend.is_loaded():
+                backend.unload_model()
+            else:
+                return {"message": f"Model {model_name} is not loaded"}
+        elif model_type == "whisper":
+            whisper_model = transcribe.get_whisper_model()
+            if whisper_model.is_loaded() and whisper_model.model_size == model_size:
+                transcribe.unload_whisper_model()
+            else:
+                return {"message": f"Model {model_name} is not loaded"}
+
+        return {"message": f"Model {model_name} unloaded successfully"}
+    except Exception as e:
+        raise HTTPException(status_code=500, detail=str(e))
+
+
 @app.get("/models/progress/{model_name}")
 async def get_model_progress(model_name: str):
     """Get model download progress via Server-Sent Events."""
@@ -1533,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
     
@@ -2010,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