From 6ecc5192d954f230d26853508130dd8511b1dcba Mon Sep 17 00:00:00 2001 From: Michael Hansen Date: Tue, 16 May 2023 17:06:57 -0500 Subject: [PATCH] Initial commit --- .projectile | 2 + augment.py | 88 +++++++++++++++++++ generate_samples.py | 209 ++++++++++++++++++++++++++++++++++++++++++++ requirements.txt | 4 + 4 files changed, 303 insertions(+) create mode 100644 .projectile create mode 100644 augment.py create mode 100755 generate_samples.py create mode 100644 requirements.txt diff --git a/.projectile b/.projectile new file mode 100644 index 0000000..ce66c37 --- /dev/null +++ b/.projectile @@ -0,0 +1,2 @@ +- /.venv/ +- /.mypy_cache/ diff --git a/augment.py b/augment.py new file mode 100644 index 0000000..834d223 --- /dev/null +++ b/augment.py @@ -0,0 +1,88 @@ +#!/usr/bin/env python3 +import argparse +import audioop +import sys +import wave +from pathlib import Path + +import numpy as np +from audiomentations import Compose, ApplyImpulseResponse, Gain + +_DIR = Path(__file__).parent + + +def main() -> None: + parser = argparse.ArgumentParser() + parser.add_argument("input_dir") + parser.add_argument("output_dir") + parser.add_argument("--sample-rate", type=int) + args = parser.parse_args() + + impulses = list((_DIR / "impulses").glob("*.wav")) + + augment = Compose( + transforms=[ + Gain(min_gain_in_db=-12, max_gain_in_db=0), + ApplyImpulseResponse(impulses), + ] + ) + + input_dir = Path(args.input_dir) + output_dir = Path(args.output_dir) + output_dir.mkdir(parents=True, exist_ok=True) + + for input_wav in input_dir.glob("*.wav"): + output_wav = output_dir / (input_wav.relative_to(input_dir)) + output_wav.parent.mkdir(parents=True, exist_ok=True) + + with wave.open(str(input_wav), "rb") as input_wav_file, wave.open( + str(output_wav), "wb" + ) as output_wav_file: + assert input_wav_file.getsampwidth() == 2 + assert input_wav_file.getnchannels() == 1 + + input_audio = ( + np.frombuffer( + input_wav_file.readframes(input_wav_file.getnframes()), + dtype=np.int16, + ).astype(np.float32) + / 32767.0 + ) + + output_audio = augment( + input_audio, sample_rate=input_wav_file.getframerate() + ) + output_wav_file.setframerate( + args.sample_rate or input_wav_file.getframerate() + ) + output_wav_file.setsampwidth(2) + output_wav_file.setnchannels(1) + + output_audio_16 = audio_float_to_int16(output_audio) + if args.sample_rate != input_wav_file.getframerate(): + output_audio_16, _state = audioop.ratecv( + output_audio_16, + 2, + 1, + input_wav_file.getframerate(), + args.sample_rate, + None, + ) + + output_wav_file.writeframes(output_audio_16) + + print(output_wav) + + +def audio_float_to_int16( + audio: np.ndarray, max_wav_value: float = 32767.0 +) -> np.ndarray: + # Don't normalize + audio_norm = audio * max_wav_value + audio_norm = np.clip(audio_norm, -max_wav_value, max_wav_value) + audio_norm = audio_norm.astype("int16") + return audio_norm + + +if __name__ == "__main__": + main() diff --git a/generate_samples.py b/generate_samples.py new file mode 100755 index 0000000..40258ec --- /dev/null +++ b/generate_samples.py @@ -0,0 +1,209 @@ +#!/usr/bin/env python3 +import argparse +import itertools as it +import json +import unicodedata +import wave +from pathlib import Path + +import numpy as np +import torch + +from espeak_phonemizer import Phonemizer +from piper_train.vits import commons + +_DIR = Path(__file__).parent + + +def main() -> None: + 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-high.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] + ) + parser.add_argument("--noise-scales", nargs="+", type=float, default=[0.667]) + parser.add_argument("--noise-scale-ws", nargs="+", type=float, default=[0.8]) + parser.add_argument("--output-dir", default="output") + args = parser.parse_args() + + model_path = Path(args.model) + model = torch.load(model_path) + model.eval() + + if torch.cuda.is_available(): + model.cuda() + + output_dir = Path(args.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"] + + phonemizer = Phonemizer(voice) + phonemes_str = phonemizer.phonemize(args.text) + phonemes = list(unicodedata.normalize("NFD", phonemes_str)) + + id_map = config["phoneme_id_map"] + phoneme_ids = list(id_map["^"]) + for phoneme in phonemes: + p_ids = id_map.get(phoneme) + if p_ids is not None: + phoneme_ids.extend(p_ids) + phoneme_ids.extend(id_map["_"]) + + phoneme_ids.extend(id_map["$"]) + + max_len = None + + sample_idx = 0 + is_done = False + settings_iter = it.cycle( + it.product( + args.slerp_weights, + args.length_scales, + args.noise_scales, + args.noise_scale_ws, + ) + ) + + speakers_iter = it.product(range(num_speakers), range(num_speakers)) + speakers_batch = list(it.islice(speakers_iter, 0, args.batch_size)) + 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]) + + x = torch.LongTensor(phoneme_ids).repeat((batch_size, 1)) + x_lengths = torch.LongTensor([len(phoneme_ids)]).repeat(batch_size) + + 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.cpu().numpy() + + audio_int16 = audio_float_to_int16(audio) + for audio_idx in range(batch_size): + wav_path = output_dir / f"{sample_idx}.wav" + with wave.open(str(wav_path), "wb") as wav_file: + wav_file.setframerate(sample_rate) + wav_file.setsampwidth(2) + wav_file.setnchannels(1) + wav_file.writeframes(audio_int16[audio_idx]) + + print(wav_path) + + sample_idx += 1 + if sample_idx >= args.max_samples: + is_done = True + break + + # Next batch + speakers_batch = list(it.islice(speakers_iter, 0, args.batch_size)) + + +def slerp(v1, v2, t, DOT_THR=0.9995, zdim=-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 + + +if __name__ == "__main__": + main() diff --git a/requirements.txt b/requirements.txt new file mode 100644 index 0000000..e4f92ed --- /dev/null +++ b/requirements.txt @@ -0,0 +1,4 @@ +espeak-phonemizer +torch +numpy +audiomentations