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,3 @@
+import os
+
+os.environ["POCKET_TTS_ERROR_WITHOUT_EOS"] = "1"
@@ -0,0 +1,180 @@
+"""Integration tests for the CLI generate command using real implementation."""
+
+from typer.testing import CliRunner
+
+from pocket_tts.data.audio import audio_read
+from pocket_tts.default_parameters import DEFAULT_VARIANT
+from pocket_tts.main import cli_app
+
+other_voice = "https://huggingface.co/kyutai/tts-voices/resolve/main/expresso/ex01-ex02_default_001_channel1_168s.wav"
+
+runner = CliRunner()
+
+
+def test_generate_basic_usage(tmp_path):
+    """Test basic generate command with default parameters."""
+    output_file = tmp_path / "test_output.wav"
+
+    result = runner.invoke(
+        cli_app,
+        ["generate", "--text", "Hello world, this is a test.", "--output-path", str(output_file)],
+    )
+
+    assert result.exit_code == 0
+    # Verify output file was created and contains audio
+    assert output_file.exists()
+    assert output_file.stat().st_size > 0
+
+    # Verify it's a valid audio file
+    audio, sample_rate = audio_read(str(output_file))
+    assert audio.shape[0] == 1  # Mono channel
+    assert audio.shape[1] > 0  # Has audio samples
+    assert sample_rate == 24000  # Expected sample rate
+
+
+def test_generate_with_custom_voice(tmp_path):
+    """Test generate command with custom voice prompt."""
+    output_file = tmp_path / "custom_voice_test.wav"
+
+    result = runner.invoke(
+        cli_app,
+        [
+            "generate",
+            "--text",
+            "Testing custom voice.",
+            "--voice",
+            other_voice,
+            "--output-path",
+            str(output_file),
+        ],
+    )
+
+    assert result.exit_code == 0
+    assert output_file.exists()
+
+    # Verify audio content
+    audio, sample_rate = audio_read(str(output_file))
+    assert audio.shape[0] == 1  # Mono channel
+    assert audio.shape[1] > 0  # Has audio samples
+    assert sample_rate == 24000
+
+
+def test_generate_with_custom_parameters(tmp_path):
+    """Test generate command with custom generation parameters."""
+    output_file = tmp_path / "custom_params_test.wav"
+
+    result = runner.invoke(
+        cli_app,
+        [
+            "generate",
+            "--text",
+            "Testing custom parameters.",
+            "--variant",
+            DEFAULT_VARIANT,
+            "--temperature",
+            "0.8",
+            "--lsd-decode-steps",
+            "2",
+            "--eos-threshold",
+            "-3.0",
+            "--frames-after-eos",
+            "7",
+            "--output-path",
+            str(output_file),
+        ],
+    )
+
+    assert result.exit_code == 0
+    assert output_file.exists()
+
+    audio, sample_rate = audio_read(str(output_file))
+    assert audio.shape[0] == 1  # Mono channel
+    assert audio.shape[1] > 0  # Has audio samples
+    assert sample_rate == 24000
+
+
+def test_generate_verbose_mode(tmp_path):
+    """Test generate command with verbose logging."""
+    output_file = tmp_path / "verbose_test.wav"
+
+    result = runner.invoke(
+        cli_app,
+        ["generate", "--text", "Testing verbose mode.", "-q", "--output-path", str(output_file)],
+    )
+
+    assert result.exit_code == 0
+    assert output_file.exists()
+
+
+def test_generate_default_text(tmp_path):
+    """Test generate command with default text when no text provided."""
+    output_file = tmp_path / "default_text_test.wav"
+
+    result = runner.invoke(cli_app, ["generate", "--output-path", str(output_file)])
+
+    assert result.exit_code == 0
+    assert output_file.exists()
+
+    audio, sample_rate = audio_read(str(output_file))
+    assert audio.shape[0] == 1  # Mono channel
+    assert audio.shape[1] > 0  # Has audio samples
+    assert sample_rate == 24000
+
+
+def test_generate_long_text(tmp_path):
+    """Test generate command with longer text."""
+    long_text = "This is a longer text to test the TTS system. " * 5
+    output_file = tmp_path / "long_text_test.wav"
+
+    result = runner.invoke(
+        cli_app, ["generate", "--text", long_text, "--output-path", str(output_file)]
+    )
+
+    assert result.exit_code == 0
+    assert output_file.exists()
+
+    audio, sample_rate = audio_read(str(output_file))
+    assert audio.shape[0] == 1  # Mono channel
+    assert audio.shape[1] > 0  # Has audio samples
+    assert sample_rate == 24000
+    # Longer text should produce longer audio
+    assert audio.shape[1] > 24000 * 10  # At least 10 second of audio
+
+
+def test_generate_command_help():
+    """Test that generate command help is displayed correctly."""
+    result = runner.invoke(cli_app, ["generate", "--help"])
+
+    assert result.exit_code == 0
+    assert "Generate speech using Kyutai Pocket TTS" in result.output
+    assert "--text" in result.output
+    assert "--voice" in result.output
+    assert "--output-path" in result.output
+    assert "--variant" in result.output
+    assert "--temperature" in result.output
+    assert "--device" in result.output
+
+
+def test_generate_multiple_runs(tmp_path):
+    """Test multiple consecutive generate commands."""
+    for i in range(3):
+        output_file = tmp_path / f"test_run_{i + 1}.wav"
+
+        result = runner.invoke(
+            cli_app,
+            [
+                "generate",
+                "--text",
+                f"This is test run number {i + 1}.",
+                "--output-path",
+                str(output_file),
+            ],
+        )
+
+        assert result.exit_code == 0
+        assert output_file.exists()
+
+        audio, sample_rate = audio_read(str(output_file))
+        assert audio.shape[0] == 1  # Mono channel
+        assert audio.shape[1] > 0  # Has audio samples
+        assert sample_rate == 24000
@@ -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())