Gabriel de Marmiesse
c90fc8ced6
The french 6 layers variant wasn't included because it's not good enough yet. Use the `french_24l` variant for now.
162 lines
4.4 KiB
Python
162 lines
4.4 KiB
Python
"""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("cosette")
|
|
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("marius")
|
|
|
|
# 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(
|
|
language="english_2026-01", 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()
|
|
# hack to make it work without auth
|
|
model.has_voice_cloning = True
|
|
|
|
# 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("marius")
|
|
|
|
# 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("fantine")
|
|
# 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("alba"),
|
|
"serious": model.get_state_for_audio_prompt("marius"),
|
|
}
|
|
|
|
# Generate with different voices
|
|
casual_audio = model.generate_audio(voices["casual"], "Hey there!")
|
|
funny_audio = model.generate_audio(voices["serious"], "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("azelma")
|
|
# 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())
|