deploy: switch to chatterbox requirements @ 98aec56

steps/_tts_models.py

@@ -1,34 +0,0 @@
-"""Process-cached loaders for TTS models (Chatterbox).
-
-ZeroGPU best practice: load weights to `cuda` outside `@spaces.GPU` scopes
-(via CUDA emulation) so the time-budgeted GPU calls only contain inference.
-On Mac/CPU these fall back to MPS or CPU.
-
-Callers should treat returned models as singletons — never call
-`del model` or `torch.cuda.empty_cache()` on them between pipeline steps.
-"""
-from __future__ import annotations
-
-import torch
-
-
-_CHATTERBOX = None
-
-
-def _select_device() -> str:
-    if torch.cuda.is_available():
-        return "cuda"
-    if hasattr(torch.backends, "mps") and torch.backends.mps.is_available():
-        return "mps"
-    return "cpu"
-
-
-def get_chatterbox():
-    global _CHATTERBOX
-    if _CHATTERBOX is None:
-        from chatterbox.mtl_tts import ChatterboxMultilingualTTS
-
-        device = _select_device()
-        print(f"[tts] Loading Chatterbox Multilingual on {device}...")
-        _CHATTERBOX = ChatterboxMultilingualTTS.from_pretrained(device)
-    return _CHATTERBOX

steps/s4_preview.py

@@ -20,12 +20,6 @@ import torchaudio
 
 TTS_ENGINE = os.getenv("TTS_ENGINE", "chatterbox").lower()
 
-# Conditional imports based on TTS_ENGINE
-if TTS_ENGINE == "chatterbox":
-    from steps._tts_models import get_chatterbox
-else:
-    get_chatterbox = None
-
 import spaces
 
 
@@ -90,21 +84,44 @@ def _clip_audio(path: str, max_sec: float = 10.0) -> str:
     return path
 
 
-@spaces.GPU(duration=30)
-def _gpu_preview_chatterbox_segment(
-    model,
-    text: str,
+@spaces.GPU(duration=60)
+def _gpu_preview_chatterbox_batch(
+    segments: list[dict],
+    ref_audio_clipped: str,
     language_id: str,
-    ref_audio_path: str,
+    output_dir: str,
 ):
-    return model.generate(
-        text[:300],
-        language_id=language_id,
-        audio_prompt_path=ref_audio_path,
-        exaggeration=0.5,
-        temperature=0.8,
-        cfg_weight=0.5,
-    )
+    """Load + run Chatterbox preview synthesis inside one GPU scope."""
+    from chatterbox.mtl_tts import ChatterboxMultilingualTTS
+
+    print("   [preview] Loading Chatterbox in GPU scope...")
+    model = ChatterboxMultilingualTTS.from_pretrained("cuda")
+    part_paths = []
+    total = len(segments)
+
+    for i, seg in enumerate(segments):
+        text = seg.get("tts_text", seg.get("translated_text", seg["text"]))
+        out_path = os.path.join(output_dir, f"cb_prev_{i:04d}.wav")
+
+        print(f"   [preview] Chatterbox: Synthesising segment {i+1}/{total}...")
+        wav = model.generate(
+            text[:300],
+            language_id=language_id,
+            audio_prompt_path=ref_audio_clipped,
+            exaggeration=0.5,
+            temperature=0.8,
+            cfg_weight=0.5,
+        )
+        torchaudio.save(
+            out_path,
+            wav.detach().cpu(),
+            model.sr,
+            encoding="PCM_S",
+            bits_per_sample=16,
+        )
+        part_paths.append(out_path)
+
+    return part_paths
 
 
 # ── Chatterbox Multilingual preview ──────────────────────────
@@ -116,12 +133,28 @@ def _preview_chatterbox(
 ):
     """Generate a stitched preview WAV using Chatterbox Multilingual."""
     try:
-        yield "   [preview] Preparing Chatterbox Multilingual...\n"
-        model = get_chatterbox()
-
         # Clip reference audio to max 10 seconds to prevent weird noise/artifacts
         ref_audio_clipped = _clip_audio(reference_audio_path, max_sec=10.0)
 
+        device = _get_device()
+        if device == "cuda":
+            yield "   [preview] Preparing Chatterbox batch preview (device=cuda)...\n"
+            part_paths = _gpu_preview_chatterbox_batch(
+                segments=segments,
+                ref_audio_clipped=ref_audio_clipped,
+                language_id=language_id,
+                output_dir=output_dir,
+            )
+            stitched = os.path.join(output_dir, "preview_chatterbox.wav")
+            _stitch_wavs(part_paths, stitched)
+            yield "   ✓ Chatterbox preview complete\n"
+            return stitched
+
+        yield f"   [preview] Preparing Chatterbox Multilingual (device={device})...\n"
+        from chatterbox.mtl_tts import ChatterboxMultilingualTTS
+
+        model = ChatterboxMultilingualTTS.from_pretrained(device)
+
         part_paths = []
         total = len(segments)
         for i, seg in enumerate(segments):
@@ -129,11 +162,13 @@ def _preview_chatterbox(
             text = seg.get("tts_text", seg.get("translated_text", seg["text"]))
             out_path = os.path.join(output_dir, f"cb_prev_{i:04d}.wav")
 
-            wav = _gpu_preview_chatterbox_segment(
-                model=model,
-                text=text,
+            wav = model.generate(
+                text[:300],
                 language_id=language_id,
-                ref_audio_path=ref_audio_clipped,
+                audio_prompt_path=ref_audio_clipped,
+                exaggeration=0.5,
+                temperature=0.8,
+                cfg_weight=0.5,
             )
             torchaudio.save(out_path, wav, model.sr, encoding="PCM_S", bits_per_sample=16)
             part_paths.append(out_path)

steps/s4_tts.py

@@ -20,31 +20,56 @@ from tqdm import tqdm
 
 TTS_ENGINE = os.getenv("TTS_ENGINE", "chatterbox").lower()
 
-# Conditional imports based on TTS_ENGINE
-if TTS_ENGINE == "chatterbox":
-    from steps._tts_models import get_chatterbox
-else:
-    # OmniVoice mode - chatterbox imports not needed
-    get_chatterbox = None
-
 import spaces
 
 
-@spaces.GPU(duration=60)
-def _gpu_chatterbox_generate(
-    model,
-    text: str,
+@spaces.GPU(duration=120)
+def _gpu_chatterbox_full_batch(
+    segments: list[dict],
+    ref_audio_clipped: str,
     language_id: str,
-    ref_audio_path: str,
+    output_dir: str,
 ):
-    return model.generate(
-        text[:300],
-        language_id=language_id,
-        audio_prompt_path=ref_audio_path,
-        exaggeration=0.5,
-        temperature=0.8,
-        cfg_weight=0.5,
-    )
+    """
+    Load + run Chatterbox inside a single GPU-decorated scope.
+
+    ZeroGPU only intercepts CUDA init while the decorated function is active,
+    so constructing the CUDA model here avoids low-level torch CUDA init errors.
+    """
+    from chatterbox.mtl_tts import ChatterboxMultilingualTTS
+
+    print("   [s4] Loading Chatterbox in GPU scope...")
+    model = ChatterboxMultilingualTTS.from_pretrained("cuda")
+    results = []
+    total = len(segments)
+
+    for i, seg in enumerate(segments):
+        text = seg.get("tts_text", seg.get("translated_text", seg["text"]))
+        out_path = os.path.join(output_dir, f"seg_{i:04d}.wav")
+        orig_dur = seg["end"] - seg["start"]
+
+        print(f"   [s4] Chatterbox: Synthesising segment {i+1}/{total}...")
+        wav = model.generate(
+            text[:300],
+            language_id=language_id,
+            audio_prompt_path=ref_audio_clipped,
+            exaggeration=0.5,
+            temperature=0.8,
+            cfg_weight=0.5,
+        )
+
+        wav = _trim_trailing_noise(wav, model.sr)
+        wav = _trim_to_duration(wav, model.sr, orig_dur)
+        torchaudio.save(
+            out_path,
+            wav.detach().cpu(),
+            model.sr,
+            encoding="PCM_S",
+            bits_per_sample=16,
+        )
+        results.append({**seg, "tts_path": out_path})
+
+    return results
 
 
 # ── Chatterbox Multilingual ─────────────────────────────────
@@ -54,12 +79,27 @@ def _synthesise_chatterbox(
     language_id: str,
     output_dir: str,
 ):
-    yield "   [s4] Preparing Chatterbox Multilingual TTS...\n"
-    model = get_chatterbox()
-
     # Clip reference audio to max 10 seconds to prevent weird noise/artifacts
     ref_audio_clipped = _clip_audio(reference_audio_path, max_sec=15.0)
 
+    device = _get_device()
+    if device == "cuda":
+        yield "   [s4] Preparing Chatterbox batch processing (device=cuda)...\n"
+        results = _gpu_chatterbox_full_batch(
+            segments=segments,
+            ref_audio_clipped=ref_audio_clipped,
+            language_id=language_id,
+            output_dir=output_dir,
+        )
+        yield f"   [s4] Chatterbox TTS complete — {len(results)} segments synthesised ✓\n"
+        yield {"__TTS_RESULT__": results}
+        return
+
+    yield f"   [s4] Preparing Chatterbox Multilingual TTS (device={device})...\n"
+    from chatterbox.mtl_tts import ChatterboxMultilingualTTS
+
+    model = ChatterboxMultilingualTTS.from_pretrained(device)
+
     results = []
     total = len(segments)
     for i, seg in enumerate(segments):
@@ -71,11 +111,13 @@ def _synthesise_chatterbox(
         max_tokens = min(1000, max(150, int(orig_dur * 75 * 1.5)))
         _ = max_tokens
 
-        wav = _gpu_chatterbox_generate(
-            model=model,
-            text=text,
+        wav = model.generate(
+            text[:300],
             language_id=language_id,
-            ref_audio_path=ref_audio_clipped,
+            audio_prompt_path=ref_audio_clipped,
+            exaggeration=0.5,
+            temperature=0.8,
+            cfg_weight=0.5,
         )
 
         wav = _trim_trailing_noise(wav, model.sr)