diff --git a/.gitignore b/.gitignore new file mode 100644 index 0000000..2927924 --- /dev/null +++ b/.gitignore @@ -0,0 +1,11 @@ +.DS_Store +.idea +*.log +tmp/ + +.venv/ +__pycache__/ +.mypy_cache/ +*.pt + +output/ diff --git a/LICENSE.md b/LICENSE.md new file mode 100644 index 0000000..0875822 --- /dev/null +++ b/LICENSE.md @@ -0,0 +1,21 @@ +MIT License + +Copyright (c) 2023 Michael Hansen + +Permission is hereby granted, free of charge, to any person obtaining a copy +of this software and associated documentation files (the "Software"), to deal +in the Software without restriction, including without limitation the rights +to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +copies of the Software, and to permit persons to whom the Software is +furnished to do so, subject to the following conditions: + +The above copyright notice and this permission notice shall be included in all +copies or substantial portions of the Software. + +THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE +SOFTWARE. diff --git a/README.md b/README.md new file mode 100644 index 0000000..055a64e --- /dev/null +++ b/README.md @@ -0,0 +1,60 @@ +# Piper Sample Generator + + +## Install + +Create a virtual environment and install the requirements: + +``` sh +git clone https://github.com/rhasspy/piper-sample-generator.git +cd piper-sample-generator/ + +python3 -m venv .venv +source .venv/bin/activate +python3 -m pip install --upgrade pip +python3 -m pip install -r requirements.txt +``` + +Download the LibriTTS generator: + +``` sh +wget -O models/en-us-libritts-high.pt '' +``` + + +## Run + +Generate a small set of samples: + +``` sh +python3 generate_samples.py 'okay, piper.' --max-samples 10 --output-dir okay_piper/ +``` + +Check the `okay_piper/` directory for 10 WAV files (named `0.wav` to `9.wav`). + +Generation can be much faster and more efficient if you have a GPU available and PyTorch is configured to use it. In this case, increase the batch size: + +``` sh +python3 generate_samples.py 'okay, piper.' --max-samples 100 --batch-size 10 --output-dir okay_piper/ +``` + +On an NVidia 2080 Ti with 11GB, a batch size of 100 was possible (generating approximately 100 samples per second). + +See `--help` for more options, including adjust the `--length-scales` (speaking speeds) and `--slerp-weights` (speaker blending) which are cycled per batch. + +### Augmentation + +Once you have samples generating, you can augment them using [audiomentation](https://iver56.github.io/audiomentations/): + +``` sh +python3 augment.py --sample-rate 16000 okay_piper/ okay_piper_augmented/ +``` + +This will do several things to each sample: + +1. Randomly decrease the volume + * The original samples are normalized, so different volume levels are needed +2. Randomly [apply an impulse response](https://iver56.github.io/audiomentations/waveform_transforms/apply_impulse_response/) using the files in `impulses/` + * Change the acoustics of the sample to sound like the speaker was in a room with echo or using a poor quality microphone +3. Resample to 16Khz for training (e.g., [openWakeWord](https://github.com/dscripka/openWakeWord)) + diff --git a/generate_samples.py b/generate_samples.py index 40258ec..1b0d2f0 100755 --- a/generate_samples.py +++ b/generate_samples.py @@ -2,6 +2,7 @@ import argparse import itertools as it import json +import logging import unicodedata import wave from pathlib import Path @@ -13,6 +14,7 @@ from espeak_phonemizer import Phonemizer from piper_train.vits import commons _DIR = Path(__file__).parent +_LOGGER = logging.getLogger(__name__) def main() -> None: @@ -29,13 +31,17 @@ def main() -> None: parser.add_argument("--noise-scale-ws", nargs="+", type=float, default=[0.8]) parser.add_argument("--output-dir", default="output") args = parser.parse_args() + logging.basicConfig(level=logging.DEBUG) + _LOGGER.debug("Loading %s", args.model) model_path = Path(args.model) model = torch.load(model_path) model.eval() + _LOGGER.info("Successfully loaded %s", args.model) if torch.cuda.is_available(): model.cuda() + _LOGGER.debug("CUDA available, using GPU") output_dir = Path(args.output_dir) output_dir.mkdir(parents=True, exist_ok=True) @@ -51,6 +57,7 @@ def main() -> None: phonemizer = Phonemizer(voice) phonemes_str = phonemizer.phonemize(args.text) phonemes = list(unicodedata.normalize("NFD", phonemes_str)) + _LOGGER.debug("Phonemes: %s", phonemes) id_map = config["phoneme_id_map"] phoneme_ids = list(id_map["^"]) @@ -61,6 +68,7 @@ def main() -> None: phoneme_ids.extend(id_map["_"]) phoneme_ids.extend(id_map["$"]) + _LOGGER.debug("Phonemes ids: %s", phoneme_ids) max_len = None @@ -77,6 +85,7 @@ def main() -> None: speakers_iter = it.product(range(num_speakers), range(num_speakers)) speakers_batch = list(it.islice(speakers_iter, 0, args.batch_size)) + batch_idx = 0 while speakers_batch: if is_done: break @@ -147,7 +156,11 @@ def main() -> None: break # Next batch + _LOGGER.debug("Batch %s complete", batch_idx + 1) speakers_batch = list(it.islice(speakers_iter, 0, args.batch_size)) + batch_idx += 1 + + _LOGGER.info("Done") def slerp(v1, v2, t, DOT_THR=0.9995, zdim=-1): diff --git a/impulses/Accoustic2_Impulse.wav b/impulses/Accoustic2_Impulse.wav new file mode 100644 index 0000000..fce6e97 Binary files /dev/null and b/impulses/Accoustic2_Impulse.wav differ diff --git a/impulses/Blatty Plate.wav b/impulses/Blatty Plate.wav new file mode 100644 index 0000000..7f940e6 Binary files /dev/null and b/impulses/Blatty Plate.wav differ diff --git a/impulses/Concrete Room.wav b/impulses/Concrete Room.wav new file mode 100644 index 0000000..8e6edcd Binary files /dev/null and b/impulses/Concrete Room.wav differ diff --git a/impulses/Derlon Sanctuary.wav b/impulses/Derlon Sanctuary.wav new file mode 100644 index 0000000..3d3eb6e Binary files /dev/null and b/impulses/Derlon Sanctuary.wav differ diff --git a/impulses/Fat Bass.wav b/impulses/Fat Bass.wav new file mode 100644 index 0000000..7721afc Binary files /dev/null and b/impulses/Fat Bass.wav differ diff --git a/impulses/Reverse Gate.wav b/impulses/Reverse Gate.wav new file mode 100644 index 0000000..fe4a7bc Binary files /dev/null and b/impulses/Reverse Gate.wav differ diff --git a/impulses/Symphonic.wav b/impulses/Symphonic.wav new file mode 100644 index 0000000..c50125b Binary files /dev/null and b/impulses/Symphonic.wav differ diff --git a/impulses/ir_bathroom1.wav b/impulses/ir_bathroom1.wav new file mode 100644 index 0000000..c1d4f65 Binary files /dev/null and b/impulses/ir_bathroom1.wav differ diff --git a/models/.gitkeep b/models/.gitkeep new file mode 100644 index 0000000..e69de29 diff --git a/models/en-us-libritts-high.pt.json b/models/en-us-libritts-high.pt.json new file mode 100644 index 0000000..357af0e --- /dev/null +++ b/models/en-us-libritts-high.pt.json @@ -0,0 +1,1314 @@ +{ + "audio": { + "sample_rate": 22050 + }, + "espeak": { + "voice": "en-us" + }, + "inference": { + "noise_scale": 0.667, + "length_scale": 1, + "noise_w": 0.8 + }, + "phoneme_map": {}, + "phoneme_id_map": { + "_": [ + 0 + ], + "^": [ + 1 + ], + "$": [ + 2 + ], + " ": [ + 3 + ], + "!": [ + 4 + ], + "'": [ + 5 + ], + "(": [ + 6 + ], + ")": [ + 7 + ], + ",": [ + 8 + ], + "-": [ + 9 + ], + ".": [ + 10 + ], + ":": [ + 11 + ], + ";": [ + 12 + ], + "?": [ + 13 + ], + "a": [ + 14 + ], + "b": [ + 15 + ], + "c": [ + 16 + ], + "d": [ + 17 + ], + "e": [ + 18 + ], + "f": [ + 19 + ], + "h": [ + 20 + ], + "i": [ + 21 + ], + "j": [ + 22 + ], + "k": [ + 23 + ], + "l": [ + 24 + ], + "m": [ + 25 + ], + "n": [ + 26 + ], + "o": [ + 27 + ], + "p": [ + 28 + ], + "q": [ + 29 + ], + "r": [ + 30 + ], + "s": [ + 31 + ], + "t": [ + 32 + ], + "u": [ + 33 + ], + "v": [ + 34 + ], + "w": [ + 35 + ], + "x": [ + 36 + ], + "y": [ + 37 + ], + "z": [ + 38 + ], + "æ": [ + 39 + ], + "ç": [ + 40 + ], + "ð": [ + 41 + ], + "ø": [ + 42 + ], + "ħ": [ + 43 + ], + "ŋ": [ + 44 + ], + "œ": [ + 45 + ], + "ǀ": [ + 46 + ], + "ǁ": [ + 47 + ], + "ǂ": [ + 48 + ], + "ǃ": [ + 49 + ], + "ɐ": [ + 50 + ], + "ɑ": [ + 51 + ], + "ɒ": [ + 52 + ], + "ɓ": [ + 53 + ], + "ɔ": [ + 54 + ], + "ɕ": [ + 55 + ], + "ɖ": [ + 56 + ], + "ɗ": [ + 57 + ], + "ɘ": [ + 58 + ], + "ə": [ + 59 + ], + "ɚ": [ + 60 + ], + "ɛ": [ + 61 + ], + "ɜ": [ + 62 + ], + "ɞ": [ + 63 + ], + "ɟ": [ + 64 + ], + "ɠ": [ + 65 + ], + "ɡ": [ + 66 + ], + "ɢ": [ + 67 + ], + "ɣ": [ + 68 + ], + "ɤ": [ + 69 + ], + "ɥ": [ + 70 + ], + "ɦ": [ + 71 + ], + "ɧ": [ + 72 + ], + "ɨ": [ + 73 + ], + "ɪ": [ + 74 + ], + "ɫ": [ + 75 + ], + "ɬ": [ + 76 + ], + "ɭ": [ + 77 + ], + "ɮ": [ + 78 + ], + "ɯ": [ + 79 + ], + "ɰ": [ + 80 + ], + "ɱ": [ + 81 + ], + "ɲ": [ + 82 + ], + "ɳ": [ + 83 + ], + "ɴ": [ + 84 + ], + "ɵ": [ + 85 + ], + "ɶ": [ + 86 + ], + "ɸ": [ + 87 + ], + "ɹ": [ + 88 + ], + "ɺ": [ + 89 + ], + "ɻ": [ + 90 + ], + "ɽ": [ + 91 + ], + "ɾ": [ + 92 + ], + "ʀ": [ + 93 + ], + "ʁ": [ + 94 + ], + "ʂ": [ + 95 + ], + "ʃ": [ + 96 + ], + "ʄ": [ + 97 + ], + "ʈ": [ + 98 + ], + "ʉ": [ + 99 + ], + "ʊ": [ + 100 + ], + "ʋ": [ + 101 + ], + "ʌ": [ + 102 + ], + "ʍ": [ + 103 + ], + "ʎ": [ + 104 + ], + "ʏ": [ + 105 + ], + "ʐ": [ + 106 + ], + "ʑ": [ + 107 + ], + "ʒ": [ + 108 + ], + "ʔ": [ + 109 + ], + "ʕ": [ + 110 + ], + "ʘ": [ + 111 + ], + "ʙ": [ + 112 + ], + "ʛ": [ + 113 + ], + "ʜ": [ + 114 + ], + "ʝ": [ + 115 + ], + "ʟ": [ + 116 + ], + "ʡ": [ + 117 + ], + "ʢ": [ + 118 + ], + "ʲ": [ + 119 + ], + "ˈ": [ + 120 + ], + "ˌ": [ + 121 + ], + "ː": [ + 122 + ], + "ˑ": [ + 123 + ], + "˞": [ + 124 + ], + "β": [ + 125 + ], + "θ": [ + 126 + ], + "χ": [ + 127 + ], + "ᵻ": [ + 128 + ], + "ⱱ": [ + 129 + ] + }, + "num_symbols": 130, + "num_speakers": 904, + "speaker_id_map": { + "3922": 0, + "8699": 1, + "4535": 2, + "6701": 3, + "3638": 4, + "922": 5, + "2531": 6, + "1638": 7, + "8848": 8, + "6544": 9, + "3615": 10, + "318": 11, + "6104": 12, + "1382": 13, + "5400": 14, + "5712": 15, + "2769": 16, + "2573": 17, + "1463": 18, + "6458": 19, + "3274": 20, + "4356": 21, + "8498": 22, + "5570": 23, + "176": 24, + "339": 25, + "28": 26, + "5909": 27, + "3869": 28, + "4899": 29, + "64": 30, + "3368": 31, + "3307": 32, + "5618": 33, + "3370": 34, + "7704": 35, + "8506": 36, + "8410": 37, + "6904": 38, + "5655": 39, + "2204": 40, + "501": 41, + "7314": 42, + "1027": 43, + "5054": 44, + "534": 45, + "2853": 46, + "5935": 47, + "2404": 48, + "7874": 49, + "816": 50, + "2053": 51, + "8066": 52, + "16": 53, + "4586": 54, + "1923": 55, + "2592": 56, + "1265": 57, + "6189": 58, + "100": 59, + "6371": 60, + "4957": 61, + "4116": 62, + "3003": 63, + "7739": 64, + "1752": 65, + "5717": 66, + "5012": 67, + "5062": 68, + "7481": 69, + "4595": 70, + "2299": 71, + "7188": 72, + "93": 73, + "4145": 74, + "8684": 75, + "7594": 76, + "2598": 77, + "3540": 78, + "7717": 79, + "6426": 80, + "4148": 81, + "335": 82, + "1379": 83, + "2512": 84, + "242": 85, + "8855": 86, + "8118": 87, + "369": 88, + "6575": 89, + "6694": 90, + "8080": 91, + "1283": 92, + "7434": 93, + "5290": 94, + "1731": 95, + "2401": 96, + "459": 97, + "192": 98, + "7910": 99, + "114": 100, + "5660": 101, + "1313": 102, + "203": 103, + "7460": 104, + "207": 105, + "6497": 106, + "6696": 107, + "7766": 108, + "6233": 109, + "3185": 110, + "2010": 111, + "2056": 112, + "3717": 113, + "5802": 114, + "5622": 115, + "2156": 116, + "4243": 117, + "1422": 118, + "5039": 119, + "4110": 120, + "1093": 121, + "1776": 122, + "7995": 123, + "6877": 124, + "5635": 125, + "54": 126, + "288": 127, + "4592": 128, + "7276": 129, + "688": 130, + "8388": 131, + "8152": 132, + "8194": 133, + "7000": 134, + "8527": 135, + "5126": 136, + "3923": 137, + "1054": 138, + "3927": 139, + "5029": 140, + "4098": 141, + "1789": 142, + "56": 143, + "7240": 144, + "5538": 145, + "1903": 146, + "6538": 147, + "3380": 148, + "6643": 149, + "7495": 150, + "8718": 151, + "8050": 152, + "126": 153, + "7245": 154, + "2517": 155, + "4438": 156, + "4945": 157, + "7145": 158, + "724": 159, + "9022": 160, + "6637": 161, + "6927": 162, + "6937": 163, + "8113": 164, + "5724": 165, + "6006": 166, + "3584": 167, + "2971": 168, + "2230": 169, + "7982": 170, + "1649": 171, + "3994": 172, + "7720": 173, + "6981": 174, + "781": 175, + "4973": 176, + "6206": 177, + "2481": 178, + "3157": 179, + "1509": 180, + "510": 181, + "7540": 182, + "8887": 183, + "7120": 184, + "2882": 185, + "7128": 186, + "8142": 187, + "7229": 188, + "2787": 189, + "8820": 190, + "2368": 191, + "4331": 192, + "4967": 193, + "4427": 194, + "6054": 195, + "3728": 196, + "274": 197, + "7134": 198, + "1603": 199, + "1383": 200, + "1165": 201, + "4363": 202, + "512": 203, + "5985": 204, + "7967": 205, + "2060": 206, + "7752": 207, + "7484": 208, + "8643": 209, + "3549": 210, + "5731": 211, + "7881": 212, + "667": 213, + "6828": 214, + "5740": 215, + "3483": 216, + "718": 217, + "6341": 218, + "1913": 219, + "3228": 220, + "7247": 221, + "7705": 222, + "1018": 223, + "8193": 224, + "6098": 225, + "3989": 226, + "7828": 227, + "5876": 228, + "7754": 229, + "4719": 230, + "8011": 231, + "7939": 232, + "5975": 233, + "2004": 234, + "6139": 235, + "8183": 236, + "3482": 237, + "3361": 238, + "4289": 239, + "231": 240, + "7789": 241, + "4598": 242, + "5239": 243, + "2638": 244, + "6300": 245, + "8474": 246, + "2194": 247, + "7832": 248, + "1079": 249, + "1335": 250, + "188": 251, + "1195": 252, + "5914": 253, + "1401": 254, + "7318": 255, + "5448": 256, + "1392": 257, + "3703": 258, + "2113": 259, + "7783": 260, + "8176": 261, + "6519": 262, + "7933": 263, + "7938": 264, + "7802": 265, + "6120": 266, + "224": 267, + "209": 268, + "5656": 269, + "3032": 270, + "6965": 271, + "258": 272, + "4837": 273, + "5489": 274, + "272": 275, + "3851": 276, + "7140": 277, + "2562": 278, + "1472": 279, + "79": 280, + "2775": 281, + "3046": 282, + "2532": 283, + "8266": 284, + "6099": 285, + "4425": 286, + "5293": 287, + "7981": 288, + "2045": 289, + "920": 290, + "511": 291, + "7416": 292, + "835": 293, + "1289": 294, + "8195": 295, + "7833": 296, + "8772": 297, + "968": 298, + "1641": 299, + "7117": 300, + "1678": 301, + "5809": 302, + "8028": 303, + "500": 304, + "6505": 305, + "7868": 306, + "14": 307, + "2238": 308, + "4744": 309, + "3733": 310, + "7515": 311, + "699": 312, + "5093": 313, + "6388": 314, + "7959": 315, + "98": 316, + "3914": 317, + "5246": 318, + "2570": 319, + "8396": 320, + "3513": 321, + "882": 322, + "7994": 323, + "5968": 324, + "8591": 325, + "806": 326, + "5261": 327, + "1271": 328, + "899": 329, + "3945": 330, + "8404": 331, + "249": 332, + "3008": 333, + "7139": 334, + "6395": 335, + "6215": 336, + "6080": 337, + "4054": 338, + "7825": 339, + "6683": 340, + "8725": 341, + "3230": 342, + "4138": 343, + "6160": 344, + "666": 345, + "6510": 346, + "3551": 347, + "8075": 348, + "225": 349, + "7169": 350, + "1851": 351, + "5984": 352, + "2960": 353, + "8329": 354, + "175": 355, + "6378": 356, + "480": 357, + "7538": 358, + "479": 359, + "5519": 360, + "8534": 361, + "4856": 362, + "101": 363, + "3521": 364, + "2256": 365, + "3083": 366, + "4278": 367, + "8713": 368, + "1226": 369, + "4222": 370, + "8494": 371, + "8776": 372, + "731": 373, + "6574": 374, + "5319": 375, + "8605": 376, + "5583": 377, + "6406": 378, + "4064": 379, + "4806": 380, + "3972": 381, + "7383": 382, + "5133": 383, + "597": 384, + "1025": 385, + "7313": 386, + "5304": 387, + "8758": 388, + "1050": 389, + "6499": 390, + "6956": 391, + "770": 392, + "4108": 393, + "2774": 394, + "3864": 395, + "4490": 396, + "4848": 397, + "1826": 398, + "6294": 399, + "7949": 400, + "1446": 401, + "7867": 402, + "8163": 403, + "953": 404, + "8138": 405, + "353": 406, + "7553": 407, + "8825": 408, + "5189": 409, + "2012": 410, + "948": 411, + "205": 412, + "1535": 413, + "8008": 414, + "1112": 415, + "7926": 416, + "4039": 417, + "716": 418, + "3967": 419, + "7932": 420, + "7525": 421, + "7316": 422, + "3448": 423, + "2393": 424, + "6788": 425, + "6550": 426, + "7011": 427, + "8791": 428, + "8119": 429, + "1777": 430, + "6014": 431, + "1046": 432, + "6269": 433, + "6188": 434, + "5266": 435, + "3490": 436, + "8786": 437, + "8824": 438, + "589": 439, + "576": 440, + "1121": 441, + "1806": 442, + "7294": 443, + "3119": 444, + "2688": 445, + "1012": 446, + "4807": 447, + "7498": 448, + "3905": 449, + "7384": 450, + "2992": 451, + "30": 452, + "497": 453, + "227": 454, + "4226": 455, + "5007": 456, + "1066": 457, + "8222": 458, + "7688": 459, + "6865": 460, + "6286": 461, + "8225": 462, + "3224": 463, + "8635": 464, + "1348": 465, + "3645": 466, + "1961": 467, + "8190": 468, + "6032": 469, + "7286": 470, + "5389": 471, + "3105": 472, + "1028": 473, + "6038": 474, + "764": 475, + "7437": 476, + "6555": 477, + "8875": 478, + "2074": 479, + "7809": 480, + "2240": 481, + "2827": 482, + "5386": 483, + "6763": 484, + "3009": 485, + "6339": 486, + "1825": 487, + "7569": 488, + "359": 489, + "7956": 490, + "2137": 491, + "8677": 492, + "4434": 493, + "329": 494, + "3289": 495, + "4290": 496, + "2999": 497, + "2427": 498, + "637": 499, + "2229": 500, + "1874": 501, + "3446": 502, + "9023": 503, + "3114": 504, + "6235": 505, + "4860": 506, + "4519": 507, + "561": 508, + "70": 509, + "4800": 510, + "2294": 511, + "6115": 512, + "2582": 513, + "8464": 514, + "5139": 515, + "6918": 516, + "337": 517, + "5810": 518, + "8401": 519, + "303": 520, + "5206": 521, + "2589": 522, + "7061": 523, + "2269": 524, + "2758": 525, + "3389": 526, + "4629": 527, + "707": 528, + "5606": 529, + "1513": 530, + "2473": 531, + "664": 532, + "5092": 533, + "5154": 534, + "6288": 535, + "6308": 536, + "4731": 537, + "3328": 538, + "7816": 539, + "3221": 540, + "8687": 541, + "7030": 542, + "476": 543, + "4257": 544, + "5918": 545, + "6317": 546, + "204": 547, + "8006": 548, + "6895": 549, + "1264": 550, + "2494": 551, + "112": 552, + "1859": 553, + "398": 554, + "1052": 555, + "3294": 556, + "1460": 557, + "8573": 558, + "5684": 559, + "8421": 560, + "5883": 561, + "7297": 562, + "246": 563, + "8057": 564, + "3835": 565, + "1748": 566, + "3816": 567, + "3357": 568, + "1053": 569, + "409": 570, + "868": 571, + "3118": 572, + "7520": 573, + "6686": 574, + "1241": 575, + "5190": 576, + "166": 577, + "1482": 578, + "5604": 579, + "1212": 580, + "2741": 581, + "1259": 582, + "984": 583, + "6492": 584, + "6167": 585, + "296": 586, + "6567": 587, + "6924": 588, + "2272": 589, + "7085": 590, + "345": 591, + "2388": 592, + "1705": 593, + "1343": 594, + "7241": 595, + "451": 596, + "5401": 597, + "6446": 598, + "612": 599, + "594": 600, + "7555": 601, + "7069": 602, + "2577": 603, + "5333": 604, + "8742": 605, + "6727": 606, + "1571": 607, + "4734": 608, + "7258": 609, + "3977": 610, + "373": 611, + "5723": 612, + "1365": 613, + "7285": 614, + "580": 615, + "836": 616, + "6782": 617, + "3654": 618, + "1974": 619, + "6258": 620, + "925": 621, + "949": 622, + "2790": 623, + "698": 624, + "6373": 625, + "2785": 626, + "1222": 627, + "2751": 628, + "3825": 629, + "5115": 630, + "1827": 631, + "3171": 632, + "119": 633, + "850": 634, + "3258": 635, + "7909": 636, + "1322": 637, + "8097": 638, + "22": 639, + "7478": 640, + "1349": 641, + "4854": 642, + "2929": 643, + "7335": 644, + "5868": 645, + "454": 646, + "7945": 647, + "2654": 648, + "3493": 649, + "1060": 650, + "8545": 651, + "6509": 652, + "5002": 653, + "7732": 654, + "3082": 655, + "1779": 656, + "2709": 657, + "7398": 658, + "8879": 659, + "639": 660, + "598": 661, + "5672": 662, + "6553": 663, + "4111": 664, + "1417": 665, + "7991": 666, + "380": 667, + "8459": 668, + "8347": 669, + "1769": 670, + "2673": 671, + "3330": 672, + "7051": 673, + "1337": 674, + "4057": 675, + "4839": 676, + "6060": 677, + "7095": 678, + "278": 679, + "1445": 680, + "6518": 681, + "2364": 682, + "1958": 683, + "548": 684, + "4010": 685, + "3072": 686, + "6993": 687, + "8575": 688, + "2149": 689, + "240": 690, + "2920": 691, + "5588": 692, + "1885": 693, + "6082": 694, + "9026": 695, + "340": 696, + "159": 697, + "7730": 698, + "7962": 699, + "1987": 700, + "3876": 701, + "8771": 702, + "5123": 703, + "3866": 704, + "3546": 705, + "7777": 706, + "115": 707, + "5337": 708, + "475": 709, + "1724": 710, + "6359": 711, + "4260": 712, + "2110": 713, + "1845": 714, + "4335": 715, + "4133": 716, + "783": 717, + "8479": 718, + "1448": 719, + "1160": 720, + "7647": 721, + "2618": 722, + "3630": 723, + "4013": 724, + "5242": 725, + "7957": 726, + "3852": 727, + "3889": 728, + "1387": 729, + "439": 730, + "1425": 731, + "2061": 732, + "7395": 733, + "7837": 734, + "5147": 735, + "2319": 736, + "3781": 737, + "1311": 738, + "4733": 739, + "8705": 740, + "3094": 741, + "2823": 742, + "1914": 743, + "954": 744, + "4381": 745, + "4044": 746, + "593": 747, + "8300": 748, + "7558": 749, + "6494": 750, + "6330": 751, + "5940": 752, + "7126": 753, + "1061": 754, + "6352": 755, + "5186": 756, + "1944": 757, + "2285": 758, + "6673": 759, + "5746": 760, + "208": 761, + "492": 762, + "216": 763, + "979": 764, + "1668": 765, + "6620": 766, + "711": 767, + "7733": 768, + "8619": 769, + "5157": 770, + "829": 771, + "3180": 772, + "3979": 773, + "1556": 774, + "3379": 775, + "5727": 776, + "596": 777, + "2127": 778, + "581": 779, + "2652": 780, + "2628": 781, + "1849": 782, + "4238": 783, + "606": 784, + "1224": 785, + "1629": 786, + "1413": 787, + "957": 788, + "8592": 789, + "2254": 790, + "1323": 791, + "122": 792, + "2093": 793, + "1100": 794, + "81": 795, + "323": 796, + "815": 797, + "2581": 798, + "543": 799, + "6037": 800, + "2397": 801, + "5513": 802, + "4495": 803, + "5776": 804, + "17": 805, + "4590": 806, + "8228": 807, + "708": 808, + "3792": 809, + "3790": 810, + "7090": 811, + "1943": 812, + "4246": 813, + "559": 814, + "3738": 815, + "2167": 816, + "1933": 817, + "2162": 818, + "549": 819, + "3025": 820, + "1182": 821, + "4358": 822, + "636": 823, + "986": 824, + "8490": 825, + "3340": 826, + "90": 827, + "1487": 828, + "1639": 829, + "1547": 830, + "4152": 831, + "1498": 832, + "1740": 833, + "6157": 834, + "217": 835, + "2201": 836, + "362": 837, + "2146": 838, + "1801": 839, + "5063": 840, + "7339": 841, + "663": 842, + "38": 843, + "1336": 844, + "3215": 845, + "210": 846, + "6075": 847, + "55": 848, + "2411": 849, + "7445": 850, + "5767": 851, + "2812": 852, + "472": 853, + "803": 854, + "4236": 855, + "7665": 856, + "1607": 857, + "1316": 858, + "7475": 859, + "3001": 860, + "1473": 861, + "3537": 862, + "3070": 863, + "1390": 864, + "1290": 865, + "2499": 866, + "154": 867, + "7518": 868, + "408": 869, + "1811": 870, + "1734": 871, + "7342": 872, + "8722": 873, + "1754": 874, + "7657": 875, + "583": 876, + "830": 877, + "6690": 878, + "1552": 879, + "2498": 880, + "1296": 881, + "3686": 882, + "157": 883, + "487": 884, + "6119": 885, + "4926": 886, + "4846": 887, + "1536": 888, + "2674": 889, + "1645": 890, + "3187": 891, + "1058": 892, + "2039": 893, + "4071": 894, + "4433": 895, + "1175": 896, + "434": 897, + "1001": 898, + "2816": 899, + "820": 900, + "2696": 901, + "4681": 902, + "2085": 903 + } +} diff --git a/piper_train/VERSION b/piper_train/VERSION new file mode 100644 index 0000000..6e8bf73 --- /dev/null +++ b/piper_train/VERSION @@ -0,0 +1 @@ +0.1.0 diff --git a/piper_train/__init__.py b/piper_train/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/piper_train/__main__.py b/piper_train/__main__.py new file mode 100644 index 0000000..a2ee0d3 --- /dev/null +++ b/piper_train/__main__.py @@ -0,0 +1,95 @@ +import argparse +import json +import logging +from pathlib import Path + +import torch +from pytorch_lightning import Trainer +from pytorch_lightning.callbacks import ModelCheckpoint + +from .vits.lightning import VitsModel + +_LOGGER = logging.getLogger(__package__) + + +def main(): + logging.basicConfig(level=logging.DEBUG) + + parser = argparse.ArgumentParser() + parser.add_argument( + "--dataset-dir", required=True, help="Path to pre-processed dataset directory" + ) + parser.add_argument( + "--checkpoint-epochs", + type=int, + help="Save checkpoint every N epochs (default: 1)", + ) + parser.add_argument( + "--quality", + default="medium", + choices=("x-low", "medium", "high"), + help="Quality/size of model (default: medium)", + ) + Trainer.add_argparse_args(parser) + VitsModel.add_model_specific_args(parser) + parser.add_argument("--seed", type=int, default=1234) + args = parser.parse_args() + _LOGGER.debug(args) + + args.dataset_dir = Path(args.dataset_dir) + if not args.default_root_dir: + args.default_root_dir = args.dataset_dir + + torch.backends.cudnn.benchmark = True + torch.manual_seed(args.seed) + + config_path = args.dataset_dir / "config.json" + dataset_path = args.dataset_dir / "dataset.jsonl" + + with open(config_path, "r", encoding="utf-8") as config_file: + # See preprocess.py for format + config = json.load(config_file) + num_symbols = int(config["num_symbols"]) + num_speakers = int(config["num_speakers"]) + sample_rate = int(config["audio"]["sample_rate"]) + + trainer = Trainer.from_argparse_args(args) + if args.checkpoint_epochs is not None: + trainer.callbacks = [ModelCheckpoint(every_n_epochs=args.checkpoint_epochs)] + _LOGGER.debug( + "Checkpoints will be saved every %s epoch(s)", args.checkpoint_epochs + ) + + dict_args = vars(args) + if args.quality == "x-low": + dict_args["hidden_channels"] = 96 + dict_args["inter_channels"] = 96 + dict_args["filter_channels"] = 384 + elif args.quality == "high": + dict_args["resblock"] = "1" + dict_args["resblock_kernel_sizes"] = (3, 7, 11) + dict_args["resblock_dilation_sizes"] = ( + (1, 3, 5), + (1, 3, 5), + (1, 3, 5), + ) + dict_args["upsample_rates"] = (8, 8, 2, 2) + dict_args["upsample_initial_channel"] = 512 + dict_args["upsample_kernel_sizes"] = (16, 16, 4, 4) + + model = VitsModel( + num_symbols=num_symbols, + num_speakers=num_speakers, + sample_rate=sample_rate, + dataset=[dataset_path], + **dict_args + ) + + trainer.fit(model) + + +# ----------------------------------------------------------------------------- + + +if __name__ == "__main__": + main() diff --git a/piper_train/_resources.py b/piper_train/_resources.py new file mode 100644 index 0000000..874f72f --- /dev/null +++ b/piper_train/_resources.py @@ -0,0 +1,19 @@ +"""Shared access to package resources""" +import os +import typing +from pathlib import Path + +try: + import importlib.resources + + files = importlib.resources.files +except (ImportError, AttributeError): + # Backport for Python < 3.9 + import importlib_resources # type: ignore + + files = importlib_resources.files + +_PACKAGE = "piper_train" +_DIR = Path(typing.cast(os.PathLike, files(_PACKAGE))) + +__version__ = (_DIR / "VERSION").read_text(encoding="utf-8").strip() diff --git a/piper_train/norm_audio/__init__.py b/piper_train/norm_audio/__init__.py new file mode 100644 index 0000000..dc024d4 --- /dev/null +++ b/piper_train/norm_audio/__init__.py @@ -0,0 +1,92 @@ +from hashlib import sha256 +from pathlib import Path +from typing import Optional, Tuple, Union + +import librosa +import torch + +from piper_train.vits.mel_processing import spectrogram_torch + +from .trim import trim_silence +from .vad import SileroVoiceActivityDetector + +_DIR = Path(__file__).parent + + +def make_silence_detector() -> SileroVoiceActivityDetector: + silence_model = _DIR / "models" / "silero_vad.onnx" + return SileroVoiceActivityDetector(silence_model) + + +def cache_norm_audio( + audio_path: Union[str, Path], + cache_dir: Union[str, Path], + detector: SileroVoiceActivityDetector, + sample_rate: int, + silence_threshold: float = 0.2, + silence_samples_per_chunk: int = 480, + silence_keep_chunks_before: int = 2, + silence_keep_chunks_after: int = 2, + filter_length: int = 1024, + window_length: int = 1024, + hop_length: int = 256, + ignore_cache: bool = False, +) -> Tuple[Path, Path]: + audio_path = Path(audio_path).absolute() + cache_dir = Path(cache_dir) + + # Cache id is the SHA256 of the full audio path + audio_cache_id = sha256(str(audio_path).encode()).hexdigest() + + audio_norm_path = cache_dir / f"{audio_cache_id}.pt" + audio_spec_path = cache_dir / f"{audio_cache_id}.spec.pt" + + # Normalize audio + audio_norm_tensor: Optional[torch.FloatTensor] = None + if ignore_cache or (not audio_norm_path.exists()): + # Trim silence first. + # + # The VAD model works on 16khz, so we determine the portion of audio + # to keep and then just load that with librosa. + vad_sample_rate = 16000 + audio_16khz, _sr = librosa.load(path=audio_path, sr=vad_sample_rate) + + offset_sec, duration_sec = trim_silence( + audio_16khz, + detector, + threshold=silence_threshold, + samples_per_chunk=silence_samples_per_chunk, + sample_rate=vad_sample_rate, + keep_chunks_before=silence_keep_chunks_before, + keep_chunks_after=silence_keep_chunks_after, + ) + + # NOTE: audio is already in [-1, 1] coming from librosa + audio_norm_array, _sr = librosa.load( + path=audio_path, + sr=sample_rate, + offset=offset_sec, + duration=duration_sec, + ) + + # Save to cache directory + audio_norm_tensor = torch.FloatTensor(audio_norm_array).unsqueeze(0) + torch.save(audio_norm_tensor, audio_norm_path) + + # Compute spectrogram + if ignore_cache or (not audio_spec_path.exists()): + if audio_norm_tensor is None: + # Load pre-cached normalized audio + audio_norm_tensor = torch.load(audio_norm_path) + + audio_spec_tensor = spectrogram_torch( + y=audio_norm_tensor, + n_fft=filter_length, + sampling_rate=sample_rate, + hop_size=hop_length, + win_size=window_length, + center=False, + ).squeeze(0) + torch.save(audio_spec_tensor, audio_spec_path) + + return audio_norm_path, audio_spec_path diff --git a/piper_train/norm_audio/trim.py b/piper_train/norm_audio/trim.py new file mode 100644 index 0000000..b076e86 --- /dev/null +++ b/piper_train/norm_audio/trim.py @@ -0,0 +1,54 @@ +from typing import Optional, Tuple + +import numpy as np + +from .vad import SileroVoiceActivityDetector + + +def trim_silence( + audio_array: np.ndarray, + detector: SileroVoiceActivityDetector, + threshold: float = 0.2, + samples_per_chunk=480, + sample_rate=16000, + keep_chunks_before: int = 2, + keep_chunks_after: int = 2, +) -> Tuple[float, Optional[float]]: + """Returns the offset/duration of trimmed audio in seconds""" + offset_sec: float = 0.0 + duration_sec: Optional[float] = None + first_chunk: Optional[int] = None + last_chunk: Optional[int] = None + seconds_per_chunk: float = samples_per_chunk / sample_rate + + chunk = audio_array[:samples_per_chunk] + audio_array = audio_array[samples_per_chunk:] + chunk_idx: int = 0 + + # Determine main block of speech + while len(audio_array) > 0: + prob = detector(chunk, sample_rate=sample_rate) + is_speech = prob >= threshold + + if is_speech: + if first_chunk is None: + # First speech + first_chunk = chunk_idx + else: + # Last speech so far + last_chunk = chunk_idx + + chunk = audio_array[:samples_per_chunk] + audio_array = audio_array[samples_per_chunk:] + chunk_idx += 1 + + if (first_chunk is not None) and (last_chunk is not None): + first_chunk = max(0, first_chunk - keep_chunks_before) + last_chunk = min(chunk_idx, last_chunk + keep_chunks_after) + + # Compute offset/duration + offset_sec = first_chunk * seconds_per_chunk + last_sec = (last_chunk + 1) * seconds_per_chunk + duration_sec = last_sec - offset_sec + + return offset_sec, duration_sec diff --git a/piper_train/norm_audio/vad.py b/piper_train/norm_audio/vad.py new file mode 100644 index 0000000..c9de88f --- /dev/null +++ b/piper_train/norm_audio/vad.py @@ -0,0 +1,54 @@ +import typing +from pathlib import Path + +import numpy as np +import onnxruntime + + +class SileroVoiceActivityDetector: + """Detects speech/silence using Silero VAD. + + https://github.com/snakers4/silero-vad + """ + + def __init__(self, onnx_path: typing.Union[str, Path]): + onnx_path = str(onnx_path) + + self.session = onnxruntime.InferenceSession(onnx_path) + self.session.intra_op_num_threads = 1 + self.session.inter_op_num_threads = 1 + + self._h = np.zeros((2, 1, 64)).astype("float32") + self._c = np.zeros((2, 1, 64)).astype("float32") + + def __call__(self, audio_array: np.ndarray, sample_rate: int = 16000): + """Return probability of speech in audio [0-1]. + + Audio must be 16Khz 16-bit mono PCM. + """ + if len(audio_array.shape) == 1: + # Add batch dimension + audio_array = np.expand_dims(audio_array, 0) + + if len(audio_array.shape) > 2: + raise ValueError( + f"Too many dimensions for input audio chunk {audio_array.shape}" + ) + + if audio_array.shape[0] > 1: + raise ValueError("Onnx model does not support batching") + + if sample_rate != 16000: + raise ValueError("Only 16Khz audio is supported") + + ort_inputs = { + "input": audio_array.astype(np.float32), + "h0": self._h, + "c0": self._c, + } + ort_outs = self.session.run(None, ort_inputs) + out, self._h, self._c = ort_outs + + out = out.squeeze(2)[:, 1] # make output type match JIT analog + + return out diff --git a/piper_train/vits/__init__.py b/piper_train/vits/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/piper_train/vits/attentions.py b/piper_train/vits/attentions.py new file mode 100644 index 0000000..20b7a9e --- /dev/null +++ b/piper_train/vits/attentions.py @@ -0,0 +1,427 @@ +import math +import typing + +import torch +from torch import nn +from torch.nn import functional as F + +from .commons import subsequent_mask +from .modules import LayerNorm + + +class Encoder(nn.Module): + def __init__( + self, + hidden_channels: int, + filter_channels: int, + n_heads: int, + n_layers: int, + kernel_size: int = 1, + p_dropout: float = 0.0, + window_size: int = 4, + **kwargs + ): + super().__init__() + self.hidden_channels = hidden_channels + self.filter_channels = filter_channels + self.n_heads = n_heads + self.n_layers = n_layers + self.kernel_size = kernel_size + self.p_dropout = p_dropout + self.window_size = window_size + + self.drop = nn.Dropout(p_dropout) + self.attn_layers = nn.ModuleList() + self.norm_layers_1 = nn.ModuleList() + self.ffn_layers = nn.ModuleList() + self.norm_layers_2 = nn.ModuleList() + for i in range(self.n_layers): + self.attn_layers.append( + MultiHeadAttention( + hidden_channels, + hidden_channels, + n_heads, + p_dropout=p_dropout, + window_size=window_size, + ) + ) + self.norm_layers_1.append(LayerNorm(hidden_channels)) + self.ffn_layers.append( + FFN( + hidden_channels, + hidden_channels, + filter_channels, + kernel_size, + p_dropout=p_dropout, + ) + ) + self.norm_layers_2.append(LayerNorm(hidden_channels)) + + def forward(self, x, x_mask): + attn_mask = x_mask.unsqueeze(2) * x_mask.unsqueeze(-1) + x = x * x_mask + for attn_layer, norm_layer_1, ffn_layer, norm_layer_2 in zip( + self.attn_layers, self.norm_layers_1, self.ffn_layers, self.norm_layers_2 + ): + y = attn_layer(x, x, attn_mask) + y = self.drop(y) + x = norm_layer_1(x + y) + + y = ffn_layer(x, x_mask) + y = self.drop(y) + x = norm_layer_2(x + y) + x = x * x_mask + return x + + +class Decoder(nn.Module): + def __init__( + self, + hidden_channels: int, + filter_channels: int, + n_heads: int, + n_layers: int, + kernel_size: int = 1, + p_dropout: float = 0.0, + proximal_bias: bool = False, + proximal_init: bool = True, + **kwargs + ): + super().__init__() + self.hidden_channels = hidden_channels + self.filter_channels = filter_channels + self.n_heads = n_heads + self.n_layers = n_layers + self.kernel_size = kernel_size + self.p_dropout = p_dropout + self.proximal_bias = proximal_bias + self.proximal_init = proximal_init + + self.drop = nn.Dropout(p_dropout) + self.self_attn_layers = nn.ModuleList() + self.norm_layers_0 = nn.ModuleList() + self.encdec_attn_layers = nn.ModuleList() + self.norm_layers_1 = nn.ModuleList() + self.ffn_layers = nn.ModuleList() + self.norm_layers_2 = nn.ModuleList() + for i in range(self.n_layers): + self.self_attn_layers.append( + MultiHeadAttention( + hidden_channels, + hidden_channels, + n_heads, + p_dropout=p_dropout, + proximal_bias=proximal_bias, + proximal_init=proximal_init, + ) + ) + self.norm_layers_0.append(LayerNorm(hidden_channels)) + self.encdec_attn_layers.append( + MultiHeadAttention( + hidden_channels, hidden_channels, n_heads, p_dropout=p_dropout + ) + ) + self.norm_layers_1.append(LayerNorm(hidden_channels)) + self.ffn_layers.append( + FFN( + hidden_channels, + hidden_channels, + filter_channels, + kernel_size, + p_dropout=p_dropout, + causal=True, + ) + ) + self.norm_layers_2.append(LayerNorm(hidden_channels)) + + def forward(self, x, x_mask, h, h_mask): + """ + x: decoder input + h: encoder output + """ + self_attn_mask = subsequent_mask(x_mask.size(2)).type_as(x) + encdec_attn_mask = h_mask.unsqueeze(2) * x_mask.unsqueeze(-1) + x = x * x_mask + for i in range(self.n_layers): + y = self.self_attn_layers[i](x, x, self_attn_mask) + y = self.drop(y) + x = self.norm_layers_0[i](x + y) + + y = self.encdec_attn_layers[i](x, h, encdec_attn_mask) + y = self.drop(y) + x = self.norm_layers_1[i](x + y) + + y = self.ffn_layers[i](x, x_mask) + y = self.drop(y) + x = self.norm_layers_2[i](x + y) + x = x * x_mask + return x + + +class MultiHeadAttention(nn.Module): + def __init__( + self, + channels: int, + out_channels: int, + n_heads: int, + p_dropout: float = 0.0, + window_size: typing.Optional[int] = None, + heads_share: bool = True, + block_length: typing.Optional[int] = None, + proximal_bias: bool = False, + proximal_init: bool = False, + ): + super().__init__() + assert channels % n_heads == 0 + + self.channels = channels + self.out_channels = out_channels + self.n_heads = n_heads + self.p_dropout = p_dropout + self.window_size = window_size + self.heads_share = heads_share + self.block_length = block_length + self.proximal_bias = proximal_bias + self.proximal_init = proximal_init + self.attn = torch.zeros(1) + + self.k_channels = channels // n_heads + self.conv_q = nn.Conv1d(channels, channels, 1) + self.conv_k = nn.Conv1d(channels, channels, 1) + self.conv_v = nn.Conv1d(channels, channels, 1) + self.conv_o = nn.Conv1d(channels, out_channels, 1) + self.drop = nn.Dropout(p_dropout) + + if window_size is not None: + n_heads_rel = 1 if heads_share else n_heads + rel_stddev = self.k_channels**-0.5 + self.emb_rel_k = nn.Parameter( + torch.randn(n_heads_rel, window_size * 2 + 1, self.k_channels) + * rel_stddev + ) + self.emb_rel_v = nn.Parameter( + torch.randn(n_heads_rel, window_size * 2 + 1, self.k_channels) + * rel_stddev + ) + + nn.init.xavier_uniform_(self.conv_q.weight) + nn.init.xavier_uniform_(self.conv_k.weight) + nn.init.xavier_uniform_(self.conv_v.weight) + if proximal_init: + with torch.no_grad(): + self.conv_k.weight.copy_(self.conv_q.weight) + self.conv_k.bias.copy_(self.conv_q.bias) + + def forward(self, x, c, attn_mask=None): + q = self.conv_q(x) + k = self.conv_k(c) + v = self.conv_v(c) + + x, self.attn = self.attention(q, k, v, mask=attn_mask) + + x = self.conv_o(x) + return x + + def attention(self, query, key, value, mask=None): + # reshape [b, d, t] -> [b, n_h, t, d_k] + b, d, t_s, t_t = (key.size(0), key.size(1), key.size(2), query.size(2)) + query = query.view(b, self.n_heads, self.k_channels, t_t).transpose(2, 3) + key = key.view(b, self.n_heads, self.k_channels, t_s).transpose(2, 3) + value = value.view(b, self.n_heads, self.k_channels, t_s).transpose(2, 3) + + scores = torch.matmul(query / math.sqrt(self.k_channels), key.transpose(-2, -1)) + if self.window_size is not None: + assert ( + t_s == t_t + ), "Relative attention is only available for self-attention." + key_relative_embeddings = self._get_relative_embeddings(self.emb_rel_k, t_s) + rel_logits = self._matmul_with_relative_keys( + query / math.sqrt(self.k_channels), key_relative_embeddings + ) + scores_local = self._relative_position_to_absolute_position(rel_logits) + scores = scores + scores_local + if self.proximal_bias: + assert t_s == t_t, "Proximal bias is only available for self-attention." + scores = scores + self._attention_bias_proximal(t_s).type_as(scores) + if mask is not None: + scores = scores.masked_fill(mask == 0, -1e4) + if self.block_length is not None: + assert ( + t_s == t_t + ), "Local attention is only available for self-attention." + block_mask = ( + torch.ones_like(scores) + .triu(-self.block_length) + .tril(self.block_length) + ) + scores = scores.masked_fill(block_mask == 0, -1e4) + p_attn = F.softmax(scores, dim=-1) # [b, n_h, t_t, t_s] + p_attn = self.drop(p_attn) + output = torch.matmul(p_attn, value) + if self.window_size is not None: + relative_weights = self._absolute_position_to_relative_position(p_attn) + value_relative_embeddings = self._get_relative_embeddings( + self.emb_rel_v, t_s + ) + output = output + self._matmul_with_relative_values( + relative_weights, value_relative_embeddings + ) + output = ( + output.transpose(2, 3).contiguous().view(b, d, t_t) + ) # [b, n_h, t_t, d_k] -> [b, d, t_t] + return output, p_attn + + def _matmul_with_relative_values(self, x, y): + """ + x: [b, h, l, m] + y: [h or 1, m, d] + ret: [b, h, l, d] + """ + ret = torch.matmul(x, y.unsqueeze(0)) + return ret + + def _matmul_with_relative_keys(self, x, y): + """ + x: [b, h, l, d] + y: [h or 1, m, d] + ret: [b, h, l, m] + """ + ret = torch.matmul(x, y.unsqueeze(0).transpose(-2, -1)) + return ret + + def _get_relative_embeddings(self, relative_embeddings, length: int): + # max_relative_position = 2 * self.window_size + 1 + # Pad first before slice to avoid using cond ops. + pad_length = max(length - (self.window_size + 1), 0) + slice_start_position = max((self.window_size + 1) - length, 0) + slice_end_position = slice_start_position + 2 * length - 1 + if pad_length > 0: + padded_relative_embeddings = F.pad( + relative_embeddings, + # convert_pad_shape([[0, 0], [pad_length, pad_length], [0, 0]]), + (0, 0, pad_length, pad_length, 0, 0), + ) + else: + padded_relative_embeddings = relative_embeddings + used_relative_embeddings = padded_relative_embeddings[ + :, slice_start_position:slice_end_position + ] + return used_relative_embeddings + + def _relative_position_to_absolute_position(self, x): + """ + x: [b, h, l, 2*l-1] + ret: [b, h, l, l] + """ + batch, heads, length, _ = x.size() + + # Concat columns of pad to shift from relative to absolute indexing. + # x = F.pad(x, convert_pad_shape([[0, 0], [0, 0], [0, 0], [0, 1]])) + x = F.pad(x, (0, 1, 0, 0, 0, 0, 0, 0)) + + # Concat extra elements so to add up to shape (len+1, 2*len-1). + x_flat = x.view([batch, heads, length * 2 * length]) + # x_flat = F.pad(x_flat, convert_pad_shape([[0, 0], [0, 0], [0, length - 1]])) + x_flat = F.pad(x_flat, (0, length - 1, 0, 0, 0, 0)) + + # Reshape and slice out the padded elements. + x_final = x_flat.view([batch, heads, length + 1, (2 * length) - 1])[ + :, :, :length, length - 1 : + ] + return x_final + + def _absolute_position_to_relative_position(self, x): + """ + x: [b, h, l, l] + ret: [b, h, l, 2*l-1] + """ + batch, heads, length, _ = x.size() + + # padd along column + # x = F.pad(x, convert_pad_shape([[0, 0], [0, 0], [0, 0], [0, length - 1]])) + x = F.pad(x, (0, length - 1, 0, 0, 0, 0, 0, 0)) + x_flat = x.view([batch, heads, (length * length) + (length * (length - 1))]) + # add 0's in the beginning that will skew the elements after reshape + # x_flat = F.pad(x_flat, convert_pad_shape([[0, 0], [0, 0], [length, 0]])) + x_flat = F.pad(x_flat, (length, 0, 0, 0, 0, 0)) + x_final = x_flat.view([batch, heads, length, 2 * length])[:, :, :, 1:] + return x_final + + def _attention_bias_proximal(self, length: int): + """Bias for self-attention to encourage attention to close positions. + Args: + length: an integer scalar. + Returns: + a Tensor with shape [1, 1, length, length] + """ + r = torch.arange(length, dtype=torch.float32) + diff = torch.unsqueeze(r, 0) - torch.unsqueeze(r, 1) + return torch.unsqueeze(torch.unsqueeze(-torch.log1p(torch.abs(diff)), 0), 0) + + +class FFN(nn.Module): + def __init__( + self, + in_channels: int, + out_channels: int, + filter_channels: int, + kernel_size: int, + p_dropout: float = 0.0, + activation: str = "", + causal: bool = False, + ): + super().__init__() + self.in_channels = in_channels + self.out_channels = out_channels + self.filter_channels = filter_channels + self.kernel_size = kernel_size + self.p_dropout = p_dropout + self.activation = activation + self.causal = causal + + self.conv_1 = nn.Conv1d(in_channels, filter_channels, kernel_size) + self.conv_2 = nn.Conv1d(filter_channels, out_channels, kernel_size) + self.drop = nn.Dropout(p_dropout) + + def forward(self, x, x_mask): + if self.causal: + padding1 = self._causal_padding(x * x_mask) + else: + padding1 = self._same_padding(x * x_mask) + + x = self.conv_1(padding1) + + if self.activation == "gelu": + x = x * torch.sigmoid(1.702 * x) + else: + x = torch.relu(x) + x = self.drop(x) + + if self.causal: + padding2 = self._causal_padding(x * x_mask) + else: + padding2 = self._same_padding(x * x_mask) + + x = self.conv_2(padding2) + + return x * x_mask + + def _causal_padding(self, x): + if self.kernel_size == 1: + return x + pad_l = self.kernel_size - 1 + pad_r = 0 + # padding = [[0, 0], [0, 0], [pad_l, pad_r]] + # x = F.pad(x, convert_pad_shape(padding)) + x = F.pad(x, (pad_l, pad_r, 0, 0, 0, 0)) + return x + + def _same_padding(self, x): + if self.kernel_size == 1: + return x + pad_l = (self.kernel_size - 1) // 2 + pad_r = self.kernel_size // 2 + # padding = [[0, 0], [0, 0], [pad_l, pad_r]] + # x = F.pad(x, convert_pad_shape(padding)) + x = F.pad(x, (pad_l, pad_r, 0, 0, 0, 0)) + return x diff --git a/piper_train/vits/commons.py b/piper_train/vits/commons.py new file mode 100644 index 0000000..3f334c3 --- /dev/null +++ b/piper_train/vits/commons.py @@ -0,0 +1,147 @@ +import logging +import math +from typing import Optional + +import torch +from torch.nn import functional as F + +_LOGGER = logging.getLogger("vits.commons") + + +def init_weights(m, mean=0.0, std=0.01): + classname = m.__class__.__name__ + if classname.find("Conv") != -1: + m.weight.data.normal_(mean, std) + + +def get_padding(kernel_size, dilation=1): + return int((kernel_size * dilation - dilation) / 2) + + +def intersperse(lst, item): + result = [item] * (len(lst) * 2 + 1) + result[1::2] = lst + return result + + +def kl_divergence(m_p, logs_p, m_q, logs_q): + """KL(P||Q)""" + kl = (logs_q - logs_p) - 0.5 + kl += ( + 0.5 * (torch.exp(2.0 * logs_p) + ((m_p - m_q) ** 2)) * torch.exp(-2.0 * logs_q) + ) + return kl + + +def rand_gumbel(shape): + """Sample from the Gumbel distribution, protect from overflows.""" + uniform_samples = torch.rand(shape) * 0.99998 + 0.00001 + return -torch.log(-torch.log(uniform_samples)) + + +def rand_gumbel_like(x): + g = rand_gumbel(x.size()).to(dtype=x.dtype, device=x.device) + return g + + +def slice_segments(x, ids_str, segment_size=4): + ret = torch.zeros_like(x[:, :, :segment_size]) + for i in range(x.size(0)): + idx_str = max(0, ids_str[i]) + idx_end = idx_str + segment_size + ret[i] = x[i, :, idx_str:idx_end] + return ret + + +def rand_slice_segments(x, x_lengths=None, segment_size=4): + b, d, t = x.size() + if x_lengths is None: + x_lengths = t + ids_str_max = x_lengths - segment_size + 1 + ids_str = (torch.rand([b]).to(device=x.device) * ids_str_max).to(dtype=torch.long) + ret = slice_segments(x, ids_str, segment_size) + return ret, ids_str + + +def get_timing_signal_1d(length, channels, min_timescale=1.0, max_timescale=1.0e4): + position = torch.arange(length, dtype=torch.float) + num_timescales = channels // 2 + log_timescale_increment = math.log(float(max_timescale) / float(min_timescale)) / ( + num_timescales - 1 + ) + inv_timescales = min_timescale * torch.exp( + torch.arange(num_timescales, dtype=torch.float) * -log_timescale_increment + ) + scaled_time = position.unsqueeze(0) * inv_timescales.unsqueeze(1) + signal = torch.cat([torch.sin(scaled_time), torch.cos(scaled_time)], 0) + signal = F.pad(signal, [0, 0, 0, channels % 2]) + signal = signal.view(1, channels, length) + return signal + + +def add_timing_signal_1d(x, min_timescale=1.0, max_timescale=1.0e4): + b, channels, length = x.size() + signal = get_timing_signal_1d(length, channels, min_timescale, max_timescale) + return x + signal.to(dtype=x.dtype, device=x.device) + + +def cat_timing_signal_1d(x, min_timescale=1.0, max_timescale=1.0e4, axis=1): + b, channels, length = x.size() + signal = get_timing_signal_1d(length, channels, min_timescale, max_timescale) + return torch.cat([x, signal.to(dtype=x.dtype, device=x.device)], axis) + + +def subsequent_mask(length: int): + mask = torch.tril(torch.ones(length, length)).unsqueeze(0).unsqueeze(0) + return mask + + +@torch.jit.script +def fused_add_tanh_sigmoid_multiply(input_a, input_b, n_channels): + n_channels_int = n_channels[0] + in_act = input_a + input_b + t_act = torch.tanh(in_act[:, :n_channels_int, :]) + s_act = torch.sigmoid(in_act[:, n_channels_int:, :]) + acts = t_act * s_act + return acts + + +def sequence_mask(length, max_length: Optional[int] = None): + if max_length is None: + max_length = length.max() + x = torch.arange(max_length, dtype=length.dtype, device=length.device) + return x.unsqueeze(0) < length.unsqueeze(1) + + +def generate_path(duration, mask): + """ + duration: [b, 1, t_x] + mask: [b, 1, t_y, t_x] + """ + b, _, t_y, t_x = mask.shape + cum_duration = torch.cumsum(duration, -1) + + cum_duration_flat = cum_duration.view(b * t_x) + path = sequence_mask(cum_duration_flat, t_y).type_as(mask) + path = path.view(b, t_x, t_y) + path = path - F.pad(path, (0, 0, 1, 0, 0, 0))[:, :-1] + path = path.unsqueeze(1).transpose(2, 3) * mask + return path + + +def clip_grad_value_(parameters, clip_value, norm_type=2): + if isinstance(parameters, torch.Tensor): + parameters = [parameters] + parameters = list(filter(lambda p: p.grad is not None, parameters)) + norm_type = float(norm_type) + if clip_value is not None: + clip_value = float(clip_value) + + total_norm = 0 + for p in parameters: + param_norm = p.grad.data.norm(norm_type) + total_norm += param_norm.item() ** norm_type + if clip_value is not None: + p.grad.data.clamp_(min=-clip_value, max=clip_value) + total_norm = total_norm ** (1.0 / norm_type) + return total_norm diff --git a/piper_train/vits/config.py b/piper_train/vits/config.py new file mode 100644 index 0000000..e878f02 --- /dev/null +++ b/piper_train/vits/config.py @@ -0,0 +1,330 @@ +"""Configuration classes""" +from dataclasses import dataclass, field +from typing import Optional, Tuple + + +@dataclass +class MelAudioConfig: + filter_length: int = 1024 + hop_length: int = 256 + win_length: int = 1024 + mel_channels: int = 80 + sample_rate: int = 22050 + sample_bytes: int = 2 + channels: int = 1 + mel_fmin: float = 0.0 + mel_fmax: Optional[float] = None + + +@dataclass +class ModelAudioConfig: + resblock: str + resblock_kernel_sizes: Tuple[int, ...] + resblock_dilation_sizes: Tuple[Tuple[int, ...], ...] + upsample_rates: Tuple[int, ...] + upsample_initial_channel: int + upsample_kernel_sizes: Tuple[int, ...] + + @staticmethod + def low_quality() -> "ModelAudioConfig": + return ModelAudioConfig( + resblock="2", + resblock_kernel_sizes=(3, 5, 7), + resblock_dilation_sizes=( + (1, 2), + (2, 6), + (3, 12), + ), + upsample_rates=(8, 8, 4), + upsample_initial_channel=256, + upsample_kernel_sizes=(16, 16, 8), + ) + + @staticmethod + def high_quality() -> "ModelAudioConfig": + return ModelAudioConfig( + resblock="1", + resblock_kernel_sizes=(3, 7, 11), + resblock_dilation_sizes=( + (1, 3, 5), + (1, 3, 5), + (1, 3, 5), + ), + upsample_rates=(8, 8, 2, 2), + upsample_initial_channel=512, + upsample_kernel_sizes=(16, 16, 4, 4), + ) + + +@dataclass +class ModelConfig: + num_symbols: int + n_speakers: int + audio: ModelAudioConfig + mel: MelAudioConfig = field(default_factory=MelAudioConfig) + + inter_channels: int = 192 + hidden_channels: int = 192 + filter_channels: int = 768 + n_heads: int = 2 + n_layers: int = 6 + kernel_size: int = 3 + p_dropout: float = 0.1 + n_layers_q: int = 3 + use_spectral_norm: bool = False + gin_channels: int = 0 # single speaker + use_sdp: bool = True # StochasticDurationPredictor + segment_size: int = 8192 + + @property + def is_multispeaker(self) -> bool: + return self.n_speakers > 1 + + @property + def resblock(self) -> str: + return self.audio.resblock + + @property + def resblock_kernel_sizes(self) -> Tuple[int, ...]: + return self.audio.resblock_kernel_sizes + + @property + def resblock_dilation_sizes(self) -> Tuple[Tuple[int, ...], ...]: + return self.audio.resblock_dilation_sizes + + @property + def upsample_rates(self) -> Tuple[int, ...]: + return self.audio.upsample_rates + + @property + def upsample_initial_channel(self) -> int: + return self.audio.upsample_initial_channel + + @property + def upsample_kernel_sizes(self) -> Tuple[int, ...]: + return self.audio.upsample_kernel_sizes + + def __post_init__(self): + if self.is_multispeaker and (self.gin_channels == 0): + self.gin_channels = 512 + + +@dataclass +class TrainingConfig: + learning_rate: float = 2e-4 + betas: Tuple[float, float] = field(default=(0.8, 0.99)) + eps: float = 1e-9 + # batch_size: int = 32 + fp16_run: bool = False + lr_decay: float = 0.999875 + init_lr_ratio: float = 1.0 + warmup_epochs: int = 0 + c_mel: int = 45 + c_kl: float = 1.0 + grad_clip: Optional[float] = None + + +# @dataclass +# class PhonemesConfig(DataClassJsonMixin): +# phoneme_separator: str = " " +# """Separator between individual phonemes in CSV input""" + +# word_separator: str = "#" +# """Separator between word phonemes in CSV input (must not match phoneme_separator)""" + +# phoneme_to_id: typing.Optional[typing.Dict[str, int]] = None +# pad: typing.Optional[str] = "_" +# bos: typing.Optional[str] = None +# eos: typing.Optional[str] = None +# blank: typing.Optional[str] = "#" +# blank_word: typing.Optional[str] = None +# blank_between: typing.Union[str, BlankBetween] = BlankBetween.WORDS +# blank_at_start: bool = True +# blank_at_end: bool = True +# simple_punctuation: bool = True +# punctuation_map: typing.Optional[typing.Dict[str, str]] = None +# separate: typing.Optional[typing.List[str]] = None +# separate_graphemes: bool = False +# separate_tones: bool = False +# tone_before: bool = False +# phoneme_map: typing.Optional[typing.Dict[str, typing.List[str]]] = None +# auto_bos_eos: bool = False +# minor_break: typing.Optional[str] = IPA.BREAK_MINOR.value +# major_break: typing.Optional[str] = IPA.BREAK_MAJOR.value +# break_phonemes_into_graphemes: bool = False +# break_phonemes_into_codepoints: bool = False +# drop_stress: bool = False +# symbols: typing.Optional[typing.List[str]] = None + +# def split_word_phonemes(self, phonemes_str: str) -> typing.List[typing.List[str]]: +# """Split phonemes string into a list of lists (outer is words, inner is individual phonemes in each word)""" +# return [ +# word_phonemes_str.split(self.phoneme_separator) +# if self.phoneme_separator +# else list(word_phonemes_str) +# for word_phonemes_str in phonemes_str.split(self.word_separator) +# ] + +# def join_word_phonemes(self, word_phonemes: typing.List[typing.List[str]]) -> str: +# """Split phonemes string into a list of lists (outer is words, inner is individual phonemes in each word)""" +# return self.word_separator.join( +# self.phoneme_separator.join(wp) for wp in word_phonemes +# ) + + +# class Phonemizer(str, Enum): +# SYMBOLS = "symbols" +# GRUUT = "gruut" +# ESPEAK = "espeak" +# EPITRAN = "epitran" + + +# class Aligner(str, Enum): +# KALDI_ALIGN = "kaldi_align" + + +# class TextCasing(str, Enum): +# LOWER = "lower" +# UPPER = "upper" + + +# class MetadataFormat(str, Enum): +# TEXT = "text" +# PHONEMES = "phonemes" +# PHONEME_IDS = "ids" + + +# @dataclass +# class DatasetConfig: +# name: str +# metadata_format: MetadataFormat = MetadataFormat.TEXT +# multispeaker: bool = False +# text_language: typing.Optional[str] = None +# audio_dir: typing.Optional[typing.Union[str, Path]] = None +# cache_dir: typing.Optional[typing.Union[str, Path]] = None + +# def get_cache_dir(self, output_dir: typing.Union[str, Path]) -> Path: +# if self.cache_dir is not None: +# cache_dir = Path(self.cache_dir) +# else: +# cache_dir = Path("cache") / self.name + +# if not cache_dir.is_absolute(): +# cache_dir = Path(output_dir) / str(cache_dir) + +# return cache_dir + + +# @dataclass +# class AlignerConfig: +# aligner: typing.Optional[Aligner] = None +# casing: typing.Optional[TextCasing] = None + + +# @dataclass +# class InferenceConfig: +# length_scale: float = 1.0 +# noise_scale: float = 0.667 +# noise_w: float = 0.8 + + +# @dataclass +# class TrainingConfig(DataClassJsonMixin): +# seed: int = 1234 +# epochs: int = 10000 +# learning_rate: float = 2e-4 +# betas: typing.Tuple[float, float] = field(default=(0.8, 0.99)) +# eps: float = 1e-9 +# batch_size: int = 32 +# fp16_run: bool = False +# lr_decay: float = 0.999875 +# segment_size: int = 8192 +# init_lr_ratio: float = 1.0 +# warmup_epochs: int = 0 +# c_mel: int = 45 +# c_kl: float = 1.0 +# grad_clip: typing.Optional[float] = None + +# min_seq_length: typing.Optional[int] = None +# max_seq_length: typing.Optional[int] = None + +# min_spec_length: typing.Optional[int] = None +# max_spec_length: typing.Optional[int] = None + +# min_speaker_utterances: typing.Optional[int] = None + +# last_epoch: int = 1 +# global_step: int = 1 +# best_loss: typing.Optional[float] = None +# audio: AudioConfig = field(default_factory=AudioConfig) +# model: ModelConfig = field(default_factory=ModelConfig) +# phonemes: PhonemesConfig = field(default_factory=PhonemesConfig) +# text_aligner: AlignerConfig = field(default_factory=AlignerConfig) +# text_language: typing.Optional[str] = None +# phonemizer: typing.Optional[Phonemizer] = None +# datasets: typing.List[DatasetConfig] = field(default_factory=list) +# inference: InferenceConfig = field(default_factory=InferenceConfig) + +# version: int = 1 +# git_commit: str = "" + +# @property +# def is_multispeaker(self): +# return self.model.is_multispeaker or any(d.multispeaker for d in self.datasets) + +# def save(self, config_file: typing.TextIO): +# """Save config as JSON to a file""" +# json.dump(self.to_dict(), config_file, indent=4) + +# def get_speaker_id(self, dataset_name: str, speaker_name: str) -> int: +# if self.speaker_id_map is None: +# self.speaker_id_map = {} + +# full_speaker_name = f"{dataset_name}_{speaker_name}" +# speaker_id = self.speaker_id_map.get(full_speaker_name) +# if speaker_id is None: +# speaker_id = len(self.speaker_id_map) +# self.speaker_id_map[full_speaker_name] = speaker_id + +# return speaker_id + +# @staticmethod +# def load(config_file: typing.TextIO) -> "TrainingConfig": +# """Load config from a JSON file""" +# return TrainingConfig.from_json(config_file.read()) + +# @staticmethod +# def load_and_merge( +# config: "TrainingConfig", +# config_files: typing.Iterable[typing.Union[str, Path, typing.TextIO]], +# ) -> "TrainingConfig": +# """Loads one or more JSON configuration files and overlays them on top of an existing config""" +# base_dict = config.to_dict() +# for maybe_config_file in config_files: +# if isinstance(maybe_config_file, (str, Path)): +# # File path +# config_file = open(maybe_config_file, "r", encoding="utf-8") +# else: +# # File object +# config_file = maybe_config_file + +# with config_file: +# # Load new config and overlay on existing config +# new_dict = json.load(config_file) +# TrainingConfig.recursive_update(base_dict, new_dict) + +# return TrainingConfig.from_dict(base_dict) + +# @staticmethod +# def recursive_update( +# base_dict: typing.Dict[typing.Any, typing.Any], +# new_dict: typing.Mapping[typing.Any, typing.Any], +# ) -> None: +# """Recursively overwrites values in base dictionary with values from new dictionary""" +# for key, value in new_dict.items(): +# if isinstance(value, collections.Mapping) and ( +# base_dict.get(key) is not None +# ): +# TrainingConfig.recursive_update(base_dict[key], value) +# else: +# base_dict[key] = value diff --git a/piper_train/vits/dataset.py b/piper_train/vits/dataset.py new file mode 100644 index 0000000..258425f --- /dev/null +++ b/piper_train/vits/dataset.py @@ -0,0 +1,214 @@ +import json +import logging +from dataclasses import dataclass +from pathlib import Path +from typing import Iterable, List, Optional, Sequence, Union + +import torch +from torch import FloatTensor, LongTensor +from torch.utils.data import Dataset + +_LOGGER = logging.getLogger("vits.dataset") + + +@dataclass +class Utterance: + phoneme_ids: List[int] + audio_norm_path: Path + audio_spec_path: Path + speaker_id: Optional[int] = None + text: Optional[str] = None + + +@dataclass +class UtteranceTensors: + phoneme_ids: LongTensor + spectrogram: FloatTensor + audio_norm: FloatTensor + speaker_id: Optional[LongTensor] = None + text: Optional[str] = None + + @property + def spec_length(self) -> int: + return self.spectrogram.size(1) + + +@dataclass +class Batch: + phoneme_ids: LongTensor + phoneme_lengths: LongTensor + spectrograms: FloatTensor + spectrogram_lengths: LongTensor + audios: FloatTensor + audio_lengths: LongTensor + speaker_ids: Optional[LongTensor] = None + + +class PiperDataset(Dataset): + """ + Dataset format: + + * phoneme_ids (required) + * audio_norm_path (required) + * audio_spec_path (required) + * text (optional) + * phonemes (optional) + * audio_path (optional) + """ + + def __init__( + self, + dataset_paths: List[Union[str, Path]], + max_phoneme_ids: Optional[int] = None, + ): + self.utterances: List[Utterance] = [] + + for dataset_path in dataset_paths: + dataset_path = Path(dataset_path) + _LOGGER.debug("Loading dataset: %s", dataset_path) + self.utterances.extend( + PiperDataset.load_dataset(dataset_path, max_phoneme_ids=max_phoneme_ids) + ) + + def __len__(self): + return len(self.utterances) + + def __getitem__(self, idx) -> UtteranceTensors: + utt = self.utterances[idx] + return UtteranceTensors( + phoneme_ids=LongTensor(utt.phoneme_ids), + audio_norm=torch.load(utt.audio_norm_path), + spectrogram=torch.load(utt.audio_spec_path), + speaker_id=LongTensor([utt.speaker_id]) + if utt.speaker_id is not None + else None, + text=utt.text, + ) + + @staticmethod + def load_dataset( + dataset_path: Path, + max_phoneme_ids: Optional[int] = None, + ) -> Iterable[Utterance]: + num_skipped = 0 + + with open(dataset_path, "r", encoding="utf-8") as dataset_file: + for line_idx, line in enumerate(dataset_file): + line = line.strip() + if not line: + continue + + try: + utt = PiperDataset.load_utterance(line) + if (max_phoneme_ids is None) or ( + len(utt.phoneme_ids) <= max_phoneme_ids + ): + yield utt + else: + num_skipped += 1 + except Exception: + _LOGGER.exception( + "Error on line %s of %s: %s", + line_idx + 1, + dataset_path, + line, + ) + + if num_skipped > 0: + _LOGGER.warning("Skipped %s utterance(s)", num_skipped) + + @staticmethod + def load_utterance(line: str) -> Utterance: + utt_dict = json.loads(line) + return Utterance( + phoneme_ids=utt_dict["phoneme_ids"], + audio_norm_path=Path(utt_dict["audio_norm_path"]), + audio_spec_path=Path(utt_dict["audio_spec_path"]), + speaker_id=utt_dict.get("speaker_id"), + text=utt_dict.get("text"), + ) + + +class UtteranceCollate: + def __init__(self, is_multispeaker: bool, segment_size: int): + self.is_multispeaker = is_multispeaker + self.segment_size = segment_size + + def __call__(self, utterances: Sequence[UtteranceTensors]) -> Batch: + num_utterances = len(utterances) + assert num_utterances > 0, "No utterances" + + max_phonemes_length = 0 + max_spec_length = 0 + max_audio_length = 0 + + num_mels = 0 + + # Determine lengths + for utt_idx, utt in enumerate(utterances): + assert utt.spectrogram is not None + assert utt.audio_norm is not None + + phoneme_length = utt.phoneme_ids.size(0) + spec_length = utt.spectrogram.size(1) + audio_length = utt.audio_norm.size(1) + + max_phonemes_length = max(max_phonemes_length, phoneme_length) + max_spec_length = max(max_spec_length, spec_length) + max_audio_length = max(max_audio_length, audio_length) + + num_mels = utt.spectrogram.size(0) + if self.is_multispeaker: + assert utt.speaker_id is not None, "Missing speaker id" + + # Audio cannot be smaller than segment size (8192) + max_audio_length = max(max_audio_length, self.segment_size) + + # Create padded tensors + phonemes_padded = LongTensor(num_utterances, max_phonemes_length) + spec_padded = FloatTensor(num_utterances, num_mels, max_spec_length) + audio_padded = FloatTensor(num_utterances, 1, max_audio_length) + + phonemes_padded.zero_() + spec_padded.zero_() + audio_padded.zero_() + + phoneme_lengths = LongTensor(num_utterances) + spec_lengths = LongTensor(num_utterances) + audio_lengths = LongTensor(num_utterances) + + speaker_ids: Optional[LongTensor] = None + if self.is_multispeaker: + speaker_ids = LongTensor(num_utterances) + + # Sort by decreasing spectrogram length + sorted_utterances = sorted( + utterances, key=lambda u: u.spectrogram.size(1), reverse=True + ) + for utt_idx, utt in enumerate(sorted_utterances): + phoneme_length = utt.phoneme_ids.size(0) + spec_length = utt.spectrogram.size(1) + audio_length = utt.audio_norm.size(1) + + phonemes_padded[utt_idx, :phoneme_length] = utt.phoneme_ids + phoneme_lengths[utt_idx] = phoneme_length + + spec_padded[utt_idx, :, :spec_length] = utt.spectrogram + spec_lengths[utt_idx] = spec_length + + audio_padded[utt_idx, :, :audio_length] = utt.audio_norm + audio_lengths[utt_idx] = audio_length + + if utt.speaker_id is not None: + assert speaker_ids is not None + speaker_ids[utt_idx] = utt.speaker_id + + return Batch( + phoneme_ids=phonemes_padded, + phoneme_lengths=phoneme_lengths, + spectrograms=spec_padded, + spectrogram_lengths=spec_lengths, + audios=audio_padded, + audio_lengths=audio_lengths, + speaker_ids=speaker_ids, + ) diff --git a/piper_train/vits/lightning.py b/piper_train/vits/lightning.py new file mode 100644 index 0000000..c6b7250 --- /dev/null +++ b/piper_train/vits/lightning.py @@ -0,0 +1,352 @@ +import logging +from pathlib import Path +from typing import List, Optional, Tuple, Union + +import pytorch_lightning as pl +import torch +from torch import autocast +from torch.nn import functional as F +from torch.utils.data import DataLoader, Dataset, random_split + +from .commons import slice_segments +from .dataset import Batch, PiperDataset, UtteranceCollate +from .losses import discriminator_loss, feature_loss, generator_loss, kl_loss +from .mel_processing import mel_spectrogram_torch, spec_to_mel_torch +from .models import MultiPeriodDiscriminator, SynthesizerTrn + +_LOGGER = logging.getLogger("vits.lightning") + + +class VitsModel(pl.LightningModule): + def __init__( + self, + num_symbols: int, + num_speakers: int, + # audio + resblock="2", + resblock_kernel_sizes=(3, 5, 7), + resblock_dilation_sizes=( + (1, 2), + (2, 6), + (3, 12), + ), + upsample_rates=(8, 8, 4), + upsample_initial_channel=256, + upsample_kernel_sizes=(16, 16, 8), + # mel + filter_length: int = 1024, + hop_length: int = 256, + win_length: int = 1024, + mel_channels: int = 80, + sample_rate: int = 22050, + sample_bytes: int = 2, + channels: int = 1, + mel_fmin: float = 0.0, + mel_fmax: Optional[float] = None, + # model + inter_channels: int = 192, + hidden_channels: int = 192, + filter_channels: int = 768, + n_heads: int = 2, + n_layers: int = 6, + kernel_size: int = 3, + p_dropout: float = 0.1, + n_layers_q: int = 3, + use_spectral_norm: bool = False, + gin_channels: int = 0, + use_sdp: bool = True, + segment_size: int = 8192, + # training + dataset: Optional[List[Union[str, Path]]] = None, + learning_rate: float = 2e-4, + betas: Tuple[float, float] = (0.8, 0.99), + eps: float = 1e-9, + batch_size: int = 1, + lr_decay: float = 0.999875, + init_lr_ratio: float = 1.0, + warmup_epochs: int = 0, + c_mel: int = 45, + c_kl: float = 1.0, + grad_clip: Optional[float] = None, + num_workers: int = 1, + seed: int = 1234, + num_test_examples: int = 5, + validation_split: float = 0.1, + max_phoneme_ids: Optional[int] = None, + **kwargs, + ): + super().__init__() + self.save_hyperparameters() + + if (self.hparams.num_speakers > 1) and (self.hparams.gin_channels <= 0): + # Default gin_channels for multi-speaker model + self.hparams.gin_channels = 512 + + # Set up models + self.model_g = SynthesizerTrn( + n_vocab=self.hparams.num_symbols, + spec_channels=self.hparams.filter_length // 2 + 1, + segment_size=self.hparams.segment_size // self.hparams.hop_length, + inter_channels=self.hparams.inter_channels, + hidden_channels=self.hparams.hidden_channels, + filter_channels=self.hparams.filter_channels, + n_heads=self.hparams.n_heads, + n_layers=self.hparams.n_layers, + kernel_size=self.hparams.kernel_size, + p_dropout=self.hparams.p_dropout, + resblock=self.hparams.resblock, + resblock_kernel_sizes=self.hparams.resblock_kernel_sizes, + resblock_dilation_sizes=self.hparams.resblock_dilation_sizes, + upsample_rates=self.hparams.upsample_rates, + upsample_initial_channel=self.hparams.upsample_initial_channel, + upsample_kernel_sizes=self.hparams.upsample_kernel_sizes, + n_speakers=self.hparams.num_speakers, + gin_channels=self.hparams.gin_channels, + use_sdp=self.hparams.use_sdp, + ) + self.model_d = MultiPeriodDiscriminator( + use_spectral_norm=self.hparams.use_spectral_norm + ) + + # Dataset splits + self._train_dataset: Optional[Dataset] = None + self._val_dataset: Optional[Dataset] = None + self._test_dataset: Optional[Dataset] = None + self._load_datasets(validation_split, num_test_examples, max_phoneme_ids) + + # State kept between training optimizers + self._y = None + self._y_hat = None + + def _load_datasets( + self, + validation_split: float, + num_test_examples: int, + max_phoneme_ids: Optional[int] = None, + ): + if self.hparams.dataset is None: + _LOGGER.debug("No dataset to load") + return + + full_dataset = PiperDataset( + self.hparams.dataset, max_phoneme_ids=max_phoneme_ids + ) + valid_set_size = int(len(full_dataset) * validation_split) + train_set_size = len(full_dataset) - valid_set_size - num_test_examples + + self._train_dataset, self._test_dataset, self._val_dataset = random_split( + full_dataset, [train_set_size, num_test_examples, valid_set_size] + ) + + def forward(self, text, text_lengths, scales, sid=None): + noise_scale = scales[0] + length_scale = scales[1] + noise_scale_w = scales[2] + audio, *_ = self.model_g.infer( + text, + text_lengths, + noise_scale=noise_scale, + length_scale=length_scale, + noise_scale_w=noise_scale_w, + sid=sid, + ) + + return audio + + def train_dataloader(self): + return DataLoader( + self._train_dataset, + collate_fn=UtteranceCollate( + is_multispeaker=self.hparams.num_speakers > 1, + segment_size=self.hparams.segment_size, + ), + num_workers=self.hparams.num_workers, + batch_size=self.hparams.batch_size, + ) + + def val_dataloader(self): + return DataLoader( + self._val_dataset, + collate_fn=UtteranceCollate( + is_multispeaker=self.hparams.num_speakers > 1, + segment_size=self.hparams.segment_size, + ), + num_workers=self.hparams.num_workers, + batch_size=self.hparams.batch_size, + ) + + def test_dataloader(self): + return DataLoader( + self._test_dataset, + collate_fn=UtteranceCollate( + is_multispeaker=self.hparams.num_speakers > 1, + segment_size=self.hparams.segment_size, + ), + num_workers=self.hparams.num_workers, + batch_size=self.hparams.batch_size, + ) + + def training_step(self, batch: Batch, batch_idx: int, optimizer_idx: int): + if optimizer_idx == 0: + return self.training_step_g(batch) + + if optimizer_idx == 1: + return self.training_step_d(batch) + + def training_step_g(self, batch: Batch): + x, x_lengths, y, _, spec, spec_lengths, speaker_ids = ( + batch.phoneme_ids, + batch.phoneme_lengths, + batch.audios, + batch.audio_lengths, + batch.spectrograms, + batch.spectrogram_lengths, + batch.speaker_ids if batch.speaker_ids is not None else None, + ) + ( + y_hat, + l_length, + _attn, + ids_slice, + _x_mask, + z_mask, + (_z, z_p, m_p, logs_p, _m_q, logs_q), + ) = self.model_g(x, x_lengths, spec, spec_lengths, speaker_ids) + self._y_hat = y_hat + + mel = spec_to_mel_torch( + spec, + self.hparams.filter_length, + self.hparams.mel_channels, + self.hparams.sample_rate, + self.hparams.mel_fmin, + self.hparams.mel_fmax, + ) + y_mel = slice_segments( + mel, + ids_slice, + self.hparams.segment_size // self.hparams.hop_length, + ) + y_hat_mel = mel_spectrogram_torch( + y_hat.squeeze(1), + self.hparams.filter_length, + self.hparams.mel_channels, + self.hparams.sample_rate, + self.hparams.hop_length, + self.hparams.win_length, + self.hparams.mel_fmin, + self.hparams.mel_fmax, + ) + y = slice_segments( + y, + ids_slice * self.hparams.hop_length, + self.hparams.segment_size, + ) # slice + + # Save for training_step_d + self._y = y + + _y_d_hat_r, y_d_hat_g, fmap_r, fmap_g = self.model_d(y, y_hat) + + with autocast(self.device.type, enabled=False): + # Generator loss + loss_dur = torch.sum(l_length.float()) + loss_mel = F.l1_loss(y_mel, y_hat_mel) * self.hparams.c_mel + loss_kl = kl_loss(z_p, logs_q, m_p, logs_p, z_mask) * self.hparams.c_kl + + loss_fm = feature_loss(fmap_r, fmap_g) + loss_gen, _losses_gen = generator_loss(y_d_hat_g) + loss_gen_all = loss_gen + loss_fm + loss_mel + loss_dur + loss_kl + + self.log("loss_gen_all", loss_gen_all) + + return loss_gen_all + + def training_step_d(self, batch: Batch): + # From training_step_g + y = self._y + y_hat = self._y_hat + y_d_hat_r, y_d_hat_g, _, _ = self.model_d(y, y_hat.detach()) + + with autocast(self.device.type, enabled=False): + # Discriminator + loss_disc, _losses_disc_r, _losses_disc_g = discriminator_loss( + y_d_hat_r, y_d_hat_g + ) + loss_disc_all = loss_disc + + self.log("loss_disc_all", loss_disc_all) + + return loss_disc_all + + def validation_step(self, batch: Batch, batch_idx: int): + val_loss = self.training_step_g(batch) + self.training_step_d(batch) + self.log("val_loss", val_loss) + + # Generate audio examples + for utt_idx, test_utt in enumerate(self._test_dataset): + text = test_utt.phoneme_ids.unsqueeze(0).to(self.device) + text_lengths = torch.LongTensor([len(test_utt.phoneme_ids)]).to(self.device) + scales = [0.667, 1.0, 0.8] + sid = ( + test_utt.speaker_id.to(self.device) + if test_utt.speaker_id is not None + else None + ) + test_audio = self(text, text_lengths, scales, sid=sid).detach() + + # Scale to make louder in [-1, 1] + test_audio = test_audio * (1.0 / max(0.01, abs(test_audio.max()))) + + tag = test_utt.text or str(utt_idx) + self.logger.experiment.add_audio( + tag, test_audio, sample_rate=self.hparams.sample_rate + ) + + return val_loss + + def configure_optimizers(self): + optimizers = [ + torch.optim.AdamW( + self.model_g.parameters(), + lr=self.hparams.learning_rate, + betas=self.hparams.betas, + eps=self.hparams.eps, + ), + torch.optim.AdamW( + self.model_d.parameters(), + lr=self.hparams.learning_rate, + betas=self.hparams.betas, + eps=self.hparams.eps, + ), + ] + schedulers = [ + torch.optim.lr_scheduler.ExponentialLR( + optimizers[0], gamma=self.hparams.lr_decay + ), + torch.optim.lr_scheduler.ExponentialLR( + optimizers[1], gamma=self.hparams.lr_decay + ), + ] + + return optimizers, schedulers + + @staticmethod + def add_model_specific_args(parent_parser): + parser = parent_parser.add_argument_group("VitsModel") + parser.add_argument("--batch-size", type=int, required=True) + parser.add_argument("--validation-split", type=float, default=0.1) + parser.add_argument("--num-test-examples", type=int, default=5) + parser.add_argument( + "--max-phoneme-ids", + type=int, + help="Exclude utterances with phoneme id lists longer than this", + ) + # + parser.add_argument("--hidden-channels", type=int, default=192) + parser.add_argument("--inter-channels", type=int, default=192) + parser.add_argument("--filter-channels", type=int, default=768) + parser.add_argument("--n-layers", type=int, default=6) + parser.add_argument("--n-heads", type=int, default=2) + # + return parent_parser diff --git a/piper_train/vits/losses.py b/piper_train/vits/losses.py new file mode 100644 index 0000000..ec327e5 --- /dev/null +++ b/piper_train/vits/losses.py @@ -0,0 +1,58 @@ +import torch + + +def feature_loss(fmap_r, fmap_g): + loss = 0 + for dr, dg in zip(fmap_r, fmap_g): + for rl, gl in zip(dr, dg): + rl = rl.float().detach() + gl = gl.float() + loss += torch.mean(torch.abs(rl - gl)) + + return loss * 2 + + +def discriminator_loss(disc_real_outputs, disc_generated_outputs): + loss = 0 + r_losses = [] + g_losses = [] + for dr, dg in zip(disc_real_outputs, disc_generated_outputs): + dr = dr.float() + dg = dg.float() + r_loss = torch.mean((1 - dr) ** 2) + g_loss = torch.mean(dg**2) + loss += r_loss + g_loss + r_losses.append(r_loss.item()) + g_losses.append(g_loss.item()) + + return loss, r_losses, g_losses + + +def generator_loss(disc_outputs): + loss = 0 + gen_losses = [] + for dg in disc_outputs: + dg = dg.float() + l_dg = torch.mean((1 - dg) ** 2) + gen_losses.append(l_dg) + loss += l_dg + + return loss, gen_losses + + +def kl_loss(z_p, logs_q, m_p, logs_p, z_mask): + """ + z_p, logs_q: [b, h, t_t] + m_p, logs_p: [b, h, t_t] + """ + z_p = z_p.float() + logs_q = logs_q.float() + m_p = m_p.float() + logs_p = logs_p.float() + z_mask = z_mask.float() + + kl = logs_p - logs_q - 0.5 + kl += 0.5 * ((z_p - m_p) ** 2) * torch.exp(-2.0 * logs_p) + kl = torch.sum(kl * z_mask) + l_kl = kl / torch.sum(z_mask) + return l_kl diff --git a/piper_train/vits/mel_processing.py b/piper_train/vits/mel_processing.py new file mode 100644 index 0000000..72e81e2 --- /dev/null +++ b/piper_train/vits/mel_processing.py @@ -0,0 +1,139 @@ +import torch +import torch.utils.data +from librosa.filters import mel as librosa_mel_fn + +MAX_WAV_VALUE = 32768.0 + + +def dynamic_range_compression_torch(x, C=1, clip_val=1e-5): + """ + PARAMS + ------ + C: compression factor + """ + return torch.log(torch.clamp(x, min=clip_val) * C) + + +def dynamic_range_decompression_torch(x, C=1): + """ + PARAMS + ------ + C: compression factor used to compress + """ + return torch.exp(x) / C + + +def spectral_normalize_torch(magnitudes): + output = dynamic_range_compression_torch(magnitudes) + return output + + +def spectral_de_normalize_torch(magnitudes): + output = dynamic_range_decompression_torch(magnitudes) + return output + + +mel_basis = {} +hann_window = {} + + +def spectrogram_torch(y, n_fft, sampling_rate, hop_size, win_size, center=False): + if torch.min(y) < -1.0: + print("min value is ", torch.min(y)) + if torch.max(y) > 1.0: + print("max value is ", torch.max(y)) + + global hann_window + dtype_device = str(y.dtype) + "_" + str(y.device) + wnsize_dtype_device = str(win_size) + "_" + dtype_device + if wnsize_dtype_device not in hann_window: + hann_window[wnsize_dtype_device] = torch.hann_window(win_size).type_as(y) + + y = torch.nn.functional.pad( + y.unsqueeze(1), + (int((n_fft - hop_size) / 2), int((n_fft - hop_size) / 2)), + mode="reflect", + ) + y = y.squeeze(1) + + spec = torch.view_as_real( + torch.stft( + y, + n_fft, + hop_length=hop_size, + win_length=win_size, + window=hann_window[wnsize_dtype_device], + center=center, + pad_mode="reflect", + normalized=False, + onesided=True, + return_complex=True, + ) + ) + + spec = torch.sqrt(spec.pow(2).sum(-1) + 1e-6) + + return spec + + +def spec_to_mel_torch(spec, n_fft, num_mels, sampling_rate, fmin, fmax): + global mel_basis + dtype_device = str(spec.dtype) + "_" + str(spec.device) + fmax_dtype_device = str(fmax) + "_" + dtype_device + if fmax_dtype_device not in mel_basis: + mel = librosa_mel_fn( + sr=sampling_rate, n_fft=n_fft, n_mels=num_mels, fmin=fmin, fmax=fmax + ) + mel_basis[fmax_dtype_device] = torch.from_numpy(mel).type_as(spec) + spec = torch.matmul(mel_basis[fmax_dtype_device], spec) + spec = spectral_normalize_torch(spec) + return spec + + +def mel_spectrogram_torch( + y, n_fft, num_mels, sampling_rate, hop_size, win_size, fmin, fmax, center=False +): + if torch.min(y) < -1.0: + print("min value is ", torch.min(y)) + if torch.max(y) > 1.0: + print("max value is ", torch.max(y)) + + global mel_basis, hann_window + dtype_device = str(y.dtype) + "_" + str(y.device) + fmax_dtype_device = str(fmax) + "_" + dtype_device + wnsize_dtype_device = str(win_size) + "_" + dtype_device + if fmax_dtype_device not in mel_basis: + mel = librosa_mel_fn( + sr=sampling_rate, n_fft=n_fft, n_mels=num_mels, fmin=fmin, fmax=fmax + ) + mel_basis[fmax_dtype_device] = torch.from_numpy(mel).type_as(y) + if wnsize_dtype_device not in hann_window: + hann_window[wnsize_dtype_device] = torch.hann_window(win_size).type_as(y) + + y = torch.nn.functional.pad( + y.unsqueeze(1), + (int((n_fft - hop_size) / 2), int((n_fft - hop_size) / 2)), + mode="reflect", + ) + y = y.squeeze(1) + spec = torch.view_as_real( + torch.stft( + y, + n_fft, + hop_length=hop_size, + win_length=win_size, + window=hann_window[wnsize_dtype_device], + center=center, + pad_mode="reflect", + normalized=False, + onesided=True, + return_complex=True, + ) + ) + + spec = torch.sqrt(spec.pow(2).sum(-1) + 1e-6) + + spec = torch.matmul(mel_basis[fmax_dtype_device], spec) + spec = spectral_normalize_torch(spec) + + return spec diff --git a/piper_train/vits/models.py b/piper_train/vits/models.py new file mode 100644 index 0000000..3845a99 --- /dev/null +++ b/piper_train/vits/models.py @@ -0,0 +1,732 @@ +import math +import typing + +import torch +from torch import nn +from torch.nn import Conv1d, Conv2d, ConvTranspose1d +from torch.nn import functional as F +from torch.nn.utils import remove_weight_norm, spectral_norm, weight_norm + +from . import attentions, commons, modules +from .commons import get_padding, init_weights + + +class StochasticDurationPredictor(nn.Module): + def __init__( + self, + in_channels: int, + filter_channels: int, + kernel_size: int, + p_dropout: float, + n_flows: int = 4, + gin_channels: int = 0, + ): + super().__init__() + filter_channels = in_channels # it needs to be removed from future version. + self.in_channels = in_channels + self.filter_channels = filter_channels + self.kernel_size = kernel_size + self.p_dropout = p_dropout + self.n_flows = n_flows + self.gin_channels = gin_channels + + self.log_flow = modules.Log() + self.flows = nn.ModuleList() + self.flows.append(modules.ElementwiseAffine(2)) + for i in range(n_flows): + self.flows.append( + modules.ConvFlow(2, filter_channels, kernel_size, n_layers=3) + ) + self.flows.append(modules.Flip()) + + self.post_pre = nn.Conv1d(1, filter_channels, 1) + self.post_proj = nn.Conv1d(filter_channels, filter_channels, 1) + self.post_convs = modules.DDSConv( + filter_channels, kernel_size, n_layers=3, p_dropout=p_dropout + ) + self.post_flows = nn.ModuleList() + self.post_flows.append(modules.ElementwiseAffine(2)) + for i in range(4): + self.post_flows.append( + modules.ConvFlow(2, filter_channels, kernel_size, n_layers=3) + ) + self.post_flows.append(modules.Flip()) + + self.pre = nn.Conv1d(in_channels, filter_channels, 1) + self.proj = nn.Conv1d(filter_channels, filter_channels, 1) + self.convs = modules.DDSConv( + filter_channels, kernel_size, n_layers=3, p_dropout=p_dropout + ) + if gin_channels != 0: + self.cond = nn.Conv1d(gin_channels, filter_channels, 1) + + def forward(self, x, x_mask, w=None, g=None, reverse=False, noise_scale=1.0): + x = torch.detach(x) + x = self.pre(x) + if g is not None: + g = torch.detach(g) + x = x + self.cond(g) + x = self.convs(x, x_mask) + x = self.proj(x) * x_mask + + if not reverse: + flows = self.flows + assert w is not None + + logdet_tot_q = 0 + h_w = self.post_pre(w) + h_w = self.post_convs(h_w, x_mask) + h_w = self.post_proj(h_w) * x_mask + e_q = torch.randn(w.size(0), 2, w.size(2)).type_as(x) * x_mask + z_q = e_q + for flow in self.post_flows: + z_q, logdet_q = flow(z_q, x_mask, g=(x + h_w)) + logdet_tot_q += logdet_q + z_u, z1 = torch.split(z_q, [1, 1], 1) + u = torch.sigmoid(z_u) * x_mask + z0 = (w - u) * x_mask + logdet_tot_q += torch.sum( + (F.logsigmoid(z_u) + F.logsigmoid(-z_u)) * x_mask, [1, 2] + ) + logq = ( + torch.sum(-0.5 * (math.log(2 * math.pi) + (e_q**2)) * x_mask, [1, 2]) + - logdet_tot_q + ) + + logdet_tot = 0 + z0, logdet = self.log_flow(z0, x_mask) + logdet_tot += logdet + z = torch.cat([z0, z1], 1) + for flow in flows: + z, logdet = flow(z, x_mask, g=x, reverse=reverse) + logdet_tot = logdet_tot + logdet + nll = ( + torch.sum(0.5 * (math.log(2 * math.pi) + (z**2)) * x_mask, [1, 2]) + - logdet_tot + ) + return nll + logq # [b] + else: + flows = list(reversed(self.flows)) + flows = flows[:-2] + [flows[-1]] # remove a useless vflow + z = torch.randn(x.size(0), 2, x.size(2)).type_as(x) * noise_scale + + for flow in flows: + z = flow(z, x_mask, g=x, reverse=reverse) + z0, z1 = torch.split(z, [1, 1], 1) + logw = z0 + return logw + + +class DurationPredictor(nn.Module): + def __init__( + self, + in_channels: int, + filter_channels: int, + kernel_size: int, + p_dropout: float, + gin_channels: int = 0, + ): + super().__init__() + + self.in_channels = in_channels + self.filter_channels = filter_channels + self.kernel_size = kernel_size + self.p_dropout = p_dropout + self.gin_channels = gin_channels + + self.drop = nn.Dropout(p_dropout) + self.conv_1 = nn.Conv1d( + in_channels, filter_channels, kernel_size, padding=kernel_size // 2 + ) + self.norm_1 = modules.LayerNorm(filter_channels) + self.conv_2 = nn.Conv1d( + filter_channels, filter_channels, kernel_size, padding=kernel_size // 2 + ) + self.norm_2 = modules.LayerNorm(filter_channels) + self.proj = nn.Conv1d(filter_channels, 1, 1) + + if gin_channels != 0: + self.cond = nn.Conv1d(gin_channels, in_channels, 1) + + def forward(self, x, x_mask, g=None): + x = torch.detach(x) + if g is not None: + g = torch.detach(g) + x = x + self.cond(g) + x = self.conv_1(x * x_mask) + x = torch.relu(x) + x = self.norm_1(x) + x = self.drop(x) + x = self.conv_2(x * x_mask) + x = torch.relu(x) + x = self.norm_2(x) + x = self.drop(x) + x = self.proj(x * x_mask) + return x * x_mask + + +class TextEncoder(nn.Module): + def __init__( + self, + n_vocab: int, + out_channels: int, + hidden_channels: int, + filter_channels: int, + n_heads: int, + n_layers: int, + kernel_size: int, + p_dropout: float, + ): + super().__init__() + self.n_vocab = n_vocab + self.out_channels = out_channels + self.hidden_channels = hidden_channels + self.filter_channels = filter_channels + self.n_heads = n_heads + self.n_layers = n_layers + self.kernel_size = kernel_size + self.p_dropout = p_dropout + + self.emb = nn.Embedding(n_vocab, hidden_channels) + nn.init.normal_(self.emb.weight, 0.0, hidden_channels**-0.5) + + self.encoder = attentions.Encoder( + hidden_channels, filter_channels, n_heads, n_layers, kernel_size, p_dropout + ) + self.proj = nn.Conv1d(hidden_channels, out_channels * 2, 1) + + def forward(self, x, x_lengths): + x = self.emb(x) * math.sqrt(self.hidden_channels) # [b, t, h] + x = torch.transpose(x, 1, -1) # [b, h, t] + x_mask = torch.unsqueeze( + commons.sequence_mask(x_lengths, x.size(2)), 1 + ).type_as(x) + + x = self.encoder(x * x_mask, x_mask) + stats = self.proj(x) * x_mask + + m, logs = torch.split(stats, self.out_channels, dim=1) + return x, m, logs, x_mask + + +class ResidualCouplingBlock(nn.Module): + def __init__( + self, + channels: int, + hidden_channels: int, + kernel_size: int, + dilation_rate: int, + n_layers: int, + n_flows: int = 4, + gin_channels: int = 0, + ): + super().__init__() + self.channels = channels + self.hidden_channels = hidden_channels + self.kernel_size = kernel_size + self.dilation_rate = dilation_rate + self.n_layers = n_layers + self.n_flows = n_flows + self.gin_channels = gin_channels + + self.flows = nn.ModuleList() + for i in range(n_flows): + self.flows.append( + modules.ResidualCouplingLayer( + channels, + hidden_channels, + kernel_size, + dilation_rate, + n_layers, + gin_channels=gin_channels, + mean_only=True, + ) + ) + self.flows.append(modules.Flip()) + + def forward(self, x, x_mask, g=None, reverse=False): + if not reverse: + for flow in self.flows: + x, _ = flow(x, x_mask, g=g, reverse=reverse) + else: + for flow in reversed(self.flows): + x = flow(x, x_mask, g=g, reverse=reverse) + return x + + +class PosteriorEncoder(nn.Module): + def __init__( + self, + in_channels: int, + out_channels: int, + hidden_channels: int, + kernel_size: int, + dilation_rate: int, + n_layers: int, + gin_channels: int = 0, + ): + super().__init__() + self.in_channels = in_channels + self.out_channels = out_channels + self.hidden_channels = hidden_channels + self.kernel_size = kernel_size + self.dilation_rate = dilation_rate + self.n_layers = n_layers + self.gin_channels = gin_channels + + self.pre = nn.Conv1d(in_channels, hidden_channels, 1) + self.enc = modules.WN( + hidden_channels, + kernel_size, + dilation_rate, + n_layers, + gin_channels=gin_channels, + ) + self.proj = nn.Conv1d(hidden_channels, out_channels * 2, 1) + + def forward(self, x, x_lengths, g=None): + x_mask = torch.unsqueeze( + commons.sequence_mask(x_lengths, x.size(2)), 1 + ).type_as(x) + x = self.pre(x) * x_mask + x = self.enc(x, x_mask, g=g) + stats = self.proj(x) * x_mask + m, logs = torch.split(stats, self.out_channels, dim=1) + z = (m + torch.randn_like(m) * torch.exp(logs)) * x_mask + return z, m, logs, x_mask + + +class Generator(torch.nn.Module): + def __init__( + self, + initial_channel: int, + resblock: typing.Optional[str], + resblock_kernel_sizes: typing.Tuple[int, ...], + resblock_dilation_sizes: typing.Tuple[typing.Tuple[int, ...], ...], + upsample_rates: typing.Tuple[int, ...], + upsample_initial_channel: int, + upsample_kernel_sizes: typing.Tuple[int, ...], + gin_channels: int = 0, + ): + super(Generator, self).__init__() + self.LRELU_SLOPE = 0.1 + self.num_kernels = len(resblock_kernel_sizes) + self.num_upsamples = len(upsample_rates) + self.conv_pre = Conv1d( + initial_channel, upsample_initial_channel, 7, 1, padding=3 + ) + resblock_module = modules.ResBlock1 if resblock == "1" else modules.ResBlock2 + + self.ups = nn.ModuleList() + for i, (u, k) in enumerate(zip(upsample_rates, upsample_kernel_sizes)): + self.ups.append( + weight_norm( + ConvTranspose1d( + upsample_initial_channel // (2**i), + upsample_initial_channel // (2 ** (i + 1)), + k, + u, + padding=(k - u) // 2, + ) + ) + ) + + self.resblocks = nn.ModuleList() + for i in range(len(self.ups)): + ch = upsample_initial_channel // (2 ** (i + 1)) + for j, (k, d) in enumerate( + zip(resblock_kernel_sizes, resblock_dilation_sizes) + ): + self.resblocks.append(resblock_module(ch, k, d)) + + self.conv_post = Conv1d(ch, 1, 7, 1, padding=3, bias=False) + self.ups.apply(init_weights) + + if gin_channels != 0: + self.cond = nn.Conv1d(gin_channels, upsample_initial_channel, 1) + + def forward(self, x, g=None): + x = self.conv_pre(x) + if g is not None: + x = x + self.cond(g) + + for i, up in enumerate(self.ups): + x = F.leaky_relu(x, self.LRELU_SLOPE) + x = up(x) + xs = torch.zeros(1) + for j, resblock in enumerate(self.resblocks): + index = j - (i * self.num_kernels) + if index == 0: + xs = resblock(x) + elif (index > 0) and (index < self.num_kernels): + xs += resblock(x) + x = xs / self.num_kernels + x = F.leaky_relu(x) + x = self.conv_post(x) + x = torch.tanh(x) + + return x + + def remove_weight_norm(self): + print("Removing weight norm...") + for l in self.ups: + remove_weight_norm(l) + for l in self.resblocks: + l.remove_weight_norm() + + +class DiscriminatorP(torch.nn.Module): + def __init__( + self, + period: int, + kernel_size: int = 5, + stride: int = 3, + use_spectral_norm: bool = False, + ): + super(DiscriminatorP, self).__init__() + self.LRELU_SLOPE = 0.1 + self.period = period + self.use_spectral_norm = use_spectral_norm + norm_f = weight_norm if not use_spectral_norm else spectral_norm + self.convs = nn.ModuleList( + [ + norm_f( + Conv2d( + 1, + 32, + (kernel_size, 1), + (stride, 1), + padding=(get_padding(kernel_size, 1), 0), + ) + ), + norm_f( + Conv2d( + 32, + 128, + (kernel_size, 1), + (stride, 1), + padding=(get_padding(kernel_size, 1), 0), + ) + ), + norm_f( + Conv2d( + 128, + 512, + (kernel_size, 1), + (stride, 1), + padding=(get_padding(kernel_size, 1), 0), + ) + ), + norm_f( + Conv2d( + 512, + 1024, + (kernel_size, 1), + (stride, 1), + padding=(get_padding(kernel_size, 1), 0), + ) + ), + norm_f( + Conv2d( + 1024, + 1024, + (kernel_size, 1), + 1, + padding=(get_padding(kernel_size, 1), 0), + ) + ), + ] + ) + self.conv_post = norm_f(Conv2d(1024, 1, (3, 1), 1, padding=(1, 0))) + + def forward(self, x): + fmap = [] + + # 1d to 2d + b, c, t = x.shape + if t % self.period != 0: # pad first + n_pad = self.period - (t % self.period) + x = F.pad(x, (0, n_pad), "reflect") + t = t + n_pad + x = x.view(b, c, t // self.period, self.period) + + for l in self.convs: + x = l(x) + x = F.leaky_relu(x, self.LRELU_SLOPE) + fmap.append(x) + x = self.conv_post(x) + fmap.append(x) + x = torch.flatten(x, 1, -1) + + return x, fmap + + +class DiscriminatorS(torch.nn.Module): + def __init__(self, use_spectral_norm=False): + super(DiscriminatorS, self).__init__() + self.LRELU_SLOPE = 0.1 + norm_f = spectral_norm if use_spectral_norm else weight_norm + self.convs = nn.ModuleList( + [ + norm_f(Conv1d(1, 16, 15, 1, padding=7)), + norm_f(Conv1d(16, 64, 41, 4, groups=4, padding=20)), + norm_f(Conv1d(64, 256, 41, 4, groups=16, padding=20)), + norm_f(Conv1d(256, 1024, 41, 4, groups=64, padding=20)), + norm_f(Conv1d(1024, 1024, 41, 4, groups=256, padding=20)), + norm_f(Conv1d(1024, 1024, 5, 1, padding=2)), + ] + ) + self.conv_post = norm_f(Conv1d(1024, 1, 3, 1, padding=1)) + + def forward(self, x): + fmap = [] + + for l in self.convs: + x = l(x) + x = F.leaky_relu(x, self.LRELU_SLOPE) + fmap.append(x) + x = self.conv_post(x) + fmap.append(x) + x = torch.flatten(x, 1, -1) + + return x, fmap + + +class MultiPeriodDiscriminator(torch.nn.Module): + def __init__(self, use_spectral_norm=False): + super(MultiPeriodDiscriminator, self).__init__() + periods = [2, 3, 5, 7, 11] + + discs = [DiscriminatorS(use_spectral_norm=use_spectral_norm)] + discs = discs + [ + DiscriminatorP(i, use_spectral_norm=use_spectral_norm) for i in periods + ] + self.discriminators = nn.ModuleList(discs) + + def forward(self, y, y_hat): + y_d_rs = [] + y_d_gs = [] + fmap_rs = [] + fmap_gs = [] + for i, d in enumerate(self.discriminators): + y_d_r, fmap_r = d(y) + y_d_g, fmap_g = d(y_hat) + y_d_rs.append(y_d_r) + y_d_gs.append(y_d_g) + fmap_rs.append(fmap_r) + fmap_gs.append(fmap_g) + + return y_d_rs, y_d_gs, fmap_rs, fmap_gs + + +class SynthesizerTrn(nn.Module): + """ + Synthesizer for Training + """ + + def __init__( + self, + n_vocab: int, + spec_channels: int, + segment_size: int, + inter_channels: int, + hidden_channels: int, + filter_channels: int, + n_heads: int, + n_layers: int, + kernel_size: int, + p_dropout: float, + resblock: str, + resblock_kernel_sizes: typing.Tuple[int, ...], + resblock_dilation_sizes: typing.Tuple[typing.Tuple[int, ...], ...], + upsample_rates: typing.Tuple[int, ...], + upsample_initial_channel: int, + upsample_kernel_sizes: typing.Tuple[int, ...], + n_speakers: int = 1, + gin_channels: int = 0, + use_sdp: bool = True, + ): + + super().__init__() + self.n_vocab = n_vocab + self.spec_channels = spec_channels + self.inter_channels = inter_channels + self.hidden_channels = hidden_channels + self.filter_channels = filter_channels + self.n_heads = n_heads + self.n_layers = n_layers + self.kernel_size = kernel_size + self.p_dropout = p_dropout + self.resblock = resblock + self.resblock_kernel_sizes = resblock_kernel_sizes + self.resblock_dilation_sizes = resblock_dilation_sizes + self.upsample_rates = upsample_rates + self.upsample_initial_channel = upsample_initial_channel + self.upsample_kernel_sizes = upsample_kernel_sizes + self.segment_size = segment_size + self.n_speakers = n_speakers + self.gin_channels = gin_channels + + self.use_sdp = use_sdp + + self.enc_p = TextEncoder( + n_vocab, + inter_channels, + hidden_channels, + filter_channels, + n_heads, + n_layers, + kernel_size, + p_dropout, + ) + self.dec = Generator( + inter_channels, + resblock, + resblock_kernel_sizes, + resblock_dilation_sizes, + upsample_rates, + upsample_initial_channel, + upsample_kernel_sizes, + gin_channels=gin_channels, + ) + self.enc_q = PosteriorEncoder( + spec_channels, + inter_channels, + hidden_channels, + 5, + 1, + 16, + gin_channels=gin_channels, + ) + self.flow = ResidualCouplingBlock( + inter_channels, hidden_channels, 5, 1, 4, gin_channels=gin_channels + ) + + if use_sdp: + self.dp = StochasticDurationPredictor( + hidden_channels, 192, 3, 0.5, 4, gin_channels=gin_channels + ) + else: + self.dp = DurationPredictor( + hidden_channels, 256, 3, 0.5, gin_channels=gin_channels + ) + + if n_speakers > 1: + self.emb_g = nn.Embedding(n_speakers, gin_channels) + + def forward(self, x, x_lengths, y, y_lengths, sid=None): + + x, m_p, logs_p, x_mask = self.enc_p(x, x_lengths) + if self.n_speakers > 1: + g = self.emb_g(sid).unsqueeze(-1) # [b, h, 1] + else: + g = None + + z, m_q, logs_q, y_mask = self.enc_q(y, y_lengths, g=g) + z_p = self.flow(z, y_mask, g=g) + + with torch.no_grad(): + # negative cross-entropy + s_p_sq_r = torch.exp(-2 * logs_p) # [b, d, t] + neg_cent1 = torch.sum( + -0.5 * math.log(2 * math.pi) - logs_p, [1], keepdim=True + ) # [b, 1, t_s] + neg_cent2 = torch.matmul( + -0.5 * (z_p**2).transpose(1, 2), s_p_sq_r + ) # [b, t_t, d] x [b, d, t_s] = [b, t_t, t_s] + neg_cent3 = torch.matmul( + z_p.transpose(1, 2), (m_p * s_p_sq_r) + ) # [b, t_t, d] x [b, d, t_s] = [b, t_t, t_s] + neg_cent4 = torch.sum( + -0.5 * (m_p**2) * s_p_sq_r, [1], keepdim=True + ) # [b, 1, t_s] + neg_cent = neg_cent1 + neg_cent2 + neg_cent3 + neg_cent4 + + attn_mask = torch.unsqueeze(x_mask, 2) * torch.unsqueeze(y_mask, -1) + attn = ( + monotonic_align.maximum_path(neg_cent, attn_mask.squeeze(1)) + .unsqueeze(1) + .detach() + ) + + w = attn.sum(2) + if self.use_sdp: + l_length = self.dp(x, x_mask, w, g=g) + l_length = l_length / torch.sum(x_mask) + else: + logw_ = torch.log(w + 1e-6) * x_mask + logw = self.dp(x, x_mask, g=g) + l_length = torch.sum((logw - logw_) ** 2, [1, 2]) / torch.sum( + x_mask + ) # for averaging + + # expand prior + m_p = torch.matmul(attn.squeeze(1), m_p.transpose(1, 2)).transpose(1, 2) + logs_p = torch.matmul(attn.squeeze(1), logs_p.transpose(1, 2)).transpose(1, 2) + + z_slice, ids_slice = commons.rand_slice_segments( + z, y_lengths, self.segment_size + ) + o = self.dec(z_slice, g=g) + return ( + o, + l_length, + attn, + ids_slice, + x_mask, + y_mask, + (z, z_p, m_p, logs_p, m_q, logs_q), + ) + + def infer( + self, + x, + x_lengths, + sid=None, + noise_scale=0.667, + length_scale=1, + noise_scale_w=0.8, + max_len=None, + ): + x, m_p, logs_p, x_mask = self.enc_p(x, x_lengths) + if self.n_speakers > 1: + assert sid is not None, "Missing speaker id" + g = self.emb_g(sid).unsqueeze(-1) # [b, h, 1] + else: + g = None + + if self.use_sdp: + logw = self.dp(x, x_mask, g=g, reverse=True, noise_scale=noise_scale_w) + else: + logw = self.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.transpose(1, 2)).transpose( + 1, 2 + ) # [b, t', t], [b, t, d] -> [b, d, t'] + logs_p = torch.matmul(attn.squeeze(1), logs_p.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 = self.flow(z_p, y_mask, g=g, reverse=True) + o = self.dec((z * y_mask)[:, :, :max_len], g=g) + + return o, attn, y_mask, (z, z_p, m_p, logs_p) + + def voice_conversion(self, y, y_lengths, sid_src, sid_tgt): + assert self.n_speakers > 1, "n_speakers have to be larger than 1." + g_src = self.emb_g(sid_src).unsqueeze(-1) + g_tgt = self.emb_g(sid_tgt).unsqueeze(-1) + z, m_q, logs_q, y_mask = self.enc_q(y, y_lengths, g=g_src) + z_p = self.flow(z, y_mask, g=g_src) + z_hat = self.flow(z_p, y_mask, g=g_tgt, reverse=True) + o_hat = self.dec(z_hat * y_mask, g=g_tgt) + return o_hat, y_mask, (z, z_p, z_hat) diff --git a/piper_train/vits/modules.py b/piper_train/vits/modules.py new file mode 100644 index 0000000..656e497 --- /dev/null +++ b/piper_train/vits/modules.py @@ -0,0 +1,527 @@ +import math +import typing + +import torch +from torch import nn +from torch.nn import Conv1d +from torch.nn import functional as F +from torch.nn.utils import remove_weight_norm, weight_norm + +from .commons import fused_add_tanh_sigmoid_multiply, get_padding, init_weights +from .transforms import piecewise_rational_quadratic_transform + + +class LayerNorm(nn.Module): + def __init__(self, channels: int, eps: float = 1e-5): + super().__init__() + self.channels = channels + self.eps = eps + + self.gamma = nn.Parameter(torch.ones(channels)) + self.beta = nn.Parameter(torch.zeros(channels)) + + def forward(self, x): + x = x.transpose(1, -1) + x = F.layer_norm(x, (self.channels,), self.gamma, self.beta, self.eps) + return x.transpose(1, -1) + + +class ConvReluNorm(nn.Module): + def __init__( + self, + in_channels: int, + hidden_channels: int, + out_channels: int, + kernel_size: int, + n_layers: int, + p_dropout: float, + ): + super().__init__() + self.in_channels = in_channels + self.hidden_channels = hidden_channels + self.out_channels = out_channels + self.kernel_size = kernel_size + self.n_layers = n_layers + self.p_dropout = p_dropout + assert n_layers > 1, "Number of layers should be larger than 0." + + self.conv_layers = nn.ModuleList() + self.norm_layers = nn.ModuleList() + self.conv_layers.append( + nn.Conv1d( + in_channels, hidden_channels, kernel_size, padding=kernel_size // 2 + ) + ) + self.norm_layers.append(LayerNorm(hidden_channels)) + self.relu_drop = nn.Sequential(nn.ReLU(), nn.Dropout(p_dropout)) + for _ in range(n_layers - 1): + self.conv_layers.append( + nn.Conv1d( + hidden_channels, + hidden_channels, + kernel_size, + padding=kernel_size // 2, + ) + ) + self.norm_layers.append(LayerNorm(hidden_channels)) + self.proj = nn.Conv1d(hidden_channels, out_channels, 1) + self.proj.weight.data.zero_() + self.proj.bias.data.zero_() + + def forward(self, x, x_mask): + x_org = x + for i in range(self.n_layers): + x = self.conv_layers[i](x * x_mask) + x = self.norm_layers[i](x) + x = self.relu_drop(x) + x = x_org + self.proj(x) + return x * x_mask + + +class DDSConv(nn.Module): + """ + Dialted and Depth-Separable Convolution + """ + + def __init__( + self, channels: int, kernel_size: int, n_layers: int, p_dropout: float = 0.0 + ): + super().__init__() + self.channels = channels + self.kernel_size = kernel_size + self.n_layers = n_layers + self.p_dropout = p_dropout + + self.drop = nn.Dropout(p_dropout) + self.convs_sep = nn.ModuleList() + self.convs_1x1 = nn.ModuleList() + self.norms_1 = nn.ModuleList() + self.norms_2 = nn.ModuleList() + for i in range(n_layers): + dilation = kernel_size**i + padding = (kernel_size * dilation - dilation) // 2 + self.convs_sep.append( + nn.Conv1d( + channels, + channels, + kernel_size, + groups=channels, + dilation=dilation, + padding=padding, + ) + ) + self.convs_1x1.append(nn.Conv1d(channels, channels, 1)) + self.norms_1.append(LayerNorm(channels)) + self.norms_2.append(LayerNorm(channels)) + + def forward(self, x, x_mask, g=None): + if g is not None: + x = x + g + for i in range(self.n_layers): + y = self.convs_sep[i](x * x_mask) + y = self.norms_1[i](y) + y = F.gelu(y) + y = self.convs_1x1[i](y) + y = self.norms_2[i](y) + y = F.gelu(y) + y = self.drop(y) + x = x + y + return x * x_mask + + +class WN(torch.nn.Module): + def __init__( + self, + hidden_channels: int, + kernel_size: int, + dilation_rate: int, + n_layers: int, + gin_channels: int = 0, + p_dropout: float = 0, + ): + super().__init__() + assert kernel_size % 2 == 1 + self.hidden_channels = hidden_channels + self.kernel_size = (kernel_size,) + self.dilation_rate = dilation_rate + self.n_layers = n_layers + self.gin_channels = gin_channels + self.p_dropout = p_dropout + + self.in_layers = torch.nn.ModuleList() + self.res_skip_layers = torch.nn.ModuleList() + self.drop = nn.Dropout(p_dropout) + + if gin_channels != 0: + cond_layer = torch.nn.Conv1d( + gin_channels, 2 * hidden_channels * n_layers, 1 + ) + self.cond_layer = torch.nn.utils.weight_norm(cond_layer, name="weight") + + for i in range(n_layers): + dilation = dilation_rate**i + padding = int((kernel_size * dilation - dilation) / 2) + in_layer = torch.nn.Conv1d( + hidden_channels, + 2 * hidden_channels, + kernel_size, + dilation=dilation, + padding=padding, + ) + in_layer = torch.nn.utils.weight_norm(in_layer, name="weight") + self.in_layers.append(in_layer) + + # last one is not necessary + if i < n_layers - 1: + res_skip_channels = 2 * hidden_channels + else: + res_skip_channels = hidden_channels + + res_skip_layer = torch.nn.Conv1d(hidden_channels, res_skip_channels, 1) + res_skip_layer = torch.nn.utils.weight_norm(res_skip_layer, name="weight") + self.res_skip_layers.append(res_skip_layer) + + def forward(self, x, x_mask, g=None, **kwargs): + output = torch.zeros_like(x) + n_channels_tensor = torch.IntTensor([self.hidden_channels]) + + if g is not None: + g = self.cond_layer(g) + + for i in range(self.n_layers): + x_in = self.in_layers[i](x) + if g is not None: + cond_offset = i * 2 * self.hidden_channels + g_l = g[:, cond_offset : cond_offset + 2 * self.hidden_channels, :] + else: + g_l = torch.zeros_like(x_in) + + acts = fused_add_tanh_sigmoid_multiply(x_in, g_l, n_channels_tensor) + acts = self.drop(acts) + + res_skip_acts = self.res_skip_layers[i](acts) + if i < self.n_layers - 1: + res_acts = res_skip_acts[:, : self.hidden_channels, :] + x = (x + res_acts) * x_mask + output = output + res_skip_acts[:, self.hidden_channels :, :] + else: + output = output + res_skip_acts + return output * x_mask + + def remove_weight_norm(self): + if self.gin_channels != 0: + torch.nn.utils.remove_weight_norm(self.cond_layer) + for l in self.in_layers: + torch.nn.utils.remove_weight_norm(l) + for l in self.res_skip_layers: + torch.nn.utils.remove_weight_norm(l) + + +class ResBlock1(torch.nn.Module): + def __init__( + self, + channels: int, + kernel_size: int = 3, + dilation: typing.Tuple[int] = (1, 3, 5), + ): + super(ResBlock1, self).__init__() + self.LRELU_SLOPE = 0.1 + self.convs1 = nn.ModuleList( + [ + weight_norm( + Conv1d( + channels, + channels, + kernel_size, + 1, + dilation=dilation[0], + padding=get_padding(kernel_size, dilation[0]), + ) + ), + weight_norm( + Conv1d( + channels, + channels, + kernel_size, + 1, + dilation=dilation[1], + padding=get_padding(kernel_size, dilation[1]), + ) + ), + weight_norm( + Conv1d( + channels, + channels, + kernel_size, + 1, + dilation=dilation[2], + padding=get_padding(kernel_size, dilation[2]), + ) + ), + ] + ) + self.convs1.apply(init_weights) + + self.convs2 = nn.ModuleList( + [ + weight_norm( + Conv1d( + channels, + channels, + kernel_size, + 1, + dilation=1, + padding=get_padding(kernel_size, 1), + ) + ), + weight_norm( + Conv1d( + channels, + channels, + kernel_size, + 1, + dilation=1, + padding=get_padding(kernel_size, 1), + ) + ), + weight_norm( + Conv1d( + channels, + channels, + kernel_size, + 1, + dilation=1, + padding=get_padding(kernel_size, 1), + ) + ), + ] + ) + self.convs2.apply(init_weights) + + def forward(self, x, x_mask=None): + for c1, c2 in zip(self.convs1, self.convs2): + xt = F.leaky_relu(x, self.LRELU_SLOPE) + if x_mask is not None: + xt = xt * x_mask + xt = c1(xt) + xt = F.leaky_relu(xt, self.LRELU_SLOPE) + if x_mask is not None: + xt = xt * x_mask + xt = c2(xt) + x = xt + x + if x_mask is not None: + x = x * x_mask + return x + + def remove_weight_norm(self): + for l in self.convs1: + remove_weight_norm(l) + for l in self.convs2: + remove_weight_norm(l) + + +class ResBlock2(torch.nn.Module): + def __init__( + self, channels: int, kernel_size: int = 3, dilation: typing.Tuple[int] = (1, 3) + ): + super(ResBlock2, self).__init__() + self.LRELU_SLOPE = 0.1 + self.convs = nn.ModuleList( + [ + weight_norm( + Conv1d( + channels, + channels, + kernel_size, + 1, + dilation=dilation[0], + padding=get_padding(kernel_size, dilation[0]), + ) + ), + weight_norm( + Conv1d( + channels, + channels, + kernel_size, + 1, + dilation=dilation[1], + padding=get_padding(kernel_size, dilation[1]), + ) + ), + ] + ) + self.convs.apply(init_weights) + + def forward(self, x, x_mask=None): + for c in self.convs: + xt = F.leaky_relu(x, self.LRELU_SLOPE) + if x_mask is not None: + xt = xt * x_mask + xt = c(xt) + x = xt + x + if x_mask is not None: + x = x * x_mask + return x + + def remove_weight_norm(self): + for l in self.convs: + remove_weight_norm(l) + + +class Log(nn.Module): + def forward( + self, x: torch.Tensor, x_mask: torch.Tensor, reverse: bool = False, **kwargs + ): + if not reverse: + y = torch.log(torch.clamp_min(x, 1e-5)) * x_mask + logdet = torch.sum(-y, [1, 2]) + return y, logdet + else: + x = torch.exp(x) * x_mask + return x + + +class Flip(nn.Module): + def forward(self, x: torch.Tensor, *args, reverse: bool = False, **kwargs): + x = torch.flip(x, [1]) + if not reverse: + logdet = torch.zeros(x.size(0)).type_as(x) + return x, logdet + else: + return x + + +class ElementwiseAffine(nn.Module): + def __init__(self, channels: int): + super().__init__() + self.channels = channels + self.m = nn.Parameter(torch.zeros(channels, 1)) + self.logs = nn.Parameter(torch.zeros(channels, 1)) + + def forward(self, x, x_mask, reverse=False, **kwargs): + if not reverse: + y = self.m + torch.exp(self.logs) * x + y = y * x_mask + logdet = torch.sum(self.logs * x_mask, [1, 2]) + return y, logdet + else: + x = (x - self.m) * torch.exp(-self.logs) * x_mask + return x + + +class ResidualCouplingLayer(nn.Module): + def __init__( + self, + channels: int, + hidden_channels: int, + kernel_size: int, + dilation_rate: int, + n_layers: int, + p_dropout: float = 0, + gin_channels: int = 0, + mean_only: bool = False, + ): + assert channels % 2 == 0, "channels should be divisible by 2" + super().__init__() + self.channels = channels + self.hidden_channels = hidden_channels + self.kernel_size = kernel_size + self.dilation_rate = dilation_rate + self.n_layers = n_layers + self.half_channels = channels // 2 + self.mean_only = mean_only + + self.pre = nn.Conv1d(self.half_channels, hidden_channels, 1) + self.enc = WN( + hidden_channels, + kernel_size, + dilation_rate, + n_layers, + p_dropout=p_dropout, + gin_channels=gin_channels, + ) + self.post = nn.Conv1d(hidden_channels, self.half_channels * (2 - mean_only), 1) + self.post.weight.data.zero_() + self.post.bias.data.zero_() + + def forward(self, x, x_mask, g=None, reverse=False): + x0, x1 = torch.split(x, [self.half_channels] * 2, 1) + h = self.pre(x0) * x_mask + h = self.enc(h, x_mask, g=g) + stats = self.post(h) * x_mask + if not self.mean_only: + m, logs = torch.split(stats, [self.half_channels] * 2, 1) + else: + m = stats + logs = torch.zeros_like(m) + + if not reverse: + x1 = m + x1 * torch.exp(logs) * x_mask + x = torch.cat([x0, x1], 1) + logdet = torch.sum(logs, [1, 2]) + return x, logdet + else: + x1 = (x1 - m) * torch.exp(-logs) * x_mask + x = torch.cat([x0, x1], 1) + return x + + +class ConvFlow(nn.Module): + def __init__( + self, + in_channels: int, + filter_channels: int, + kernel_size: int, + n_layers: int, + num_bins: int = 10, + tail_bound: float = 5.0, + ): + super().__init__() + self.in_channels = in_channels + self.filter_channels = filter_channels + self.kernel_size = kernel_size + self.n_layers = n_layers + self.num_bins = num_bins + self.tail_bound = tail_bound + self.half_channels = in_channels // 2 + + self.pre = nn.Conv1d(self.half_channels, filter_channels, 1) + self.convs = DDSConv(filter_channels, kernel_size, n_layers, p_dropout=0.0) + self.proj = nn.Conv1d( + filter_channels, self.half_channels * (num_bins * 3 - 1), 1 + ) + self.proj.weight.data.zero_() + self.proj.bias.data.zero_() + + def forward(self, x, x_mask, g=None, reverse=False): + x0, x1 = torch.split(x, [self.half_channels] * 2, 1) + h = self.pre(x0) + h = self.convs(h, x_mask, g=g) + h = self.proj(h) * x_mask + + b, c, t = x0.shape + h = h.reshape(b, c, -1, t).permute(0, 1, 3, 2) # [b, cx?, t] -> [b, c, t, ?] + + unnormalized_widths = h[..., : self.num_bins] / math.sqrt(self.filter_channels) + unnormalized_heights = h[..., self.num_bins : 2 * self.num_bins] / math.sqrt( + self.filter_channels + ) + unnormalized_derivatives = h[..., 2 * self.num_bins :] + + x1, logabsdet = piecewise_rational_quadratic_transform( + x1, + unnormalized_widths, + unnormalized_heights, + unnormalized_derivatives, + inverse=reverse, + tails="linear", + tail_bound=self.tail_bound, + ) + + x = torch.cat([x0, x1], 1) * x_mask + + logdet = torch.sum(logabsdet * x_mask, [1, 2]) + if not reverse: + return x, logdet + else: + return x diff --git a/piper_train/vits/transforms.py b/piper_train/vits/transforms.py new file mode 100644 index 0000000..329862f --- /dev/null +++ b/piper_train/vits/transforms.py @@ -0,0 +1,212 @@ +import numpy as np +import torch +from torch.nn import functional as F + +DEFAULT_MIN_BIN_WIDTH = 1e-3 +DEFAULT_MIN_BIN_HEIGHT = 1e-3 +DEFAULT_MIN_DERIVATIVE = 1e-3 + + +def piecewise_rational_quadratic_transform( + inputs, + unnormalized_widths, + unnormalized_heights, + unnormalized_derivatives, + inverse=False, + tails=None, + tail_bound=1.0, + min_bin_width=DEFAULT_MIN_BIN_WIDTH, + min_bin_height=DEFAULT_MIN_BIN_HEIGHT, + min_derivative=DEFAULT_MIN_DERIVATIVE, +): + + if tails is None: + spline_fn = rational_quadratic_spline + spline_kwargs = {} + else: + spline_fn = unconstrained_rational_quadratic_spline + spline_kwargs = {"tails": tails, "tail_bound": tail_bound} + + outputs, logabsdet = spline_fn( + inputs=inputs, + unnormalized_widths=unnormalized_widths, + unnormalized_heights=unnormalized_heights, + unnormalized_derivatives=unnormalized_derivatives, + inverse=inverse, + min_bin_width=min_bin_width, + min_bin_height=min_bin_height, + min_derivative=min_derivative, + **spline_kwargs + ) + return outputs, logabsdet + + +def searchsorted(bin_locations, inputs, eps=1e-6): + # bin_locations[..., -1] += eps + bin_locations[..., bin_locations.size(-1) - 1] += eps + return torch.sum(inputs[..., None] >= bin_locations, dim=-1) - 1 + + +def unconstrained_rational_quadratic_spline( + inputs, + unnormalized_widths, + unnormalized_heights, + unnormalized_derivatives, + inverse=False, + tails="linear", + tail_bound=1.0, + min_bin_width=DEFAULT_MIN_BIN_WIDTH, + min_bin_height=DEFAULT_MIN_BIN_HEIGHT, + min_derivative=DEFAULT_MIN_DERIVATIVE, +): + inside_interval_mask = (inputs >= -tail_bound) & (inputs <= tail_bound) + outside_interval_mask = ~inside_interval_mask + + outputs = torch.zeros_like(inputs) + logabsdet = torch.zeros_like(inputs) + + if tails == "linear": + unnormalized_derivatives = F.pad(unnormalized_derivatives, pad=(1, 1)) + constant = np.log(np.exp(1 - min_derivative) - 1) + unnormalized_derivatives[..., 0] = constant + # unnormalized_derivatives[..., -1] = constant + unnormalized_derivatives[..., unnormalized_derivatives.size(-1) - 1] = constant + + outputs[outside_interval_mask] = inputs[outside_interval_mask] + logabsdet[outside_interval_mask] = 0 + else: + raise RuntimeError("{} tails are not implemented.".format(tails)) + + ( + outputs[inside_interval_mask], + logabsdet[inside_interval_mask], + ) = rational_quadratic_spline( + inputs=inputs[inside_interval_mask], + unnormalized_widths=unnormalized_widths[inside_interval_mask, :], + unnormalized_heights=unnormalized_heights[inside_interval_mask, :], + unnormalized_derivatives=unnormalized_derivatives[inside_interval_mask, :], + inverse=inverse, + left=-tail_bound, + right=tail_bound, + bottom=-tail_bound, + top=tail_bound, + min_bin_width=min_bin_width, + min_bin_height=min_bin_height, + min_derivative=min_derivative, + ) + + return outputs, logabsdet + + +def rational_quadratic_spline( + inputs, + unnormalized_widths, + unnormalized_heights, + unnormalized_derivatives, + inverse=False, + left=0.0, + right=1.0, + bottom=0.0, + top=1.0, + min_bin_width=DEFAULT_MIN_BIN_WIDTH, + min_bin_height=DEFAULT_MIN_BIN_HEIGHT, + min_derivative=DEFAULT_MIN_DERIVATIVE, +): + # if torch.min(inputs) < left or torch.max(inputs) > right: + # raise ValueError("Input to a transform is not within its domain") + + num_bins = unnormalized_widths.shape[-1] + + # if min_bin_width * num_bins > 1.0: + # raise ValueError("Minimal bin width too large for the number of bins") + # if min_bin_height * num_bins > 1.0: + # raise ValueError("Minimal bin height too large for the number of bins") + + widths = F.softmax(unnormalized_widths, dim=-1) + widths = min_bin_width + (1 - min_bin_width * num_bins) * widths + cumwidths = torch.cumsum(widths, dim=-1) + cumwidths = F.pad(cumwidths, pad=(1, 0), mode="constant", value=0.0) + cumwidths = (right - left) * cumwidths + left + cumwidths[..., 0] = left + # cumwidths[..., -1] = right + cumwidths[..., cumwidths.size(-1) - 1] = right + widths = cumwidths[..., 1:] - cumwidths[..., :-1] + + derivatives = min_derivative + F.softplus(unnormalized_derivatives) + + heights = F.softmax(unnormalized_heights, dim=-1) + heights = min_bin_height + (1 - min_bin_height * num_bins) * heights + cumheights = torch.cumsum(heights, dim=-1) + cumheights = F.pad(cumheights, pad=(1, 0), mode="constant", value=0.0) + cumheights = (top - bottom) * cumheights + bottom + cumheights[..., 0] = bottom + # cumheights[..., -1] = top + cumheights[..., cumheights.size(-1) - 1] = top + heights = cumheights[..., 1:] - cumheights[..., :-1] + + if inverse: + bin_idx = searchsorted(cumheights, inputs)[..., None] + else: + bin_idx = searchsorted(cumwidths, inputs)[..., None] + + input_cumwidths = cumwidths.gather(-1, bin_idx)[..., 0] + input_bin_widths = widths.gather(-1, bin_idx)[..., 0] + + input_cumheights = cumheights.gather(-1, bin_idx)[..., 0] + delta = heights / widths + input_delta = delta.gather(-1, bin_idx)[..., 0] + + input_derivatives = derivatives.gather(-1, bin_idx)[..., 0] + input_derivatives_plus_one = derivatives[..., 1:].gather(-1, bin_idx)[..., 0] + + input_heights = heights.gather(-1, bin_idx)[..., 0] + + if inverse: + a = (inputs - input_cumheights) * ( + input_derivatives + input_derivatives_plus_one - 2 * input_delta + ) + input_heights * (input_delta - input_derivatives) + b = input_heights * input_derivatives - (inputs - input_cumheights) * ( + input_derivatives + input_derivatives_plus_one - 2 * input_delta + ) + c = -input_delta * (inputs - input_cumheights) + + discriminant = b.pow(2) - 4 * a * c + assert (discriminant >= 0).all(), discriminant + + root = (2 * c) / (-b - torch.sqrt(discriminant)) + outputs = root * input_bin_widths + input_cumwidths + + theta_one_minus_theta = root * (1 - root) + denominator = input_delta + ( + (input_derivatives + input_derivatives_plus_one - 2 * input_delta) + * theta_one_minus_theta + ) + derivative_numerator = input_delta.pow(2) * ( + input_derivatives_plus_one * root.pow(2) + + 2 * input_delta * theta_one_minus_theta + + input_derivatives * (1 - root).pow(2) + ) + logabsdet = torch.log(derivative_numerator) - 2 * torch.log(denominator) + + return outputs, -logabsdet + + theta = (inputs - input_cumwidths) / input_bin_widths + theta_one_minus_theta = theta * (1 - theta) + + numerator = input_heights * ( + input_delta * theta.pow(2) + input_derivatives * theta_one_minus_theta + ) + denominator = input_delta + ( + (input_derivatives + input_derivatives_plus_one - 2 * input_delta) + * theta_one_minus_theta + ) + outputs = input_cumheights + numerator / denominator + + derivative_numerator = input_delta.pow(2) * ( + input_derivatives_plus_one * theta.pow(2) + + 2 * input_delta * theta_one_minus_theta + + input_derivatives * (1 - theta).pow(2) + ) + logabsdet = torch.log(derivative_numerator) - 2 * torch.log(denominator) + + return outputs, logabsdet diff --git a/piper_train/vits/utils.py b/piper_train/vits/utils.py new file mode 100644 index 0000000..84d7ecc --- /dev/null +++ b/piper_train/vits/utils.py @@ -0,0 +1,16 @@ +import numpy as np +import torch + + +def to_gpu(x: torch.Tensor) -> torch.Tensor: + return x.contiguous().cuda(non_blocking=True) + + +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 diff --git a/piper_train/vits/wavfile.py b/piper_train/vits/wavfile.py new file mode 100644 index 0000000..38a71c9 --- /dev/null +++ b/piper_train/vits/wavfile.py @@ -0,0 +1,860 @@ +""" +Module to read / write wav files using NumPy arrays + +Functions +--------- +`read`: Return the sample rate (in samples/sec) and data from a WAV file. + +`write`: Write a NumPy array as a WAV file. + +""" +import io +import struct +import sys +import warnings +from enum import IntEnum + +import numpy + +__all__ = ["WavFileWarning", "read", "write"] + + +class WavFileWarning(UserWarning): + pass + + +class WAVE_FORMAT(IntEnum): + """ + WAVE form wFormatTag IDs + + Complete list is in mmreg.h in Windows 10 SDK. ALAC and OPUS are the + newest additions, in v10.0.14393 2016-07 + """ + + UNKNOWN = 0x0000 + PCM = 0x0001 + ADPCM = 0x0002 + IEEE_FLOAT = 0x0003 + VSELP = 0x0004 + IBM_CVSD = 0x0005 + ALAW = 0x0006 + MULAW = 0x0007 + DTS = 0x0008 + DRM = 0x0009 + WMAVOICE9 = 0x000A + WMAVOICE10 = 0x000B + OKI_ADPCM = 0x0010 + DVI_ADPCM = 0x0011 + IMA_ADPCM = 0x0011 # Duplicate + MEDIASPACE_ADPCM = 0x0012 + SIERRA_ADPCM = 0x0013 + G723_ADPCM = 0x0014 + DIGISTD = 0x0015 + DIGIFIX = 0x0016 + DIALOGIC_OKI_ADPCM = 0x0017 + MEDIAVISION_ADPCM = 0x0018 + CU_CODEC = 0x0019 + HP_DYN_VOICE = 0x001A + YAMAHA_ADPCM = 0x0020 + SONARC = 0x0021 + DSPGROUP_TRUESPEECH = 0x0022 + ECHOSC1 = 0x0023 + AUDIOFILE_AF36 = 0x0024 + APTX = 0x0025 + AUDIOFILE_AF10 = 0x0026 + PROSODY_1612 = 0x0027 + LRC = 0x0028 + DOLBY_AC2 = 0x0030 + GSM610 = 0x0031 + MSNAUDIO = 0x0032 + ANTEX_ADPCME = 0x0033 + CONTROL_RES_VQLPC = 0x0034 + DIGIREAL = 0x0035 + DIGIADPCM = 0x0036 + CONTROL_RES_CR10 = 0x0037 + NMS_VBXADPCM = 0x0038 + CS_IMAADPCM = 0x0039 + ECHOSC3 = 0x003A + ROCKWELL_ADPCM = 0x003B + ROCKWELL_DIGITALK = 0x003C + XEBEC = 0x003D + G721_ADPCM = 0x0040 + G728_CELP = 0x0041 + MSG723 = 0x0042 + INTEL_G723_1 = 0x0043 + INTEL_G729 = 0x0044 + SHARP_G726 = 0x0045 + MPEG = 0x0050 + RT24 = 0x0052 + PAC = 0x0053 + MPEGLAYER3 = 0x0055 + LUCENT_G723 = 0x0059 + CIRRUS = 0x0060 + ESPCM = 0x0061 + VOXWARE = 0x0062 + CANOPUS_ATRAC = 0x0063 + G726_ADPCM = 0x0064 + G722_ADPCM = 0x0065 + DSAT = 0x0066 + DSAT_DISPLAY = 0x0067 + VOXWARE_BYTE_ALIGNED = 0x0069 + VOXWARE_AC8 = 0x0070 + VOXWARE_AC10 = 0x0071 + VOXWARE_AC16 = 0x0072 + VOXWARE_AC20 = 0x0073 + VOXWARE_RT24 = 0x0074 + VOXWARE_RT29 = 0x0075 + VOXWARE_RT29HW = 0x0076 + VOXWARE_VR12 = 0x0077 + VOXWARE_VR18 = 0x0078 + VOXWARE_TQ40 = 0x0079 + VOXWARE_SC3 = 0x007A + VOXWARE_SC3_1 = 0x007B + SOFTSOUND = 0x0080 + VOXWARE_TQ60 = 0x0081 + MSRT24 = 0x0082 + G729A = 0x0083 + MVI_MVI2 = 0x0084 + DF_G726 = 0x0085 + DF_GSM610 = 0x0086 + ISIAUDIO = 0x0088 + ONLIVE = 0x0089 + MULTITUDE_FT_SX20 = 0x008A + INFOCOM_ITS_G721_ADPCM = 0x008B + CONVEDIA_G729 = 0x008C + CONGRUENCY = 0x008D + SBC24 = 0x0091 + DOLBY_AC3_SPDIF = 0x0092 + MEDIASONIC_G723 = 0x0093 + PROSODY_8KBPS = 0x0094 + ZYXEL_ADPCM = 0x0097 + PHILIPS_LPCBB = 0x0098 + PACKED = 0x0099 + MALDEN_PHONYTALK = 0x00A0 + RACAL_RECORDER_GSM = 0x00A1 + RACAL_RECORDER_G720_A = 0x00A2 + RACAL_RECORDER_G723_1 = 0x00A3 + RACAL_RECORDER_TETRA_ACELP = 0x00A4 + NEC_AAC = 0x00B0 + RAW_AAC1 = 0x00FF + RHETOREX_ADPCM = 0x0100 + IRAT = 0x0101 + VIVO_G723 = 0x0111 + VIVO_SIREN = 0x0112 + PHILIPS_CELP = 0x0120 + PHILIPS_GRUNDIG = 0x0121 + DIGITAL_G723 = 0x0123 + SANYO_LD_ADPCM = 0x0125 + SIPROLAB_ACEPLNET = 0x0130 + SIPROLAB_ACELP4800 = 0x0131 + SIPROLAB_ACELP8V3 = 0x0132 + SIPROLAB_G729 = 0x0133 + SIPROLAB_G729A = 0x0134 + SIPROLAB_KELVIN = 0x0135 + VOICEAGE_AMR = 0x0136 + G726ADPCM = 0x0140 + DICTAPHONE_CELP68 = 0x0141 + DICTAPHONE_CELP54 = 0x0142 + QUALCOMM_PUREVOICE = 0x0150 + QUALCOMM_HALFRATE = 0x0151 + TUBGSM = 0x0155 + MSAUDIO1 = 0x0160 + WMAUDIO2 = 0x0161 + WMAUDIO3 = 0x0162 + WMAUDIO_LOSSLESS = 0x0163 + WMASPDIF = 0x0164 + UNISYS_NAP_ADPCM = 0x0170 + UNISYS_NAP_ULAW = 0x0171 + UNISYS_NAP_ALAW = 0x0172 + UNISYS_NAP_16K = 0x0173 + SYCOM_ACM_SYC008 = 0x0174 + SYCOM_ACM_SYC701_G726L = 0x0175 + SYCOM_ACM_SYC701_CELP54 = 0x0176 + SYCOM_ACM_SYC701_CELP68 = 0x0177 + KNOWLEDGE_ADVENTURE_ADPCM = 0x0178 + FRAUNHOFER_IIS_MPEG2_AAC = 0x0180 + DTS_DS = 0x0190 + CREATIVE_ADPCM = 0x0200 + CREATIVE_FASTSPEECH8 = 0x0202 + CREATIVE_FASTSPEECH10 = 0x0203 + UHER_ADPCM = 0x0210 + ULEAD_DV_AUDIO = 0x0215 + ULEAD_DV_AUDIO_1 = 0x0216 + QUARTERDECK = 0x0220 + ILINK_VC = 0x0230 + RAW_SPORT = 0x0240 + ESST_AC3 = 0x0241 + GENERIC_PASSTHRU = 0x0249 + IPI_HSX = 0x0250 + IPI_RPELP = 0x0251 + CS2 = 0x0260 + SONY_SCX = 0x0270 + SONY_SCY = 0x0271 + SONY_ATRAC3 = 0x0272 + SONY_SPC = 0x0273 + TELUM_AUDIO = 0x0280 + TELUM_IA_AUDIO = 0x0281 + NORCOM_VOICE_SYSTEMS_ADPCM = 0x0285 + FM_TOWNS_SND = 0x0300 + MICRONAS = 0x0350 + MICRONAS_CELP833 = 0x0351 + BTV_DIGITAL = 0x0400 + INTEL_MUSIC_CODER = 0x0401 + INDEO_AUDIO = 0x0402 + QDESIGN_MUSIC = 0x0450 + ON2_VP7_AUDIO = 0x0500 + ON2_VP6_AUDIO = 0x0501 + VME_VMPCM = 0x0680 + TPC = 0x0681 + LIGHTWAVE_LOSSLESS = 0x08AE + OLIGSM = 0x1000 + OLIADPCM = 0x1001 + OLICELP = 0x1002 + OLISBC = 0x1003 + OLIOPR = 0x1004 + LH_CODEC = 0x1100 + LH_CODEC_CELP = 0x1101 + LH_CODEC_SBC8 = 0x1102 + LH_CODEC_SBC12 = 0x1103 + LH_CODEC_SBC16 = 0x1104 + NORRIS = 0x1400 + ISIAUDIO_2 = 0x1401 + SOUNDSPACE_MUSICOMPRESS = 0x1500 + MPEG_ADTS_AAC = 0x1600 + MPEG_RAW_AAC = 0x1601 + MPEG_LOAS = 0x1602 + NOKIA_MPEG_ADTS_AAC = 0x1608 + NOKIA_MPEG_RAW_AAC = 0x1609 + VODAFONE_MPEG_ADTS_AAC = 0x160A + VODAFONE_MPEG_RAW_AAC = 0x160B + MPEG_HEAAC = 0x1610 + VOXWARE_RT24_SPEECH = 0x181C + SONICFOUNDRY_LOSSLESS = 0x1971 + INNINGS_TELECOM_ADPCM = 0x1979 + LUCENT_SX8300P = 0x1C07 + LUCENT_SX5363S = 0x1C0C + CUSEEME = 0x1F03 + NTCSOFT_ALF2CM_ACM = 0x1FC4 + DVM = 0x2000 + DTS2 = 0x2001 + MAKEAVIS = 0x3313 + DIVIO_MPEG4_AAC = 0x4143 + NOKIA_ADAPTIVE_MULTIRATE = 0x4201 + DIVIO_G726 = 0x4243 + LEAD_SPEECH = 0x434C + LEAD_VORBIS = 0x564C + WAVPACK_AUDIO = 0x5756 + OGG_VORBIS_MODE_1 = 0x674F + OGG_VORBIS_MODE_2 = 0x6750 + OGG_VORBIS_MODE_3 = 0x6751 + OGG_VORBIS_MODE_1_PLUS = 0x676F + OGG_VORBIS_MODE_2_PLUS = 0x6770 + OGG_VORBIS_MODE_3_PLUS = 0x6771 + ALAC = 0x6C61 + _3COM_NBX = 0x7000 # Can't have leading digit + OPUS = 0x704F + FAAD_AAC = 0x706D + AMR_NB = 0x7361 + AMR_WB = 0x7362 + AMR_WP = 0x7363 + GSM_AMR_CBR = 0x7A21 + GSM_AMR_VBR_SID = 0x7A22 + COMVERSE_INFOSYS_G723_1 = 0xA100 + COMVERSE_INFOSYS_AVQSBC = 0xA101 + COMVERSE_INFOSYS_SBC = 0xA102 + SYMBOL_G729_A = 0xA103 + VOICEAGE_AMR_WB = 0xA104 + INGENIENT_G726 = 0xA105 + MPEG4_AAC = 0xA106 + ENCORE_G726 = 0xA107 + ZOLL_ASAO = 0xA108 + SPEEX_VOICE = 0xA109 + VIANIX_MASC = 0xA10A + WM9_SPECTRUM_ANALYZER = 0xA10B + WMF_SPECTRUM_ANAYZER = 0xA10C + GSM_610 = 0xA10D + GSM_620 = 0xA10E + GSM_660 = 0xA10F + GSM_690 = 0xA110 + GSM_ADAPTIVE_MULTIRATE_WB = 0xA111 + POLYCOM_G722 = 0xA112 + POLYCOM_G728 = 0xA113 + POLYCOM_G729_A = 0xA114 + POLYCOM_SIREN = 0xA115 + GLOBAL_IP_ILBC = 0xA116 + RADIOTIME_TIME_SHIFT_RADIO = 0xA117 + NICE_ACA = 0xA118 + NICE_ADPCM = 0xA119 + VOCORD_G721 = 0xA11A + VOCORD_G726 = 0xA11B + VOCORD_G722_1 = 0xA11C + VOCORD_G728 = 0xA11D + VOCORD_G729 = 0xA11E + VOCORD_G729_A = 0xA11F + VOCORD_G723_1 = 0xA120 + VOCORD_LBC = 0xA121 + NICE_G728 = 0xA122 + FRACE_TELECOM_G729 = 0xA123 + CODIAN = 0xA124 + FLAC = 0xF1AC + EXTENSIBLE = 0xFFFE + DEVELOPMENT = 0xFFFF + + +KNOWN_WAVE_FORMATS = {WAVE_FORMAT.PCM, WAVE_FORMAT.IEEE_FLOAT} + + +def _raise_bad_format(format_tag): + try: + format_name = WAVE_FORMAT(format_tag).name + except ValueError: + format_name = f"{format_tag:#06x}" + raise ValueError( + f"Unknown wave file format: {format_name}. Supported " + "formats: " + ", ".join(x.name for x in KNOWN_WAVE_FORMATS) + ) + + +def _read_fmt_chunk(fid, is_big_endian): + """ + Returns + ------- + size : int + size of format subchunk in bytes (minus 8 for "fmt " and itself) + format_tag : int + PCM, float, or compressed format + channels : int + number of channels + fs : int + sampling frequency in samples per second + bytes_per_second : int + overall byte rate for the file + block_align : int + bytes per sample, including all channels + bit_depth : int + bits per sample + + Notes + ----- + Assumes file pointer is immediately after the 'fmt ' id + """ + if is_big_endian: + fmt = ">" + else: + fmt = "<" + + size = struct.unpack(fmt + "I", fid.read(4))[0] + + if size < 16: + raise ValueError("Binary structure of wave file is not compliant") + + res = struct.unpack(fmt + "HHIIHH", fid.read(16)) + bytes_read = 16 + + format_tag, channels, fs, bytes_per_second, block_align, bit_depth = res + + if format_tag == WAVE_FORMAT.EXTENSIBLE and size >= (16 + 2): + ext_chunk_size = struct.unpack(fmt + "H", fid.read(2))[0] + bytes_read += 2 + if ext_chunk_size >= 22: + extensible_chunk_data = fid.read(22) + bytes_read += 22 + raw_guid = extensible_chunk_data[2 + 4 : 2 + 4 + 16] + # GUID template {XXXXXXXX-0000-0010-8000-00AA00389B71} (RFC-2361) + # MS GUID byte order: first three groups are native byte order, + # rest is Big Endian + if is_big_endian: + tail = b"\x00\x00\x00\x10\x80\x00\x00\xAA\x00\x38\x9B\x71" + else: + tail = b"\x00\x00\x10\x00\x80\x00\x00\xAA\x00\x38\x9B\x71" + if raw_guid.endswith(tail): + format_tag = struct.unpack(fmt + "I", raw_guid[:4])[0] + else: + raise ValueError("Binary structure of wave file is not compliant") + + if format_tag not in KNOWN_WAVE_FORMATS: + _raise_bad_format(format_tag) + + # move file pointer to next chunk + if size > bytes_read: + fid.read(size - bytes_read) + + # fmt should always be 16, 18 or 40, but handle it just in case + _handle_pad_byte(fid, size) + + return (size, format_tag, channels, fs, bytes_per_second, block_align, bit_depth) + + +def _read_data_chunk( + fid, format_tag, channels, bit_depth, is_big_endian, block_align, mmap=False +): + """ + Notes + ----- + Assumes file pointer is immediately after the 'data' id + + It's possible to not use all available bits in a container, or to store + samples in a container bigger than necessary, so bytes_per_sample uses + the actual reported container size (nBlockAlign / nChannels). Real-world + examples: + + Adobe Audition's "24-bit packed int (type 1, 20-bit)" + + nChannels = 2, nBlockAlign = 6, wBitsPerSample = 20 + + http://www-mmsp.ece.mcgill.ca/Documents/AudioFormats/WAVE/Samples/AFsp/M1F1-int12-AFsp.wav + is: + + nChannels = 2, nBlockAlign = 4, wBitsPerSample = 12 + + http://www-mmsp.ece.mcgill.ca/Documents/AudioFormats/WAVE/Docs/multichaudP.pdf + gives an example of: + + nChannels = 2, nBlockAlign = 8, wBitsPerSample = 20 + """ + if is_big_endian: + fmt = ">" + else: + fmt = "<" + + # Size of the data subchunk in bytes + size = struct.unpack(fmt + "I", fid.read(4))[0] + + # Number of bytes per sample (sample container size) + bytes_per_sample = block_align // channels + n_samples = size // bytes_per_sample + + if format_tag == WAVE_FORMAT.PCM: + if 1 <= bit_depth <= 8: + dtype = "u1" # WAV of 8-bit integer or less are unsigned + elif bytes_per_sample in {3, 5, 6, 7}: + # No compatible dtype. Load as raw bytes for reshaping later. + dtype = "V1" + elif bit_depth <= 64: + # Remaining bit depths can map directly to signed numpy dtypes + dtype = f"{fmt}i{bytes_per_sample}" + else: + raise ValueError( + "Unsupported bit depth: the WAV file " + f"has {bit_depth}-bit integer data." + ) + elif format_tag == WAVE_FORMAT.IEEE_FLOAT: + if bit_depth in {32, 64}: + dtype = f"{fmt}f{bytes_per_sample}" + else: + raise ValueError( + "Unsupported bit depth: the WAV file " + f"has {bit_depth}-bit floating-point data." + ) + else: + _raise_bad_format(format_tag) + + start = fid.tell() + if not mmap: + try: + count = size if dtype == "V1" else n_samples + data = numpy.fromfile(fid, dtype=dtype, count=count) + except io.UnsupportedOperation: # not a C-like file + fid.seek(start, 0) # just in case it seeked, though it shouldn't + data = numpy.frombuffer(fid.read(size), dtype=dtype) + + if dtype == "V1": + # Rearrange raw bytes into smallest compatible numpy dtype + dt = f"{fmt}i4" if bytes_per_sample == 3 else f"{fmt}i8" + a = numpy.zeros( + (len(data) // bytes_per_sample, numpy.dtype(dt).itemsize), dtype="V1" + ) + if is_big_endian: + a[:, :bytes_per_sample] = data.reshape((-1, bytes_per_sample)) + else: + a[:, -bytes_per_sample:] = data.reshape((-1, bytes_per_sample)) + data = a.view(dt).reshape(a.shape[:-1]) + else: + if bytes_per_sample in {1, 2, 4, 8}: + start = fid.tell() + data = numpy.memmap( + fid, dtype=dtype, mode="c", offset=start, shape=(n_samples,) + ) + fid.seek(start + size) + else: + raise ValueError( + "mmap=True not compatible with " + f"{bytes_per_sample}-byte container size." + ) + + _handle_pad_byte(fid, size) + + if channels > 1: + data = data.reshape(-1, channels) + return data + + +def _skip_unknown_chunk(fid, is_big_endian): + if is_big_endian: + fmt = ">I" + else: + fmt = ">> from os.path import dirname, join as pjoin + >>> from scipy.io import wavfile + >>> import scipy.io + + Get the filename for an example .wav file from the tests/data directory. + + >>> data_dir = pjoin(dirname(scipy.io.__file__), 'tests', 'data') + >>> wav_fname = pjoin(data_dir, 'test-44100Hz-2ch-32bit-float-be.wav') + + Load the .wav file contents. + + >>> samplerate, data = wavfile.read(wav_fname) + >>> print(f"number of channels = {data.shape[1]}") + number of channels = 2 + >>> length = data.shape[0] / samplerate + >>> print(f"length = {length}s") + length = 0.01s + + Plot the waveform. + + >>> import matplotlib.pyplot as plt + >>> import numpy as np + >>> time = np.linspace(0., length, data.shape[0]) + >>> plt.plot(time, data[:, 0], label="Left channel") + >>> plt.plot(time, data[:, 1], label="Right channel") + >>> plt.legend() + >>> plt.xlabel("Time [s]") + >>> plt.ylabel("Amplitude") + >>> plt.show() + + """ + if hasattr(filename, "read"): + fid = filename + mmap = False + else: + # pylint: disable=consider-using-with + fid = open(filename, "rb") + + try: + file_size, is_big_endian = _read_riff_chunk(fid) + fmt_chunk_received = False + data_chunk_received = False + while fid.tell() < file_size: + # read the next chunk + chunk_id = fid.read(4) + + if not chunk_id: + if data_chunk_received: + # End of file but data successfully read + warnings.warn( + f"Reached EOF prematurely; finished at {fid.tell()} bytes, " + "expected {file_size} bytes from header.", + WavFileWarning, + stacklevel=2, + ) + break + + raise ValueError("Unexpected end of file.") + if len(chunk_id) < 4: + msg = f"Incomplete chunk ID: {repr(chunk_id)}" + # If we have the data, ignore the broken chunk + if fmt_chunk_received and data_chunk_received: + warnings.warn(msg + ", ignoring it.", WavFileWarning, stacklevel=2) + else: + raise ValueError(msg) + + if chunk_id == b"fmt ": + fmt_chunk_received = True + fmt_chunk = _read_fmt_chunk(fid, is_big_endian) + format_tag, channels, fs = fmt_chunk[1:4] + bit_depth = fmt_chunk[6] + block_align = fmt_chunk[5] + elif chunk_id == b"fact": + _skip_unknown_chunk(fid, is_big_endian) + elif chunk_id == b"data": + data_chunk_received = True + if not fmt_chunk_received: + raise ValueError("No fmt chunk before data") + data = _read_data_chunk( + fid, + format_tag, + channels, + bit_depth, + is_big_endian, + block_align, + mmap, + ) + elif chunk_id == b"LIST": + # Someday this could be handled properly but for now skip it + _skip_unknown_chunk(fid, is_big_endian) + elif chunk_id in {b"JUNK", b"Fake"}: + # Skip alignment chunks without warning + _skip_unknown_chunk(fid, is_big_endian) + else: + warnings.warn( + "Chunk (non-data) not understood, skipping it.", + WavFileWarning, + stacklevel=2, + ) + _skip_unknown_chunk(fid, is_big_endian) + finally: + if not hasattr(filename, "read"): + fid.close() + else: + fid.seek(0) + + return fs, data + + +def write(filename, rate, data): + """ + Write a NumPy array as a WAV file. + + Parameters + ---------- + filename : string or open file handle + Output wav file. + rate : int + The sample rate (in samples/sec). + data : ndarray + A 1-D or 2-D NumPy array of either integer or float data-type. + + Notes + ----- + * Writes a simple uncompressed WAV file. + * To write multiple-channels, use a 2-D array of shape + (Nsamples, Nchannels). + * The bits-per-sample and PCM/float will be determined by the data-type. + + Common data types: [1]_ + + ===================== =========== =========== ============= + WAV format Min Max NumPy dtype + ===================== =========== =========== ============= + 32-bit floating-point -1.0 +1.0 float32 + 32-bit PCM -2147483648 +2147483647 int32 + 16-bit PCM -32768 +32767 int16 + 8-bit PCM 0 255 uint8 + ===================== =========== =========== ============= + + Note that 8-bit PCM is unsigned. + + References + ---------- + .. [1] IBM Corporation and Microsoft Corporation, "Multimedia Programming + Interface and Data Specifications 1.0", section "Data Format of the + Samples", August 1991 + http://www.tactilemedia.com/info/MCI_Control_Info.html + + Examples + -------- + Create a 100Hz sine wave, sampled at 44100Hz. + Write to 16-bit PCM, Mono. + + >>> from scipy.io.wavfile import write + >>> samplerate = 44100; fs = 100 + >>> t = np.linspace(0., 1., samplerate) + >>> amplitude = np.iinfo(np.int16).max + >>> data = amplitude * np.sin(2. * np.pi * fs * t) + >>> write("example.wav", samplerate, data.astype(np.int16)) + + """ + if hasattr(filename, "write"): + fid = filename + else: + # pylint: disable=consider-using-with + fid = open(filename, "wb") + + fs = rate + + try: + dkind = data.dtype.kind + if not ( + dkind == "i" or dkind == "f" or (dkind == "u" and data.dtype.itemsize == 1) + ): + raise ValueError(f"Unsupported data type '{data.dtype}'") + + header_data = b"" + + header_data += b"RIFF" + header_data += b"\x00\x00\x00\x00" + header_data += b"WAVE" + + # fmt chunk + header_data += b"fmt " + if dkind == "f": + format_tag = WAVE_FORMAT.IEEE_FLOAT + else: + format_tag = WAVE_FORMAT.PCM + if data.ndim == 1: + channels = 1 + else: + channels = data.shape[1] + bit_depth = data.dtype.itemsize * 8 + bytes_per_second = fs * (bit_depth // 8) * channels + block_align = channels * (bit_depth // 8) + + fmt_chunk_data = struct.pack( + " 0xFFFFFFFF: + raise ValueError("Data exceeds wave file size limit") + + fid.write(header_data) + + # data chunk + fid.write(b"data") + fid.write(struct.pack("" or ( + data.dtype.byteorder == "=" and sys.byteorder == "big" + ): + data = data.byteswap() + _array_tofile(fid, data) + + # Determine file size and place it in correct + # position at start of the file. + size = fid.tell() + fid.seek(4) + fid.write(struct.pack("