Initial commit

.projectile

@@ -0,0 +1,2 @@
+- /.venv/
+- /.mypy_cache/

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()

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()

requirements.txt

@@ -0,0 +1,4 @@
+espeak-phonemizer
+torch
+numpy
+audiomentations