Copy the files in the new repo for public release

2026-01-12 11:25:16 +01:00
commit f5fa328a58
54 changed files with 6411 additions and 0 deletions
@@ -0,0 +1,172 @@
+"""We also add the imports in the functions to make sure we didn't forget them."""
+# ruff: noqa: F841
+
+import pytest
+
+
+def test_readme_example():
+    import scipy.io.wavfile
+
+    from pocket_tts import TTSModel
+
+    tts_model = TTSModel.load_model()
+    voice_state = tts_model.get_state_for_audio_prompt(
+        "hf://kyutai/tts-voices/alba-mackenna/casual.wav"
+    )
+    audio = tts_model.generate_audio(voice_state, "Hello world, this is a test.")
+    # Audio is a torch tensor containing PCM data.
+    scipy.io.wavfile.write("output.wav", tts_model.sample_rate, audio.numpy())
+
+
+def test_quick_start():
+    import scipy.io.wavfile
+
+    from pocket_tts import TTSModel
+
+    # Load the model
+    tts_model = TTSModel.load_model()
+
+    # Get voice state from an audio file
+    voice_state = tts_model.get_state_for_audio_prompt(
+        "hf://kyutai/tts-voices/alba-mackenna/casual.wav"
+    )
+
+    # Generate audio
+    audio = tts_model.generate_audio(voice_state, "Hello world, this is a test.")
+
+    # Save to file
+    scipy.io.wavfile.write("output.wav", tts_model.sample_rate, audio.numpy())
+
+
+def test_load_model():
+    from pocket_tts import TTSModel
+
+    # Load with default settings
+    model = TTSModel.load_model()
+
+    # Load with custom parameters
+    model = TTSModel.load_model(
+        variant="b6369a24", temp=0.5, lsd_decode_steps=5, eos_threshold=-3.0
+    )
+
+
+def test_device():
+    from pocket_tts import TTSModel
+
+    model = TTSModel.load_model()
+    print(f"Model running on: {model.device}")
+
+
+def test_sample_rate():
+    from pocket_tts import TTSModel
+
+    model = TTSModel.load_model()
+    print(f"Sample rate: {model.sample_rate} Hz")
+
+
+@pytest.fixture
+def make_my_voice_file():
+    import requests
+
+    url = "https://huggingface.co/kyutai/tts-voices/resolve/main/expresso/ex01-ex02_default_001_channel1_168s.wav"
+    response = requests.get(url)
+    with open("my_voice.wav", "wb") as f:
+        f.write(response.content)
+
+
+@pytest.mark.usefixtures("make_my_voice_file")
+def test_get_state_for_audio_prompt():
+    from pocket_tts import TTSModel
+
+    model = TTSModel.load_model()
+    # From HuggingFace URL
+    voice_state = model.get_state_for_audio_prompt(
+        "hf://kyutai/tts-voices/alba-mackenna/casual.wav"
+    )
+
+    # From local file
+    voice_state = model.get_state_for_audio_prompt("./my_voice.wav")
+
+    # From HTTP URL
+    voice_state = model.get_state_for_audio_prompt(
+        "https://huggingface.co/kyutai/tts-voices/resolve"
+        "/main/expresso/ex01-ex02_default_001_channel1_168s.wav"
+    )
+
+
+def test_generate_audio():
+    from pocket_tts import TTSModel
+
+    model = TTSModel.load_model()
+
+    voice_state = model.get_state_for_audio_prompt(
+        "hf://kyutai/tts-voices/alba-mackenna/casual.wav"
+    )
+
+    # Generate audio
+    audio = model.generate_audio(voice_state, "Hello world!", frames_after_eos=2, copy_state=True)
+
+    print(f"Generated audio shape: {audio.shape}")
+    print(f"Audio duration: {audio.shape[-1] / model.sample_rate:.2f} seconds")
+
+
+def test_generate_audio_stream():
+    from pocket_tts import TTSModel
+
+    model = TTSModel.load_model()
+
+    voice_state = model.get_state_for_audio_prompt(
+        "hf://kyutai/tts-voices/alba-mackenna/casual.wav"
+    )
+    # Stream generation
+    for chunk in model.generate_audio_stream(voice_state, "Long text content..."):
+        # Process each chunk as it's generated
+        print(f"Generated chunk: {chunk.shape[0]} samples")
+        # Could save chunks to file or play in real-time
+
+
+def test_voice_management():
+    from pocket_tts import TTSModel
+
+    model = TTSModel.load_model()
+    # Preload multiple voices
+    voices = {
+        "casual": model.get_state_for_audio_prompt(
+            "hf://kyutai/tts-voices/alba-mackenna/casual.wav"
+        ),
+        "funny": model.get_state_for_audio_prompt(
+            "https://huggingface.co/kyutai/tts-voices/resolve/main/expresso/ex01-ex02_default_001_channel1_168s.wav"
+        ),
+    }
+
+    # Generate with different voices
+    casual_audio = model.generate_audio(voices["casual"], "Hey there!")
+    funny_audio = model.generate_audio(voices["funny"], "Good morning.")
+
+
+def test_batch_processing():
+    import scipy.io.wavfile
+    import torch
+
+    from pocket_tts import TTSModel
+
+    model = TTSModel.load_model()
+
+    voice_state = model.get_state_for_audio_prompt(
+        "hf://kyutai/tts-voices/alba-mackenna/casual.wav"
+    )
+    # Process multiple texts efficiently by re-using the same voice state
+    texts = [
+        "First sentence to generate.",
+        "Second sentence to generate.",
+        "Third sentence to generate.",
+    ]
+
+    audios = []
+    for text in texts:
+        audio = model.generate_audio(voice_state, text)
+        audios.append(audio)
+
+    # Concatenate all audio
+    full_audio = torch.cat(audios, dim=0)
+    scipy.io.wavfile.write("batch_output.wav", model.sample_rate, full_audio.numpy())