Fix ZeroGPU and private model loading

app.py

@@ -6,6 +6,8 @@ import gradio as gr
 import spaces
 
 MODEL = None
+# ZeroGPU supports CUDA placement at module load time via CUDA emulation.
+TARGET_DEVICE = "cuda"
 
 DEFAULT_CONFIG = {
     "audio": 'https://storage.googleapis.com/chatterbox-demo-samples/mtl-v3-single-language-prompts/es-latam/es_mx_f1.wav',
@@ -32,9 +34,15 @@ def default_text_for_ui():
 def get_or_load_model():
     global MODEL
     if MODEL is None:
-        print("Model not loaded, initializing on CPU...")
-        MODEL = ChatterboxTTS.from_pretrained("cpu")
-        print("Model loaded.")
+        print(f"Model not loaded, initializing on {TARGET_DEVICE}...")
+        try:
+            MODEL = ChatterboxTTS.from_pretrained(TARGET_DEVICE)
+        except Exception as exc:
+            if TARGET_DEVICE != "cuda":
+                raise
+            print(f"CUDA model initialization failed, falling back to CPU: {exc}")
+            MODEL = ChatterboxTTS.from_pretrained("cpu")
+        print(f"Model loaded on {MODEL.device}.")
     return MODEL
 
 
@@ -57,9 +65,8 @@ def generate_tts_audio(
     cfgw_input: float = 0.5,
 ):
     """Generate speech from text with optional reference audio styling."""
-    device = "cuda" if torch.cuda.is_available() else "cpu"
     current_model = get_or_load_model()
-    current_model.to(device)
+    device = current_model.device
     if seed_num_input != 0:
         set_seed(int(seed_num_input), device)
     chosen_prompt = audio_prompt_path_input or default_audio_for_ui()
@@ -77,6 +84,9 @@ def generate_tts_audio(
     return (current_model.sr, wav.squeeze(0).cpu().numpy())
 
 
+get_or_load_model()
+
+
 with gr.Blocks() as demo:
     gr.Markdown(
         """

chatterbox/src/chatterbox/models/t3/inference/t3_hf_backend.py

@@ -23,14 +23,12 @@ class T3HuggingfaceBackend(LlamaPreTrainedModel, GenerationMixin):
         speech_head,
         latents_queue=None,
         logits_queue=None,
-        alignment_stream_analyzer: 'AlignmentStreamAnalyzer'=None,
     ):
         super().__init__(config)
         self.model = llama
         self.speech_enc = speech_enc
         self.speech_head = speech_head
         self._added_cond = False
-        self.alignment_stream_analyzer = alignment_stream_analyzer
 
     @torch.inference_mode()
     def prepare_inputs_for_generation(
@@ -105,9 +103,6 @@ class T3HuggingfaceBackend(LlamaPreTrainedModel, GenerationMixin):
         logits = self.speech_head(hidden_states)
         # assert inputs_embeds.size(0) == 1 # (disabled for CFG)
 
-        # NOTE: hallucination handler may modify logits to force emit an EOS token
-        # logits = self.alignment_stream_analyzer.step(logits)
-
         return CausalLMOutputWithCrossAttentions(
             logits=logits,
             past_key_values=tfmr_out.past_key_values,

chatterbox/src/chatterbox/models/t3/t3.py

@@ -16,7 +16,6 @@ from .modules.cond_enc import T3CondEnc, T3Cond
 from .modules.t3_config import T3ConfigMultilingual
 from .llama_configs import LLAMA_CONFIGS
 from .inference.t3_hf_backend import T3HuggingfaceBackend
-from .inference.alignment_stream_analyzer import AlignmentStreamAnalyzer
 
 
 logger = logging.getLogger(__name__)
@@ -255,21 +254,11 @@ class T3(nn.Module):
         # TODO? synchronize the expensive compile function
         # with self.compile_lock:
         if not self.compiled:
-            alignment_stream_analyzer = AlignmentStreamAnalyzer(
-                self.tfmr,
-                None,
-                text_tokens_slice=(len_cond, len_cond + text_tokens.size(-1)),
-                alignment_layer_idx=9, # TODO: hparam or something?
-                eos_idx=self.hp.stop_speech_token,
-            )
-            assert alignment_stream_analyzer.eos_idx == self.hp.stop_speech_token
-
             patched_model = T3HuggingfaceBackend(
                 config=self.cfg,
                 llama=self.tfmr,
                 speech_enc=self.speech_emb,
                 speech_head=self.speech_head,
-                alignment_stream_analyzer=alignment_stream_analyzer,
             )
             self.patched_model = patched_model
             self.compiled = True
@@ -317,7 +306,7 @@ class T3(nn.Module):
             inputs_embeds=inputs_embeds,
             past_key_values=None,
             use_cache=True,
-            output_attentions=True,
+            output_attentions=False,
             output_hidden_states=True,
             return_dict=True,
         )
@@ -333,11 +322,6 @@ class T3(nn.Module):
             cfg = torch.as_tensor(cfg_weight, device=cond.device, dtype=cond.dtype)
             logits = cond + cfg * (cond - uncond)
             
-            if self.patched_model.alignment_stream_analyzer is not None:
-                if logits.dim() == 1:            # guard in case something upstream squeezed
-                    logits = logits.unsqueeze(0) # (1, V)
-                logits = self.patched_model.alignment_stream_analyzer.step(logits)  # (1, V)
-
             # Apply repetition penalty
             ids_for_proc = generated_ids[:1, ...]   # batch = 1
             logits = repetition_penalty_processor(ids_for_proc, logits)  # expects (B,V)
@@ -372,7 +356,7 @@ class T3(nn.Module):
             output = self.patched_model(
                 inputs_embeds=next_token_embed,
                 past_key_values=past,
-                output_attentions=True,
+                output_attentions=False,
                 output_hidden_states=True,
                 return_dict=True,
             )

chatterbox/src/chatterbox/tts.py

@@ -11,7 +11,7 @@ from huggingface_hub import snapshot_download, hf_hub_download
 
 from .models.t3 import T3
 from .models.t3.modules.t3_config import T3ConfigMultilingual
-from .models.s3tokenizer import S3_SR, drop_invalid_tokens
+from .models.s3tokenizer import S3_SR, S3_TOKEN_RATE, drop_invalid_tokens
 from .models.s3gen import S3GEN_SR, S3Gen
 from .models.tokenizers import MTLTokenizer
 from .models.voice_encoder import VoiceEncoder
@@ -312,5 +312,12 @@ class ChatterboxTTS:
                 ref_dict=self.conds.gen,
             )
             wav = wav.squeeze(0).detach().cpu().numpy()
+
+            # Drop the final speech token's audio: it is emitted just before
+            # EOS with degraded attention and decodes to ~40 ms of noise.
+            n_tokens = int(speech_tokens.shape[-1])
+            st_len = max(1, n_tokens - 1)
+            wav = wav[: st_len * (S3GEN_SR // S3_TOKEN_RATE)]
+
             watermarked_wav = self.watermarker.apply_watermark(wav, sample_rate=self.sr)
         return torch.from_numpy(watermarked_wav).unsqueeze(0)

requirements.txt

@@ -1,9 +1,12 @@
+--extra-index-url https://download.pytorch.org/whl/cu128
+torch==2.8.0
+torchaudio==2.8.0
+
 gradio
 numpy==1.26.0
 resampy==0.4.3
 librosa==0.10.0
 s3tokenizer
-torchaudio<2.8
 
 transformers==4.46.3
 diffusers==0.29.0
@@ -11,4 +14,4 @@ omegaconf==2.3.0
 resemble-perth==1.0.1
 silero-vad==5.1.2
 conformer==0.3.2
-safetensors
+safetensors