Upload 2 files

app.py

@@ -1,135 +1,160 @@
-import os
-import io
-import torch
-import torchaudio
-import numpy as np
-import gradio as gr
-import pocket_tts
-from pocket_tts import TTSModel
-
-# 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!")
-
-# Official Kyutai Pocket-TTS voice list
-VOICES = [
-    "alba",
-    "marius",
-    "jaime",
-    "stuart",
-    "kelly",
-    "leo",
-    "carla",
-    "serena"
-]
-
-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
-
-    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: {e}")
-        try:
-            import librosa
-            return librosa.effects.time_stretch(audio_np.astype(np.float32), rate=speed)
-        except Exception:
-            return audio_np
-
-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.")
-
-    clean_text = text.strip()
-    clean_voice = voice.lower().replace("pocket-", "").replace("-", "_").strip()
-    if clean_voice not in VOICES:
-        clean_voice = "alba"
-
-    speed_factor = max(0.5, min(2.0, float(speed) if speed else 1.0))
-
-    # 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())
-
-    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
-    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_int16_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__":
+import os
+import io
+import spaces
+import torch
+import torchaudio
+import numpy as np
+import gradio as gr
+from pocket_tts import TTSModel
+
+# 1. Load the Pocket-TTS model at startup (CPU-only, ~100M params)
+print("Loading Kyutai Pocket-TTS model...")
+tts_model = TTSModel.load_model()
+print("Pocket-TTS model loaded successfully!")
+
+# Official Kyutai Pocket-TTS English voice list
+# See: https://huggingface.co/kyutai/tts-voices
+VOICES = [
+    "alba",
+    "anna",
+    "azelma",
+    "bill_boerst",
+    "caro_davy",
+    "charles",
+    "cosette",
+    "eponine",
+    "eve",
+    "fantine",
+    "george",
+    "jane",
+    "jean",
+    "javert",
+    "marius",
+    "mary",
+    "michael",
+    "paul",
+    "peter_yearsley",
+    "stuart_bell",
+    "vera",
+]
+
+# Pre-cache voice states at startup for faster inference
+print("Pre-caching voice states...")
+voice_states = {}
+for voice_name in VOICES:
+    try:
+        voice_states[voice_name] = tts_model.get_state_for_audio_prompt(voice_name)
+        print(f"  Cached voice: {voice_name}")
+    except Exception as e:
+        print(f"  Warning: Could not cache voice '{voice_name}': {e}")
+print("Voice states cached!")
+
+
+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
+
+    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: {e}")
+        try:
+            import librosa
+            return librosa.effects.time_stretch(audio_np.astype(np.float32), rate=speed)
+        except Exception:
+            return audio_np
+
+
+@spaces.GPU
+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.")
+
+    clean_text = text.strip()
+    clean_voice = voice.lower().strip()
+    if clean_voice not in VOICES:
+        clean_voice = "alba"
+
+    speed_factor = max(0.5, min(2.0, float(speed) if speed else 1.0))
+
+    # 1. Get the cached voice state, or load it on demand
+    if clean_voice in voice_states:
+        voice_state = voice_states[clean_voice]
+    else:
+        voice_state = tts_model.get_state_for_audio_prompt(clean_voice)
+
+    # 2. Generate audio using the official API
+    audio_tensor = tts_model.generate_audio(voice_state, clean_text)
+
+    # Convert to numpy
+    audio_np = audio_tensor.numpy().astype(np.float32)
+    sample_rate = tts_model.sample_rate
+
+    # 3. 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_int16_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.queue().launch(server_name="0.0.0.0", server_port=7860)

requirements.txt

@@ -1,8 +1,9 @@
-pocket-tts
-torch
-torchaudio
-gradio>=4.0.0
-numpy
-scipy
-librosa
-soundfile
+pocket-tts
+torch
+torchaudio
+gradio>=4.0.0
+numpy
+scipy
+librosa
+soundfile
+spaces