mirror of
https://github.com/storytold/storyteller-ml.git
synced 2026-10-12 23:48:57 +00:00
Brandon Thomas
925 lines
36 KiB
Plaintext
925 lines
36 KiB
Plaintext
{
|
|
"cells": [
|
|
{
|
|
"cell_type": "markdown",
|
|
"source": [
|
|
"# TalkNet Training (offline)\n",
|
|
"\n",
|
|
"To train a 22KHz TalkNet, run the cells below and follow the instructions.\n",
|
|
"\n",
|
|
"The notebook will automatically resume training any models from the last saved checkpoint. If you're resuming from a new session, always re-run step 1 first."
|
|
],
|
|
"metadata": {
|
|
"id": "Gss5Ox_RNiba"
|
|
}
|
|
},
|
|
{
|
|
"cell_type": "markdown",
|
|
"source": [
|
|
"**Step 1:** Configure training data paths. Create the following and change the paths below:\n",
|
|
"* A dataset of .wav files, packaged as a .zip or .tar file\n",
|
|
"* Training and validation filelists, in LJSpeech format with relative paths (note: ARPABET transcripts are not supported)\n",
|
|
"* An output path for checkpoints"
|
|
],
|
|
"metadata": {}
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": null,
|
|
"source": [
|
|
"### CHANGE THESE ###\n",
|
|
"\n",
|
|
"dataset = \"/absolute/path_to_dataset.zip\"\n",
|
|
"train_filelist = \"/absolute/train_filelist.txt\"\n",
|
|
"val_filelist = \"/absolute/val_filelist.txt\"\n",
|
|
"output_dir = \"/absolute/name_of_character\"\n",
|
|
"\n",
|
|
"### ------------ ###\n",
|
|
"\n",
|
|
"import os\n",
|
|
"import gdown\n",
|
|
"import platform\n",
|
|
"\n",
|
|
"assert os.path.exists(dataset), \"Cannot find dataset\"\n",
|
|
"assert os.path.exists(train_filelist), \"Cannot find training filelist\"\n",
|
|
"assert os.path.exists(val_filelist), \"Cannot find validation filelist\"\n",
|
|
"if not os.path.exists(output_dir):\n",
|
|
" os.makedirs(output_dir)\n",
|
|
"\n",
|
|
"if os.getcwd() != output_dir:\n",
|
|
" cwd = os.getcwd()\n",
|
|
"else:\n",
|
|
" os.chdir(cwd)\n",
|
|
" \n",
|
|
"# Download pre-trained models\n",
|
|
"if not os.path.exists(\"talknet_spect.nemo\"):\n",
|
|
" print(\"Downloading pre-trained models...\")\n",
|
|
" zip_path = \"TalkNet_Pretrained.zip\"\n",
|
|
" if not os.path.exists(zip_path) or os.stat(zip_path).st_size < 100:\n",
|
|
" d = 'https://drive.google.com/uc?id='\n",
|
|
" gdown.download(d+\"19wSym9mNEnmzLS9XdPlfNAW9_u-mP1hR\", zip_path, quiet=False)\n",
|
|
" print(platform.system())\n",
|
|
" if \"Windows\" in platform.system():\n",
|
|
" !tar -xf {zip_path}\n",
|
|
" else:\n",
|
|
" !unzip -qo {zip_path}\n",
|
|
" os.remove(zip_path)\n",
|
|
"\n",
|
|
"print(\"OK\")"
|
|
],
|
|
"outputs": [],
|
|
"metadata": {
|
|
"cellView": "form",
|
|
"id": "nfSawDUD5tqv"
|
|
}
|
|
},
|
|
{
|
|
"cell_type": "markdown",
|
|
"source": [
|
|
"**Step 2:** Dataset processing, part 1.\n",
|
|
"\n",
|
|
"If this step fails, try the following:\n",
|
|
"* Make sure your filelists are correct. They should have relative \n",
|
|
"paths that match the contents of the archive."
|
|
],
|
|
"metadata": {}
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": null,
|
|
"source": [
|
|
"import os\n",
|
|
"import shutil\n",
|
|
"import sys\n",
|
|
"import json\n",
|
|
"import nemo\n",
|
|
"import torch\n",
|
|
"import torchaudio\n",
|
|
"import numpy as np\n",
|
|
"#from pysptk import sptk\n",
|
|
"from pathlib import Path\n",
|
|
"from tqdm.notebook import tqdm\n",
|
|
"import ffmpeg\n",
|
|
"\n",
|
|
"def fix_transcripts(inpath):\n",
|
|
" found_arpabet = False\n",
|
|
" found_grapheme = False\n",
|
|
" with open(inpath, \"r\", encoding=\"utf8\") as f:\n",
|
|
" lines = f.readlines()\n",
|
|
" with open(inpath, \"w\", encoding=\"utf8\") as f:\n",
|
|
" for l in lines:\n",
|
|
" if l.strip() == \"\":\n",
|
|
" continue\n",
|
|
" if \"{\" in l:\n",
|
|
" if not found_arpabet:\n",
|
|
" print(\"Warning: Skipping ARPABET lines (not supported).\")\n",
|
|
" found_arpabet = True\n",
|
|
" else:\n",
|
|
" f.write(l)\n",
|
|
" found_grapheme = True\n",
|
|
" assert found_grapheme, \"No non-ARPABET lines found in \" + inpath\n",
|
|
"\n",
|
|
"def generate_json(inpath, outpath):\n",
|
|
" output = \"\"\n",
|
|
" sample_rate = 22050\n",
|
|
" with open(inpath, \"r\", encoding=\"utf8\") as f:\n",
|
|
" for l in f.readlines():\n",
|
|
" lpath = l.split(\"|\")[0].strip()\n",
|
|
" if lpath[:5] != \"wavs/\":\n",
|
|
" lpath = \"wavs/\" + lpath\n",
|
|
" lpath = os.path.join(output_dir, \"_training_data\", lpath)\n",
|
|
" size = os.stat(lpath).st_size\n",
|
|
" x = {\n",
|
|
" \"audio_filepath\": lpath,\n",
|
|
" \"duration\": size / (sample_rate * 2),\n",
|
|
" \"text\": l.split(\"|\")[1].strip(),\n",
|
|
" }\n",
|
|
" output += json.dumps(x) + \"\\n\"\n",
|
|
" with open(outpath, \"w\", encoding=\"utf8\") as w:\n",
|
|
" w.write(output)\n",
|
|
"\n",
|
|
"def convert_to_22k(inpath):\n",
|
|
" if inpath.strip()[-4:].lower() != \".wav\":\n",
|
|
" print(\"Warning: \" + inpath.strip() + \" is not a .wav file!\")\n",
|
|
" return\n",
|
|
" ffmpeg.input(inpath).output(\n",
|
|
" inpath + \"_22k.wav\",\n",
|
|
" ar=\"22050\",\n",
|
|
" ac=\"1\",\n",
|
|
" acodec=\"pcm_s16le\",\n",
|
|
" map_metadata=\"-1\",\n",
|
|
" fflags=\"+bitexact\",\n",
|
|
" ).overwrite_output().run(quiet=True)\n",
|
|
" os.remove(inpath)\n",
|
|
" os.rename(inpath + \"_22k.wav\", inpath)\n",
|
|
"\n",
|
|
"# Extract dataset\n",
|
|
"os.chdir(output_dir)\n",
|
|
"\n",
|
|
"if os.path.exists(\"_training_data\"):\n",
|
|
" shutil.rmtree(\"_training_data\")\n",
|
|
"os.mkdir(\"_training_data\")\n",
|
|
"os.chdir(\"_training_data\")\n",
|
|
"os.mkdir(\"wavs\")\n",
|
|
"os.chdir(\"wavs\")\n",
|
|
"if dataset[-4:] == \".zip\":\n",
|
|
" if \"Windows\" in platform.system():\n",
|
|
" !tar -xf \"{dataset}\"\n",
|
|
" else:\n",
|
|
" !unzip -qo \"{dataset}\"\n",
|
|
"elif dataset[-4:] == \".tar\":\n",
|
|
" !tar -xf \"{dataset}\"\n",
|
|
"else:\n",
|
|
" raise Exception(\"Unknown extension for dataset\")\n",
|
|
"if os.path.exists(os.path.join(output_dir, \"_training_data\", \"wavs\", \"wavs\")):\n",
|
|
" shutil.move(\n",
|
|
" os.path.join(output_dir, \"_training_data\", \"wavs\", \"wavs\"), \n",
|
|
" os.path.join(output_dir, \"_training_data\", \"tempwavs\")\n",
|
|
" )\n",
|
|
" shutil.rmtree(os.path.join(output_dir, \"_training_data\", \"wavs\"))\n",
|
|
" shutil.move(\n",
|
|
" os.path.join(output_dir, \"_training_data\", \"tempwavs\"), \n",
|
|
" os.path.join(output_dir, \"_training_data\", \"wavs\")\n",
|
|
" )\n",
|
|
"# Filelist for preprocessing\n",
|
|
"os.chdir(output_dir)\n",
|
|
"shutil.copy(train_filelist, \"trainfiles.txt\")\n",
|
|
"shutil.copy(val_filelist, \"valfiles.txt\")\n",
|
|
"fix_transcripts(\"trainfiles.txt\")\n",
|
|
"fix_transcripts(\"valfiles.txt\")\n",
|
|
"seen_files = []\n",
|
|
"with open(\"trainfiles.txt\") as f:\n",
|
|
" t = f.read().split(\"\\n\")\n",
|
|
"with open(\"valfiles.txt\") as f:\n",
|
|
" v = f.read().split(\"\\n\")\n",
|
|
" all_filelist = t[:] + v[:]\n",
|
|
"with open(\"allfiles.txt\", \"w\") as f:\n",
|
|
" for x in all_filelist:\n",
|
|
" if x.strip() == \"\":\n",
|
|
" continue\n",
|
|
" if x.split(\"|\")[0] not in seen_files:\n",
|
|
" seen_files.append(x.split(\"|\")[0])\n",
|
|
" f.write(x.strip() + \"\\n\")\n",
|
|
"\n",
|
|
"# Ensure audio is 22k\n",
|
|
"print(\"Converting audio...\")\n",
|
|
"for r, _, f in os.walk(os.path.join(\"_training_data\", \"wavs\")):\n",
|
|
" for name in tqdm(f):\n",
|
|
" convert_to_22k(os.path.join(r, name))\n",
|
|
"\n",
|
|
"# Convert to JSON\n",
|
|
"generate_json(\"trainfiles.txt\", \"trainfiles.json\")\n",
|
|
"generate_json(\"valfiles.txt\", \"valfiles.json\")\n",
|
|
"generate_json(\"allfiles.txt\", \"allfiles.json\")\n",
|
|
"\n",
|
|
"print(\"OK\")"
|
|
],
|
|
"outputs": [],
|
|
"metadata": {
|
|
"cellView": "form",
|
|
"id": "bxFr3Fdi_kOC"
|
|
}
|
|
},
|
|
{
|
|
"cell_type": "markdown",
|
|
"source": [
|
|
"**Step 3:** Dataset processing, part 2. This takes a while, but\n",
|
|
"you only have to run this once per dataset (results are saved to Drive).\n",
|
|
"\n",
|
|
"If this step fails, try the following:\n",
|
|
"* Make sure your dataset only contains WAV files."
|
|
],
|
|
"metadata": {}
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": null,
|
|
"source": [
|
|
"# Extract phoneme duration\n",
|
|
"\n",
|
|
"import json\n",
|
|
"from nemo.collections.asr.models import EncDecCTCModel\n",
|
|
"asr_model = EncDecCTCModel.from_pretrained(model_name=\"asr_talknet_aligner\").cpu().eval()\n",
|
|
"\n",
|
|
"def forward_extractor(tokens, log_probs, blank):\n",
|
|
" \"\"\"Computes states f and p.\"\"\"\n",
|
|
" n, m = len(tokens), log_probs.shape[0]\n",
|
|
" # `f[s, t]` -- max sum of log probs for `s` first codes\n",
|
|
" # with `t` first timesteps with ending in `tokens[s]`.\n",
|
|
" f = np.empty((n + 1, m + 1), dtype=float)\n",
|
|
" f.fill(-(10 ** 9))\n",
|
|
" p = np.empty((n + 1, m + 1), dtype=int)\n",
|
|
" f[0, 0] = 0.0 # Start\n",
|
|
" for s in range(1, n + 1):\n",
|
|
" c = tokens[s - 1]\n",
|
|
" for t in range((s + 1) // 2, m + 1):\n",
|
|
" f[s, t] = log_probs[t - 1, c]\n",
|
|
" # Option #1: prev char is equal to current one.\n",
|
|
" if s == 1 or c == blank or c == tokens[s - 3]:\n",
|
|
" options = f[s : (s - 2 if s > 1 else None) : -1, t - 1]\n",
|
|
" else: # Is not equal to current one.\n",
|
|
" options = f[s : (s - 3 if s > 2 else None) : -1, t - 1]\n",
|
|
" f[s, t] += np.max(options)\n",
|
|
" p[s, t] = np.argmax(options)\n",
|
|
" return f, p\n",
|
|
"\n",
|
|
"\n",
|
|
"def backward_extractor(f, p):\n",
|
|
" \"\"\"Computes durs from f and p.\"\"\"\n",
|
|
" n, m = f.shape\n",
|
|
" n -= 1\n",
|
|
" m -= 1\n",
|
|
" durs = np.zeros(n, dtype=int)\n",
|
|
" if f[-1, -1] >= f[-2, -1]:\n",
|
|
" s, t = n, m\n",
|
|
" else:\n",
|
|
" s, t = n - 1, m\n",
|
|
" while s > 0:\n",
|
|
" durs[s - 1] += 1\n",
|
|
" s -= p[s, t]\n",
|
|
" t -= 1\n",
|
|
" assert durs.shape[0] == n\n",
|
|
" assert np.sum(durs) == m\n",
|
|
" assert np.all(durs[1::2] > 0)\n",
|
|
" return durs\n",
|
|
"\n",
|
|
"def preprocess_tokens(tokens, blank):\n",
|
|
" new_tokens = [blank]\n",
|
|
" for c in tokens:\n",
|
|
" new_tokens.extend([c, blank])\n",
|
|
" tokens = new_tokens\n",
|
|
" return tokens\n",
|
|
"\n",
|
|
"data_config = {\n",
|
|
" 'manifest_filepath': \"allfiles.json\",\n",
|
|
" 'sample_rate': 22050,\n",
|
|
" 'labels': asr_model.decoder.vocabulary,\n",
|
|
" 'batch_size': 1,\n",
|
|
"}\n",
|
|
"\n",
|
|
"parser = nemo.collections.asr.data.audio_to_text.AudioToCharWithDursF0Dataset.make_vocab(\n",
|
|
" notation='phonemes', punct=True, spaces=True, stresses=False, add_blank_at=\"last\"\n",
|
|
")\n",
|
|
"\n",
|
|
"dataset = nemo.collections.asr.data.audio_to_text._AudioTextDataset(\n",
|
|
" manifest_filepath=data_config['manifest_filepath'], sample_rate=data_config['sample_rate'], parser=parser,\n",
|
|
")\n",
|
|
"\n",
|
|
"dl = torch.utils.data.DataLoader(\n",
|
|
" dataset=dataset, batch_size=data_config['batch_size'], collate_fn=dataset.collate_fn, shuffle=False,\n",
|
|
")\n",
|
|
"\n",
|
|
"blank_id = asr_model.decoder.num_classes_with_blank - 1\n",
|
|
"\n",
|
|
"if os.path.exists(os.path.join(output_dir, \"durations.pt\")):\n",
|
|
" print(\"durations.pt already exists; skipping\")\n",
|
|
"else:\n",
|
|
" dur_data = {}\n",
|
|
" for sample_idx, test_sample in tqdm(enumerate(dl), total=len(dl)):\n",
|
|
" log_probs, _, greedy_predictions = asr_model(\n",
|
|
" input_signal=test_sample[0], input_signal_length=test_sample[1]\n",
|
|
" )\n",
|
|
"\n",
|
|
" log_probs = log_probs[0].cpu().detach().numpy()\n",
|
|
" seq_ids = test_sample[2][0].cpu().detach().numpy()\n",
|
|
"\n",
|
|
" target_tokens = preprocess_tokens(seq_ids, blank_id)\n",
|
|
"\n",
|
|
" f, p = forward_extractor(target_tokens, log_probs, blank_id)\n",
|
|
" durs = backward_extractor(f, p)\n",
|
|
"\n",
|
|
" dur_key = Path(dl.dataset.collection[sample_idx].audio_file).stem\n",
|
|
" dur_data[dur_key] = {\n",
|
|
" 'blanks': torch.tensor(durs[::2], dtype=torch.long).cpu().detach(), \n",
|
|
" 'tokens': torch.tensor(durs[1::2], dtype=torch.long).cpu().detach()\n",
|
|
" }\n",
|
|
"\n",
|
|
" del test_sample\n",
|
|
"\n",
|
|
" torch.save(dur_data, os.path.join(output_dir, \"durations.pt\"))\n",
|
|
"\n",
|
|
"#Extract F0 (pitch)\n",
|
|
"import crepe\n",
|
|
"from scipy.io import wavfile\n",
|
|
"\n",
|
|
"def crepe_f0(audio_file, hop_length=256):\n",
|
|
" sr, audio = wavfile.read(audio_file)\n",
|
|
" audio_x = np.arange(0, len(audio)) / 22050.0\n",
|
|
" time, frequency, confidence, activation = crepe.predict(audio, sr, viterbi=True)\n",
|
|
"\n",
|
|
" x = np.arange(0, len(audio), hop_length) / 22050.0\n",
|
|
" freq_interp = np.interp(x, time, frequency)\n",
|
|
" conf_interp = np.interp(x, time, confidence)\n",
|
|
" audio_interp = np.interp(x, audio_x, np.absolute(audio)) / 32768.0\n",
|
|
" weights = [0.5, 0.25, 0.25]\n",
|
|
" audio_smooth = np.convolve(audio_interp, np.array(weights)[::-1], \"same\")\n",
|
|
"\n",
|
|
" conf_threshold = 0.25\n",
|
|
" audio_threshold = 0.0005\n",
|
|
" for i in range(len(freq_interp)):\n",
|
|
" if conf_interp[i] < conf_threshold:\n",
|
|
" freq_interp[i] = 0.0\n",
|
|
" if audio_smooth[i] < audio_threshold:\n",
|
|
" freq_interp[i] = 0.0\n",
|
|
"\n",
|
|
" # Hack to make f0 and mel lengths equal\n",
|
|
" if len(audio) % hop_length == 0:\n",
|
|
" freq_interp = np.pad(freq_interp, pad_width=[0, 1])\n",
|
|
" return torch.from_numpy(freq_interp.astype(np.float32))\n",
|
|
"\n",
|
|
"if os.path.exists(os.path.join(output_dir, \"f0s.pt\")):\n",
|
|
" print(\"f0s.pt already exists; skipping\")\n",
|
|
"else:\n",
|
|
" f0_data = {}\n",
|
|
" with open(\"allfiles.json\") as f:\n",
|
|
" for i, l in enumerate(f.readlines()):\n",
|
|
" print(str(i))\n",
|
|
" audio_path = json.loads(l)[\"audio_filepath\"]\n",
|
|
" f0_data[Path(audio_path).stem] = crepe_f0(audio_path)\n",
|
|
"\n",
|
|
" # calculate f0 stats (mean & std) only for train set\n",
|
|
" with open(\"trainfiles.json\") as f:\n",
|
|
" train_ids = {Path(json.loads(l)[\"audio_filepath\"]).stem for l in f}\n",
|
|
" all_f0 = torch.cat([f0[f0 >= 1e-5] for f0_id, f0 in f0_data.items() if f0_id in train_ids])\n",
|
|
"\n",
|
|
" F0_MEAN, F0_STD = all_f0.mean().item(), all_f0.std().item() \n",
|
|
" print(\"F0_MEAN: \" + str(F0_MEAN) + \", F0_STD: \" + str(F0_STD))\n",
|
|
" torch.save(f0_data, os.path.join(output_dir, \"f0s.pt\"))\n",
|
|
" with open(os.path.join(output_dir, \"f0_info.json\"), \"w\") as f:\n",
|
|
" f.write(json.dumps({\"FO_MEAN\": F0_MEAN, \"F0_STD\": F0_STD}))\n",
|
|
"\n",
|
|
"print(\"OK\")"
|
|
],
|
|
"outputs": [],
|
|
"metadata": {
|
|
"cellView": "form",
|
|
"id": "sos9vsxPkIN7"
|
|
}
|
|
},
|
|
{
|
|
"cell_type": "markdown",
|
|
"source": [
|
|
"**Step 4:** Train duration predictor.\n",
|
|
"\n",
|
|
"If CUDA runs out of memory, try the following:\n",
|
|
"* Click on Kernel -> Restart, re-run step 1, and try again.\n",
|
|
"* If that doesn't help, reduce the batch size (default 64)."
|
|
],
|
|
"metadata": {}
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": null,
|
|
"source": [
|
|
"batch_size = 64\n",
|
|
"\n",
|
|
"import os\n",
|
|
"from hydra.experimental import compose, initialize\n",
|
|
"from hydra.core.global_hydra import GlobalHydra\n",
|
|
"from omegaconf import OmegaConf\n",
|
|
"import pytorch_lightning as pl\n",
|
|
"from nemo.collections.common.callbacks import LogEpochTimeCallback\n",
|
|
"from nemo.collections.tts.models import TalkNetDursModel\n",
|
|
"from nemo.core.config import hydra_runner\n",
|
|
"from nemo.utils.exp_manager import exp_manager\n",
|
|
"\n",
|
|
"epochs = 20\n",
|
|
"learning_rate = 1e-3\n",
|
|
"min_learning_rate = 3e-6\n",
|
|
"load_checkpoints = True\n",
|
|
"\n",
|
|
"def train(cfg):\n",
|
|
" cfg.sample_rate = 22050\n",
|
|
" cfg.train_dataset = os.path.join(output_dir, \"trainfiles.json\")\n",
|
|
" cfg.validation_datasets = os.path.join(output_dir, \"valfiles.json\")\n",
|
|
" cfg.durs_file = os.path.join(output_dir, \"durations.pt\")\n",
|
|
" cfg.f0_file = os.path.join(output_dir, \"f0s.pt\")\n",
|
|
" cfg.trainer.accelerator = \"dp\"\n",
|
|
" cfg.trainer.max_epochs = epochs\n",
|
|
" cfg.trainer.check_val_every_n_epoch = 5\n",
|
|
" cfg.model.train_ds.dataloader_params.batch_size = batch_size\n",
|
|
" cfg.model.validation_ds.dataloader_params.batch_size = batch_size\n",
|
|
" cfg.model.optim.lr = learning_rate\n",
|
|
" cfg.model.optim.sched.min_lr = min_learning_rate\n",
|
|
" cfg.exp_manager.exp_dir = output_dir\n",
|
|
"\n",
|
|
" # Find checkpoints\n",
|
|
" ckpt_path = \"\"\n",
|
|
" if load_checkpoints:\n",
|
|
" path0 = os.path.join(output_dir, \"TalkNetDurs\")\n",
|
|
" if os.path.exists(path0):\n",
|
|
" path1 = sorted(os.listdir(path0))\n",
|
|
" for i in range(len(path1)):\n",
|
|
" path2 = os.path.join(path0, path1[-(1+i)], \"checkpoints\")\n",
|
|
" if os.path.exists(path2):\n",
|
|
" match = [x for x in os.listdir(path2) if \"last.ckpt\" in x]\n",
|
|
" if len(match) > 0:\n",
|
|
" ckpt_path = os.path.join(path2, match[0])\n",
|
|
" print(\"Resuming training from \" + match[0])\n",
|
|
" break\n",
|
|
" \n",
|
|
" if ckpt_path != \"\":\n",
|
|
" trainer = pl.Trainer(**cfg.trainer, resume_from_checkpoint = ckpt_path)\n",
|
|
" model = TalkNetDursModel(cfg=cfg.model, trainer=trainer)\n",
|
|
" else:\n",
|
|
" warmstart_path = os.path.join(cwd, \"talknet_durs.nemo\")\n",
|
|
" trainer = pl.Trainer(**cfg.trainer)\n",
|
|
" model = TalkNetDursModel.restore_from(warmstart_path, override_config_path=cfg)\n",
|
|
" model.set_trainer(trainer)\n",
|
|
" model.setup_training_data(cfg.model.train_ds)\n",
|
|
" model.setup_validation_data(cfg.model.validation_ds)\n",
|
|
" model.setup_optimization(cfg.model.optim)\n",
|
|
" print(\"Warm-starting from \" + warmstart_path)\n",
|
|
" exp_manager(trainer, cfg.get('exp_manager', None))\n",
|
|
" trainer.callbacks.extend([pl.callbacks.LearningRateMonitor(), LogEpochTimeCallback()]) # noqa\n",
|
|
" trainer.fit(model)\n",
|
|
"\n",
|
|
"os.chdir(cwd)\n",
|
|
"GlobalHydra().clear()\n",
|
|
"initialize(config_path=\"conf\")\n",
|
|
"cfg = compose(config_name=\"talknet-durs\")\n",
|
|
"train(cfg)\n"
|
|
],
|
|
"outputs": [],
|
|
"metadata": {
|
|
"cellView": "form",
|
|
"id": "nM7-bMpKO7U2"
|
|
}
|
|
},
|
|
{
|
|
"cell_type": "markdown",
|
|
"source": [
|
|
"**Step 5:** Train pitch predictor.\n",
|
|
"\n",
|
|
"If CUDA runs out of memory, try the following:\n",
|
|
"* Click on Kernel -> Restart, re-run step 1, and try again.\n",
|
|
"* If that doesn't help, reduce the batch size (default 64)."
|
|
],
|
|
"metadata": {}
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": null,
|
|
"source": [
|
|
"batch_size = 64\n",
|
|
"epochs = 50\n",
|
|
"\n",
|
|
"import json\n",
|
|
"\n",
|
|
"with open(os.path.join(output_dir, \"f0_info.json\"), \"r\") as f:\n",
|
|
" f0_info = json.load(f)\n",
|
|
" f0_mean = f0_info[\"FO_MEAN\"]\n",
|
|
" f0_std = f0_info[\"F0_STD\"]\n",
|
|
"\n",
|
|
"learning_rate = 1e-3\n",
|
|
"min_learning_rate = 3e-6\n",
|
|
"load_checkpoints = True\n",
|
|
"\n",
|
|
"import os\n",
|
|
"from hydra.experimental import compose, initialize\n",
|
|
"from hydra.core.global_hydra import GlobalHydra\n",
|
|
"from omegaconf import OmegaConf\n",
|
|
"import pytorch_lightning as pl\n",
|
|
"from nemo.collections.common.callbacks import LogEpochTimeCallback\n",
|
|
"from nemo.collections.tts.models import TalkNetPitchModel\n",
|
|
"from nemo.core.config import hydra_runner\n",
|
|
"from nemo.utils.exp_manager import exp_manager\n",
|
|
"\n",
|
|
"def train(cfg):\n",
|
|
" cfg.sample_rate = 22050\n",
|
|
" cfg.train_dataset = os.path.join(output_dir, \"trainfiles.json\")\n",
|
|
" cfg.validation_datasets = os.path.join(output_dir, \"valfiles.json\")\n",
|
|
" cfg.durs_file = os.path.join(output_dir, \"durations.pt\")\n",
|
|
" cfg.f0_file = os.path.join(output_dir, \"f0s.pt\")\n",
|
|
" cfg.trainer.accelerator = \"dp\"\n",
|
|
" cfg.trainer.max_epochs = epochs\n",
|
|
" cfg.trainer.check_val_every_n_epoch = 5\n",
|
|
" cfg.model.f0_mean=f0_mean\n",
|
|
" cfg.model.f0_std=f0_std\n",
|
|
" cfg.model.train_ds.dataloader_params.batch_size = batch_size\n",
|
|
" cfg.model.validation_ds.dataloader_params.batch_size = batch_size\n",
|
|
" cfg.model.optim.lr = learning_rate\n",
|
|
" cfg.model.optim.sched.min_lr = min_learning_rate\n",
|
|
" cfg.exp_manager.exp_dir = output_dir\n",
|
|
"\n",
|
|
" # Find checkpoints\n",
|
|
" ckpt_path = \"\"\n",
|
|
" if load_checkpoints:\n",
|
|
" path0 = os.path.join(output_dir, \"TalkNetPitch\")\n",
|
|
" if os.path.exists(path0):\n",
|
|
" path1 = sorted(os.listdir(path0))\n",
|
|
" for i in range(len(path1)):\n",
|
|
" path2 = os.path.join(path0, path1[-(1+i)], \"checkpoints\")\n",
|
|
" if os.path.exists(path2):\n",
|
|
" match = [x for x in os.listdir(path2) if \"last.ckpt\" in x]\n",
|
|
" if len(match) > 0:\n",
|
|
" ckpt_path = os.path.join(path2, match[0])\n",
|
|
" print(\"Resuming training from \" + match[0])\n",
|
|
" break\n",
|
|
" \n",
|
|
" if ckpt_path != \"\":\n",
|
|
" trainer = pl.Trainer(**cfg.trainer, resume_from_checkpoint = ckpt_path)\n",
|
|
" model = TalkNetPitchModel(cfg=cfg.model, trainer=trainer)\n",
|
|
" else:\n",
|
|
" warmstart_path = os.path.join(cwd, \"talknet_pitch.nemo\")\n",
|
|
" trainer = pl.Trainer(**cfg.trainer)\n",
|
|
" model = TalkNetPitchModel.restore_from(warmstart_path, override_config_path=cfg)\n",
|
|
" model.set_trainer(trainer)\n",
|
|
" model.setup_training_data(cfg.model.train_ds)\n",
|
|
" model.setup_validation_data(cfg.model.validation_ds)\n",
|
|
" model.setup_optimization(cfg.model.optim)\n",
|
|
" print(\"Warm-starting from \" + warmstart_path)\n",
|
|
" exp_manager(trainer, cfg.get('exp_manager', None))\n",
|
|
" trainer.callbacks.extend([pl.callbacks.LearningRateMonitor(), LogEpochTimeCallback()]) # noqa\n",
|
|
" trainer.fit(model)\n",
|
|
"\n",
|
|
"os.chdir(cwd)\n",
|
|
"GlobalHydra().clear()\n",
|
|
"initialize(config_path=\"conf\")\n",
|
|
"cfg = compose(config_name=\"talknet-pitch\")\n",
|
|
"train(cfg)"
|
|
],
|
|
"outputs": [],
|
|
"metadata": {
|
|
"cellView": "form",
|
|
"id": "JLfm00NuJfon"
|
|
}
|
|
},
|
|
{
|
|
"cell_type": "markdown",
|
|
"source": [
|
|
"**Step 6:** Train spectrogram generator. 200+ epochs are recommended. \n",
|
|
"\n",
|
|
"This is the slowest of the three models to train, and the hardest to\n",
|
|
"get good results from. If your character sounds noisy or robotic,\n",
|
|
"try improving the dataset, or adjusting the epochs and learning rate.\n",
|
|
"\n",
|
|
"If CUDA runs out of memory, try the following:\n",
|
|
"* Click on Kernel -> Restart, re-run step 1, and try again.\n",
|
|
"* If that doesn't help, reduce the batch size (default 32)."
|
|
],
|
|
"metadata": {}
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": null,
|
|
"source": [
|
|
"epochs = 200\n",
|
|
"batch_size = 32\n",
|
|
"\n",
|
|
"# Advanced settings. You can probably leave these at their defaults (1e-3, 3e-6, empty, checked).\n",
|
|
"learning_rate = 1e-3\n",
|
|
"min_learning_rate = 3e-6\n",
|
|
"pretrained_path = \"\"\n",
|
|
"load_checkpoints = True\n",
|
|
"\n",
|
|
"import os\n",
|
|
"from hydra.experimental import compose, initialize\n",
|
|
"from hydra.core.global_hydra import GlobalHydra\n",
|
|
"from omegaconf import OmegaConf\n",
|
|
"import pytorch_lightning as pl\n",
|
|
"from nemo.collections.common.callbacks import LogEpochTimeCallback\n",
|
|
"from nemo.collections.tts.models import TalkNetSpectModel\n",
|
|
"from nemo.core.config import hydra_runner\n",
|
|
"from nemo.utils.exp_manager import exp_manager\n",
|
|
"\n",
|
|
"def train(cfg):\n",
|
|
" cfg.sample_rate = 22050\n",
|
|
" cfg.train_dataset = os.path.join(output_dir, \"trainfiles.json\")\n",
|
|
" cfg.validation_datasets = os.path.join(output_dir, \"valfiles.json\")\n",
|
|
" cfg.durs_file = os.path.join(output_dir, \"durations.pt\")\n",
|
|
" cfg.f0_file = os.path.join(output_dir, \"f0s.pt\")\n",
|
|
" cfg.trainer.accelerator = \"dp\"\n",
|
|
" cfg.trainer.max_epochs = epochs\n",
|
|
" cfg.trainer.check_val_every_n_epoch = 5\n",
|
|
" cfg.model.train_ds.dataloader_params.batch_size = batch_size\n",
|
|
" cfg.model.validation_ds.dataloader_params.batch_size = batch_size\n",
|
|
" cfg.model.optim.lr = learning_rate\n",
|
|
" cfg.model.optim.sched.min_lr = min_learning_rate\n",
|
|
" cfg.exp_manager.exp_dir = output_dir\n",
|
|
"\n",
|
|
" # Find checkpoints\n",
|
|
" ckpt_path = \"\"\n",
|
|
" if load_checkpoints:\n",
|
|
" path0 = os.path.join(output_dir, \"TalkNetSpect\")\n",
|
|
" if os.path.exists(path0):\n",
|
|
" path1 = sorted(os.listdir(path0))\n",
|
|
" for i in range(len(path1)):\n",
|
|
" path2 = os.path.join(path0, path1[-(1+i)], \"checkpoints\")\n",
|
|
" if os.path.exists(path2):\n",
|
|
" match = [x for x in os.listdir(path2) if \"last.ckpt\" in x]\n",
|
|
" if len(match) > 0:\n",
|
|
" ckpt_path = os.path.join(path2, match[0])\n",
|
|
" print(\"Resuming training from \" + match[0])\n",
|
|
" break\n",
|
|
" \n",
|
|
" if ckpt_path != \"\":\n",
|
|
" trainer = pl.Trainer(**cfg.trainer, resume_from_checkpoint = ckpt_path)\n",
|
|
" model = TalkNetSpectModel(cfg=cfg.model, trainer=trainer)\n",
|
|
" else:\n",
|
|
" if pretrained_path != \"\":\n",
|
|
" warmstart_path = pretrained_path\n",
|
|
" else:\n",
|
|
" warmstart_path = os.path.join(cwd, \"talknet_spect.nemo\")\n",
|
|
" trainer = pl.Trainer(**cfg.trainer)\n",
|
|
" model = TalkNetSpectModel.restore_from(warmstart_path, override_config_path=cfg)\n",
|
|
" model.set_trainer(trainer)\n",
|
|
" model.setup_training_data(cfg.model.train_ds)\n",
|
|
" model.setup_validation_data(cfg.model.validation_ds)\n",
|
|
" model.setup_optimization(cfg.model.optim)\n",
|
|
" print(\"Warm-starting from \" + warmstart_path)\n",
|
|
" exp_manager(trainer, cfg.get('exp_manager', None))\n",
|
|
" trainer.callbacks.extend([pl.callbacks.LearningRateMonitor(), LogEpochTimeCallback()]) # noqa\n",
|
|
" trainer.fit(model)\n",
|
|
"\n",
|
|
"os.chdir(cwd)\n",
|
|
"GlobalHydra().clear()\n",
|
|
"initialize(config_path=\"conf\")\n",
|
|
"cfg = compose(config_name=\"talknet-spect\")\n",
|
|
"train(cfg)"
|
|
],
|
|
"outputs": [],
|
|
"metadata": {
|
|
"cellView": "form",
|
|
"id": "N9hh4WPHbCcn"
|
|
}
|
|
},
|
|
{
|
|
"cell_type": "markdown",
|
|
"source": [
|
|
"**Step 7:** Generate GTA spectrograms. This will help HiFi-GAN learn what your TalkNet model sounds like.\n",
|
|
"\n",
|
|
"If this step fails, make sure you've finished training the spectrogram generator."
|
|
],
|
|
"metadata": {}
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": null,
|
|
"source": [
|
|
"import sys\n",
|
|
"import os\n",
|
|
"import torch\n",
|
|
"import numpy as np\n",
|
|
"from tqdm import tqdm\n",
|
|
"from nemo.collections.tts.models import TalkNetSpectModel\n",
|
|
"import shutil\n",
|
|
"\n",
|
|
"def fix_paths(inpath):\n",
|
|
" output = \"\"\n",
|
|
" with open(inpath, \"r\", encoding=\"utf8\") as f:\n",
|
|
" for l in f.readlines():\n",
|
|
" if l[:5].lower() != \"wavs/\":\n",
|
|
" output += \"wavs/\" + l\n",
|
|
" else:\n",
|
|
" output += l\n",
|
|
" with open(inpath, \"w\", encoding=\"utf8\") as w:\n",
|
|
" w.write(output)\n",
|
|
"\n",
|
|
"shutil.copyfile(train_filelist, os.path.join(cwd, \"hifi-gan\", \"training.txt\"))\n",
|
|
"shutil.copyfile(val_filelist, os.path.join(cwd, \"hifi-gan\", \"validation.txt\"))\n",
|
|
"fix_paths(os.path.join(cwd, \"hifi-gan\", \"training.txt\"))\n",
|
|
"fix_paths(os.path.join(cwd, \"hifi-gan\", \"validation.txt\"))\n",
|
|
"fix_paths(os.path.join(output_dir, \"allfiles.txt\"))\n",
|
|
"\n",
|
|
"os.chdir(cwd)\n",
|
|
"outdir = os.path.join(output_dir, \"_training_data\", \"wavs\")\n",
|
|
"if not os.path.exists(outdir):\n",
|
|
" os.mkdir(outdir)\n",
|
|
"\n",
|
|
"model_path = \"\"\n",
|
|
"path0 = os.path.join(output_dir, \"TalkNetSpect\")\n",
|
|
"if os.path.exists(path0):\n",
|
|
" path1 = sorted(os.listdir(path0))\n",
|
|
" for i in range(len(path1)):\n",
|
|
" path2 = os.path.join(path0, path1[-(1+i)], \"checkpoints\")\n",
|
|
" if os.path.exists(path2):\n",
|
|
" match = [x for x in os.listdir(path2) if \"TalkNetSpect.nemo\" in x]\n",
|
|
" if len(match) > 0:\n",
|
|
" model_path = os.path.join(path2, match[0])\n",
|
|
" break\n",
|
|
"assert model_path != \"\", \"TalkNetSpect.nemo not found\"\n",
|
|
"\n",
|
|
"dur_path = os.path.join(output_dir, \"durations.pt\")\n",
|
|
"f0_path = os.path.join(output_dir, \"f0s.pt\")\n",
|
|
"\n",
|
|
"model = TalkNetSpectModel.restore_from(model_path)\n",
|
|
"model.eval()\n",
|
|
"with open(os.path.join(output_dir, \"allfiles.txt\"), \"r\", encoding=\"utf-8\") as f:\n",
|
|
" dataset = f.readlines()\n",
|
|
"durs = torch.load(dur_path)\n",
|
|
"f0s = torch.load(f0_path)\n",
|
|
"\n",
|
|
"for x in tqdm(dataset):\n",
|
|
" x_name = os.path.splitext(os.path.basename(x.split(\"|\")[0].strip()))[0]\n",
|
|
" x_tokens = model.parse(text=x.split(\"|\")[1].strip())\n",
|
|
" x_durs = (\n",
|
|
" torch.stack(\n",
|
|
" (\n",
|
|
" durs[x_name][\"blanks\"],\n",
|
|
" torch.cat((durs[x_name][\"tokens\"], torch.zeros(1).int())),\n",
|
|
" ),\n",
|
|
" dim=1,\n",
|
|
" )\n",
|
|
" .view(-1)[:-1]\n",
|
|
" .view(1, -1)\n",
|
|
" .to(\"cuda:0\")\n",
|
|
" )\n",
|
|
" x_f0s = f0s[x_name].view(1, -1).to(\"cuda:0\")\n",
|
|
" x_spect = model.force_spectrogram(tokens=x_tokens, durs=x_durs, f0=x_f0s)\n",
|
|
" rel_path = os.path.splitext(x.split(\"|\")[0].strip())[0][5:]\n",
|
|
" abs_dir = os.path.join(outdir, os.path.dirname(rel_path))\n",
|
|
" if abs_dir != \"\" and not os.path.exists(abs_dir):\n",
|
|
" os.makedirs(abs_dir, exist_ok=True)\n",
|
|
" np.save(os.path.join(outdir, rel_path + \".npy\"), x_spect.detach().cpu().numpy())\n"
|
|
],
|
|
"outputs": [],
|
|
"metadata": {
|
|
"cellView": "form",
|
|
"id": "Ajfyfz2p9Ior"
|
|
}
|
|
},
|
|
{
|
|
"cell_type": "markdown",
|
|
"source": [
|
|
"**Step 8:** Train HiFi-GAN. 2,000+ steps are recommended.\n",
|
|
"Stop this cell to finish training the model.\n",
|
|
"\n",
|
|
"* Click on Kernel -> Restart, re-run step 1, and try again.\n",
|
|
"If this step still fails to start, make sure step 7 finished successfully.\n",
|
|
"\n",
|
|
"Note: If the training process starts at step 2500000, delete the HiFiGAN folder and try again."
|
|
],
|
|
"metadata": {}
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": null,
|
|
"source": [
|
|
"import gdown\n",
|
|
"d = 'https://drive.google.com/uc?id='\n",
|
|
"\n",
|
|
"os.chdir(os.path.join(cwd, \"hifi-gan\"))\n",
|
|
"\n",
|
|
"if not os.path.exists(os.path.join(output_dir, \"HiFiGAN\")):\n",
|
|
" os.makedirs(os.path.join(output_dir, \"HiFiGAN\"))\n",
|
|
"if not os.path.exists(os.path.join(output_dir, \"HiFiGAN\", \"do_00000000\")):\n",
|
|
" print(\"Downloading universal model...\")\n",
|
|
" gdown.download(d+\"1qpgI41wNXFcH-iKq1Y42JlBC9j0je8PW\", os.path.join(output_dir, \"HiFiGAN\", \"g_00000000\"), quiet=False)\n",
|
|
" gdown.download(d+\"1O63eHZR9t1haCdRHQcEgMfMNxiOciSru\", os.path.join(output_dir, \"HiFiGAN\", \"do_00000000\"), quiet=False)\n",
|
|
" start_from_universal = \"--warm_start True \"\n",
|
|
"else:\n",
|
|
" start_from_universal = \"\"\n",
|
|
"\n",
|
|
"hifi_train = os.path.join(cwd, \"hifi-gan\", \"training.txt\")\n",
|
|
"hifi_val = os.path.join(cwd, \"hifi-gan\", \"validation.txt\")\n",
|
|
"hifi_wavs = os.path.join(output_dir, \"_training_data\")\n",
|
|
"\n",
|
|
"!python train.py --fine_tuning True --config config_v1b.json \\\n",
|
|
"{start_from_universal} \\\n",
|
|
"--checkpoint_interval 250 --checkpoint_path \"{os.path.join(output_dir, 'HiFiGAN')}\" \\\n",
|
|
"--input_training_file \"{hifi_train}\" \\\n",
|
|
"--input_validation_file \"{hifi_val}\" \\\n",
|
|
"--input_wavs_dir \"{hifi_wavs}\"\n"
|
|
],
|
|
"outputs": [],
|
|
"metadata": {
|
|
"cellView": "form",
|
|
"id": "yVBjGhRB9hUJ"
|
|
}
|
|
},
|
|
{
|
|
"cell_type": "markdown",
|
|
"source": [
|
|
"**Step 9:** Package the models. They'll be saved to the output directory as [character_name]_TalkNet.zip.\n",
|
|
"\n",
|
|
"When done, upload it to Google Drive, with permissions set to \"Anyone with the link\". \n",
|
|
"You can then use it with TalkNet by selecting \"Custom model\" as your character. \n",
|
|
"\n",
|
|
"This cell will also delete the training checkpoints and logs.\n",
|
|
"That should free up roughly 2 GB of space.\n",
|
|
"If you wish to keep them, set delete_checkpoints to False."
|
|
],
|
|
"metadata": {}
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": null,
|
|
"source": [
|
|
"character_name = \"Character\"\n",
|
|
"delete_checkpoints = True\n",
|
|
"\n",
|
|
"import shutil\n",
|
|
"from zipfile import ZipFile\n",
|
|
"\n",
|
|
"def find_talknet(model_dir):\n",
|
|
" ckpt_path = \"\"\n",
|
|
" path0 = os.path.join(output_dir, model_dir)\n",
|
|
" if os.path.exists(path0):\n",
|
|
" path1 = sorted(os.listdir(path0))\n",
|
|
" for i in range(len(path1)):\n",
|
|
" path2 = os.path.join(path0, path1[-(1+i)], \"checkpoints\")\n",
|
|
" if os.path.exists(path2):\n",
|
|
" match = [x for x in os.listdir(path2) if \".nemo\" in x]\n",
|
|
" if len(match) > 0:\n",
|
|
" ckpt_path = os.path.join(path2, match[0])\n",
|
|
" break\n",
|
|
" assert ckpt_path != \"\", \"Couldn't find \" + model_dir\n",
|
|
" return ckpt_path\n",
|
|
"\n",
|
|
"durs_path = find_talknet(\"TalkNetDurs\")\n",
|
|
"pitch_path = find_talknet(\"TalkNetPitch\")\n",
|
|
"spect_path = find_talknet(\"TalkNetSpect\")\n",
|
|
"assert os.path.exists(os.path.join(output_dir, \"HiFiGAN\", \"g_00000000\")), \"Couldn't find HiFi-GAN\"\n",
|
|
"\n",
|
|
"zip = ZipFile(os.path.join(output_dir, character_name + \"_TalkNet.zip\"), 'w')\n",
|
|
"zip.write(durs_path, \"TalkNetDurs.nemo\")\n",
|
|
"zip.write(pitch_path, \"TalkNetPitch.nemo\")\n",
|
|
"zip.write(spect_path, \"TalkNetSpect.nemo\")\n",
|
|
"zip.write(os.path.join(output_dir, \"HiFiGAN\", \"g_00000000\"), \"hifiganmodel\")\n",
|
|
"zip.write(os.path.join(output_dir, \"HiFiGAN\", \"config.json\"), \"config.json\")\n",
|
|
"zip.write(os.path.join(output_dir, \"f0_info.json\"), \"f0_info.json\")\n",
|
|
"zip.close()\n",
|
|
"print(\"Archived model to \" + os.path.join(output_dir, character_name + \"_TalkNet.zip\"))\n",
|
|
"\n",
|
|
"if delete_checkpoints:\n",
|
|
" shutil.rmtree((os.path.join(output_dir, \"TalkNetDurs\")))\n",
|
|
" shutil.rmtree((os.path.join(output_dir, \"TalkNetPitch\")))\n",
|
|
" shutil.rmtree((os.path.join(output_dir, \"TalkNetSpect\")))\n",
|
|
" shutil.rmtree((os.path.join(output_dir, \"HiFiGAN\")))\n"
|
|
],
|
|
"outputs": [],
|
|
"metadata": {
|
|
"cellView": "form",
|
|
"id": "5OtwfrTT-blU"
|
|
}
|
|
}
|
|
],
|
|
"metadata": {
|
|
"accelerator": "GPU",
|
|
"colab": {
|
|
"collapsed_sections": [],
|
|
"name": "TalkNet_Training_Offline.ipynb",
|
|
"provenance": []
|
|
},
|
|
"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.10"
|
|
}
|
|
},
|
|
"nbformat": 4,
|
|
"nbformat_minor": 2
|
|
} |