Update app.py

app.py

@@ -3,16 +3,16 @@ import io
 import torch
 import torchaudio
 import numpy as np
-import scipy.io.wavfile as wavfile
 import gradio as gr
+import pocket_tts
 from pocket_tts import TTSModel
 
-# Load the official Kyutai Pocket-TTS model on CUDA GPU if available, else CPU
-device = "cuda" if torch.cuda.is_available() else "cpu"
-print(f"Loading Pocket-TTS model on device: {device}...")
-model = TTSModel.from_pretrained("kyutai/pocket-tts-v0.1", device=device)
+# 1. Load the Pocket-TTS model using the official pocket_tts API
+print("Loading Kyutai Pocket-TTS model...")
+model = TTSModel.load_model()
+print("Pocket-TTS model loaded successfully!")
 
-# Available voices in Kyutai Pocket-TTS
+# Official Kyutai Pocket-TTS voice list
 VOICES = [
     "alba",
     "marius",
@@ -44,7 +44,7 @@ def change_speed_pitch_preserved(audio_np: np.ndarray, sample_rate: int, speed:
         )
         return stretched_tensor.squeeze(0).numpy()
     except Exception as e:
-        print(f"SoX tempo stretch fallback to phase vocoder: {e}")
+        print(f"SoX tempo stretch fallback: {e}")
         try:
             import librosa
             return librosa.effects.time_stretch(audio_np.astype(np.float32), rate=speed)
@@ -63,19 +63,23 @@ def synthesize(text: str, voice: str, speed: float = 1.0):
     if clean_voice not in VOICES:
         clean_voice = "alba"
 
-    # Clamp speed between 0.5x and 2.0x
     speed_factor = max(0.5, min(2.0, float(speed) if speed else 1.0))
 
-    # 1. Generate audio using Kyutai Pocket-TTS model
-    with torch.no_grad():
-        audio_tensor = model.generate(text=clean_text, voice=clean_voice)
+    # 1. Prepare voice state & stream chunks
+    voice_state = model.get_voice_state(clean_voice)
+    generation_state = model.get_state_for_audio_generation(clean_text, voice_state)
+    
+    audio_chunks = []
+    for chunk in model.generate_audio_stream(generation_state):
+        if isinstance(chunk, torch.Tensor):
+            audio_chunks.append(chunk.detach().cpu().float().numpy().squeeze())
+        else:
+            audio_chunks.append(np.array(chunk, dtype=np.float32).squeeze())
 
-    # Convert to 1D float numpy array
-    if isinstance(audio_tensor, torch.Tensor):
-        audio_np = audio_tensor.cpu().float().numpy().squeeze()
-    else:
-        audio_np = np.array(audio_tensor, dtype=np.float32).squeeze()
+    if not audio_chunks:
+        raise gr.Error("No audio was generated by the model.")
 
+    audio_np = np.concatenate(audio_chunks)
     sample_rate = getattr(model, "sample_rate", 24000)
 
     # 2. Adjust speed with pitch preservation
@@ -87,7 +91,7 @@ def synthesize(text: str, voice: str, speed: float = 1.0):
     if max_val > 0:
         audio_np = (audio_np / max_val) * 0.95
 
-    # Return in Gradio (sample_rate, numpy_array) format
+    # Return in Gradio (sample_rate, numpy_int16_array) format
     int16_audio = (audio_np * 32767).astype(np.int16)
     return (sample_rate, int16_audio)