Update app.py

app.py

@@ -1,70 +1,131 @@
+import os
 import io
 import torch
-import scipy.io.wavfile
+import torchaudio
+import numpy as np
+import scipy.io.wavfile as wavfile
 import gradio as gr
-import spaces
 from pocket_tts import TTSModel
 
-print("Initializing Kyutai Pocket-TTS for ZeroGPU...")
-model = None
-
-def get_model():
-    global model
-    if model is None:
-        model = TTSModel.load_model()
-    return model
-
-VOICE_LIST = [
-    "stuart_bell", "alba", "fantine", "bill_boerst",
-    "caro_davy", "peter_yearsley", "azelma", "cosette",
-    "eponine", "javert", "marius", "jean", "jane",
-    "anna", "charles", "eve", "george", "mary",
-    "michael", "paul", "vera"
+# 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)
+
+# Available voices in Kyutai Pocket-TTS
+VOICES = [
+    "alba",
+    "marius",
+    "jaime",
+    "stuart",
+    "kelly",
+    "leo",
+    "carla",
+    "serena"
 ]
 
-# Cache voice states so we don't recompute voice prompt embeddings every request
-voice_state_cache = {}
+def change_speed_pitch_preserved(audio_np: np.ndarray, sample_rate: int, speed: float) -> np.ndarray:
+    """
+    Adjusts speech speed cleanly while preserving natural human voice pitch and formants.
+    """
+    if abs(speed - 1.0) < 0.02:
+        return audio_np
 
-@spaces.GPU
-def generate_tts(text: str, voice: str):
-    clean_text = text.strip()
-    if not clean_text:
-        return None
-    
-    m = get_model()
-    clean_voice = voice.replace("pocket-", "").replace("-", "_")
-    
-    # Get or compute voice state for the voice prompt
-    if clean_voice not in voice_state_cache:
+    try:
+        # High quality DSP time-stretching using torchaudio / sox tempo
+        tensor = torch.from_numpy(audio_np).float()
+        if tensor.dim() == 1:
+            tensor = tensor.unsqueeze(0)
+
+        # 'tempo -s' applies SoX's speech-optimized WSOLA with phase alignment
+        effects = [["tempo", "-s", str(speed)]]
+        stretched_tensor, _ = torchaudio.sox_effects.apply_effects_tensor(
+            tensor, sample_rate, effects
+        )
+        return stretched_tensor.squeeze(0).numpy()
+    except Exception as e:
+        print(f"SoX tempo stretch fallback to phase vocoder: {e}")
         try:
-            voice_state_cache[clean_voice] = m.get_state_for_audio_prompt(clean_voice)
+            import librosa
+            return librosa.effects.time_stretch(audio_np.astype(np.float32), rate=speed)
         except Exception:
-            # If string voice identifier is passed directly or with fallback
-            voice_state_cache[clean_voice] = m.get_state_for_audio_prompt(f"voices/{clean_voice}.wav") if hasattr(m, "get_state_for_audio_prompt") else {}
+            return audio_np
 
-    model_state = voice_state_cache[clean_voice]
+def synthesize(text: str, voice: str, speed: float = 1.0):
+    """
+    Generates audio from Pocket-TTS and applies clean server-side speed adjustment.
+    """
+    if not text or not text.strip():
+        raise gr.Error("Text prompt cannot be empty.")
 
-    # Generate audio
-    audio = m.generate_audio(model_state, clean_text)
-    
-    if isinstance(audio, torch.Tensor):
-        audio_np = audio.detach().cpu().numpy()
+    clean_text = text.strip()
+    clean_voice = voice.lower().replace("pocket-", "").replace("-", "_").strip()
+    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)
+
+    # Convert to 1D float numpy array
+    if isinstance(audio_tensor, torch.Tensor):
+        audio_np = audio_tensor.cpu().float().numpy().squeeze()
     else:
-        audio_np = audio
-        
-    sample_rate = getattr(m, "sample_rate", 24000)
-    return (sample_rate, audio_np)
-
-demo = gr.Interface(
-    fn=generate_tts,
-    inputs=[
-        gr.Textbox(label="Text to speak", value="Hello! This is Kyutai Pocket TTS running smoothly."),
-        gr.Dropdown(choices=VOICE_LIST, value="stuart_bell", label="Voice"),
-    ],
-    outputs=gr.Audio(label="Generated Audio", type="numpy"),
-    title="Kyutai Pocket-TTS ZeroGPU API",
-    api_name="predict"
-)
+        audio_np = np.array(audio_tensor, dtype=np.float32).squeeze()
+
+    sample_rate = getattr(model, "sample_rate", 24000)
+
+    # 2. Adjust speed with pitch preservation
+    if abs(speed_factor - 1.0) >= 0.02:
+        audio_np = change_speed_pitch_preserved(audio_np, sample_rate, speed_factor)
+
+    # Normalize audio to prevent clipping
+    max_val = np.max(np.abs(audio_np))
+    if max_val > 0:
+        audio_np = (audio_np / max_val) * 0.95
+
+    # Return in Gradio (sample_rate, numpy_array) format
+    int16_audio = (audio_np * 32767).astype(np.int16)
+    return (sample_rate, int16_audio)
+
+# --- Gradio UI & API Interface ---
+with gr.Blocks(title="Kyutai Pocket-TTS Server") as demo:
+    gr.Markdown("# 🎙️ Kyutai Pocket-TTS Server with Speed Control")
+    
+    with gr.Row():
+        with gr.Column():
+            text_input = gr.Textbox(
+                label="Text to Synthesize",
+                placeholder="Enter text to speak...",
+                lines=4,
+                value="The quick brown fox jumps over the lazy dog."
+            )
+            voice_input = gr.Dropdown(
+                label="Voice",
+                choices=VOICES,
+                value="alba"
+            )
+            speed_slider = gr.Slider(
+                label="Speed Multiplier",
+                minimum=0.5,
+                maximum=2.0,
+                step=0.05,
+                value=1.0
+            )
+            generate_btn = gr.Button("Generate Speech", variant="primary")
+            
+        with gr.Column():
+            audio_output = gr.Audio(label="Synthesized Audio", type="numpy")
+
+    generate_btn.click(
+        fn=synthesize,
+        inputs=[text_input, voice_input, speed_slider],
+        outputs=audio_output,
+        api_name="predict"
+    )
 
 if __name__ == "__main__":
-    demo.launch(server_name="0.0.0.0", server_port=7860)
+    demo.queue().launch(server_name="0.0.0.0", server_port=7860)