Spaces:
Sleeping
Sleeping
| import gradio as gr | |
| import io | |
| import wave | |
| import numpy as np | |
| # Lazy imports for optional dependencies | |
| try: | |
| import torch # type: ignore | |
| except Exception: # pragma: no cover | |
| torch = None # type: ignore | |
| try: | |
| from pocket_tts import TTSModel # type: ignore | |
| except Exception: # pragma: no cover | |
| TTSModel = None # type: ignore | |
| # Global state for lazy initialization | |
| _POCKET_STATE = { | |
| "initialized": False, | |
| "model": None, | |
| "voice_states": {}, | |
| "sample_rate": 24000, | |
| } | |
| def _get_available_voices() -> dict[str, str]: | |
| """Get available voices from the local ./voices/ directory. | |
| Scans ./voices/ directory for audio files (WAV, MP3, etc.) | |
| """ | |
| import os | |
| voices_dir = os.path.join(os.path.dirname(__file__), "voices") | |
| local_voices = {} | |
| if os.path.exists(voices_dir): | |
| for f in os.listdir(voices_dir): | |
| # Support common audio formats | |
| if f.lower().endswith(('.wav', '.mp3', '.flac', '.ogg', '.m4a')): | |
| voice_name = os.path.splitext(f)[0] | |
| local_voices[voice_name] = os.path.join(voices_dir, f) | |
| if local_voices: | |
| print(f"Found {len(local_voices)} local voice(s): {list(local_voices.keys())}") | |
| else: | |
| print("WARNING: No voices found in voices/ directory") | |
| return local_voices | |
| # Scan voices at import time | |
| PRESET_VOICES = _get_available_voices() | |
| def _init_pocket( | |
| temp: float = 0.7, | |
| lsd_decode_steps: int = 1, | |
| noise_clamp: float | None = None, | |
| eos_threshold: float = -4.0, | |
| ) -> None: | |
| """Lazy initialization of the Pocket TTS model.""" | |
| if _POCKET_STATE["initialized"]: | |
| return | |
| if TTSModel is None: | |
| raise gr.Error( | |
| "pocket-tts is not installed. Please install with: pip install pocket-tts" | |
| ) | |
| if torch is None: | |
| raise gr.Error("PyTorch is not installed. Please install torch>=2.5.0") | |
| print("Initializing Pocket TTS...") | |
| # Log in to HuggingFace if token is available (enables voice cloning on Spaces) | |
| import os | |
| hf_token = os.environ.get("HF_TOKEN") | |
| if hf_token: | |
| print("HF_TOKEN found, using for authentication") | |
| # Auto-detect device: CPU by default, CUDA if available | |
| # Note: The pocket-tts docs mention GPU doesn't provide speedup for this model | |
| device = "cuda" if torch.cuda.is_available() else "cpu" | |
| print(f"Using device: {device}") | |
| try: | |
| model = TTSModel.load_model( | |
| temp=float(temp), | |
| lsd_decode_steps=int(lsd_decode_steps), | |
| noise_clamp=float(noise_clamp) if noise_clamp is not None else None, | |
| eos_threshold=float(eos_threshold), | |
| ) | |
| _POCKET_STATE.update({ | |
| "initialized": True, | |
| "model": model, | |
| "sample_rate": model.sample_rate, | |
| }) | |
| print(f"Pocket TTS initialized. Sample rate: {model.sample_rate} Hz") | |
| # Auto-create missing embeddings if voice cloning is available | |
| if model.has_voice_cloning: | |
| _create_missing_embeddings(model) | |
| else: | |
| print("Voice cloning not available - using pre-computed embeddings only") | |
| except Exception as e: | |
| raise gr.Error(f"Failed to initialize Pocket TTS model: {str(e)}") | |
| def _create_missing_embeddings(model) -> None: | |
| """Create embeddings for any voices that have audio files but no embedding.""" | |
| import os | |
| from pocket_tts.data.audio import audio_read | |
| from pocket_tts.data.audio_utils import convert_audio | |
| import safetensors.torch | |
| voices_dir = os.path.join(os.path.dirname(__file__), "voices") | |
| embeddings_dir = os.path.join(os.path.dirname(__file__), "embeddings") | |
| if not os.path.exists(voices_dir): | |
| return | |
| os.makedirs(embeddings_dir, exist_ok=True) | |
| audio_extensions = ('.wav', '.mp3', '.flac', '.ogg', '.m4a') | |
| for voice_name, voice_path in PRESET_VOICES.items(): | |
| embedding_path = os.path.join(embeddings_dir, f"{voice_name}.safetensors") | |
| # Skip if embedding already exists or no local file | |
| if os.path.exists(embedding_path) or voice_path is None: | |
| continue | |
| # Skip fallback HuggingFace voices | |
| if voice_path.startswith("hf://"): | |
| continue | |
| print(f"Creating embedding for '{voice_name}'...") | |
| try: | |
| # Convert to WAV if needed | |
| audio_path = voice_path | |
| if not voice_path.lower().endswith('.wav'): | |
| from pydub import AudioSegment | |
| import tempfile | |
| audio = AudioSegment.from_file(voice_path) | |
| temp_wav = tempfile.NamedTemporaryFile(suffix='.wav', delete=False) | |
| audio.export(temp_wav.name, format='wav') | |
| audio_path = temp_wav.name | |
| # Read and encode audio | |
| audio, sr = audio_read(audio_path) | |
| audio_tensor = convert_audio(audio, sr, model.config.mimi.sample_rate, 1) | |
| with torch.no_grad(): | |
| audio_prompt = model._encode_audio(audio_tensor.unsqueeze(0).to(model.device)) | |
| # Save embedding | |
| safetensors.torch.save_file( | |
| {"audio_prompt": audio_prompt.cpu()}, | |
| embedding_path | |
| ) | |
| print(f" Saved: {embedding_path}") | |
| except Exception as e: | |
| print(f" Error creating embedding for {voice_name}: {e}") | |
| def _convert_to_wav(audio_path: str) -> str: | |
| """Convert audio file to WAV format if needed. | |
| Returns the path to a WAV file (original if already WAV, or converted temp file). | |
| Uses pydub for MP3 (requires ffmpeg), soundfile for other formats. | |
| """ | |
| import tempfile | |
| # Check if already WAV | |
| if audio_path.lower().endswith('.wav'): | |
| return audio_path | |
| print(f"Converting {audio_path} to WAV format...") | |
| # Create temp file path | |
| import os | |
| tmp_fd, wav_path = tempfile.mkstemp(suffix=".wav") | |
| os.close(tmp_fd) | |
| # Try pydub first (better MP3 support via ffmpeg) | |
| try: | |
| from pydub import AudioSegment | |
| audio = AudioSegment.from_file(audio_path) | |
| audio.export(wav_path, format="wav") | |
| print(f"Converted via pydub to: {wav_path}") | |
| return wav_path | |
| except ImportError: | |
| pass # pydub not installed, try soundfile | |
| except Exception as e: | |
| print(f"pydub conversion failed: {e}, trying soundfile...") | |
| # Fall back to soundfile | |
| try: | |
| import soundfile as sf | |
| audio_data, sample_rate = sf.read(audio_path) | |
| sf.write(wav_path, audio_data, sample_rate) | |
| print(f"Converted via soundfile to: {wav_path}") | |
| return wav_path | |
| except Exception as e: | |
| raise gr.Error(f"Failed to convert audio file: {str(e)}. Please upload a WAV file directly or install pydub+ffmpeg for MP3 support.") | |
| def _get_voice_state(voice_name: str | None, custom_audio_path: str | None): | |
| """Get or create voice state for generation. | |
| Args: | |
| voice_name: Name of preset voice (alba, marius, etc.) | |
| custom_audio_path: Path to custom audio file for voice cloning | |
| Returns: | |
| Voice state dict for the model | |
| """ | |
| model = _POCKET_STATE["model"] | |
| # Custom audio takes priority | |
| if custom_audio_path: | |
| print(f"Loading custom voice from: {custom_audio_path}") | |
| # Convert to WAV if needed | |
| wav_path = _convert_to_wav(custom_audio_path) | |
| return model.get_state_for_audio_prompt(wav_path) | |
| # Use preset voice | |
| if not voice_name or voice_name not in PRESET_VOICES: | |
| # Default to first available voice | |
| voice_name = list(PRESET_VOICES.keys())[0] if PRESET_VOICES else None | |
| if not voice_name: | |
| raise gr.Error("No voices available. Add audio files to the voices/ directory.") | |
| # Check cache | |
| if voice_name in _POCKET_STATE["voice_states"]: | |
| return _POCKET_STATE["voice_states"][voice_name] | |
| # Check for pre-computed embedding first (no voice cloning needed) | |
| import os | |
| embeddings_dir = os.path.join(os.path.dirname(__file__), "embeddings") | |
| embedding_path = os.path.join(embeddings_dir, f"{voice_name}.safetensors") | |
| if os.path.exists(embedding_path): | |
| print(f"Loading pre-computed embedding for '{voice_name}' from: {embedding_path}") | |
| import safetensors.torch | |
| from pocket_tts.modules.stateful_module import init_states | |
| # Load the audio prompt embedding | |
| state_dict = safetensors.torch.load_file(embedding_path) | |
| audio_prompt = state_dict["audio_prompt"].to(model.device) | |
| # Create fresh model state and condition it with the audio prompt | |
| # (same logic as model.get_state_for_audio_prompt uses internally) | |
| voice_state = init_states(model.flow_lm, batch_size=1, sequence_length=1000) | |
| model._run_flow_lm_and_increment_step(model_state=voice_state, audio_conditioning=audio_prompt) | |
| # Detach all tensors to make them leaf tensors (required for deepcopy) | |
| def detach_tensors(obj): | |
| if isinstance(obj, torch.Tensor): | |
| return obj.detach().clone() | |
| elif isinstance(obj, dict): | |
| return {k: detach_tensors(v) for k, v in obj.items()} | |
| else: | |
| return obj | |
| voice_state = detach_tensors(voice_state) | |
| _POCKET_STATE["voice_states"][voice_name] = voice_state | |
| return voice_state | |
| # Fall back to voice cloning (requires auth) | |
| if not model.has_voice_cloning: | |
| raise gr.Error( | |
| f"No embedding found for voice '{voice_name}'. " | |
| f"Voice cloning is not available (requires HF auth). " | |
| f"Run the app locally first to create embeddings." | |
| ) | |
| voice_path = PRESET_VOICES[voice_name] | |
| print(f"Loading preset voice '{voice_name}' from: {voice_path}") | |
| # Convert to WAV if needed (local files may be MP3, etc.) | |
| wav_path = _convert_to_wav(voice_path) | |
| voice_state = model.get_state_for_audio_prompt(wav_path) | |
| _POCKET_STATE["voice_states"][voice_name] = voice_state | |
| return voice_state | |
| def _audio_np_to_int16(audio_np: np.ndarray) -> np.ndarray: | |
| """Convert float audio array to int16.""" | |
| audio_clipped = np.clip(audio_np, -1.0, 1.0) | |
| return (audio_clipped * 32767.0).astype(np.int16) | |
| def _wav_bytes_from_int16(audio_int16: np.ndarray, sample_rate: int) -> bytes: | |
| """Create WAV bytes from int16 audio array.""" | |
| buffer = io.BytesIO() | |
| with wave.open(buffer, "wb") as wf: | |
| wf.setnchannels(1) | |
| wf.setsampwidth(2) | |
| wf.setframerate(sample_rate) | |
| wf.writeframes(audio_int16.tobytes()) | |
| return buffer.getvalue() | |
| def _split_into_sentences(text: str) -> list[str]: | |
| """Split text into sentences for chunk-by-chunk generation. | |
| Uses simple punctuation-based splitting for natural speech chunks. | |
| """ | |
| import re | |
| # Split on sentence-ending punctuation, keeping the punctuation | |
| # Handle common patterns: . ! ? and combinations like "..." or "?!" | |
| sentences = re.split(r'(?<=[.!?])\s+', text.strip()) | |
| # Filter out empty strings and strip whitespace | |
| return [s.strip() for s in sentences if s.strip()] | |
| def pocket_tts_stream( | |
| text: str, | |
| voice: str, | |
| custom_audio, | |
| temperature: float, | |
| lsd_decode_steps: int, | |
| noise_clamp: float | None, | |
| eos_threshold: float, | |
| frames_after_eos: int, | |
| ): | |
| """Generate speech with sentence-level streaming. | |
| Splits text into sentences and yields complete audio for each sentence, | |
| matching Kokoro's smooth streaming pattern. | |
| """ | |
| if not text or not text.strip(): | |
| raise gr.Error("Please enter text to synthesize.") | |
| # Initialize model with current parameters | |
| _init_pocket( | |
| temp=temperature, | |
| lsd_decode_steps=lsd_decode_steps, | |
| noise_clamp=noise_clamp if noise_clamp and noise_clamp > 0 else None, | |
| eos_threshold=eos_threshold, | |
| ) | |
| model = _POCKET_STATE["model"] | |
| sample_rate = _POCKET_STATE["sample_rate"] | |
| # Get voice state | |
| custom_path = custom_audio if custom_audio else None | |
| voice_state = _get_voice_state(voice, custom_path) | |
| # Split text into sentences for natural chunking | |
| sentences = _split_into_sentences(text) | |
| if not sentences: | |
| raise gr.Error("No valid sentences found in text.") | |
| produced_any = False | |
| # Buffer for initial audio - wait for ~5 seconds before yielding first chunk | |
| # This prevents stuttering from short first sentences | |
| min_initial_samples = int(sample_rate * 5) # 5 seconds of audio | |
| audio_buffer = [] | |
| buffer_samples = 0 | |
| initial_buffer_yielded = False | |
| try: | |
| for idx, sentence in enumerate(sentences): | |
| # Generate complete audio for this sentence (non-streaming per sentence) | |
| audio = model.generate_audio( | |
| voice_state, | |
| sentence, | |
| frames_after_eos=frames_after_eos if frames_after_eos > 0 else None, | |
| copy_state=True, | |
| ) | |
| produced_any = True | |
| # Convert tensor to numpy | |
| audio_np = audio.cpu().numpy() if hasattr(audio, 'cpu') else audio | |
| if not initial_buffer_yielded: | |
| # Accumulate in buffer until we have enough audio | |
| audio_buffer.append(audio_np) | |
| buffer_samples += len(audio_np) | |
| # Check if we have enough or this is the last sentence | |
| if buffer_samples >= min_initial_samples or idx == len(sentences) - 1: | |
| # Yield the accumulated buffer | |
| combined = np.concatenate(audio_buffer, axis=0) | |
| audio_int16 = _audio_np_to_int16(combined) | |
| yield _wav_bytes_from_int16(audio_int16, sample_rate) | |
| audio_buffer = [] | |
| buffer_samples = 0 | |
| initial_buffer_yielded = True | |
| else: | |
| # After initial buffer, yield each sentence immediately | |
| audio_int16 = _audio_np_to_int16(audio_np) | |
| yield _wav_bytes_from_int16(audio_int16, sample_rate) | |
| except gr.Error: | |
| raise | |
| except Exception as e: | |
| raise gr.Error(f"Error during speech generation: {str(e)[:200]}...") | |
| if not produced_any: | |
| raise gr.Error("No audio was generated.") | |
| def generate_tts( | |
| text: str, | |
| voice: str, | |
| custom_audio, | |
| temperature: float, | |
| lsd_decode_steps: int, | |
| noise_clamp: float, | |
| eos_threshold: float, | |
| frames_after_eos: int, | |
| ): | |
| """Main streaming dispatcher for Pocket TTS.""" | |
| yield from pocket_tts_stream( | |
| text, | |
| voice, | |
| custom_audio, | |
| temperature, | |
| lsd_decode_steps, | |
| noise_clamp, | |
| eos_threshold, | |
| frames_after_eos, | |
| ) | |
| # --- Gradio UI --- | |
| with gr.Blocks() as demo: | |
| gr.HTML( | |
| "<h1 style='text-align: center;'>Pocket-TTS</h1>" | |
| ) | |
| device_info = gr.Markdown( | |
| "<p style='text-align: center;'>Powered by kyutai/pocket-tts | Running on CPU | Voices cloned from Kokoro-82M</p>" | |
| ) | |
| def update_device_info(): | |
| device = "CUDA" if torch.cuda.is_available() else "CPU" | |
| return f"<p style='text-align: center;'>Powered by kyutai/pocket-tts | Running on {device} | Voices cloned from Kokoro-82M</p>" | |
| demo.load(update_device_info, outputs=device_info) | |
| with gr.Row(): | |
| with gr.Column(): | |
| # Text input | |
| text_input = gr.Textbox( | |
| label="Input Text", | |
| placeholder="Enter the text you want to convert to speech here...", | |
| lines=5, | |
| value="The quick brown fox jumps over the lazy dog. I am already far north of London, and as I walk in the streets of Petersburgh, I feel a cold northern breeze play upon my cheeks, which braces my nerves and fills me with delight. Do you understand this feeling? This breeze, which has traveled from the regions towards which I am advancing, gives me a foretaste of those icy climes. Inspirited by this wind of promise, my daydreams become more fervent and vivid.", | |
| ) | |
| # Voice selection | |
| with gr.Group(): | |
| gr.Markdown("### Voice Selection") | |
| gr.Markdown("Select a preset voice OR upload your own WAV file for voice cloning.") | |
| voice_dropdown = gr.Dropdown( | |
| choices=list(PRESET_VOICES.keys()), | |
| label="Preset Voice", | |
| value=list(PRESET_VOICES.keys())[0] if PRESET_VOICES else None, | |
| info="Select a pre-loaded voice. Ignored if custom audio is uploaded.", | |
| ) | |
| gr.Markdown("--- OR ---") | |
| ref_audio_input = gr.Audio( | |
| label="Custom Voice (WAV)", | |
| type="filepath", | |
| sources=["upload", "microphone"], | |
| ) | |
| generate_btn = gr.Button( | |
| "Generate Speech", | |
| variant="primary", | |
| ) | |
| with gr.Column(): | |
| audio_output = gr.Audio( | |
| label="Generated Speech", | |
| streaming=True, | |
| autoplay=True, | |
| ) | |
| with gr.Accordion("Advanced Options", open=False): | |
| temp_slider = gr.Slider( | |
| minimum=0.1, | |
| maximum=1.5, | |
| value=0.7, | |
| step=0.05, | |
| label="Temperature", | |
| info="Controls randomness. Higher = more varied, lower = more consistent.", | |
| ) | |
| lsd_steps_slider = gr.Slider( | |
| minimum=1, | |
| maximum=10, | |
| value=1, | |
| step=1, | |
| label="LSD Decode Steps", | |
| info="Number of generation steps. Higher = potentially better quality but slower.", | |
| ) | |
| noise_clamp_slider = gr.Slider( | |
| minimum=0.0, | |
| maximum=5.0, | |
| value=0.0, | |
| step=0.1, | |
| label="Noise Clamp", | |
| info="Maximum value for noise sampling. 0 = disabled.", | |
| ) | |
| eos_threshold_slider = gr.Slider( | |
| minimum=-10.0, | |
| maximum=0.0, | |
| value=-4.0, | |
| step=0.5, | |
| label="EOS Threshold", | |
| info="Threshold for end-of-sequence detection. More negative = longer audio.", | |
| ) | |
| frames_after_eos_slider = gr.Slider( | |
| minimum=0, | |
| maximum=10, | |
| value=2, | |
| step=1, | |
| label="Frames After EOS", | |
| info="Additional frames to generate after EOS detection.", | |
| ) | |
| # Connect inputs | |
| generate_inputs = [ | |
| text_input, | |
| voice_dropdown, | |
| ref_audio_input, | |
| temp_slider, | |
| lsd_steps_slider, | |
| noise_clamp_slider, | |
| eos_threshold_slider, | |
| frames_after_eos_slider, | |
| ] | |
| generate_btn.click( | |
| fn=generate_tts, | |
| inputs=generate_inputs, | |
| outputs=audio_output, | |
| api_name="generate_speech", | |
| ) | |
| text_input.submit( | |
| fn=generate_tts, | |
| inputs=generate_inputs, | |
| outputs=audio_output, | |
| api_name="generate_speech_enter", | |
| ) | |
| if __name__ == "__main__": | |
| demo.queue().launch(debug=True, theme="Nymbo/Nymbo_Theme") | |