pipelines/no_midi_preparation.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "\n",
    "# Prepare your own dataset for DiffSinger (MIDI-less version)\n",
    "\n",
    "## 1 Overview\n",
    "\n",
    "This Jupyter Notebook will guide you to prepare your own dataset for DiffSinger with 44.1 kHz sampling rate.\n",
    "Please read and follow the guidance carefully, take actions when there are notice for <font color=\"red\">manual action</font> and pay attention to blocks marked with <font color=\"red\">optional step</font>.\n",
    "\n",
    "### 1.1 Introduction to this pipeline and MIDI-less version\n",
    "\n",
    "This pipeline does not support customized phoneme dictionaries. It uses the [opencpop strict pinyin dictionary](../dictionaries/opencpop-extension.txt) by default.\n",
    "\n",
    "MIDI-less version is a simplified version of DiffSinger where MIDI layers, word layers and slur layers are removed from the data labels. The model uses raw phoneme sequence with durations as input, and applies pitch embedding directly from the ground truth. Predictors for phoneme durations and pitch curve are also removed. Below are some limitations and advantages of the MIDI-less version:\n",
    "\n",
    "- The model will not predict phoneme durations and f0 sequence by itself. You must specify `ph_dur` and `f0_seq` at inference time.\n",
    "- Performance of pitch control will be better than MIDI-A version, because MIDI keys are misleading information for the diffusion decoder when f0 sequence is already embedded.\n",
    "- MIDIs and slurs does not need to be labeled, thus the labeling work is easier than other versions.\n",
    "- More varieties of data can be used as training materials, even including speech.\n",
    "\n",
    "### 1.2 Install dependencies\n",
    "\n",
    "Please run the following cell the first time you start this notebook.\n",
    "\n",
    "**Note**: You should ensure you are in a Conda environment with Python 3.8 or 3.9 before you install dependencies of this pipeline.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false,
    "jupyter": {
     "outputs_hidden": false
    }
   },
   "outputs": [],
   "source": [
    "%conda install -c conda-forge montreal-forced-aligner==2.0.6 --yes\n",
    "%pip install -r requirements.txt\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 1.3 Initializing environment\n",
    "\n",
    "Please run the following cell every time you start this notebook.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false,
    "jupyter": {
     "outputs_hidden": false
    }
   },
   "outputs": [],
   "source": [
    "import glob\n",
    "import os\n",
    "import shutil\n",
    "import sys\n",
    "\n",
    "import librosa\n",
    "import matplotlib.pyplot as plt\n",
    "import numpy as np\n",
    "import parselmouth as pm\n",
    "import soundfile\n",
    "import textgrid as tg\n",
    "import tqdm\n",
    "\n",
    "\n",
    "def length(src: str):\n",
    "    if os.path.isfile(src) and src.endswith('.wav'):\n",
    "        return librosa.get_duration(filename=src) / 3600.\n",
    "    elif os.path.isdir(src):\n",
    "        total = 0\n",
    "        for ch in [os.path.join(src, c) for c in os.listdir(src)]:\n",
    "            total += length(ch)\n",
    "        return total\n",
    "    return 0\n",
    "\n",
    "\n",
    "print('Environment initialized successfully.')\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 2 Raw recordings and audio slicing\n",
    "\n",
    "### 2.1 Choose raw recordings\n",
    "\n",
    "Your recordings must meet the following conditions:\n",
    "\n",
    "1. They must be in one single folder. Files in sub-folders will be ignored.\n",
    "2. They must be in WAV format.\n",
    "3. They must have a sampling rate higher than 32 kHz.\n",
    "4. They should be clean, unaccompanied voices with no significant noise or reverb.\n",
    "5. They should contain only voices from one single human.\n",
    "\n",
    "<font color=\"#66ccff\">NOTICE</font>: Before you train a model, you must obtain permission from the copyright holder of the dataset and make sure the provider is fully aware that you will train a model from their data, that you will or will not distribute the synthesized voices and model weights, and the potential risks of this kind of activity.\n",
    "\n",
    "> **Tips for building multiple datasets to train combined models**\n",
    ">\n",
    "> If you have multiple speakers/singers, or you have multiple styles/timbres from one person, please make one dataset for each speaker/singer/style/timbre separately. Then you are able to configure parameters and settings for preprocessing and training combined models from these datasets following instructions in **Section 5** of this pipeline.\n",
    "\n",
    "<font color=\"red\">Optional step</font>: The raw data must be sliced into parts of about 5-15 seconds. If you want to do this yourself, please skip to section 2.3. Otherwise, please edit paths in the following cell before you run it.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false,
    "jupyter": {
     "outputs_hidden": false
    }
   },
   "outputs": [],
   "source": [
    "########################################\n",
    "\n",
    "# Configuration for data paths\n",
    "raw_path = r'path/to/your/raw/recordings'  # Path to your raw, unsliced recordings\n",
    "\n",
    "########################################\n",
    "\n",
    "assert os.path.exists(raw_path) and os.path.isdir(raw_path), 'The chosen path does not exist or is not a directory.'\n",
    "print('Raw recording path:', raw_path)\n",
    "print()\n",
    "print('===== Recording List =====')\n",
    "raw_filelist = glob.glob(f'{raw_path}/*.wav', recursive=True)\n",
    "raw_length = length(raw_path)\n",
    "if len(raw_filelist) > 5:\n",
    "    print('\\n'.join(raw_filelist[:5] + [f'... ({len(raw_filelist) - 5} more)']))\n",
    "else:\n",
    "    print('\\n'.join(raw_filelist))\n",
    "print()\n",
    "print(f'Found {len(raw_filelist)} valid recordings with total length of {round(raw_length, 2)} hours.')\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 2.2 Audio slicing\n",
    "\n",
    "We provide an audio slicer which automatically cuts recordings into short pieces.\n",
    "\n",
    "The audio slicer is based on silence detection and has several arguments that have to be specified. You should modify these arguments according to your data.\n",
    "\n",
    "For more details of each argument, see its [GitHub repository](https://github.com/openvpi/audio-slicer).\n",
    "\n",
    "Please edit paths and arguments in the following cell before you run it.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false,
    "jupyter": {
     "outputs_hidden": false
    }
   },
   "outputs": [],
   "source": [
    "########################################\n",
    "\n",
    "# Configuration for data paths\n",
    "sliced_path = r'path/to/your/sliced/recordings'  # Path to hold the sliced segments of your recordings\n",
    "\n",
    "# Slicer arguments\n",
    "db_threshold_ = -40.\n",
    "min_length_ = 5000\n",
    "min_interval_ = 300\n",
    "hop_size_ = 10\n",
    "max_sil_kept_ = 500\n",
    "\n",
    "########################################\n",
    "\n",
    "assert 'raw_path' in locals().keys(), 'Raw path of your recordings has not been specified.'\n",
    "assert not os.path.exists(sliced_path) or os.path.isdir(sliced_path), 'The chosen path is not a directory.'\n",
    "os.makedirs(sliced_path, exist_ok=True)\n",
    "print('Sliced recording path:', sliced_path)\n",
    "\n",
    "from utils.slicer2 import Slicer\n",
    "\n",
    "for file in tqdm.tqdm(raw_filelist):\n",
    "    y, sr = librosa.load(file, sr=None, mono=True)\n",
    "    slicer = Slicer(\n",
    "        sr=sr,\n",
    "        threshold=db_threshold_,\n",
    "        min_length=min_length_,\n",
    "        min_interval=min_interval_,\n",
    "        hop_size=hop_size_,\n",
    "        max_sil_kept=max_sil_kept_\n",
    "    )\n",
    "    chunks = slicer.slice(y)\n",
    "    for i, chunk in enumerate(chunks):\n",
    "        soundfile.write(os.path.join(sliced_path, f'%s_slice_%04d.wav' % (os.path.basename(file).rsplit('.', maxsplit=1)[0], i)), chunk, sr)\n",
    "\n",
    "print()\n",
    "print('===== Segment List =====')\n",
    "sliced_filelist = glob.glob(f'{sliced_path}/*.wav', recursive=True)\n",
    "sliced_length = length(sliced_path)\n",
    "if len(sliced_filelist) > 5:\n",
    "    print('\\n'.join(sliced_filelist[:5] + [f'... ({len(sliced_filelist) - 5} more)']))\n",
    "else:\n",
    "    print('\\n'.join(sliced_filelist))\n",
    "print()\n",
    "print(f'Sliced your recordings into {len(sliced_filelist)} segments with total length of {round(sliced_length, 2)} hours.')\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 2.3 Validating recording segments\n",
    "\n",
    "In this section, we validate your recording segments.\n",
    "\n",
    "<font color=\"red\">Optional step</font>: If you skipped section 2.2, please specify the path to your sliced recordings in the following cell and run it. Otherwise, skip this cell.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false,
    "jupyter": {
     "outputs_hidden": false
    }
   },
   "outputs": [],
   "source": [
    "########################################\n",
    "\n",
    "# Configuration for data paths\n",
    "sliced_path = r'path/to/your/sliced/recordings'  # Path to your sliced segments of recordings\n",
    "\n",
    "########################################\n",
    "\n",
    "assert os.path.exists(sliced_path) and os.path.isdir(sliced_path), 'The chosen path does not exist or is not a directory.'\n",
    "\n",
    "print('Sliced recording path:', sliced_path)\n",
    "print()\n",
    "print('===== Segment List =====')\n",
    "sliced_filelist = glob.glob(f'{sliced_path}/*.wav', recursive=True)\n",
    "sliced_length = length(sliced_path)\n",
    "if len(sliced_filelist) > 5:\n",
    "    print('\\n'.join(sliced_filelist[:5] + [f'... ({len(sliced_filelist) - 5} more)']))\n",
    "else:\n",
    "    print('\\n'.join(sliced_filelist))\n",
    "print()\n",
    "print(f'Found {len(sliced_filelist)} valid segments with total length of {round(sliced_length, 2)} hours.')\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Run the following cell to check if there are segments with an unexpected length (less than 2 seconds or more than 20 seconds).\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false,
    "jupyter": {
     "outputs_hidden": false
    }
   },
   "outputs": [],
   "source": [
    "reported = False\n",
    "for file in tqdm.tqdm(sliced_filelist):\n",
    "    wave_seconds = librosa.get_duration(filename=file)\n",
    "    if wave_seconds < 2.:\n",
    "        reported = True\n",
    "        print(f'Too short! \\'{file}\\' has a length of {round(wave_seconds, 1)} seconds!')\n",
    "    if wave_seconds > 20.:\n",
    "        reported = True\n",
    "        print(f'Too long! \\'{file}\\' has a length of {round(wave_seconds, 1)} seconds!')\n",
    "if not reported:\n",
    "    print('Congratulations! All segments have proper length.')\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<font color=\"red\">Manual action</font>: please consider removing segments too short and manually slicing segments to long, as reported above.\n",
    "\n",
    "Move on when this is done or there are no segments reported.\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 3 Label your segments\n",
    "\n",
    "### 3.1 Label syllable sequence\n",
    "\n",
    "All segments should have their transcriptions (or lyrics) annotated. Run the following cell to see the example segment (from Opencpop dataset) and its corresponding annotation.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false,
    "jupyter": {
     "outputs_hidden": false
    }
   },
   "outputs": [],
   "source": [
    "from IPython.display import Audio\n",
    "\n",
    "# noinspection PyTypeChecker\n",
    "display(Audio(filename='assets/2001000001.wav'))\n",
    "with open('assets/2001000001.lab', 'r') as f:\n",
    "    print(f.read())\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<font color=\"red\">Manual action</font>: now your task is to annotation transcriptions for each segment like the example shown above.\n",
    "\n",
    "Each segment should have one annotation file with the same filename as it and `.lab` extension, and placed in the same directory. In the annotation file, you should write all syllables sung or spoken in this segment. Syllables should be split by space, and only syllables that appears in the dictionary are allowed. In addition, all phonemes in the dictionary should be covered in the annotations.\n",
    "\n",
    "**Special notes**: `AP` and `SP` should not appear in the annotation.\n",
    "\n",
    "**News**:  We developed [MinLabel](https://github.com/SineStriker/qsynthesis-revenge/tree/main/src/Test/MinLabel), a simple yet efficient tool to help finishing this step. You can download the binary executable for Windows [here](https://huggingface.co/datasets/fox7005/tool/resolve/main/MinLabel%200.0.1.6.zip).\n",
    "\n",
    "<font color=\"red\">Optional step</font>: if you want us to help you create all empty `lab` files (instead of creating them yourself), please run the following cell.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false,
    "jupyter": {
     "outputs_hidden": false
    }
   },
   "outputs": [],
   "source": [
    "for file in tqdm.tqdm(sliced_filelist):\n",
    "    filename = os.path.basename(file)\n",
    "    name_without_ext = filename.rsplit('.', maxsplit=1)[0]\n",
    "    annotation = os.path.join(sliced_path, f'{name_without_ext}.lab')\n",
    "    if not os.path.exists(annotation):\n",
    "        with open(annotation, 'a'):\n",
    "            ...\n",
    "print('Creating missing lab files done.')\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Run the following cell to see if all segments are annotated and all annotations are valid. If there are failed checks, please fix them and run again.\n",
    "\n",
    "A summary of your phoneme coverage will be generated. If there are some phonemes that have extremely few occurrences (for example, less than 20), it is highly recommended to add more recordings to cover these phonemes.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false,
    "jupyter": {
     "outputs_hidden": false
    }
   },
   "outputs": [],
   "source": [
    "import utils.distribution as dist\n",
    "\n",
    "# Load dictionary\n",
    "dict_path = '../dictionaries/opencpop-extension.txt'\n",
    "with open(dict_path, 'r', encoding='utf8') as f:\n",
    "    rules = [ln.strip().split('\\t') for ln in f.readlines()]\n",
    "dictionary = {}\n",
    "phoneme_set = set()\n",
    "for r in rules:\n",
    "    phonemes = r[1].split()\n",
    "    dictionary[r[0]] = phonemes\n",
    "    phoneme_set.update(phonemes)\n",
    "\n",
    "# Run checks\n",
    "check_failed = False\n",
    "covered = set()\n",
    "phoneme_map = {}\n",
    "for ph in sorted(phoneme_set):\n",
    "    phoneme_map[ph] = 0\n",
    "\n",
    "segment_pairs = []\n",
    "\n",
    "for file in tqdm.tqdm(sliced_filelist):\n",
    "    filename = os.path.basename(file)\n",
    "    name_without_ext = filename.rsplit('.', maxsplit=1)[0]\n",
    "    annotation = os.path.join(sliced_path, f'{name_without_ext}.lab')\n",
    "    if not os.path.exists(annotation):\n",
    "        print(f'No annotation found for \\'{filename}\\'!')\n",
    "        check_failed = True\n",
    "        continue\n",
    "    with open(annotation, 'r', encoding='utf8') as f:\n",
    "        syllables = f.read().strip().split()\n",
    "    if not syllables:\n",
    "        print(f'Annotation file \\'{annotation}\\' is empty!')\n",
    "        check_failed = True\n",
    "    else:\n",
    "        oov = []\n",
    "        for s in syllables:\n",
    "            if s not in dictionary:\n",
    "                oov.append(s)\n",
    "            else:\n",
    "                for ph in dictionary[s]:\n",
    "                    phoneme_map[ph] += 1\n",
    "                covered.update(dictionary[s])\n",
    "        if oov:\n",
    "            print(f'Syllable(s) {oov} not allowed in annotation file \\'{annotation}\\'')\n",
    "            check_failed = True\n",
    "\n",
    "# Phoneme coverage\n",
    "uncovered = phoneme_set - covered\n",
    "if uncovered:\n",
    "    print(f'The following phonemes are not covered!')\n",
    "    print(sorted(uncovered))\n",
    "    print('Please add more recordings to cover these phonemes.')\n",
    "    check_failed = True\n",
    "\n",
    "if not check_failed:\n",
    "    print('Congratulations! All annotations are well prepared.')\n",
    "    print('Here is a summary of your phoneme coverage.')\n",
    "\n",
    "phoneme_list = sorted(phoneme_set)\n",
    "phoneme_counts = [phoneme_map[ph] for ph in phoneme_list]\n",
    "dist.draw_distribution(\n",
    "    title='Phoneme Distribution Summary',\n",
    "    x_label='Phoneme',\n",
    "    y_label='Number of occurrences',\n",
    "    items=phoneme_list,\n",
    "    values=phoneme_counts\n",
    ")\n",
    "phoneme_summary = os.path.join(sliced_path, 'phoneme_distribution.jpg')\n",
    "plt.savefig(fname=phoneme_summary,\n",
    "            bbox_inches='tight',\n",
    "            pad_inches=0.25)\n",
    "plt.show()\n",
    "print(f'Summary saved to \\'{phoneme_summary}\\'.')\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 3.2 Forced alignment\n",
    "\n",
    "Given the transcriptions of each segment, we are able to align the phoneme sequence to its corresponding audio, thus obtaining position and duration information of each phoneme.\n",
    "\n",
    "We use [Montreal Forced Aligner](https://github.com/MontrealCorpusTools/Montreal-Forced-Aligner) to do forced phoneme alignment.\n",
    "\n",
    "To run MFA alignment, please first run the following cell to resample all recordings to 16 kHz. The resampled recordings and copies of the phoneme labels will be saved at `pipelines/segments/`. Also, the folder `pipelines/textgrids/` will be created for temporarily storing aligned TextGrids.\n",
    "\n",
    "<font color=\"yellow\">WARNING</font>: This will overwrite all files in `pipelines/segments/` and `pipelines/textgrids/`.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false,
    "jupyter": {
     "outputs_hidden": false
    }
   },
   "outputs": [],
   "source": [
    "segments_dir = 'segments'\n",
    "textgrids_dir = 'textgrids'\n",
    "if os.path.exists(segments_dir):\n",
    "    shutil.rmtree(segments_dir)\n",
    "os.makedirs(segments_dir)\n",
    "if os.path.exists(textgrids_dir):\n",
    "    shutil.rmtree(textgrids_dir)\n",
    "os.makedirs(textgrids_dir)\n",
    "samplerate = 16000\n",
    "for file in tqdm.tqdm(sliced_filelist):\n",
    "    y, _ = librosa.load(file, sr=samplerate, mono=True)\n",
    "    filename = os.path.basename(file)\n",
    "    soundfile.write(os.path.join(segments_dir, filename), y, samplerate, subtype='PCM_16')\n",
    "    name_without_ext = filename.rsplit('.', maxsplit=1)[0]\n",
    "    annotation = os.path.join(sliced_path, f'{name_without_ext}.lab')\n",
    "    shutil.copy(annotation, segments_dir)\n",
    "print('Resampling and copying done.')\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Run the following cell to download the MFA pretrained model and perform forced alignment. You may edit the --beam argument of MFA according to your dataset: Longer segments and worse quality requires larger beams.\n",
    "\n",
    "This cell also checks if alignments for all segments are sucessfully generated. If the checks fail, there are probably severe errors in your labels, or your segments are too long, or you did not use a proper --beam value.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false,
    "jupyter": {
     "outputs_hidden": false
    }
   },
   "outputs": [],
   "source": [
    "import requests\n",
    "\n",
    "mfa_zip = f'assets/mfa-opencpop-extension.zip'\n",
    "mfa_uri = 'https://huggingface.co/datasets/fox7005/tool/resolve/main/mfa-opencpop-extension.zip'\n",
    "if not os.path.exists(mfa_zip):\n",
    "    # Download\n",
    "    print('Model not found, downloading...')\n",
    "    with open(mfa_zip, 'wb') as f:\n",
    "        f.write(requests.get(mfa_uri).content)\n",
    "    print('Done.')\n",
    "else:\n",
    "    pass\n",
    "\n",
    "segments_dir = 'segments'\n",
    "textgrids_dir = 'textgrids'\n",
    "os.makedirs(textgrids_dir, exist_ok=True)\n",
    "\n",
    "import montreal_forced_aligner\n",
    "\n",
    "sys.argv = [\n",
    "    'mfa',\n",
    "    'align',\n",
    "    segments_dir,\n",
    "    dict_path,\n",
    "    mfa_zip,\n",
    "    textgrids_dir,\n",
    "    '--beam',\n",
    "    '100',  # Edit --beam here.\n",
    "    '--clean',\n",
    "    '--overwrite'\n",
    "]\n",
    "montreal_forced_aligner.command_line.mfa.main()\n",
    "\n",
    "print('Checking aligments...')\n",
    "missing = []\n",
    "for wavfile in tqdm.tqdm(sliced_filelist):\n",
    "    name = os.path.basename(wavfile).rsplit('.', maxsplit=1)[0]\n",
    "    tgfile = os.path.join(textgrids_dir, f'{name}.TextGrid')\n",
    "    if not os.path.exists(tgfile):\n",
    "        missing.append(tgfile)\n",
    "if len(missing) > 0:\n",
    "    print('These TextGrids are missing! There are possible severe errors in labels of those corresponding segments. '\n",
    "          'If you do believe there are no errors, consider increase the \\'--beam\\' argument for MFA.')\n",
    "    for fn in missing:\n",
    "        print(f' - {fn}')\n",
    "else:\n",
    "    print('All alignments have been successfully generated. Please move on.')\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 3.3 Optimize and finish the TextGrids\n",
    "\n",
    "In this section, we run some scripts to reduce errors for long utterances and detect `AP`s which have not been labeled before. The optimized TextGrids can be saved for future use if you specify a backup directory in the following cell.\n",
    "\n",
    "Edit the path and adjust arguments according to your needs in the following cell before you run it. Optimized results will be saved at `pipelines/textgrids/revised/`.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false,
    "jupyter": {
     "outputs_hidden": false
    }
   },
   "outputs": [],
   "source": [
    "########################################\n",
    "\n",
    "# Configuration for voice arguments based on your dataset\n",
    "f0_min = 40.  # Minimum value of pitch\n",
    "f0_max = 1100.  # Maximum value of pitch\n",
    "br_len = 0.1  # Minimum length of aspiration in seconds\n",
    "br_db = -60.  # Threshold of RMS in dB for detecting aspiration\n",
    "br_centroid = 2000.  # Threshold of spectral centroid in Hz for detecting aspiration\n",
    "\n",
    "# Other arguments, do not edit unless you understand them\n",
    "time_step = 0.005  # Time step for feature extraction\n",
    "min_space = 0.04  # Minimum length of space in seconds\n",
    "voicing_thresh_vowel = 0.45  # Threshold of voicing for fixing long utterances\n",
    "voicing_thresh_breath = 0.6  # Threshold of voicing for detecting aspiration\n",
    "br_win_sz = 0.05  # Size of sliding window in seconds for detecting aspiration\n",
    "\n",
    "########################################\n",
    "\n",
    "# import utils.tg_optimizer as optimizer\n",
    "\n",
    "textgrids_revised_dir = 'textgrids/revised'\n",
    "os.makedirs(textgrids_revised_dir, exist_ok=True)\n",
    "for wavfile in tqdm.tqdm(sliced_filelist):\n",
    "    name = os.path.basename(wavfile).rsplit('.', maxsplit=1)[0]\n",
    "    textgrid = tg.TextGrid()\n",
    "    textgrid.read(os.path.join(textgrids_dir, f'{name}.TextGrid'))\n",
    "    words = textgrid[0]\n",
    "    phones = textgrid[1]\n",
    "    sound = pm.Sound(wavfile)\n",
    "    f0_voicing_breath = sound.to_pitch_ac(\n",
    "        time_step=time_step,\n",
    "        voicing_threshold=voicing_thresh_breath,\n",
    "        pitch_floor=f0_min,\n",
    "        pitch_ceiling=f0_max,\n",
    "    ).selected_array['frequency']\n",
    "    f0_voicing_vowel = sound.to_pitch_ac(\n",
    "        time_step=time_step,\n",
    "        voicing_threshold=voicing_thresh_vowel,\n",
    "        pitch_floor=f0_min,\n",
    "        pitch_ceiling=f0_max,\n",
    "    ).selected_array['frequency']\n",
    "    y, sr = librosa.load(wavfile, sr=24000, mono=True)\n",
    "    hop_size = int(time_step * sr)\n",
    "    spectral_centroid = librosa.feature.spectral_centroid(y=y, sr=sr, n_fft=2048, hop_length=hop_size).squeeze(0)\n",
    "\n",
    "    # Fix long utterances\n",
    "    i = j = 0\n",
    "    while i < len(words):\n",
    "        word = words[i]\n",
    "        phone = phones[j]\n",
    "        if word.mark is not None and word.mark != '':\n",
    "            i += 1\n",
    "            j += len(dictionary[word.mark])\n",
    "            continue\n",
    "        if i == 0:\n",
    "            i += 1\n",
    "            j += 1\n",
    "            continue\n",
    "        prev_word = words[i - 1]\n",
    "        prev_phone = phones[j - 1]\n",
    "        # Extend length of long utterances\n",
    "        while word.minTime < word.maxTime - time_step:\n",
    "            pos = min(f0_voicing_vowel.shape[0] - 1, int(word.minTime / time_step))\n",
    "            if f0_voicing_vowel[pos] < f0_min:\n",
    "                break\n",
    "            prev_word.maxTime += time_step\n",
    "            prev_phone.maxTime += time_step\n",
    "            word.minTime += time_step\n",
    "            phone.minTime += time_step\n",
    "        i += 1\n",
    "        j += 1\n",
    "\n",
    "    # Detect aspiration\n",
    "    i = j = 0\n",
    "    while i < len(words):\n",
    "        word = words[i]\n",
    "        phone = phones[j]\n",
    "        if word.mark is not None and word.mark != '':\n",
    "            i += 1\n",
    "            j += len(dictionary[word.mark])\n",
    "            continue\n",
    "        if word.maxTime - word.minTime < br_len:\n",
    "            i += 1\n",
    "            j += 1\n",
    "            continue\n",
    "        ap_ranges = []\n",
    "        br_start = None\n",
    "        win_pos = word.minTime\n",
    "        while win_pos + br_win_sz <= word.maxTime:\n",
    "            all_noisy = (f0_voicing_breath[int(win_pos / time_step) : int((win_pos + br_win_sz) / time_step)] < f0_min).all()\n",
    "            rms_db = 20 * np.log10(np.clip(sound.get_rms(from_time=win_pos, to_time=win_pos + br_win_sz), a_min=1e-12, a_max=1))\n",
    "            # print(win_pos, win_pos + br_win_sz, all_noisy, rms_db)\n",
    "            if all_noisy and rms_db >= br_db:\n",
    "                if br_start is None:\n",
    "                    br_start = win_pos\n",
    "            else:\n",
    "                if br_start is not None:\n",
    "                    br_end = win_pos + br_win_sz - time_step\n",
    "                    if br_end - br_start >= br_len:\n",
    "                        centroid = spectral_centroid[int(br_start / time_step) : int(br_end / time_step)].mean()\n",
    "                        if centroid >= br_centroid:\n",
    "                            ap_ranges.append((br_start, br_end))\n",
    "                    br_start = None\n",
    "                    win_pos = br_end\n",
    "            win_pos += time_step\n",
    "        if br_start is not None:\n",
    "            br_end = win_pos + br_win_sz - time_step\n",
    "            if br_end - br_start >= br_len:\n",
    "                centroid = spectral_centroid[int(br_start / time_step) : int(br_end / time_step)].mean()\n",
    "                if centroid >= br_centroid:\n",
    "                    ap_ranges.append((br_start, br_end))\n",
    "        # print(ap_ranges)\n",
    "        if len(ap_ranges) == 0:\n",
    "            i += 1\n",
    "            j += 1\n",
    "            continue\n",
    "        words.removeInterval(word)\n",
    "        phones.removeInterval(phone)\n",
    "        if word.minTime < ap_ranges[0][0]:\n",
    "            words.add(minTime=word.minTime, maxTime=ap_ranges[0][0], mark=None)\n",
    "            phones.add(minTime=phone.minTime, maxTime=ap_ranges[0][0], mark=None)\n",
    "            i += 1\n",
    "            j += 1\n",
    "        for k, ap in enumerate(ap_ranges):\n",
    "            if k > 0:\n",
    "                words.add(minTime=ap_ranges[k - 1][1], maxTime=ap[0], mark=None)\n",
    "                phones.add(minTime=ap_ranges[k - 1][1], maxTime=ap[0], mark=None)\n",
    "                i += 1\n",
    "                j += 1\n",
    "            words.add(minTime=ap[0], maxTime=min(word.maxTime, ap[1]), mark='AP')\n",
    "            phones.add(minTime=ap[0], maxTime=min(word.maxTime, ap[1]), mark='AP')\n",
    "            i += 1\n",
    "            j += 1\n",
    "        if ap_ranges[-1][1] < word.maxTime:\n",
    "            words.add(minTime=ap_ranges[-1][1], maxTime=word.maxTime, mark=None)\n",
    "            phones.add(minTime=ap_ranges[-1][1], maxTime=phone.maxTime, mark=None)\n",
    "            i += 1\n",
    "            j += 1\n",
    "\n",
    "    # Remove short spaces\n",
    "    i = j = 0\n",
    "    while i < len(words):\n",
    "        word = words[i]\n",
    "        phone = phones[j]\n",
    "        if word.mark is not None and word.mark != '':\n",
    "            i += 1\n",
    "            j += (1 if word.mark == 'AP' else len(dictionary[word.mark]))\n",
    "            continue\n",
    "        if word.maxTime - word.minTime >= min_space:\n",
    "            word.mark = 'SP'\n",
    "            phone.mark = 'SP'\n",
    "            i += 1\n",
    "            j += 1\n",
    "            continue\n",
    "        if i == 0:\n",
    "            if len(words) >= 2:\n",
    "                words[i + 1].minTime = word.minTime\n",
    "                phones[j + 1].minTime = phone.minTime\n",
    "                words.removeInterval(word)\n",
    "                phones.removeInterval(phone)\n",
    "            else:\n",
    "                break\n",
    "        elif i == len(words) - 1:\n",
    "            if len(words) >= 2:\n",
    "                words[i - 1].maxTime = word.maxTime\n",
    "                phones[j - 1].maxTime = phone.maxTime\n",
    "                words.removeInterval(word)\n",
    "                phones.removeInterval(phone)\n",
    "            else:\n",
    "                break\n",
    "        else:\n",
    "            words[i - 1].maxTime = words[i + 1].minTime = (word.minTime + word.maxTime) / 2\n",
    "            phones[j - 1].maxTime = phones[j + 1].minTime = (phone.minTime + phone.maxTime) / 2\n",
    "            words.removeInterval(word)\n",
    "            phones.removeInterval(phone)\n",
    "    textgrid.write(os.path.join(textgrids_revised_dir, f'{name}.TextGrid'))\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "`TextGrid` saved in `pipelines/textgrids/revised` can be edited via [Praat](https://github.com/praat/praat). You may examine these files and fix label errors by yourself if you want a more accurate model with higher performance. However, this is not required since manual labeling takes much time.\n",
    "\n",
    "Run the following cell to see summary of word-level pitch coverage of your dataset. (Data may not be accurate due to octave errors in pitch extraction.)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false,
    "jupyter": {
     "outputs_hidden": false
    }
   },
   "outputs": [],
   "source": [
    "import utils.distribution as dist\n",
    "\n",
    "\n",
    "def key_to_name(midi_key):\n",
    "    note_names = ['C', 'C#', 'D', 'D#', 'E', 'F', 'F#', 'G', 'G#', 'A', 'A#', 'B']\n",
    "    return note_names[midi_key % 12] + str(midi_key // 12 - 1)\n",
    "\n",
    "\n",
    "pit_map = {}\n",
    "if not f0_min in locals():\n",
    "    f0_min = 40.\n",
    "if not f0_max in locals():\n",
    "    f0_max = 1100.\n",
    "if not voicing_thresh_vowel in locals():\n",
    "    voicing_thresh_vowel = 0.45\n",
    "for wavfile in tqdm.tqdm(sliced_filelist):\n",
    "    name = os.path.basename(wavfile).rsplit('.', maxsplit=1)[0]\n",
    "    textgrid = tg.TextGrid()\n",
    "    textgrid.read(os.path.join(textgrids_revised_dir, f'{name}.TextGrid'))\n",
    "    timestep = 0.01\n",
    "    f0 = pm.Sound(wavfile).to_pitch_ac(\n",
    "        time_step=timestep,\n",
    "        voicing_threshold=voicing_thresh_vowel,\n",
    "        pitch_floor=f0_min,\n",
    "        pitch_ceiling=f0_max,\n",
    "    ).selected_array['frequency']\n",
    "    pitch = 12. * np.log2(f0 / 440.) + 69.\n",
    "    for word in textgrid[0]:\n",
    "        if word.mark in ['AP', 'SP']:\n",
    "            continue\n",
    "        if word.maxTime - word.minTime < timestep:\n",
    "            continue\n",
    "        word_pit = pitch[int(word.minTime / timestep) : int(word.maxTime / timestep)]\n",
    "        word_pit = np.extract(word_pit >= 0, word_pit)\n",
    "        if word_pit.shape[0] == 0:\n",
    "            continue\n",
    "        counts = np.bincount(word_pit.astype(np.int64))\n",
    "        midi = counts.argmax()\n",
    "        if midi in pit_map:\n",
    "            pit_map[midi] += 1\n",
    "        else:\n",
    "            pit_map[midi] = 1\n",
    "midi_keys = sorted(pit_map.keys())\n",
    "midi_keys = list(range(midi_keys[0], midi_keys[-1] + 1))\n",
    "dist.draw_distribution(\n",
    "    title='Pitch Distribution Summary',\n",
    "    x_label='Pitch',\n",
    "    y_label='Number of occurrences',\n",
    "    items=[key_to_name(k) for k in midi_keys],\n",
    "    values=[pit_map.get(k, 0) for k in midi_keys]\n",
    ")\n",
    "pitch_summary = os.path.join(sliced_path, 'pitch_distribution.jpg')\n",
    "plt.savefig(fname=pitch_summary,\n",
    "            bbox_inches='tight',\n",
    "            pad_inches=0.25)\n",
    "plt.show()\n",
    "print(f'Summary saved to \\'{pitch_summary}\\'.')\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 4 Building the final dataset\n",
    "\n",
    "Congratulations! If you have gone through all sections above with success, it means that you are now prepared for building your final dataset. There are only a few steps to go before you can run scripts to train your own model.\n",
    "\n",
    "### 4.1 Name and format your dataset\n",
    "\n",
    "Please provide a unique name for your dataset, usually the name of the singer/speaker (whether real or virtual). For example, `opencpop` will be a good name for the dataset. You can also add tags to represent dataset version, model capacity or improvements. For example, `v2` represents the version, `large` represents the capacity, and `fix_br` means you fixed breaths since your trained last model.\n",
    "\n",
    "Please edit the following cell before you run it. Remember only using letters, numbers and underlines (`_`).\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false,
    "jupyter": {
     "outputs_hidden": false
    }
   },
   "outputs": [],
   "source": [
    "########################################\n",
    "\n",
    "# Name and tags of your dataset\n",
    "dataset_name = '???'  # Required\n",
    "dataset_tags = ''  # Optional\n",
    "\n",
    "########################################\n",
    "\n",
    "import random\n",
    "import re\n",
    "\n",
    "from textgrid import TextGrid\n",
    "\n",
    "assert dataset_name != '', 'Dataset name cannot be empty.'\n",
    "assert re.search(r'[^0-9A-Za-z_]', dataset_name) is None, 'Dataset name contains invalid characters.'\n",
    "full_name = dataset_name\n",
    "if dataset_tags != '':\n",
    "    assert re.search(r'[^0-9A-Za-z_]', dataset_name) is None, 'Dataset tags contain invalid characters.'\n",
    "    full_name += f'_{dataset_tags}'\n",
    "assert not os.path.exists(f'../data/{full_name}'), f'The name \\'{full_name}\\' already exists in your \\'data\\' folder!'\n",
    "\n",
    "print('Dataset name:', dataset_name)\n",
    "if dataset_tags != '':\n",
    "    print('Tags:', dataset_tags)\n",
    "\n",
    "formatted_path = f'../data/{full_name}/raw/wavs'\n",
    "os.makedirs(formatted_path)\n",
    "transcriptions = []\n",
    "samplerate = 44100\n",
    "min_sil = int(0.1 * samplerate)\n",
    "max_sil = int(0.5 * samplerate)\n",
    "for wavfile in tqdm.tqdm(sliced_filelist):\n",
    "    name = os.path.basename(wavfile).rsplit('.', maxsplit=1)[0]\n",
    "    y, _ = librosa.load(wavfile, sr=samplerate, mono=True)\n",
    "    tg = TextGrid()\n",
    "    tg.read(os.path.join(textgrids_revised_dir, f'{name}.TextGrid'))\n",
    "    ph_seq = [ph.mark for ph in tg[1]]\n",
    "    ph_dur = [ph.maxTime - ph.minTime for ph in tg[1]]\n",
    "    if random.random() < 0.5:\n",
    "        len_sil = random.randrange(min_sil, max_sil)\n",
    "        y = np.concatenate((np.zeros((len_sil,), dtype=np.float32), y))\n",
    "        if ph_seq[0] == 'SP':\n",
    "            ph_dur[0] += len_sil / samplerate\n",
    "        else:\n",
    "            ph_seq.insert(0, 'SP')\n",
    "            ph_dur.insert(0, len_sil / samplerate)\n",
    "    if random.random() < 0.5:\n",
    "        len_sil = random.randrange(min_sil, max_sil)\n",
    "        y = np.concatenate((y, np.zeros((len_sil,), dtype=np.float32)))\n",
    "        if ph_seq[-1] == 'SP':\n",
    "            ph_dur[-1] += len_sil / samplerate\n",
    "        else:\n",
    "            ph_seq.append('SP')\n",
    "            ph_dur.append(len_sil / samplerate)\n",
    "    ph_seq = ' '.join(ph_seq)\n",
    "    ph_dur = ' '.join([str(round(d, 6)) for d in ph_dur])\n",
    "    soundfile.write(os.path.join(formatted_path, f'{name}.wav'), y, samplerate)\n",
    "    transcriptions.append(f'{name}|啊|{ph_seq}|rest|0|{ph_dur}|0')\n",
    "with open(f'../data/{full_name}/raw/transcriptions.txt', 'w', encoding='utf8') as f:\n",
    "    print('\\n'.join(transcriptions), file=f)\n",
    "print(f'All wavs and transcriptions saved at \\'data/{full_name}/raw/\\'.')\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now that the dataset and transcriptions have been saved, you can run the following cell to clean up all temporary files generated by pipelines above.\n",
    "\n",
    "<font color=\"yellow\">WARNING</font>: This will remove `pipelines/segments/` and `pipelines/textgrids/` folders. You should specify a directory in the following cell to back up your TextGrids if you want them for future use.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false,
    "jupyter": {
     "outputs_hidden": false
    }
   },
   "outputs": [],
   "source": [
    "########################################\n",
    "\n",
    "# Optional path to back up your TextGrids\n",
    "textgrids_backup_path = r''  # If left empty, the TextGrids will not be backed up\n",
    "\n",
    "########################################\n",
    "\n",
    "assert textgrids_backup_path == '' or not os.path.exists(textgrids_backup_path) or os.path.isdir(textgrids_backup_path), 'The backup path is not a directory.'\n",
    "\n",
    "if textgrids_backup_path != '':\n",
    "    os.makedirs(textgrids_backup_path, exist_ok=True)\n",
    "    for tg in tqdm.tqdm(glob.glob(f'{textgrids_revised_dir}/*.TextGrid')):\n",
    "        filename = os.path.basename(tg)\n",
    "        shutil.copy(tg, os.path.join(textgrids_backup_path, filename))\n",
    "\n",
    "shutil.rmtree(segments_dir)\n",
    "shutil.rmtree(textgrids_dir)\n",
    "print('Cleaning up done.')\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "tags": []
   },
   "source": [
    "### 4.2 Configuring parameters\n",
    "\n",
    "Here you can configure some parameters for preprocessing, training and the neural networks. Read the explanations below, edit parameters according to your preference, and run the following cell to generate the configuration file.\n",
    "\n",
    "> You do not need to actually run this cell if you are building multiple datasets to train combined models. In this case, please see instructions in **Section 5**.\n",
    "\n",
    "#### 4.2.1 The neural networks\n",
    "\n",
    "These parameters control the capacity and structure of the neural networks.\n",
    "\n",
    "##### `residual_channels` and `residual_layers`\n",
    "\n",
    "These two hyperparameters refer to the width and the depth of the diffusion decoder network. Generally speaking, `384x20` represents a `base` model capacity and `512x20` represents a `large` model capacity. `384x30` is also a reasonable choice. Larger models consumes more GPU memory and runs slower at training and inference time, but they produce better results.\n",
    "\n",
    "GPU memory required for training:\n",
    "\n",
    "Base model - at least 6 GB (12 GB recommended)\n",
    "Large model - at least 12 GB (24 GB recommended)\n",
    "\n",
    "##### `f0_embed_type`\n",
    "\n",
    "There are two ways of f0 embedding: `discrete` and `continuous`. The discrete embedding takes 256 bins from 50Hz to 1100Hz and embeds the discrete f0 bins with `torch.nn.Embedding`, while the continuous embedding converts f0 to continuous mel frequency and embeds these values with `torch.nn.Linear`.\n",
    "\n",
    "The discrete embedding has been tested for long and is guaranteed to be stable. The continuous embedding has not been widely tested, but some improvements at the extreme edges of the pitch range were observed.\n",
    "\n",
    "#### 4.2.2 Data augmentation\n",
    "\n",
    "Data augmentation improves the performance or functionalities of your model, **but may boost the size of your training data**.\n",
    "\n",
    "##### `random_pitch_shifting`\n",
    "\n",
    "Once enabled, the pitch of your data will be randomly shifted without keeping the formant when preprocessing. Besides, the number of keys that each piece of data is shifted will be embedded into the networks. This will broaden the range of pitch and allows you to control the frame-level key shift values (like the GEN parameter in VOCALOID) at inference time.\n",
    "\n",
    "This type of augmentation accepts the following arguments:\n",
    "- `range` controls the range of pitch keys that will be randomly shifted.\n",
    "- `scale` controls the amount of data that the augmentation will be applied to.\n",
    "\n",
    "$ D_{augmentation} \\approx (1 + scale) \\cdot D_{original} $\n",
    "\n",
    "##### `fixed_pitch_shifting`\n",
    "\n",
    "Once enabled, the pitch of your data will be shifted several numbers of keys. These data with pitch shifting will be marked as different speakers from the original data, **thus making the model a multi-speaker combined model**. This will also broaden the range of pitch (maybe slightly better than random pitch shifting). **This augmentation is not compatible with random pitch shifting.**\n",
    "\n",
    "This type of augmentation accepts the following arguments:\n",
    "- `targets` controls the number of targets of pitch shifting and the number of keys that will be shifted towards each target.\n",
    "- `scale` controls the amount of data that **each target** of the augmentation will be applied to.\n",
    "\n",
    "$ D_{augmentation} \\approx (1 + N \\cdot scale) \\cdot D_{original} $ , where $ N $ represents the number of targets.\n",
    "\n",
    "##### `random_time_stretching`\n",
    "\n",
    "Once enabled, the speed of your data will be randomly changed when preprocessing. The ratio of the speed change will be embedded into the networks, which allows you to control the frame-level speed or velocity (similar to but much more flexible than the VEL parameter in VOCALOID) at inference time. In other words, by applying global time stretching at training time, you gain the ability to apply local time stretching at inference time. This can be used to adjust the texture of consonants and the ratio of different parts of vowels. **Some audio segments will be longer after this augmentation is applied. Please be careful of your batch size and your GPU memory usage.**\n",
    "\n",
    "This type of augmentation accepts the following arguments:\n",
    "- `range` controls the range of the speed changing ratio.\n",
    "- `domain` determines in which domain the speed ratio will follow a uniform distribution in: `log` or `linear`\n",
    "- `scale` controls the amount of data that the augmentation will be applied to.\n",
    "\n",
    "$ D_{augmentation} \\approx (1 + scale \\cdot \\frac{1}{b - a} \\cdot \\int_{a}^{b} f(x) dx) \\cdot D_{original} $ , where $ a, b $ represents the range of the speed ratio, $ f(x) $ represents the PDF of the speed ratio.\n",
    "\n",
    "---\n",
    "> When there are more than two types of augmentation enabled, a cascade and joint augmentation scaling algorithm is applied. Briefly speaking, the following rules will be satisfied after applying and combining multiple types of augmentation:\n",
    "> 1. The number of data pieces applied with the $ k $th augmentation will be $ scale_{k} $ times than those not applied with the $ k $th augmentation.\n",
    "> 2. The number of data pieces applied with at least one type of augmentation will be $ \\sum_{i = 1}^{n} scale_{i} $ times than those not applied with any augmentation (purely raw data).\n",
    "\n",
    "#### 4.2.3 Training and validating\n",
    "\n",
    "##### `test_prefixes`\n",
    "\n",
    "All files with name prefixes specified in this list will be put into the test set. Each time when a checkpoint is saved, the program will first run inference on the test set and put the result on the TensorBoard. Thus, you can listen to these demos and judge the quality of your model. If you leave it empty, test cases will be randomly selected.\n",
    "\n",
    "##### `max_tokens` and `max_sentences`\n",
    "\n",
    "These two parameters jointly determine the batch size at training time, the former representing maximum number of frames in one batch and the latter limiting the maximum batch size. Larger batches consumes more GPU memory at training time. This value can be adjusted according to your GPU memory. Remember not to set this value too low because the model may not converge with small batches.\n",
    "\n",
    "##### `lr`, `decay_steps`, `gamma`\n",
    "\n",
    "The learning rate starts at `lr`, decays with the rate `gamma` at every `decay_steps` during training. If you decreased your batch size, you may consider using a smaller learning rate and more decay steps, or larger gamma.\n",
    "\n",
    "##### `val_check_interval`, `num_ckpt_keep` and `max_updates`\n",
    "\n",
    "These three values refer to the training steps between validating and saving checkpoints, the number of the most recent checkpoints reserved, and the maximum training steps. With default batch size and 5 hours of training data, 250k ~ 350k training steps is reasonable. If you decrease the batch size, you may increase the training steps.\n",
    "\n",
    "##### `permanent_ckpt_start` and `permanent_ckpt_interval`\n",
    "\n",
    "These two values help you save permanent checkpoints during training. Normally, old checkpoints will be removed and only the newest `num_ckpt_keep` checkpoints will be kept. With these two values, the program will save a permanent checkpoints each `permanent_ckpt_interval` steps after the training reaches `permenant_ckpt_start` steps. Permanent checkpoints will never be removed and can be used for comparison between different training steps and as backups in case of over-fitting.\n",
    "\n",
    "To enable permanent checkpoints, please ensure these values are positive multiples of `val_check_interval`. To disable permanent checkpoints, please set `permanent_ckpt_interval` to `-1`.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false,
    "jupyter": {
     "outputs_hidden": false
    }
   },
   "outputs": [],
   "source": [
    "########################################\n",
    "\n",
    "# The neural networks\n",
    "residual_channels = 512\n",
    "residual_layers = 20\n",
    "f0_embed_type = 'continuous'\n",
    "\n",
    "# Data augmentation\n",
    "random_pitch_shifting = {\n",
    "    'enabled': False,\n",
    "    'range': [-5., 5.],\n",
    "    'scale': 2.\n",
    "}\n",
    "fixed_pitch_shifting = {\n",
    "    'enabled': False,\n",
    "    'targets': [-5., 5.],\n",
    "    'scale': 0.75\n",
    "}\n",
    "random_time_stretching = {\n",
    "    'enabled': False,\n",
    "    'range': [0.5, 2.],\n",
    "    'domain': 'log',  # or linear\n",
    "    'scale': 2.\n",
    "}\n",
    "\n",
    "# Training and validating\n",
    "test_prefixes = [\n",
    "\n",
    "]\n",
    "\n",
    "max_tokens = 80000\n",
    "max_sentences = 48\n",
    "\n",
    "lr = 0.0004\n",
    "decay_steps = 50000\n",
    "gamma = 0.5\n",
    "\n",
    "val_check_interval = 2000\n",
    "num_ckpt_keep = 5\n",
    "max_updates = 320000\n",
    "permanent_ckpt_start = 120000\n",
    "permanent_ckpt_interval = 40000\n",
    "\n",
    "########################################\n",
    "\n",
    "from copy import deepcopy\n",
    "import datetime\n",
    "import random\n",
    "\n",
    "import yaml\n",
    "\n",
    "training_cases = [os.path.basename(w).rsplit('.', maxsplit=1)[0] for w in sliced_filelist]\n",
    "valid_test_cases = []\n",
    "if len(test_prefixes) > 0:\n",
    "    for prefix in deepcopy(test_prefixes):\n",
    "        if prefix in training_cases:\n",
    "            valid_test_cases.append(training_cases)\n",
    "            test_prefixes.remove(prefix)\n",
    "            training_cases.remove(prefix)\n",
    "    i = 0\n",
    "    while i < len(training_cases):\n",
    "        for prefix in test_prefixes:\n",
    "            if training_cases[i].startswith(prefix):\n",
    "                valid_test_cases.append(training_cases[i])\n",
    "                training_cases.pop(i)\n",
    "                i -= 1\n",
    "                break\n",
    "        i += 1\n",
    "else:\n",
    "    test_prefixes += sorted(random.sample(training_cases, 10))\n",
    "\n",
    "configs = {\n",
    "    'base_config': ['configs/acoustic/nomidi.yaml'],\n",
    "    'speakers': [dataset_name],\n",
    "    'raw_data_dir': [f'data/{full_name}/raw'],\n",
    "    'binary_data_dir': f'data/{full_name}/binary',\n",
    "    'residual_channels': residual_channels,\n",
    "    'residual_layers': residual_layers,\n",
    "    'f0_embed_type': f0_embed_type,\n",
    "    'test_prefixes': test_prefixes,\n",
    "    'max_tokens': max_tokens,\n",
    "    'max_sentences': max_sentences,\n",
    "    'lr': lr,\n",
    "    'decay_steps': decay_steps,\n",
    "    'gamma': gamma,\n",
    "    'val_check_interval': val_check_interval,\n",
    "    'num_valid_plots': min(10, len(test_prefixes)),\n",
    "    'num_ckpt_keep': num_ckpt_keep,\n",
    "    'max_updates': max_updates,\n",
    "    'permanent_ckpt_start': permanent_ckpt_start,\n",
    "    'permanent_ckpt_interval': permanent_ckpt_interval\n",
    "}\n",
    "\n",
    "augmentation_args = {}\n",
    "if random_pitch_shifting['enabled']:\n",
    "    augmentation_args['random_pitch_shifting'] = {\n",
    "        'range': random_pitch_shifting['range'],\n",
    "        'scale': random_pitch_shifting['scale']\n",
    "    }\n",
    "    configs['use_key_shift_embed'] = True\n",
    "if fixed_pitch_shifting['enabled']:\n",
    "    augmentation_args['fixed_pitch_shifting'] = {\n",
    "        'targets': fixed_pitch_shifting['targets'],\n",
    "        'scale': fixed_pitch_shifting['scale']\n",
    "    }\n",
    "    configs['use_spk_id'] = True\n",
    "    configs['num_spk'] = 1 + len(fixed_pitch_shifting['targets'])\n",
    "if random_time_stretching['enabled']:\n",
    "    augmentation_args['random_time_stretching'] = {\n",
    "        'range': random_time_stretching['range'],\n",
    "        'domain': random_time_stretching['domain'],\n",
    "        'scale': random_time_stretching['scale']\n",
    "    }\n",
    "    configs['use_speed_embed'] = True\n",
    "configs['augmentation_args'] = augmentation_args\n",
    "\n",
    "with open(f'../data/{full_name}/config.yaml', 'w', encoding='utf8') as f:\n",
    "    yaml.dump(configs, f, sort_keys=False, allow_unicode=True)\n",
    "\n",
    "date = datetime.datetime.now().strftime('%m%d')\n",
    "exp_name = f'{date}_{dataset_name}_ds1000'\n",
    "if dataset_tags != '':\n",
    "    exp_name += f'_{dataset_tags}'\n",
    "print('Congratulations! All steps have been done and you are now prepared to train your own model.\\n'\n",
    "      'Before you start, please read and follow instructions in the repository README.\\n'\n",
    "      'Here are the commands for you to copy that you can run preprocessing and training:\\n')\n",
    "\n",
    "print('============ Linux ============\\n'\n",
    "      'export PYTHONPATH=.\\n'\n",
    "      'export CUDA_VISIBLE_DEVICES=0\\n'\n",
    "      f'python data_gen/binarize.py --config data/{full_name}/config.yaml\\n'\n",
    "      f'python run.py --config data/{full_name}/config.yaml --exp_name {exp_name} --reset\\n')\n",
    "\n",
    "print('===== Windows (PowerShell) =====\\n'\n",
    "      '$env:PYTHONPATH=\".\"\\n'\n",
    "      '$env:CUDA_VISIBLE_DEVICES=0\\n'\n",
    "      f'python data_gen/binarize.py --config data/{full_name}/config.yaml\\n'\n",
    "      f'python run.py --config data/{full_name}/config.yaml --exp_name {exp_name} --reset\\n')\n",
    "\n",
    "print('===== Windows (Command Prompt) =====\\n'\n",
    "      'set PYTHONPATH=.\\n'\n",
    "      'set CUDA_VISIBLE_DEVICES=0\\n'\n",
    "      f'python data_gen/binarize.py --config data/{full_name}/config.yaml\\n'\n",
    "      f'python run.py --config data/{full_name}/config.yaml --exp_name {exp_name} --reset\\n')\n",
    "\n",
    "print(f'If you want to train your model on another machine (like a remote GPU), please copy the whole \\'data/{full_name}/\\' folder.')\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 5 (Additional section) Configuring multiple datasets for combined models\n",
    "\n",
    "If you have multiple datasets, you can train a combined model with them for better performance and case coverage. In addition, you gain ability to switch between different speakers/singers/styles/timbres, or even mix any of them in any proportions, static or dynamic over time, via one single model!\n",
    "\n",
    "This section will guide you to create configuration files for this kind of models, so that you can run joint preprocessing and training from multiple datasets.\n",
    "\n",
    "<font color=\"#66ccff\">NOTICE</font>: Before you train a combined model, you must obtain permission from copyright holders of each dataset and make sure all these providers are fully aware that you will train a combined model from their data, that you will or will not distribute the synthesized voices and model weights, that their voices will or will not be mixed, and the potential risks of this kind of activity.\n",
    "\n",
    "### 5.1 Selecting datasets\n",
    "\n",
    "In the following cell, you can choose the datasets and name your combined model.\n",
    "\n",
    "##### `model_name` and `model_tags`\n",
    "\n",
    "Similar to `dataset_name` and `dataset_tags` described in Section 4.1, but here you are able to use `+` in your model name. For example, `female_triplet` and `alice+bob` are nice names for a combined model.\n",
    "\n",
    "##### `datasets`\n",
    "\n",
    "Selection of datasets. You specify datasets using their full names and set a speaker name for each dataset. Here is an example:\n",
    "\n",
    "```python\n",
    "datasets = [\n",
    "    {\n",
    "        'dataset': 'alice',\n",
    "        'speaker': 'Alice-Normal'\n",
    "    },\n",
    "    {\n",
    "        'dataset': 'alice_sweet',\n",
    "        'speaker': 'Alice-Sweet'\n",
    "    },\n",
    "    {\n",
    "        'dataset': 'bob_v2',\n",
    "        'speaker': 'Bob'\n",
    "    },\n",
    "]\n",
    "```\n",
    "You must type in full names of datasets (folder name in the `data/` directory). You may use letters, numbers, underline (`_`) and hyphens (`-`) in the speaker name.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false,
    "jupyter": {
     "outputs_hidden": false
    }
   },
   "outputs": [],
   "source": [
    "########################################\n",
    "\n",
    "# Name and tags of your combined model\n",
    "model_name = '???'  # Required\n",
    "model_tags = ''  # Optional\n",
    "\n",
    "datasets = [\n",
    "    {\n",
    "        'dataset': '???',\n",
    "        'speaker': '???'\n",
    "    },\n",
    "]\n",
    "\n",
    "########################################\n",
    "\n",
    "import re\n",
    "\n",
    "assert model_name != '', 'Model name cannot be empty.'\n",
    "assert re.search(r'[^0-9A-Za-z_\\+]', model_name) is None, 'Model name contains invalid characters.'\n",
    "model_full_name = model_name\n",
    "if model_tags != '':\n",
    "    assert re.search(r'[^0-9A-Za-z_]', model_name) is None, 'Dataset tags contain invalid characters.'\n",
    "    model_full_name += f'_{model_tags}'\n",
    "assert not os.path.exists(f'../data/{model_full_name}'), f'The name \\'{model_full_name}\\' already exists in your \\'data\\' folder!'\n",
    "\n",
    "speakers = []\n",
    "raw_data_dirs = []\n",
    "for selection in datasets:\n",
    "    assert selection['dataset'] in os.listdir('../data/'), f'Dataset \\'{selection[\"dataset\"]}\\' not found.'\n",
    "    assert os.path.exists(f'../data/{selection[\"dataset\"]}/raw/wavs'), f'Wave directory not found in dataset \\'{selection[\"dataset\"]}\\''\n",
    "    assert os.path.exists(f'../data/{selection[\"dataset\"]}/raw/transcriptions.txt'), f'Transcriptions not found in dataset \\'{selection[\"dataset\"]}\\''\n",
    "    assert re.search(r'[^0-9A-Za-z_-]', selection['speaker']) is None, 'Speaker name contains invalid characters.'\n",
    "    speakers.append(selection['speaker'])\n",
    "    raw_data_dirs.append(f'data/{selection[\"dataset\"]}/raw')\n",
    "\n",
    "print('Model name:', model_name)\n",
    "if model_tags != '':\n",
    "    print('Tags:', model_tags)\n",
    "os.makedirs(f'../data/{model_full_name}/')\n",
    "print(f'Created \\'data/{model_full_name}/\\'')\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 5.2 Configuring parameters\n",
    "\n",
    "Most parameters for combined models are the same as described in Section 4.2, except the following differences.\n",
    "\n",
    "##### `random_pitch_shifting`, `fixed_pitch_shifting` and `random_time_stretching`\n",
    "\n",
    "Please pay attention to the size of your training data!\n",
    "\n",
    "##### `test_prefixes`\n",
    "\n",
    "You can use prefixes or full names of wave file, with or without speaker id in this parameter. For example:\n",
    "- `xxx` will fully match one single filename in any of selected datasets, or if there is none, match all filenames starting with `xxx`.\n",
    "- `0:xxx` will fully match one single filename in **the first** dataset, or if there is none, match all filenames starting with `xxx` **in that dataset**.\n",
    "\n",
    "If not specified, test cases will be randomly selected from all datasets.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false,
    "jupyter": {
     "outputs_hidden": false
    }
   },
   "outputs": [],
   "source": [
    "########################################\n",
    "\n",
    "# The neural networks\n",
    "residual_channels = 512\n",
    "residual_layers = 20\n",
    "f0_embed_type = 'continuous'\n",
    "\n",
    "# Data augmentation\n",
    "random_pitch_shifting = {\n",
    "    'enabled': False,\n",
    "    'range': [-5., 5.],\n",
    "    'scale': 1.5\n",
    "}\n",
    "fixed_pitch_shifting = {\n",
    "    'enabled': False,\n",
    "    'targets': [-5., 5.],\n",
    "    'scale': 0.75\n",
    "}\n",
    "random_time_stretching = {\n",
    "    'enabled': False,\n",
    "    'range': [0.5, 2.],\n",
    "    'domain': 'log',  # or linear\n",
    "    'scale': 1.5\n",
    "}\n",
    "\n",
    "# Training and validating\n",
    "test_prefixes = [\n",
    "\n",
    "]\n",
    "\n",
    "max_tokens = 80000\n",
    "max_sentences = 48\n",
    "\n",
    "lr = 0.0004\n",
    "decay_steps = 50000\n",
    "gamma = 0.5\n",
    "\n",
    "val_check_interval = 2000\n",
    "num_ckpt_keep = 5\n",
    "max_updates = 320000\n",
    "permanent_ckpt_start = 120000\n",
    "permanent_ckpt_interval = 40000\n",
    "\n",
    "########################################\n",
    "\n",
    "import datetime\n",
    "import random\n",
    "\n",
    "import yaml\n",
    "\n",
    "training_cases = []\n",
    "for raw_data_dir in raw_data_dirs:\n",
    "    with open(os.path.join('..', raw_data_dir, 'transcriptions.txt'), 'r', encoding='utf8') as f:\n",
    "        training_cases.append([line.split('|')[0] for line in f.readlines()])\n",
    "valid_test_cases = []\n",
    "if len(test_prefixes) > 0:\n",
    "    for prefix in deepcopy(test_prefixes):\n",
    "        if prefix.contains(':'):\n",
    "            idx, fn = prefix.split(':')\n",
    "            idx = int(idx)\n",
    "            if fn in training_cases[idx]:\n",
    "                valid_test_cases.append(prefix)\n",
    "                test_prefixes.remove(prefix)\n",
    "                training_cases[idx].remove(fn)\n",
    "        else:\n",
    "            for idx, cases in training_cases:\n",
    "                if prefix in cases:\n",
    "                    valid_test_cases.append(f'{idx}:{prefix}')\n",
    "                    test_prefixes.remove(prefix)\n",
    "                    cases.remove(prefix)\n",
    "\n",
    "    for prefix in deepcopy(test_prefixes):\n",
    "        if prefix.contains(':'):\n",
    "            idx, fn = prefix.split(':')\n",
    "            idx = int(idx)\n",
    "            for case in deepcopy(training_cases[idx]):\n",
    "                if case.startswith(fn):\n",
    "                    valid_test_cases.append(prefix)\n",
    "                    training_cases[idx].remove(case)\n",
    "        else:\n",
    "            for idx, cases in training_cases:\n",
    "                for case in deepcopy(cases):\n",
    "                    if case.startswith(prefix):\n",
    "                        valid_test_cases.append(f'{idx}:{case}')\n",
    "                        cases.remove(case)\n",
    "else:\n",
    "    total = min(20, max(10, 4 * len(datasets)))\n",
    "    quotient, remainder = total // len(datasets), total % len(datasets)\n",
    "    if quotient == 0:\n",
    "        test_counts = [1] * len(datasets)\n",
    "    else:\n",
    "        test_counts = [quotient + 1] * remainder + [quotient] * (len(datasets) - remainder)\n",
    "    for i, count in enumerate(test_counts):\n",
    "        test_prefixes += [f'{i}:{n}' for n in sorted(random.sample(training_cases[i], count))]\n",
    "\n",
    "configs = {\n",
    "    'base_config': ['configs/acoustic/nomidi.yaml'],\n",
    "    'speakers': speakers,\n",
    "    'num_spk': len(speakers),\n",
    "    'use_spk_id': True,\n",
    "    'raw_data_dir': raw_data_dirs,\n",
    "    'binary_data_dir': f'data/{model_full_name}/binary',\n",
    "    'residual_channels': residual_channels,\n",
    "    'residual_layers': residual_layers,\n",
    "    'f0_embed_type': f0_embed_type,\n",
    "    'test_prefixes': test_prefixes,\n",
    "    'max_tokens': max_tokens,\n",
    "    'max_sentences': max_sentences,\n",
    "    'lr': lr,\n",
    "    'decay_steps': decay_steps,\n",
    "    'gamma': gamma,\n",
    "    'val_check_interval': val_check_interval,\n",
    "    'num_valid_plots': min(20, len(test_prefixes)),\n",
    "    'num_ckpt_keep': num_ckpt_keep,\n",
    "    'max_updates': max_updates,\n",
    "    'permanent_ckpt_start': permanent_ckpt_start,\n",
    "    'permanent_ckpt_interval': permanent_ckpt_interval\n",
    "}\n",
    "\n",
    "augmentation_args = {}\n",
    "if random_pitch_shifting['enabled']:\n",
    "    augmentation_args['random_pitch_shifting'] = {\n",
    "        'range': random_pitch_shifting['range'],\n",
    "        'scale': random_pitch_shifting['scale']\n",
    "    }\n",
    "    configs['use_key_shift_embed'] = True\n",
    "if fixed_pitch_shifting['enabled']:\n",
    "    augmentation_args['fixed_pitch_shifting'] = {\n",
    "        'targets': fixed_pitch_shifting['targets'],\n",
    "        'scale': fixed_pitch_shifting['scale']\n",
    "    }\n",
    "    configs['use_spk_id'] = True\n",
    "    configs['num_spk'] *= (1 + len(fixed_pitch_shifting['targets']))\n",
    "if random_time_stretching['enabled']:\n",
    "    augmentation_args['random_time_stretching'] = {\n",
    "        'range': random_time_stretching['range'],\n",
    "        'domain': random_time_stretching['domain'],\n",
    "        'scale': random_time_stretching['scale']\n",
    "    }\n",
    "    configs['use_speed_embed'] = True\n",
    "configs['augmentation_args'] = augmentation_args\n",
    "\n",
    "with open(f'../data/{model_full_name}/config.yaml', 'w', encoding='utf8') as f:\n",
    "    yaml.dump(configs, f, sort_keys=False, allow_unicode=True)\n",
    "\n",
    "date = datetime.datetime.now().strftime('%m%d')\n",
    "exp_name = f'{date}_{model_name}_ds1000'\n",
    "if model_tags != '':\n",
    "    exp_name += f'_{model_tags}'\n",
    "print('Congratulations! All steps have been done and you are now prepared to train your combined model.\\n'\n",
    "      'Before you start, please read and follow instructions in the repository README.\\n'\n",
    "      'Here are the commands for you to copy that you can run preprocessing and training:\\n')\n",
    "\n",
    "print('============ Linux ============\\n'\n",
    "      'export PYTHONPATH=.\\n'\n",
    "      'export CUDA_VISIBLE_DEVICES=0\\n'\n",
    "      f'python data_gen/binarize.py --config data/{model_full_name}/config.yaml\\n'\n",
    "      f'python run.py --config data/{model_full_name}/config.yaml --exp_name {exp_name} --reset\\n')\n",
    "\n",
    "print('===== Windows (PowerShell) =====\\n'\n",
    "      '$env:PYTHONPATH=\".\"\\n'\n",
    "      '$env:CUDA_VISIBLE_DEVICES=0\\n'\n",
    "      f'python data_gen/binarize.py --config data/{model_full_name}/config.yaml\\n'\n",
    "      f'python run.py --config data/{model_full_name}/config.yaml --exp_name {exp_name} --reset\\n')\n",
    "\n",
    "print('===== Windows (Command Prompt) =====\\n'\n",
    "      'set PYTHONPATH=.\\n'\n",
    "      'set CUDA_VISIBLE_DEVICES=0\\n'\n",
    "      f'python data_gen/binarize.py --config data/{model_full_name}/config.yaml\\n'\n",
    "      f'python run.py --config data/{model_full_name}/config.yaml --exp_name {exp_name} --reset\\n')\n",
    "\n",
    "print(f'To preprocess the selected datasets, please make sure these directories exist:')\n",
    "for d in raw_data_dirs:\n",
    "    print(f' - {d}')\n",
    "print()\n",
    "print(f'If you want to train your model on another machine (like a remote GPU), please copy the whole \\'data/{model_full_name}/\\' folder.')\n"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.8.15"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}