mirror of
https://github.com/rhasspy/piper-sample-generator.git
synced 2026-08-27 18:15:58 -04:00
472 lines
16 KiB
Python
Executable file
472 lines
16 KiB
Python
Executable file
#!/usr/bin/env python3
|
|
import argparse
|
|
import gc
|
|
import itertools as it
|
|
import json
|
|
import logging
|
|
import os
|
|
import wave
|
|
from pathlib import Path
|
|
from typing import Any, Dict, List, Optional, Tuple, Union
|
|
|
|
import numpy as np
|
|
import torch
|
|
import torchaudio
|
|
import webrtcvad
|
|
from piper_phonemize import phonemize_espeak
|
|
|
|
from piper_train.vits import commons
|
|
|
|
_DIR = Path(__file__).parent
|
|
_LOGGER = logging.getLogger(__name__)
|
|
logging.basicConfig(level=logging.DEBUG)
|
|
|
|
|
|
# Main generation function
|
|
def generate_samples(
|
|
text: Union[List[str], str],
|
|
output_dir: Union[str, Path],
|
|
max_samples: Optional[int] = None,
|
|
file_names: Optional[List[str]] = None,
|
|
model: Union[str, Path] = _DIR / "models" / "en_US-libritts_r-medium.pt",
|
|
batch_size: int = 1,
|
|
slerp_weights: Tuple[float, ...] = (0.5,),
|
|
length_scales: Tuple[float, ...] = (0.75, 1, 1.25),
|
|
noise_scales: Tuple[float, ...] = (0.667,),
|
|
noise_scale_ws: Tuple[float, ...] = (0.8,),
|
|
max_speakers: Optional[float] = None,
|
|
verbose: bool = False,
|
|
auto_reduce_batch_size: bool = False,
|
|
min_phoneme_count: Optional[int] = None,
|
|
**kwargs,
|
|
) -> None:
|
|
"""
|
|
Generate synthetic speech clips, saving the clips to the specified output directory.
|
|
|
|
Args:
|
|
text (List[str]): The text to convert into speech. Can be either a
|
|
a list of strings, or a path to a file with text on each line.
|
|
output_dir (str): The location to save the generated clips.
|
|
max_samples (int): The maximum number of samples to generate.
|
|
file_names (List[str]): The names to use when saving the files. Must be the same length
|
|
as the `text` argument, if a list.
|
|
model (str): The path to the STT model to use for generation.
|
|
batch_size (int): The batch size to use when generated the clips
|
|
slerp_weights (List[float]): The weights to use when mixing speakers via SLERP.
|
|
length_scales (List[float]): Controls the average duration/speed of the generated speech.
|
|
noise_scales (List[float]): A parameter for overall variability of the generated speech.
|
|
noise_scale_ws (List[float]): A parameter for the stochastic duration of words/phonemes.
|
|
max_speakers (int): The maximum speaker number to use, if the model is multi-speaker.
|
|
verbose (bool): Enable or disable more detailed logging messages (default: False).
|
|
auto_reduce_batch_size (bool): Automatically and temporarily reduce the batch size
|
|
if CUDA OOM errors are detected, and try to resume generation.
|
|
min_phoneme_count (int): If set, ensure this number of phonemes is always sent to the model.
|
|
Clip audio to extract original phrase.
|
|
|
|
Returns:
|
|
None
|
|
"""
|
|
|
|
if max_samples is None:
|
|
max_samples = len(text)
|
|
|
|
_LOGGER.debug("Loading %s", model)
|
|
model_path = Path(model)
|
|
|
|
torch_model = torch.load(model_path)
|
|
torch_model.eval()
|
|
_LOGGER.info("Successfully loaded the model")
|
|
|
|
if torch.cuda.is_available():
|
|
torch_model.cuda()
|
|
_LOGGER.debug("CUDA available, using GPU")
|
|
|
|
output_dir = Path(output_dir)
|
|
output_dir.mkdir(parents=True, exist_ok=True)
|
|
|
|
config_path = f"{model_path}.json"
|
|
with open(config_path, "r", encoding="utf-8") as config_file:
|
|
config = json.load(config_file)
|
|
|
|
voice = config["espeak"]["voice"]
|
|
sample_rate = config["audio"]["sample_rate"]
|
|
num_speakers = config["num_speakers"]
|
|
if max_speakers is not None:
|
|
num_speakers = min(num_speakers, max_speakers)
|
|
|
|
max_len = None
|
|
|
|
sample_idx = 0
|
|
is_done = False
|
|
settings_iter = it.cycle(
|
|
it.product(
|
|
slerp_weights,
|
|
length_scales,
|
|
noise_scales,
|
|
noise_scale_ws,
|
|
)
|
|
)
|
|
|
|
# Define resampler to get to 16khz (https://pytorch.org/audio/stable/tutorials/audio_resampling_tutorial.html#kaiser-best)
|
|
sample_rate = 22050
|
|
resample_rate = 16000
|
|
resampler = torchaudio.transforms.Resample(
|
|
sample_rate,
|
|
resample_rate,
|
|
lowpass_filter_width=64,
|
|
rolloff=0.9475937167399596,
|
|
resampling_method="kaiser_window",
|
|
beta=14.769656459379492,
|
|
)
|
|
|
|
speakers_iter = it.cycle(it.product(range(num_speakers), range(num_speakers)))
|
|
speakers_batch = list(it.islice(speakers_iter, 0, batch_size))
|
|
if isinstance(text, str) and os.path.exists(text):
|
|
texts = it.cycle(
|
|
[
|
|
i.strip()
|
|
for i in open(text, "r", encoding="utf-8").readlines()
|
|
if len(i.strip()) > 0
|
|
]
|
|
)
|
|
elif isinstance(text, list):
|
|
texts = it.cycle(text)
|
|
else:
|
|
texts = it.cycle([text])
|
|
|
|
if file_names:
|
|
file_names = it.cycle(file_names)
|
|
|
|
batch_idx = 0
|
|
while speakers_batch:
|
|
if is_done:
|
|
break
|
|
|
|
batch_size = len(speakers_batch)
|
|
slerp_weight, length_scale, noise_scale, noise_scale_w = next(settings_iter)
|
|
|
|
with torch.no_grad():
|
|
speaker_1 = torch.LongTensor([s[0] for s in speakers_batch])
|
|
speaker_2 = torch.LongTensor([s[1] for s in speakers_batch])
|
|
|
|
phoneme_ids_by_batch = []
|
|
clip_indexes_by_batch = []
|
|
for i in range(batch_size):
|
|
phoneme_ids, clip_phoneme_index = get_phonemes(
|
|
voice, config, next(texts), verbose, min_phoneme_count
|
|
)
|
|
phoneme_ids_by_batch.append(phoneme_ids)
|
|
clip_indexes_by_batch.append(clip_phoneme_index)
|
|
|
|
def right_pad_lists(lists):
|
|
max_length = max(len(lst) for lst in lists)
|
|
padded_lists = []
|
|
for lst in lists:
|
|
padded_l = lst + [1] * (
|
|
max_length - len(lst)
|
|
) # phoneme 1 (corresponding to '^' character seems to work best)
|
|
padded_lists.append(padded_l)
|
|
return padded_lists
|
|
|
|
phoneme_ids_by_batch = right_pad_lists(phoneme_ids_by_batch)
|
|
|
|
if auto_reduce_batch_size:
|
|
oom_error = True
|
|
counter = 1
|
|
while oom_error is True:
|
|
try:
|
|
audio, phoneme_samples = generate_audio(
|
|
torch_model,
|
|
speaker_1[0 : batch_size // counter],
|
|
speaker_2[0 : batch_size // counter],
|
|
phoneme_ids_by_batch[0 : batch_size // counter],
|
|
slerp_weight,
|
|
noise_scale,
|
|
noise_scale_w,
|
|
length_scale,
|
|
max_len,
|
|
)
|
|
oom_error = False
|
|
except torch.cuda.OutOfMemoryError:
|
|
torch.cuda.empty_cache()
|
|
gc.collect()
|
|
counter += 1 # reduce batch size to avoid OOM errors
|
|
else:
|
|
audio, phoneme_samples = generate_audio(
|
|
torch_model,
|
|
speaker_1,
|
|
speaker_2,
|
|
phoneme_ids_by_batch,
|
|
slerp_weight,
|
|
noise_scale,
|
|
noise_scale_w,
|
|
length_scale,
|
|
max_len,
|
|
)
|
|
|
|
# Clip audio when using min_phoneme_count
|
|
for i, clip_phoneme_index in enumerate(clip_indexes_by_batch):
|
|
if clip_phoneme_index is not None:
|
|
last_sample_idx = int(
|
|
phoneme_samples[i].flatten()[clip_phoneme_index:].sum().item()
|
|
)
|
|
|
|
# Fill remainder of audio with silence.
|
|
# It will be removed in the next stage.
|
|
audio[i, 0, :-last_sample_idx] = 0
|
|
|
|
# Resample audio
|
|
audio = resampler(audio.cpu()).numpy()
|
|
|
|
audio_int16 = audio_float_to_int16(audio)
|
|
for audio_idx in range(audio_int16.shape[0]):
|
|
# Use webrtcvad to trip silence from the clips
|
|
audio_data = remove_silence(audio_int16[audio_idx].flatten())[
|
|
None,
|
|
]
|
|
|
|
if isinstance(file_names, it.cycle):
|
|
wav_path = output_dir / next(file_names)
|
|
else:
|
|
wav_path = output_dir / f"{sample_idx}.wav"
|
|
|
|
wav_file: wave.Wave_write = wave.open(str(wav_path), "wb")
|
|
with wav_file:
|
|
wav_file.setframerate(resample_rate)
|
|
wav_file.setsampwidth(2)
|
|
wav_file.setnchannels(1)
|
|
wav_file.writeframes(audio_data)
|
|
|
|
sample_idx += 1
|
|
if sample_idx >= max_samples:
|
|
is_done = True
|
|
break
|
|
|
|
# print(f"Batch {batch_idx +1}/{max_samples//batch_size} complete", " "*200, end='\r')
|
|
|
|
# Next batch
|
|
_LOGGER.debug("Batch %s/%s complete", batch_idx + 1, max_samples // batch_size)
|
|
speakers_batch = list(it.islice(speakers_iter, 0, batch_size))
|
|
batch_idx += 1
|
|
|
|
_LOGGER.info("Done")
|
|
|
|
|
|
def remove_silence(
|
|
x: np.ndarray,
|
|
frame_duration: float = 0.030,
|
|
sample_rate: int = 16000,
|
|
min_start: int = 2000,
|
|
) -> np.ndarray:
|
|
"""Uses webrtc voice activity detection to remove silence from the clips"""
|
|
vad = webrtcvad.Vad(0)
|
|
if x.dtype in (np.float32, np.float64):
|
|
x = (x * 32767).astype(np.int16)
|
|
x_new = x[0:min_start].tolist()
|
|
step_size = int(sample_rate * frame_duration)
|
|
for i in range(min_start, x.shape[0] - step_size, step_size):
|
|
vad_res = vad.is_speech(x[i : i + step_size].tobytes(), sample_rate)
|
|
if vad_res:
|
|
x_new.extend(x[i : i + step_size].tolist())
|
|
return np.array(x_new).astype(np.int16)
|
|
|
|
|
|
def generate_audio(
|
|
model,
|
|
speaker_1,
|
|
speaker_2,
|
|
phoneme_ids,
|
|
slerp_weight,
|
|
noise_scale,
|
|
noise_scale_w,
|
|
length_scale,
|
|
max_len,
|
|
):
|
|
x = torch.LongTensor(phoneme_ids)
|
|
x_lengths = torch.LongTensor([len(i) for i in phoneme_ids])
|
|
|
|
if torch.cuda.is_available():
|
|
speaker_1 = speaker_1.cuda()
|
|
speaker_2 = speaker_2.cuda()
|
|
x = x.cuda()
|
|
x_lengths = x_lengths.cuda()
|
|
|
|
x, m_p_orig, logs_p_orig, x_mask = model.enc_p(x, x_lengths)
|
|
emb0 = model.emb_g(speaker_1)
|
|
emb1 = model.emb_g(speaker_2)
|
|
g = slerp(emb0, emb1, slerp_weight).unsqueeze(-1) # [b, h, 1]
|
|
|
|
if model.use_sdp:
|
|
logw = model.dp(x, x_mask, g=g, reverse=True, noise_scale=noise_scale_w)
|
|
else:
|
|
logw = model.dp(x, x_mask, g=g)
|
|
w = torch.exp(logw) * x_mask * length_scale
|
|
w_ceil = torch.ceil(w)
|
|
y_lengths = torch.clamp_min(torch.sum(w_ceil, [1, 2]), 1).long()
|
|
y_mask = torch.unsqueeze(
|
|
commons.sequence_mask(y_lengths, y_lengths.max()), 1
|
|
).type_as(x_mask)
|
|
attn_mask = torch.unsqueeze(x_mask, 2) * torch.unsqueeze(y_mask, -1)
|
|
attn = commons.generate_path(w_ceil, attn_mask)
|
|
|
|
m_p = torch.matmul(attn.squeeze(1), m_p_orig.transpose(1, 2)).transpose(
|
|
1, 2
|
|
) # [b, t', t], [b, t, d] -> [b, d, t']
|
|
logs_p = torch.matmul(attn.squeeze(1), logs_p_orig.transpose(1, 2)).transpose(
|
|
1, 2
|
|
) # [b, t', t], [b, t, d] -> [b, d, t']
|
|
|
|
z_p = m_p + torch.randn_like(m_p) * torch.exp(logs_p) * noise_scale
|
|
z = model.flow(z_p, y_mask, g=g, reverse=True)
|
|
o = model.dec((z * y_mask)[:, :, :max_len], g=g)
|
|
|
|
audio = o
|
|
phoneme_samples = w_ceil * 256 # hop length
|
|
|
|
return audio, phoneme_samples
|
|
|
|
|
|
def get_phonemes(
|
|
voice: str,
|
|
config: Dict[str, Any],
|
|
text: str,
|
|
verbose: bool = False,
|
|
min_phoneme_count: Optional[int] = None,
|
|
) -> Tuple[List[int], Optional[int]]:
|
|
# Combine all sentences
|
|
phonemes = [
|
|
p
|
|
for sentence_phonemes in phonemize_espeak(text, voice)
|
|
for p in sentence_phonemes
|
|
]
|
|
if verbose is True:
|
|
_LOGGER.debug("Phonemes: %s", phonemes)
|
|
|
|
id_map = config["phoneme_id_map"]
|
|
|
|
# Beginning of utterance
|
|
phoneme_ids = list(id_map["^"])
|
|
|
|
# Phoneme ids for just the text
|
|
text_phoneme_ids = []
|
|
|
|
for phoneme in phonemes:
|
|
p_ids = id_map.get(phoneme)
|
|
if p_ids is not None:
|
|
phoneme_ids.extend(p_ids)
|
|
text_phoneme_ids.extend(p_ids)
|
|
phoneme_ids.extend(id_map["_"])
|
|
text_phoneme_ids.extend(id_map["_"])
|
|
|
|
# Index where audio should be clipped at.
|
|
# When None, all of the audio will be used.
|
|
clip_phoneme_index: Optional[int] = None
|
|
|
|
if min_phoneme_count is not None:
|
|
# Repeat phrase until minimum phoneme count is met.
|
|
# NOTE: It is critical that the ^ and $ phonemes are not repeated here.
|
|
while (len(phoneme_ids) - 1) < min_phoneme_count:
|
|
# We will clip audio at the beginning of the last phrase
|
|
clip_phoneme_index = len(phoneme_ids) - 1
|
|
|
|
phoneme_ids.extend(text_phoneme_ids)
|
|
|
|
# End of utterance
|
|
phoneme_ids.extend(id_map["$"])
|
|
|
|
return phoneme_ids, clip_phoneme_index
|
|
|
|
|
|
def slerp(v1, v2, t: float, DOT_THR: float = 0.9995, zdim: int = -1):
|
|
"""SLERP for pytorch tensors interpolating `v1` to `v2` with scale of `t`.
|
|
|
|
`DOT_THR` determines when the vectors are too close to parallel.
|
|
If they are too close, then a regular linear interpolation is used.
|
|
|
|
`zdim` is the feature dimension over which to compute norms and find angles.
|
|
For example: if a sequence of 5 vectors is input with shape [5, 768]
|
|
Then `zdim = 1` or `zdim = -1` computes SLERP along the feature dim of 768.
|
|
|
|
Theory Reference:
|
|
https://splines.readthedocs.io/en/latest/rotation/slerp.html
|
|
PyTorch reference:
|
|
https://discuss.pytorch.org/t/help-regarding-slerp-function-for-generative-model-sampling/32475/3
|
|
Numpy reference:
|
|
https://gist.github.com/dvschultz/3af50c40df002da3b751efab1daddf2c
|
|
"""
|
|
|
|
# take the dot product between normalized vectors
|
|
v1_norm = v1 / torch.norm(v1, dim=zdim, keepdim=True)
|
|
v2_norm = v2 / torch.norm(v2, dim=zdim, keepdim=True)
|
|
dot = (v1_norm * v2_norm).sum(zdim)
|
|
|
|
# if the vectors are too close, return a simple linear interpolation
|
|
if (torch.abs(dot) > DOT_THR).any():
|
|
res = (1 - t) * v1 + t * v2
|
|
|
|
# else apply SLERP
|
|
else:
|
|
# compute the angle terms we need
|
|
theta = torch.acos(dot)
|
|
theta_t = theta * t
|
|
sin_theta = torch.sin(theta)
|
|
sin_theta_t = torch.sin(theta_t)
|
|
|
|
# compute the sine scaling terms for the vectors
|
|
s1 = torch.sin(theta - theta_t) / sin_theta
|
|
s2 = sin_theta_t / sin_theta
|
|
|
|
# interpolate the vectors
|
|
res = (s1.unsqueeze(zdim) * v1) + (s2.unsqueeze(zdim) * v2)
|
|
|
|
return res
|
|
|
|
|
|
def audio_float_to_int16(
|
|
audio: np.ndarray, max_wav_value: float = 32767.0
|
|
) -> np.ndarray:
|
|
"""Normalize audio and convert to int16 range"""
|
|
audio_norm = audio * (max_wav_value / max(0.01, np.max(np.abs(audio))))
|
|
audio_norm = np.clip(audio_norm, -max_wav_value, max_wav_value)
|
|
audio_norm = audio_norm.astype("int16")
|
|
return audio_norm
|
|
|
|
|
|
def main() -> None:
|
|
"""Main entry point."""
|
|
|
|
# Get command line arguments
|
|
parser = argparse.ArgumentParser()
|
|
parser.add_argument("text")
|
|
parser.add_argument("--max-samples", required=True, type=int)
|
|
parser.add_argument(
|
|
"--model", default=_DIR / "models" / "en_US-libritts_r-medium.pt"
|
|
)
|
|
parser.add_argument("--batch-size", type=int, default=1)
|
|
parser.add_argument("--slerp-weights", nargs="+", type=float, default=[0.5])
|
|
parser.add_argument(
|
|
"--length-scales", nargs="+", type=float, default=[1.0, 0.75, 1.25, 1.4]
|
|
)
|
|
parser.add_argument(
|
|
"--noise-scales",
|
|
nargs="+",
|
|
type=float,
|
|
default=[0.667, 0.75, 0.85, 0.9, 1.0, 1.4],
|
|
)
|
|
parser.add_argument("--noise-scale-ws", nargs="+", type=float, default=[0.8])
|
|
parser.add_argument("--output-dir", default="output")
|
|
parser.add_argument(
|
|
"--max-speakers",
|
|
type=int,
|
|
help="Maximum number of speakers to use (default: all)",
|
|
)
|
|
parser.add_argument("--min-phoneme-count", type=int)
|
|
parser.add_argument("--verbose", action="store_true")
|
|
args = parser.parse_args().__dict__
|
|
|
|
# Generate speech
|
|
generate_samples(**args)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|