diff --git a/voice_conversion/softvc/README.md b/voice_conversion/softvc/README.md new file mode 100644 index 0000000..0ae7a93 --- /dev/null +++ b/voice_conversion/softvc/README.md @@ -0,0 +1,22 @@ +FakeYou SoftVC +============== + +This work is based off research by Michael and Paul, leveraging bshall's SoftVC: + +https://bshall.github.io/soft-vc/ + +These are the repos contained herein: + +* https://github.com/bshall/hubert +* https://github.com/bshall/acoustic-model +* https://github.com/bshall/hifigan + +And this is our initial notebook: + +* [Colab Notebook](https://colab.research.google.com/drive/1g0YsQQY8fzXHmZUWAMzQO60c6OUlyKBl#scrollTo=Sf_rjYw0P83I) + + + + + + diff --git a/voice_conversion/softvc/__init__.py b/voice_conversion/softvc/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/voice_conversion/softvc/acoustic-model/.gitignore b/voice_conversion/softvc/acoustic-model/.gitignore new file mode 100644 index 0000000..0202868 --- /dev/null +++ b/voice_conversion/softvc/acoustic-model/.gitignore @@ -0,0 +1,132 @@ +# Byte-compiled / optimized / DLL files +__pycache__/ +*.py[cod] +*$py.class + +# C extensions +*.so + +# Distribution / packaging +.Python +build/ +develop-eggs/ +dist/ +downloads/ +eggs/ +.eggs/ +lib/ +lib64/ +parts/ +sdist/ +var/ +wheels/ +pip-wheel-metadata/ +share/python-wheels/ +*.egg-info/ +.installed.cfg +*.egg +MANIFEST + +# PyInstaller +# Usually these files are written by a python script from a template +# before PyInstaller builds the exe, so as to inject date/other infos into it. +*.manifest +*.spec + +# Installer logs +pip-log.txt +pip-delete-this-directory.txt + +# Unit test / coverage reports +htmlcov/ +.tox/ +.nox/ +.coverage +.coverage.* +.cache +nosetests.xml +coverage.xml +*.cover +*.py,cover +.hypothesis/ +.pytest_cache/ + +# Translations +*.mo +*.pot + +# Django stuff: +*.log +local_settings.py +db.sqlite3 +db.sqlite3-journal + +# Flask stuff: +instance/ +.webassets-cache + +# Scrapy stuff: +.scrapy + +# Sphinx documentation +docs/_build/ + +# PyBuilder +target/ + +# Jupyter Notebook +.ipynb_checkpoints + +# IPython +profile_default/ +ipython_config.py + +# pyenv +.python-version + +# pipenv +# According to pypa/pipenv#598, it is recommended to include Pipfile.lock in version control. +# However, in case of collaboration, if having platform-specific dependencies or dependencies +# having no cross-platform support, pipenv may install dependencies that don't work, or not +# install all needed dependencies. +#Pipfile.lock + +# PEP 582; used by e.g. github.com/David-OConnor/pyflow +__pypackages__/ + +# Celery stuff +celerybeat-schedule +celerybeat.pid + +# SageMath parsed files +*.sage.py + +# Environments +.env +.venv +env/ +venv/ +ENV/ +env.bak/ +venv.bak/ + +# VSCode project settings +.vscode + +# Spyder project settings +.spyderproject +.spyproject + +# Rope project settings +.ropeproject + +# mkdocs documentation +/site + +# mypy +.mypy_cache/ +.dmypy.json +dmypy.json + +# Pyre type checker +.pyre/ diff --git a/voice_conversion/softvc/acoustic-model/LICENSE b/voice_conversion/softvc/acoustic-model/LICENSE new file mode 100644 index 0000000..6eb2af0 --- /dev/null +++ b/voice_conversion/softvc/acoustic-model/LICENSE @@ -0,0 +1,21 @@ +MIT License + +Copyright (c) 2021 Benjamin van Niekerk + +Permission is hereby granted, free of charge, to any person obtaining a copy +of this software and associated documentation files (the "Software"), to deal +in the Software without restriction, including without limitation the rights +to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +copies of the Software, and to permit persons to whom the Software is +furnished to do so, subject to the following conditions: + +The above copyright notice and this permission notice shall be included in all +copies or substantial portions of the Software. + +THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE +SOFTWARE. diff --git a/voice_conversion/softvc/acoustic-model/README.md b/voice_conversion/softvc/acoustic-model/README.md new file mode 100644 index 0000000..fcc1daf --- /dev/null +++ b/voice_conversion/softvc/acoustic-model/README.md @@ -0,0 +1,180 @@ +

+ + Open In Colab + +

+ +# Acoustic-Model + +Training and inference scripts for the acoustic models in [A Comparison of Discrete and Soft Speech Units for Improved Voice Conversion](https://ieeexplore.ieee.org/abstract/document/9746484). For more details see [soft-vc](https://github.com/bshall/soft-vc). Audio samples can be found [here](https://bshall.github.io/soft-vc/). Colab demo can be found [here](https://colab.research.google.com/github/bshall/soft-vc/blob/main/soft-vc-demo.ipynb). + +
+ Soft-VC +
+
+ + Fig 1: Architecture of the voice conversion system. a) The discrete content encoder clusters audio features to produce a sequence of discrete speech units. b) The soft content encoder is trained to predict the discrete units. The acoustic model transforms the discrete/soft speech units into a target spectrogram. The vocoder converts the spectrogram into an audio waveform. + +
+ +## Example Usage + +### Programmatic Usage + +```python +import torch +import numpy as np + +# Load checkpoint (either hubert_soft or hubert_discrete) +acoustic = torch.hub.load("bshall/acoustic-model:main", "hubert_soft").cuda() + +# Load speech units +units = torch.from_numpy(np.load("path/to/units")) + +# Generate mel-spectrogram +mel = acoustic.generate(units) +``` + +### Script-Based Usage + +``` +usage: generate.py [-h] {soft,discrete} in-dir out-dir + +Generate spectrograms from input speech units (discrete or soft). + +positional arguments: + {soft,discrete} available models (HuBERT-Soft or HuBERT-Discrete) + in-dir path to the dataset directory. + out-dir path to the output directory. + +optional arguments: + -h, --help show this help message and exit +``` + +## Training + +### Step 1: Dataset Preparation + +Download and extract the [LJSpeech](https://keithito.com/LJ-Speech-Dataset/) dataset. The training script expects the following tree structure for the dataset directory: + +``` +└───wavs + ├───dev + │ ├───LJ001-0001.wav + │ ├───... + │ └───LJ050-0278.wav + └───train + ├───LJ002-0332.wav + ├───... + └───LJ047-0007.wav +``` + +The `train` and `dev` directories should contain the training and validation splits respectively. The splits used for the paper can be found [here](https://github.com/bshall/acoustic-model/releases/tag/v0.1). + +### Step 2: Extract Spectrograms + +Extract mel-spectrograms using the `mel.py` script: + +``` +usage: mels.py [-h] in-dir out-dir + +Extract mel-spectrograms for an audio dataset. + +positional arguments: + in-dir path to the dataset directory. + out-dir path to the output directory. + +optional arguments: + -h, --help show this help message and exit +``` + +for example: + +``` +python mel.py path/to/LJSpeech-1.1/wavs path/to/LJSpeech-1.1/mels +``` + +At this point the directory tree should look like: + +``` +├───mels +│ ├───... +└───wavs + ├───... +``` + +### Step 3: Extract Discrete or Soft Speech Units + +Use the HuBERT-Soft or HuBERT-Discrete content encoders to extract speech units. First clone the [content encoder repo](https://github.com/bshall/hubert) and then run `encode.py` (see the repo for details): + +``` +usage: encode.py [-h] [--extension EXTENSION] {soft,discrete} in-dir out-dir + +Encode an audio dataset. + +positional arguments: + {soft,discrete} available models (HuBERT-Soft or HuBERT-Discrete) + in-dir path to the dataset directory. + out-dir path to the output directory. + +optional arguments: + -h, --help show this help message and exit + --extension EXTENSION + extension of the audio files (defaults to .flac). +``` + +for example: + +``` +python encode.py soft path/to/LJSpeech-1.1/wavs path/to/LJSpeech-1.1/soft --extension .wav +``` + +At this point the directory tree should look like: + +``` +├───mels +│ ├───... +├───soft/discrete +│ ├───... +└───wavs + ├───... +``` + +### Step 4: Train the Acoustic-Model + +``` +usage: train.py [-h] [--resume RESUME] [--discrete] dataset-dir checkpoint-dir + +Train the acoustic model. + +positional arguments: + dataset-dir path to the data directory. + checkpoint-dir path to the checkpoint directory. + +optional arguments: + -h, --help show this help message and exit + --resume RESUME path to the checkpoint to resume from. + --discrete Use discrete units. +``` + +## Links + +- [Soft-VC repo](https://github.com/bshall/soft-vc) +- [Soft-VC paper](https://ieeexplore.ieee.org/abstract/document/9746484) +- [HuBERT content encoders](https://github.com/bshall/hubert) +- [HiFiGAN vocoder](https://github.com/bshall/hifigan) + +## Citation + +If you found this work helpful please consider citing our paper: + +``` +@inproceedings{ + soft-vc-2022, + author={van Niekerk, Benjamin and Carbonneau, Marc-André and Zaïdi, Julian and Baas, Matthew and Seuté, Hugo and Kamper, Herman}, + booktitle={ICASSP}, + title={A Comparison of Discrete and Soft Speech Units for Improved Voice Conversion}, + year={2022} +} +``` \ No newline at end of file diff --git a/voice_conversion/softvc/acoustic-model/acoustic-model.png b/voice_conversion/softvc/acoustic-model/acoustic-model.png new file mode 100644 index 0000000..41cc822 Binary files /dev/null and b/voice_conversion/softvc/acoustic-model/acoustic-model.png differ diff --git a/voice_conversion/softvc/acoustic-model/acoustic/__init__.py b/voice_conversion/softvc/acoustic-model/acoustic/__init__.py new file mode 100644 index 0000000..38186d0 --- /dev/null +++ b/voice_conversion/softvc/acoustic-model/acoustic/__init__.py @@ -0,0 +1 @@ +from .model import AcousticModel, hubert_discrete, hubert_soft diff --git a/voice_conversion/softvc/acoustic-model/acoustic/dataset.py b/voice_conversion/softvc/acoustic-model/acoustic/dataset.py new file mode 100644 index 0000000..97cdaa6 --- /dev/null +++ b/voice_conversion/softvc/acoustic-model/acoustic/dataset.py @@ -0,0 +1,55 @@ +from pathlib import Path +import numpy as np + +import torch +import torch.nn.functional as F +from torch.utils.data import Dataset +from torch.nn.utils.rnn import pad_sequence + + +class MelDataset(Dataset): + def __init__(self, root: Path, train: bool = True, discrete: bool = False): + self.discrete = discrete + self.mels_dir = root / "mels" + self.units_dir = root / "discrete" if discrete else root / "soft" + + pattern = "train/**/*.npy" if train else "dev/**/*.npy" + self.metadata = [ + path.relative_to(self.mels_dir).with_suffix("") + for path in self.mels_dir.rglob(pattern) + ] + + def __len__(self): + return len(self.metadata) + + def __getitem__(self, index): + path = self.metadata[index] + mel_path = self.mels_dir / path + units_path = self.units_dir / path + + mel = np.load(mel_path.with_suffix(".npy")).T + units = np.load(units_path.with_suffix(".npy")) + + length = 2 * units.shape[0] + + mel = torch.from_numpy(mel[:length, :]) + mel = F.pad(mel, (0, 0, 1, 0)) + units = torch.from_numpy(units) + if self.discrete: + units = units.long() + return mel, units + + def pad_collate(self, batch): + mels, units = zip(*batch) + + mels, units = list(mels), list(units) + + mels_lengths = torch.tensor([x.size(0) - 1 for x in mels]) + units_lengths = torch.tensor([x.size(0) for x in units]) + + mels = pad_sequence(mels, batch_first=True) + units = pad_sequence( + units, batch_first=True, padding_value=100 if self.discrete else 0 + ) + + return mels, mels_lengths, units, units_lengths diff --git a/voice_conversion/softvc/acoustic-model/acoustic/model.py b/voice_conversion/softvc/acoustic-model/acoustic/model.py new file mode 100644 index 0000000..00b2054 --- /dev/null +++ b/voice_conversion/softvc/acoustic-model/acoustic/model.py @@ -0,0 +1,168 @@ +import torch +import torch.nn as nn +from torch.nn.modules.utils import consume_prefix_in_state_dict_if_present + +URLS = { + "hubert-discrete": "https://github.com/bshall/acoustic-model/releases/download/v0.1/hubert-discrete-d49e1c77.pt", + "hubert-soft": "https://github.com/bshall/acoustic-model/releases/download/v0.1/hubert-soft-0321fd7e.pt", +} + + +class AcousticModel(nn.Module): + def __init__(self, discrete: bool = False, upsample: bool = True): + super().__init__() + self.encoder = Encoder(discrete, upsample) + self.decoder = Decoder() + + def forward(self, x: torch.Tensor, mels: torch.Tensor) -> torch.Tensor: + x = self.encoder(x) + return self.decoder(x, mels) + + @torch.inference_mode() + def generate(self, x: torch.Tensor) -> torch.Tensor: + x = self.encoder(x) + return self.decoder.generate(x) + + +class Encoder(nn.Module): + def __init__(self, discrete: bool = False, upsample: bool = True): + super().__init__() + self.embedding = nn.Embedding(100 + 1, 256) if discrete else None + self.prenet = PreNet(256, 256, 256) + self.convs = nn.Sequential( + nn.Conv1d(256, 512, 5, 1, 2), + nn.ReLU(), + nn.InstanceNorm1d(512), + nn.ConvTranspose1d(512, 512, 4, 2, 1) if upsample else nn.Identity(), + nn.Conv1d(512, 512, 5, 1, 2), + nn.ReLU(), + nn.InstanceNorm1d(512), + nn.Conv1d(512, 512, 5, 1, 2), + nn.ReLU(), + nn.InstanceNorm1d(512), + ) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + if self.embedding is not None: + x = self.embedding(x) + x = self.prenet(x) + x = self.convs(x.transpose(1, 2)) + return x.transpose(1, 2) + + +class Decoder(nn.Module): + def __init__(self): + super().__init__() + self.prenet = PreNet(128, 256, 256) + self.lstm1 = nn.LSTM(512 + 256, 768, batch_first=True) + self.lstm2 = nn.LSTM(768, 768, batch_first=True) + self.lstm3 = nn.LSTM(768, 768, batch_first=True) + self.proj = nn.Linear(768, 128, bias=False) + + def forward(self, x: torch.Tensor, mels: torch.Tensor) -> torch.Tensor: + mels = self.prenet(mels) + x, _ = self.lstm1(torch.cat((x, mels), dim=-1)) + res = x + x, _ = self.lstm2(x) + x = res + x + res = x + x, _ = self.lstm3(x) + x = res + x + return self.proj(x) + + @torch.inference_mode() + def generate(self, xs: torch.Tensor) -> torch.Tensor: + m = torch.zeros(xs.size(0), 128, device=xs.device) + h1 = torch.zeros(1, xs.size(0), 768, device=xs.device) + c1 = torch.zeros(1, xs.size(0), 768, device=xs.device) + h2 = torch.zeros(1, xs.size(0), 768, device=xs.device) + c2 = torch.zeros(1, xs.size(0), 768, device=xs.device) + h3 = torch.zeros(1, xs.size(0), 768, device=xs.device) + c3 = torch.zeros(1, xs.size(0), 768, device=xs.device) + + mel = [] + for x in torch.unbind(xs, dim=1): + m = self.prenet(m) + x = torch.cat((x, m), dim=1).unsqueeze(1) + x1, (h1, c1) = self.lstm1(x, (h1, c1)) + x2, (h2, c2) = self.lstm2(x1, (h2, c2)) + x = x1 + x2 + x3, (h3, c3) = self.lstm3(x, (h3, c3)) + x = x + x3 + m = self.proj(x).squeeze(1) + mel.append(m) + return torch.stack(mel, dim=1) + + +class PreNet(nn.Module): + def __init__( + self, + input_size: int, + hidden_size: int, + output_size: int, + dropout: float = 0.5, + ): + super().__init__() + self.net = nn.Sequential( + nn.Linear(input_size, hidden_size), + nn.ReLU(), + nn.Dropout(dropout), + nn.Linear(hidden_size, output_size), + nn.ReLU(), + nn.Dropout(dropout), + ) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + return self.net(x) + + +def _acoustic( + name: str, + discrete: bool, + upsample: bool, + pretrained: bool = True, + progress: bool = True, +) -> AcousticModel: + acoustic = AcousticModel(discrete, upsample) + if pretrained: + checkpoint = torch.hub.load_state_dict_from_url(URLS[name], progress=progress) + consume_prefix_in_state_dict_if_present(checkpoint["acoustic-model"], "module.") + acoustic.load_state_dict(checkpoint["acoustic-model"]) + acoustic.eval() + return acoustic + + +def hubert_discrete( + pretrained: bool = True, + progress: bool = True, +) -> AcousticModel: + r"""HuBERT-Discrete acoustic model from `"A Comparison of Discrete and Soft Speech Units for Improved Voice Conversion"`. + Args: + pretrained (bool): load pretrained weights into the model + progress (bool): show progress bar when downloading model + """ + return _acoustic( + "hubert-discrete", + discrete=True, + upsample=True, + pretrained=pretrained, + progress=progress, + ) + + +def hubert_soft( + pretrained: bool = True, + progress: bool = True, +) -> AcousticModel: + r"""HuBERT-Soft acoustic model from `"A Comparison of Discrete and Soft Speech Units for Improved Voice Conversion"`. + Args: + pretrained (bool): load pretrained weights into the model + progress (bool): show progress bar when downloading model + """ + return _acoustic( + "hubert-soft", + discrete=False, + upsample=True, + pretrained=pretrained, + progress=progress, + ) diff --git a/voice_conversion/softvc/acoustic-model/acoustic/utils.py b/voice_conversion/softvc/acoustic-model/acoustic/utils.py new file mode 100644 index 0000000..ad505fe --- /dev/null +++ b/voice_conversion/softvc/acoustic-model/acoustic/utils.py @@ -0,0 +1,99 @@ +import torch +import torch.nn.functional as F +import matplotlib + +import torchaudio.transforms as transforms + +matplotlib.use("Agg") +import matplotlib.pylab as plt + + +class Metric: + def __init__(self): + self.steps = 0 + self.value = 0 + + def update(self, value): + self.steps += 1 + self.value += (value - self.value) / self.steps + return self.value + + def reset(self): + self.steps = 0 + self.value = 0 + + +class LogMelSpectrogram(torch.nn.Module): + def __init__(self): + super().__init__() + self.melspctrogram = transforms.MelSpectrogram( + sample_rate=16000, + n_fft=1024, + win_length=1024, + hop_length=160, + center=False, + power=1.0, + norm="slaney", + onesided=True, + n_mels=128, + mel_scale="slaney", + ) + + def forward(self, wav): + padding = (1024 - 160) // 2 + wav = F.pad(wav, (padding, padding), "reflect") + mel = self.melspctrogram(wav) + logmel = torch.log(torch.clamp(mel, min=1e-5)) + return logmel + + +def save_checkpoint( + checkpoint_dir, + acoustic, + optimizer, + step, + loss, + best, + logger, +): + state = { + "acoustic-model": acoustic.state_dict(), + "optimizer": optimizer.state_dict(), + "step": step, + "loss": loss, + } + checkpoint_dir.mkdir(exist_ok=True, parents=True) + checkpoint_path = checkpoint_dir / f"model-{step}.pt" + torch.save(state, checkpoint_path) + if best: + best_path = checkpoint_dir / "model-best.pt" + torch.save(state, best_path) + logger.info(f"Saved checkpoint: {checkpoint_path.stem}") + + +def load_checkpoint( + load_path, + acoustic, + optimizer, + rank, + logger, +): + logger.info(f"Loading checkpoint from {load_path}") + checkpoint = torch.load(load_path, map_location={"cuda:0": f"cuda:{rank}"}) + acoustic.load_state_dict(checkpoint["acoustic-model"]) + if "optimizer" in checkpoint: + optimizer.load_state_dict(checkpoint["optimizer"]) + step = checkpoint.get("step", 0) + loss = checkpoint.get("loss", float("inf")) + return step, loss + + +def plot_spectrogram(spectrogram): + fig, ax = plt.subplots(figsize=(10, 2)) + im = ax.imshow(spectrogram, aspect="auto", origin="lower", interpolation="none") + plt.colorbar(im, ax=ax) + + fig.canvas.draw() + plt.close() + + return fig diff --git a/voice_conversion/softvc/acoustic-model/generate.py b/voice_conversion/softvc/acoustic-model/generate.py new file mode 100644 index 0000000..0696b6c --- /dev/null +++ b/voice_conversion/softvc/acoustic-model/generate.py @@ -0,0 +1,51 @@ +import argparse + +from pathlib import Path +import numpy as np +from tqdm import tqdm + +import torch + + +def generate(args): + print("Loading acoustic model checkpoint") + acoustic = torch.hub.load("bshall/acoustic-model:main", f"hubert_{args.model}").cuda() + + print(f"Generating from {args.in_dir} -> {args.out_dir}") + for path in tqdm(list(args.in_dir.rglob("*.npy"))): + units = np.load(path) + units_dtype = torch.long if args.model == "discrete" else torch.float + units = torch.tensor(units, dtype=units_dtype).cuda() + + with torch.inference_mode(): + mel_ = acoustic.generate(units) + mel_ = mel_.transpose(1, 2) + + out_path = args.out_dir / path.relative_to(args.in_dir) + out_path.parent.mkdir(exist_ok=True, parents=True) + np.save(out_path.with_suffix(".npy"), mel_.squeeze().cpu().numpy()) + + +if __name__ == "__main__": + parser = argparse.ArgumentParser( + description="Generate spectrograms from input speech units (discrete or soft)." + ) + parser.add_argument( + "model", + help="available models (HuBERT-Soft or HuBERT-Discrete)", + choices=["soft", "discrete"], + ) + parser.add_argument( + "in_dir", + metavar="in-dir", + help="path to the dataset directory.", + type=Path, + ) + parser.add_argument( + "out_dir", + metavar="out-dir", + help="path to the output directory.", + type=Path, + ) + args = parser.parse_args() + generate(args) diff --git a/voice_conversion/softvc/acoustic-model/hubconf.py b/voice_conversion/softvc/acoustic-model/hubconf.py new file mode 100644 index 0000000..8d3ea3a --- /dev/null +++ b/voice_conversion/softvc/acoustic-model/hubconf.py @@ -0,0 +1,3 @@ +dependencies = ["torch", "torchaudio"] + +from acoustic import hubert_discrete, hubert_soft diff --git a/voice_conversion/softvc/acoustic-model/mels.py b/voice_conversion/softvc/acoustic-model/mels.py new file mode 100644 index 0000000..7f06cc1 --- /dev/null +++ b/voice_conversion/softvc/acoustic-model/mels.py @@ -0,0 +1,64 @@ +import argparse +from pathlib import Path +from concurrent.futures import ProcessPoolExecutor +from multiprocessing import cpu_count + +import numpy as np +from tqdm import tqdm + +import torchaudio +from torchaudio.functional import resample + +from acoustic.utils import LogMelSpectrogram + + +melspectrogram = LogMelSpectrogram() + + +def process_wav(in_path, out_path): + wav, sr = torchaudio.load(in_path) + wav = resample(wav, sr, 16000) + + logmel = melspectrogram(wav.unsqueeze(0)) + + np.save(out_path.with_suffix(".npy"), logmel.squeeze().numpy()) + return out_path, logmel.shape[-1] + + +def preprocess_dataset(args): + args.out_dir.mkdir(parents=True, exist_ok=True) + + futures = [] + executor = ProcessPoolExecutor(max_workers=cpu_count()) + print(f"Extracting features for {args.in_dir}") + for in_path in args.in_dir.rglob("*.wav"): + relative_path = in_path.relative_to(args.in_dir) + out_path = args.out_dir / relative_path.with_suffix("") + out_path.parent.mkdir(parents=True, exist_ok=True) + futures.append(executor.submit(process_wav, in_path, out_path)) + + results = [future.result() for future in tqdm(futures)] + + lengths = {path.stem: length for path, length in results} + frames = sum(lengths.values()) + frame_shift_ms = 160 / 16000 + hours = frames * frame_shift_ms / 3600 + print(f"Wrote {len(lengths)} utterances, {frames} frames ({hours:.2f} hours)") + + +if __name__ == "__main__": + parser = argparse.ArgumentParser(description="Extract mel-spectrograms for an audio dataset.") + parser.add_argument( + "in_dir", + metavar="in-dir", + help="path to the dataset directory.", + type=Path, + ) + parser.add_argument( + "out_dir", + metavar="out-dir", + help="path to the output directory.", + type=Path, + ) + args = parser.parse_args() + preprocess_dataset(args) diff --git a/voice_conversion/softvc/acoustic-model/train.py b/voice_conversion/softvc/acoustic-model/train.py new file mode 100644 index 0000000..88972a3 --- /dev/null +++ b/voice_conversion/softvc/acoustic-model/train.py @@ -0,0 +1,311 @@ +import argparse +import logging +from pathlib import Path + +import torch +import torch.nn.functional as F +import torch.optim as optim +from torch.utils.data import DataLoader +from torch.utils.tensorboard import SummaryWriter +import torch.distributed as dist +from torch.utils.data.distributed import DistributedSampler +import torch.multiprocessing as mp +from torch.nn.parallel import DistributedDataParallel as DDP + +from acoustic import AcousticModel +from acoustic.dataset import MelDataset +from acoustic.utils import Metric, save_checkpoint, load_checkpoint, plot_spectrogram + + +logging.basicConfig(level=logging.DEBUG) +logger = logging.getLogger(__name__) + +######################################################################################## +# Define hyperparameters for training: +######################################################################################## + +BATCH_SIZE = 32 +LEARNING_RATE = 4e-4 +BETAS = (0.8, 0.99) +WEIGHT_DECAY = 1e-5 +STEPS = 80000 +LOG_INTERVAL = 5 +VALIDATION_INTERVAL = 1000 +CHECKPOINT_INTERVAL = 1000 +BACKEND = "nccl" +INIT_METHOD = "tcp://localhost:54321" + + +def train(rank, world_size, args): + dist.init_process_group( + BACKEND, + rank=rank, + world_size=world_size, + init_method=INIT_METHOD, + ) + + #################################################################################### + # Setup logging utilities: + #################################################################################### + + log_dir = args.checkpoint_dir / "logs" + log_dir.mkdir(exist_ok=True, parents=True) + + if rank == 0: + logger.setLevel(logging.INFO) + handler = logging.FileHandler(log_dir / f"{args.checkpoint_dir.stem}.log") + handler.setLevel(logging.INFO) + formatter = logging.Formatter( + "%(asctime)s [%(levelname)s] %(message)s", datefmt="%m/%d/%Y %I:%M:%S" + ) + handler.setFormatter(formatter) + logger.addHandler(handler) + else: + logger.setLevel(logging.ERROR) + + writer = SummaryWriter(log_dir) if rank == 0 else None + + #################################################################################### + # Initialize models and optimizer + #################################################################################### + + acoustic = AcousticModel().to(rank) + + acoustic = DDP(acoustic, device_ids=[rank]) + + optimizer = optim.AdamW( + acoustic.parameters(), + lr=LEARNING_RATE, + betas=BETAS, + weight_decay=WEIGHT_DECAY, + ) + + #################################################################################### + # Initialize datasets and dataloaders + #################################################################################### + + train_dataset = MelDataset( + root=args.dataset_dir, + train=True, + discrete=args.discrete, + ) + train_sampler = DistributedSampler(train_dataset, drop_last=True) + train_loader = DataLoader( + train_dataset, + batch_size=BATCH_SIZE, + sampler=train_sampler, + collate_fn=train_dataset.pad_collate, + num_workers=8, + pin_memory=True, + shuffle=False, + drop_last=True, + ) + + validation_dataset = MelDataset( + root=args.dataset_dir, + train=False, + ) + validation_loader = DataLoader( + validation_dataset, + batch_size=1, + shuffle=False, + num_workers=8, + pin_memory=True, + ) + + #################################################################################### + # Load checkpoint if args.resume is set + #################################################################################### + + if args.resume is not None: + global_step, best_loss = load_checkpoint( + load_path=args.resume, + acoustic=acoustic, + optimizer=optimizer, + rank=rank, + logger=logger, + ) + else: + global_step, best_loss = 0, float("inf") + + # =================================================================================# + # Start training loop + # =================================================================================# + + n_epochs = STEPS // len(train_loader) + 1 + start_epoch = global_step // len(train_loader) + 1 + + logger.info("**" * 40) + logger.info(f"PyTorch version: {torch.__version__}") + logger.info(f"CUDA version: {torch.version.cuda}") + logger.info(f"CUDNN version: {torch.backends.cudnn.version()}") + logger.info(f"CUDNN enabled: {torch.backends.cudnn.enabled}") + logger.info(f"CUDNN deterministic: {torch.backends.cudnn.deterministic}") + logger.info(f"CUDNN benchmark: {torch.backends.cudnn.benchmark}") + logger.info(f"# of GPUS: {torch.cuda.device_count()}") + logger.info(f"batch size: {BATCH_SIZE}") + logger.info(f"iterations per epoch: {len(train_loader)}") + logger.info(f"# of epochs: {n_epochs}") + logger.info(f"started at epoch: {start_epoch}") + logger.info("**" * 40 + "\n") + + average_loss = Metric() + epoch_loss = Metric() + + validation_loss = Metric() + + for epoch in range(start_epoch, n_epochs + 1): + train_sampler.set_epoch(epoch) + + acoustic.train() + epoch_loss.reset() + + for mels, mels_lengths, units, units_lengths in train_loader: + mels, mels_lengths = mels.to(rank), mels_lengths.to(rank) + units, units_lengths = units.to(rank), units_lengths.to(rank) + + ############################################################################ + # Compute training loss + ############################################################################ + + optimizer.zero_grad() + + mels_ = acoustic(units, mels[:, :-1, :]) + + loss = F.l1_loss(mels_, mels[:, 1:, :], reduction="none") + loss = torch.sum(loss, dim=(1, 2)) / (mels_.size(-1) * mels_lengths) + loss = torch.mean(loss) + + loss.backward() + optimizer.step() + + global_step += 1 + + ############################################################################ + # Update and log training metrics + ############################################################################ + + average_loss.update(loss.item()) + epoch_loss.update(loss.item()) + + if rank == 0 and global_step % LOG_INTERVAL == 0: + writer.add_scalar( + "train/loss", + average_loss.value, + global_step, + ) + average_loss.reset() + + # --------------------------------------------------------------------------# + # Start validation loop + # --------------------------------------------------------------------------# + + if global_step % VALIDATION_INTERVAL == 0: + acoustic.eval() + validation_loss.reset() + + for i, (mels, units) in enumerate(validation_loader, 1): + mels, units = mels.to(rank), units.to(rank) + + with torch.no_grad(): + mels_ = acoustic(units, mels[:, :-1, :]) + loss = F.l1_loss(mels_, mels[:, 1:, :]) + + #################################################################### + # Update validation metrics and log generated mels + #################################################################### + + validation_loss.update(loss.item()) + + if rank == 0: + writer.add_figure( + f"generated/mel_{i}", + plot_spectrogram( + mels_.squeeze().transpose(0, 1).cpu().numpy() + ), + global_step, + ) + + acoustic.train() + + ############################################################################ + # Log validation metrics + ############################ + + if rank == 0: + writer.add_scalar( + "validation/loss", + validation_loss.value, + global_step, + ) + logger.info( + f"valid -- epoch: {epoch}, loss: {validation_loss.value:.4f}" + ) + + new_best = best_loss > validation_loss.value + if new_best or global_step % CHECKPOINT_INTERVAL: + if new_best: + logger.info("-------- new best model found!") + best_loss = validation_loss.value + + if rank == 0: + save_checkpoint( + checkpoint_dir=args.checkpoint_dir, + acoustic=acoustic, + optimizer=optimizer, + step=global_step, + loss=validation_loss.value, + best=new_best, + logger=logger, + ) + + # -----------------------------------------------------------------------------# + # End validation loop + # -----------------------------------------------------------------------------# + + #################################################################################### + # Log training metrics + #################################################################################### + + logger.info(f"train -- epoch: {epoch}, loss: {epoch_loss.value:.4f}") + + # =================================================================================# + # End training loop + # ==================================================================================# + + dist.destroy_process_group() + + +if __name__ == "__main__": + parser = argparse.ArgumentParser(description="Train the acoustic model.") + parser.add_argument( + "dataset_dir", + metavar="dataset-dir", + help="path to the data directory.", + type=Path, + ) + parser.add_argument( + "checkpoint_dir", + metavar="checkpoint-dir", + help="path to the checkpoint directory.", + type=Path, + ) + parser.add_argument( + "--resume", + help="path to the checkpoint to resume from.", + type=Path, + ) + parser.add_argument( + "--discrete", + action="store_true", + help="use discrete units.", + ) + args = parser.parse_args() + + world_size = torch.cuda.device_count() + mp.spawn( + train, + args=(world_size, args), + nprocs=world_size, + join=True, + ) diff --git a/voice_conversion/softvc/hifigan/.gitignore b/voice_conversion/softvc/hifigan/.gitignore new file mode 100644 index 0000000..0202868 --- /dev/null +++ b/voice_conversion/softvc/hifigan/.gitignore @@ -0,0 +1,132 @@ +# Byte-compiled / optimized / DLL files +__pycache__/ +*.py[cod] +*$py.class + +# C extensions +*.so + +# Distribution / packaging +.Python +build/ +develop-eggs/ +dist/ +downloads/ +eggs/ +.eggs/ +lib/ +lib64/ +parts/ +sdist/ +var/ +wheels/ +pip-wheel-metadata/ +share/python-wheels/ +*.egg-info/ +.installed.cfg +*.egg +MANIFEST + +# PyInstaller +# Usually these files are written by a python script from a template +# before PyInstaller builds the exe, so as to inject date/other infos into it. +*.manifest +*.spec + +# Installer logs +pip-log.txt +pip-delete-this-directory.txt + +# Unit test / coverage reports +htmlcov/ +.tox/ +.nox/ +.coverage +.coverage.* +.cache +nosetests.xml +coverage.xml +*.cover +*.py,cover +.hypothesis/ +.pytest_cache/ + +# Translations +*.mo +*.pot + +# Django stuff: +*.log +local_settings.py +db.sqlite3 +db.sqlite3-journal + +# Flask stuff: +instance/ +.webassets-cache + +# Scrapy stuff: +.scrapy + +# Sphinx documentation +docs/_build/ + +# PyBuilder +target/ + +# Jupyter Notebook +.ipynb_checkpoints + +# IPython +profile_default/ +ipython_config.py + +# pyenv +.python-version + +# pipenv +# According to pypa/pipenv#598, it is recommended to include Pipfile.lock in version control. +# However, in case of collaboration, if having platform-specific dependencies or dependencies +# having no cross-platform support, pipenv may install dependencies that don't work, or not +# install all needed dependencies. +#Pipfile.lock + +# PEP 582; used by e.g. github.com/David-OConnor/pyflow +__pypackages__/ + +# Celery stuff +celerybeat-schedule +celerybeat.pid + +# SageMath parsed files +*.sage.py + +# Environments +.env +.venv +env/ +venv/ +ENV/ +env.bak/ +venv.bak/ + +# VSCode project settings +.vscode + +# Spyder project settings +.spyderproject +.spyproject + +# Rope project settings +.ropeproject + +# mkdocs documentation +/site + +# mypy +.mypy_cache/ +.dmypy.json +dmypy.json + +# Pyre type checker +.pyre/ diff --git a/voice_conversion/softvc/hifigan/LICENSE b/voice_conversion/softvc/hifigan/LICENSE new file mode 100644 index 0000000..6eb2af0 --- /dev/null +++ b/voice_conversion/softvc/hifigan/LICENSE @@ -0,0 +1,21 @@ +MIT License + +Copyright (c) 2021 Benjamin van Niekerk + +Permission is hereby granted, free of charge, to any person obtaining a copy +of this software and associated documentation files (the "Software"), to deal +in the Software without restriction, including without limitation the rights +to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +copies of the Software, and to permit persons to whom the Software is +furnished to do so, subject to the following conditions: + +The above copyright notice and this permission notice shall be included in all +copies or substantial portions of the Software. + +THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE +SOFTWARE. diff --git a/voice_conversion/softvc/hifigan/README.md b/voice_conversion/softvc/hifigan/README.md new file mode 100644 index 0000000..d963302 --- /dev/null +++ b/voice_conversion/softvc/hifigan/README.md @@ -0,0 +1,139 @@ +

+ + Open In Colab + +

+ +# HiFi-GAN + +Training and inference scripts for the vocoder models in [A Comparison of Discrete and Soft Speech Units for Improved Voice Conversion](https://ieeexplore.ieee.org/abstract/document/9746484). For more details see [soft-vc](https://github.com/bshall/soft-vc). Audio samples can be found [here](https://bshall.github.io/soft-vc/). Colab demo can be found [here](https://colab.research.google.com/github/bshall/soft-vc/blob/main/soft-vc-demo.ipynb). + +
+ Soft-VC +
+
+ + Fig 1: Architecture of the voice conversion system. a) The discrete content encoder clusters audio features to produce a sequence of discrete speech units. b) The soft content encoder is trained to predict the discrete units. The acoustic model transforms the discrete/soft speech units into a target spectrogram. The vocoder converts the spectrogram into an audio waveform. + +
+ +## Example Usage + +### Programmatic Usage + +```python +import torch +import numpy as np + +# Load checkpoint +hifigan = torch.hub.load("bshall/hifigan:main", "hifigan_hubert_soft").cuda() +# Load mel-spectrogram +mel = torch.from_numpy(np.load("path/to/mel")).unsqueeze(0).cuda() +# Generate +wav, sr = hifigan.generate(mel) +``` + +### Script-Based Usage + +``` +usage: generate.py [-h] {soft,discrete,base} in-dir out-dir + +Generate audio for a directory of mel-spectrogams using HiFi-GAN. + +positional arguments: + {soft,discrete,base} available models (HuBERT-Soft, HuBERT-Discrete, or + Base). + in-dir path to input directory containing the mel- + spectrograms. + out-dir path to output directory. + +optional arguments: + -h, --help show this help message and exit +``` + +## Training + +### Step 1: Dataset Preparation + +Download and extract the [LJSpeech](https://keithito.com/LJ-Speech-Dataset/) dataset. The training script expects the following tree structure for the dataset directory: + +``` +└───wavs + ├───dev + │ ├───LJ001-0001.wav + │ ├───... + │ └───LJ050-0278.wav + └───train + ├───LJ002-0332.wav + ├───... + └───LJ047-0007.wav +``` + +The `train` and `dev` directories should contain the training and validation splits respectively. The splits used for the paper can be found [here](https://github.com/bshall/hifigan/releases/tag/v0.1). + +### Step 2: Resample the Audio + +Resample the audio to 16kHz using the `resample.py` script: + +``` +usage: resample.py [-h] [--sample-rate SAMPLE_RATE] in-dir out-dir + +Resample an audio dataset. + +positional arguments: + in-dir path to the dataset directory. + out-dir path to the output directory. + +optional arguments: + -h, --help show this help message and exit + --sample-rate SAMPLE_RATE + target sample rate (default 16kHz) +``` + +for example: + +``` +python reample.py path/to/LJSpeech-1.1/ path/to/LJSpeech-Resampled/ +``` + +### Step 3: Train HifiGAN + +``` +usage: train.py [-h] [--resume RESUME] [--finetune] dataset-dir checkpoint-dir + +Train or finetune HiFi-GAN. + +positional arguments: + dataset-dir path to the preprocessed data directory + checkpoint-dir path to the checkpoint directory + +optional arguments: + -h, --help show this help message and exit + --resume RESUME path to the checkpoint to resume from + --finetune whether to finetune (note that a resume path must be given) +``` + +## Links + +- [Soft-VC repo](https://github.com/bshall/soft-vc) +- [Soft-VC paper](https://ieeexplore.ieee.org/abstract/document/9746484) +- [HuBERT content encoders](https://github.com/bshall/hubert) +- [Acoustic models](https://github.com/bshall/acoustic-model) + +## Citation + +If you found this work helpful please consider citing our paper: + +``` +@inproceedings{ + soft-vc-2022, + author={van Niekerk, Benjamin and Carbonneau, Marc-André and Zaïdi, Julian and Baas, Matthew and Seuté, Hugo and Kamper, Herman}, + booktitle={ICASSP}, + title={A Comparison of Discrete and Soft Speech Units for Improved Voice Conversion}, + year={2022} +} +``` + +## Acknowledgements +This repo is based heavily on [https://github.com/jik876/hifi-gan](https://github.com/jik876/hifi-gan). \ No newline at end of file diff --git a/voice_conversion/softvc/hifigan/generate.py b/voice_conversion/softvc/hifigan/generate.py new file mode 100644 index 0000000..22c58cf --- /dev/null +++ b/voice_conversion/softvc/hifigan/generate.py @@ -0,0 +1,50 @@ +from pathlib import Path +import numpy as np +import argparse +import torch +import torchaudio +from tqdm import tqdm + + +def generate(args): + print("Loading checkpoint") + model_name = f"hifigan_hubert_{args.model}" if args.model != "base" else "hifigan" + hifigan = torch.hub.load("bshall/hifigan:main", model_name).cuda() + + print(f"Generating audio from {args.in_dir}") + for path in tqdm(list(args.in_dir.rglob("*.npy"))): + mel = torch.from_numpy(np.load(path)) + mel = mel.unsqueeze(0).cuda() + + wav, sr = hifigan.generate(mel) + wav = wav.squeeze(0).cpu() + + out_path = args.out_dir / path.relative_to(args.in_dir) + out_path.parent.mkdir(exist_ok=True, parents=True) + torchaudio.save(out_path.with_suffix(".wav"), wav, sr) + + +if __name__ == "__main__": + parser = argparse.ArgumentParser( + description="Generate audio for a directory of mel-spectrogams using HiFi-GAN." + ) + parser.add_argument( + "model", + help="available models (HuBERT-Soft, HuBERT-Discrete, or Base).", + choices=["soft", "discrete", "base"], + ) + parser.add_argument( + "in_dir", + metavar="in-dir", + help="path to input directory containing the mel-spectrograms.", + type=Path, + ) + parser.add_argument( + "out_dir", + metavar="out-dir", + help="path to output directory.", + type=Path, + ) + args = parser.parse_args() + + generate(args) diff --git a/voice_conversion/softvc/hifigan/hifigan/__init__.py b/voice_conversion/softvc/hifigan/hifigan/__init__.py new file mode 100644 index 0000000..e594d66 --- /dev/null +++ b/voice_conversion/softvc/hifigan/hifigan/__init__.py @@ -0,0 +1 @@ +from .generator import hifigan, hifigan_hubert_discrete, hifigan_hubert_soft \ No newline at end of file diff --git a/voice_conversion/softvc/hifigan/hifigan/dataset.py b/voice_conversion/softvc/hifigan/hifigan/dataset.py new file mode 100644 index 0000000..0142cd0 --- /dev/null +++ b/voice_conversion/softvc/hifigan/hifigan/dataset.py @@ -0,0 +1,126 @@ +from pathlib import Path +import math +import random +import numpy as np +import torch +import torch.nn.functional as F + +from torch.utils.data import Dataset + +import torchaudio +import torchaudio.transforms as transforms + + +class LogMelSpectrogram(torch.nn.Module): + def __init__(self): + super().__init__() + self.melspctrogram = transforms.MelSpectrogram( + sample_rate=16000, + n_fft=1024, + win_length=1024, + hop_length=160, + center=False, + power=1.0, + norm="slaney", + onesided=True, + n_mels=128, + mel_scale="slaney", + ) + + def forward(self, wav): + wav = F.pad(wav, ((1024 - 160) // 2, (1024 - 160) // 2), "reflect") + mel = self.melspctrogram(wav) + logmel = torch.log(torch.clamp(mel, min=1e-5)) + return logmel + + +class MelDataset(Dataset): + def __init__( + self, + root: Path, + segment_length: int, + sample_rate: int, + hop_length: int, + train: bool = True, + finetune: bool = False, + ): + self.wavs_dir = root / "wavs" + self.mels_dir = root / "mels" + self.data_dir = self.wavs_dir if not finetune else self.mels_dir + + self.segment_length = segment_length + self.sample_rate = sample_rate + self.hop_length = hop_length + self.train = train + self.finetune = finetune + + suffix = ".wav" if not finetune else ".npy" + pattern = f"train/**/*{suffix}" if train else "dev/**/*{suffix}" + + self.metadata = [ + path.relative_to(self.data_dir).with_suffix("") + for path in self.data_dir.rglob(pattern) + ] + + self.logmel = LogMelSpectrogram() + + def __len__(self): + return len(self.metadata) + + def __getitem__(self, index): + path = self.metadata[index] + wav_path = self.wavs_dir / path + + info = torchaudio.info(wav_path.with_suffix(".wav")) + if info.sample_rate != self.sample_rate: + raise ValueError( + f"Sample rate {info.sample_rate} doesn't match target of {self.sample_rate}" + ) + + if self.finetune: + mel_path = self.mels_dir / path + src_logmel = torch.from_numpy(np.load(mel_path.with_suffix(".npy"))) + src_logmel = src_logmel.unsqueeze(0) + + mel_frames_per_segment = math.ceil(self.segment_length / self.hop_length) + mel_diff = src_logmel.size(-1) - mel_frames_per_segment if self.train else 0 + mel_offset = random.randint(0, max(mel_diff, 0)) + + frame_offset = self.hop_length * mel_offset + else: + frame_diff = info.num_frames - self.segment_length + frame_offset = random.randint(0, max(frame_diff, 0)) + + wav, _ = torchaudio.load( + filepath=wav_path.with_suffix(".wav"), + frame_offset=frame_offset if self.train else 0, + num_frames=self.segment_length if self.train else -1, + ) + + if wav.size(-1) < self.segment_length: + wav = F.pad(wav, (0, self.segment_length - wav.size(-1))) + + if not self.finetune and self.train: + gain = random.random() * (0.99 - 0.4) + 0.4 + flip = -1 if random.random() > 0.5 else 1 + wav = flip * gain * wav / wav.abs().max() + + tgt_logmel = self.logmel(wav.unsqueeze(0)).squeeze(0) + + if self.finetune: + if self.train: + src_logmel = src_logmel[ + :, :, mel_offset : mel_offset + mel_frames_per_segment + ] + + if src_logmel.size(-1) < mel_frames_per_segment: + src_logmel = F.pad( + src_logmel, + (0, mel_frames_per_segment - src_logmel.size(-1)), + "constant", + src_logmel.min(), + ) + else: + src_logmel = tgt_logmel.clone() + + return wav, src_logmel, tgt_logmel diff --git a/voice_conversion/softvc/hifigan/hifigan/discriminator.py b/voice_conversion/softvc/hifigan/hifigan/discriminator.py new file mode 100644 index 0000000..8f53c3d --- /dev/null +++ b/voice_conversion/softvc/hifigan/hifigan/discriminator.py @@ -0,0 +1,262 @@ +# adopted from https://github.com/jik876/hifi-gan/blob/master/models.py +import torch +import torch.nn as nn +import torch.nn.functional as F +from typing import Tuple, List + +from hifigan.utils import get_padding + + +LRELU_SLOPE = 0.1 + + +class PeriodDiscriminator(torch.nn.Module): + """HiFiGAN Period Discriminator""" + + def __init__( + self, + period: int, + kernel_size: int = 5, + stride: int = 3, + use_spectral_norm: bool = False, + ) -> None: + super().__init__() + self.period = period + norm_f = nn.utils.spectral_norm if use_spectral_norm else nn.utils.weight_norm + self.convs = nn.ModuleList( + [ + norm_f( + nn.Conv2d( + 1, + 32, + (kernel_size, 1), + (stride, 1), + padding=(get_padding(5, 1), 0), + ) + ), + norm_f( + nn.Conv2d( + 32, + 128, + (kernel_size, 1), + (stride, 1), + padding=(get_padding(5, 1), 0), + ) + ), + norm_f( + nn.Conv2d( + 128, + 512, + (kernel_size, 1), + (stride, 1), + padding=(get_padding(5, 1), 0), + ) + ), + norm_f( + nn.Conv2d( + 512, + 1024, + (kernel_size, 1), + (stride, 1), + padding=(get_padding(5, 1), 0), + ) + ), + norm_f(nn.Conv2d(1024, 1024, (kernel_size, 1), 1, padding=(2, 0))), + ] + ) + self.conv_post = norm_f(nn.Conv2d(1024, 1, (3, 1), 1, padding=(1, 0))) + + def forward(self, x: torch.Tensor) -> Tuple[torch.Tensor, List[torch.Tensor]]: + """ + Args: + x (Tensor): input waveform. + Returns: + [Tensor]: discriminator scores per sample in the batch. + [List[Tensor]]: list of features from each convolutional layer. + """ + feat = [] + + # 1d to 2d + b, c, t = x.shape + if t % self.period != 0: # pad first + n_pad = self.period - (t % self.period) + x = F.pad(x, (0, n_pad), "reflect") + t = t + n_pad + x = x.view(b, c, t // self.period, self.period) + + for l in self.convs: + x = l(x) + x = F.leaky_relu(x, LRELU_SLOPE) + feat.append(x) + x = self.conv_post(x) + feat.append(x) + x = torch.flatten(x, 1, -1) + + return x, feat + + +class MultiPeriodDiscriminator(torch.nn.Module): + """HiFiGAN Multi-Period Discriminator (MPD)""" + + def __init__(self): + super().__init__() + self.discriminators = nn.ModuleList( + [ + PeriodDiscriminator(2), + PeriodDiscriminator(3), + PeriodDiscriminator(5), + PeriodDiscriminator(7), + PeriodDiscriminator(11), + ] + ) + + def forward( + self, x: torch.Tensor + ) -> Tuple[List[torch.Tensor], List[List[torch.Tensor]]]: + """ + Args: + x (Tensor): input waveform. + Returns: + [List[Tensor]]: list of scores from each discriminator. + [List[List[Tensor]]]: list of features from each discriminator's convolutional layers. + """ + scores = [] + feats = [] + for _, d in enumerate(self.discriminators): + score, feat = d(x) + scores.append(score) + feats.append(feat) + return scores, feats + + +class ScaleDiscriminator(torch.nn.Module): + """HiFiGAN Scale Discriminator.""" + + def __init__(self, use_spectral_norm: bool = False) -> None: + super().__init__() + norm_f = nn.utils.spectral_norm if use_spectral_norm else nn.utils.weight_norm + self.convs = nn.ModuleList( + [ + norm_f(nn.Conv1d(1, 128, 15, 1, padding=7)), + norm_f(nn.Conv1d(128, 128, 41, 2, groups=4, padding=20)), + norm_f(nn.Conv1d(128, 256, 41, 2, groups=16, padding=20)), + norm_f(nn.Conv1d(256, 512, 41, 4, groups=16, padding=20)), + norm_f(nn.Conv1d(512, 1024, 41, 4, groups=16, padding=20)), + norm_f(nn.Conv1d(1024, 1024, 41, 1, groups=16, padding=20)), + norm_f(nn.Conv1d(1024, 1024, 5, 1, padding=2)), + ] + ) + self.conv_post = norm_f(nn.Conv1d(1024, 1, 3, 1, padding=1)) + + def forward(self, x: torch.Tensor) -> Tuple[torch.Tensor, List[torch.Tensor]]: + """ + Args: + x (Tensor): input waveform. + Returns: + Tensor: discriminator scores. + List[Tensor]: list of features from the convolutional layers. + """ + feat = [] + for l in self.convs: + x = l(x) + x = F.leaky_relu(x, LRELU_SLOPE) + feat.append(x) + x = self.conv_post(x) + feat.append(x) + x = torch.flatten(x, 1, -1) + return x, feat + + +class MultiScaleDiscriminator(torch.nn.Module): + """HiFiGAN Multi-Scale Discriminator.""" + + def __init__(self): + super().__init__() + self.discriminators = nn.ModuleList( + [ + ScaleDiscriminator(use_spectral_norm=True), + ScaleDiscriminator(), + ScaleDiscriminator(), + ] + ) + self.meanpools = nn.ModuleList( + [nn.AvgPool1d(4, 2, padding=2), nn.AvgPool1d(4, 2, padding=2)] + ) + + def forward( + self, x: torch.Tensor + ) -> Tuple[List[torch.Tensor], List[List[torch.Tensor]]]: + """ + Args: + x (Tensor): input waveform. + Returns: + List[Tensor]: discriminator scores. + List[List[Tensor]]: list of features from each discriminator's convolutional layers. + """ + scores = [] + feats = [] + for i, d in enumerate(self.discriminators): + if i != 0: + x = self.meanpools[i - 1](x) + score, feat = d(x) + scores.append(score) + feats.append(feat) + return scores, feats + + +class HifiganDiscriminator(nn.Module): + """HiFiGAN discriminator""" + + def __init__(self): + super().__init__() + self.mpd = MultiPeriodDiscriminator() + self.msd = MultiScaleDiscriminator() + + def forward( + self, x: torch.Tensor + ) -> Tuple[List[torch.Tensor], List[List[torch.Tensor]]]: + """ + Args: + x (Tensor): input waveform. + Returns: + List[Tensor]: discriminator scores. + List[List[Tensor]]: list of features from from each discriminator's convolutional layers. + """ + scores, feats = self.mpd(x) + scores_, feats_ = self.msd(x) + return scores + scores_, feats + feats_ + + +def feature_loss( + features_real: List[List[torch.Tensor]], features_generate: List[List[torch.Tensor]] +) -> float: + loss = 0 + for r, g in zip(features_real, features_generate): + for rl, gl in zip(r, g): + loss += torch.mean(torch.abs(rl - gl)) + return loss * 2 + + +def discriminator_loss(real, generated): + loss = 0 + real_losses = [] + generated_losses = [] + for r, g in zip(real, generated): + r_loss = torch.mean((1 - r) ** 2) + g_loss = torch.mean(g ** 2) + loss += r_loss + g_loss + real_losses.append(r_loss.item()) + generated_losses.append(g_loss.item()) + + return loss, real_losses, generated_losses + + +def generator_loss(discriminator_outputs): + loss = 0 + generator_losses = [] + for x in discriminator_outputs: + l = torch.mean((1 - x) ** 2) + generator_losses.append(l) + loss += l + + return loss, generator_losses diff --git a/voice_conversion/softvc/hifigan/hifigan/generator.py b/voice_conversion/softvc/hifigan/hifigan/generator.py new file mode 100644 index 0000000..17282c3 --- /dev/null +++ b/voice_conversion/softvc/hifigan/hifigan/generator.py @@ -0,0 +1,282 @@ +# adapted from https://github.com/jik876/hifi-gan/blob/master/models.py +import torch +import torch.nn as nn +import torch.nn.functional as F +from torch.nn.utils import remove_weight_norm, weight_norm +from torch.nn.modules.utils import consume_prefix_in_state_dict_if_present +from typing import Tuple + +from hifigan.utils import get_padding + + +URLS = { + "hifigan": "https://github.com/bshall/hifigan/releases/download/v0.1/hifigan-67926ec6.pt", + "hifigan-hubert-soft": "https://github.com/bshall/hifigan/releases/download/v0.1/hifigan-hubert-discrete-bbad3043.pt", + "hifigan-hubert-discrete": "https://github.com/bshall/hifigan/releases/download/v0.1/hifigan-hubert-soft-65f03469.pt", +} + +LRELU_SLOPE = 0.1 + + +class HifiganGenerator(torch.nn.Module): + def __init__( + self, + in_channels: int = 128, + resblock_dilation_sizes: Tuple[Tuple[int, ...], ...] = ( + (1, 3, 5), + (1, 3, 5), + (1, 3, 5), + ), + resblock_kernel_sizes: Tuple[int, ...] = (3, 7, 11), + upsample_kernel_sizes: Tuple[int, ...] = (20, 8, 4, 4), + upsample_initial_channel: int = 512, + upsample_factors: int = (10, 4, 2, 2), + inference_padding: int = 5, + sample_rate: int = 16000, + ) -> None: + r"""HiFiGAN Generator + Args: + in_channels (int): number of input channels. + resblock_dilation_sizes (Tuple[Tuple[int, ...], ...]): list of dilation values in each layer of a `ResBlock`. + resblock_kernel_sizes (Tuple[int, ...]): list of kernel sizes for each `ResBlock`. + upsample_kernel_sizes (Tuple[int, ...]): list of kernel sizes for each transposed convolution. + upsample_initial_channel (int): number of channels for the first upsampling layer. This is divided by 2 + for each consecutive upsampling layer. + upsample_factors (Tuple[int, ...]): upsampling factors (stride) for each upsampling layer. + inference_padding (int): constant padding applied to the input at inference time. + sample_rate (int): sample rate of the generated audio. + """ + super().__init__() + self.inference_padding = inference_padding + self.num_kernels = len(resblock_kernel_sizes) + self.num_upsamples = len(upsample_factors) + self.sample_rate = sample_rate + # initial upsampling layers + self.conv_pre = weight_norm( + nn.Conv1d(in_channels, upsample_initial_channel, 7, 1, padding=3) + ) + + # upsampling layers + self.ups = nn.ModuleList() + for i, (u, k) in enumerate(zip(upsample_factors, upsample_kernel_sizes)): + self.ups.append( + weight_norm( + nn.ConvTranspose1d( + upsample_initial_channel // (2 ** i), + upsample_initial_channel // (2 ** (i + 1)), + k, + u, + padding=(k - u) // 2, + ) + ) + ) + # MRF blocks + self.resblocks = nn.ModuleList() + for i in range(len(self.ups)): + ch = upsample_initial_channel // (2 ** (i + 1)) + for _, (k, d) in enumerate( + zip(resblock_kernel_sizes, resblock_dilation_sizes) + ): + self.resblocks.append(ResBlock1(ch, k, d)) + # post convolution layer + self.conv_post = weight_norm(nn.Conv1d(ch, 1, 7, 1, padding=3)) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + o = self.conv_pre(x) + for i in range(self.num_upsamples): + o = F.leaky_relu(o, LRELU_SLOPE) + o = self.ups[i](o) + z_sum = None + for j in range(self.num_kernels): + if z_sum is None: + z_sum = self.resblocks[i * self.num_kernels + j](o) + else: + z_sum += self.resblocks[i * self.num_kernels + j](o) + o = z_sum / self.num_kernels + o = F.leaky_relu(o) + o = self.conv_post(o) + o = torch.tanh(o) + return o + + @torch.no_grad() + def generate(self, x: torch.Tensor) -> torch.Tensor: + x = F.pad(x, (self.inference_padding, self.inference_padding), "replicate") + return self(x), self.sample_rate + + def remove_weight_norm(self): + print("Removing weight norm...") + for l in self.ups: + remove_weight_norm(l) + for l in self.resblocks: + l.remove_weight_norm() + remove_weight_norm(self.conv_pre) + remove_weight_norm(self.conv_post) + + +class ResBlock1(torch.nn.Module): + def __init__( + self, channels: int, kernel_size: int = 3, dilation: Tuple[int, ...] = (1, 3, 5) + ) -> None: + super().__init__() + self.convs1 = nn.ModuleList( + [ + weight_norm( + nn.Conv1d( + channels, + channels, + kernel_size, + 1, + dilation=dilation[0], + padding=get_padding(kernel_size, dilation[0]), + ) + ), + weight_norm( + nn.Conv1d( + channels, + channels, + kernel_size, + 1, + dilation=dilation[1], + padding=get_padding(kernel_size, dilation[1]), + ) + ), + weight_norm( + nn.Conv1d( + channels, + channels, + kernel_size, + 1, + dilation=dilation[2], + padding=get_padding(kernel_size, dilation[2]), + ) + ), + ] + ) + + self.convs2 = nn.ModuleList( + [ + weight_norm( + nn.Conv1d( + channels, + channels, + kernel_size, + 1, + dilation=1, + padding=get_padding(kernel_size, 1), + ) + ), + weight_norm( + nn.Conv1d( + channels, + channels, + kernel_size, + 1, + dilation=1, + padding=get_padding(kernel_size, 1), + ) + ), + weight_norm( + nn.Conv1d( + channels, + channels, + kernel_size, + 1, + dilation=1, + padding=get_padding(kernel_size, 1), + ) + ), + ] + ) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + for c1, c2 in zip(self.convs1, self.convs2): + xt = F.leaky_relu(x, LRELU_SLOPE) + xt = c1(xt) + xt = F.leaky_relu(xt, LRELU_SLOPE) + xt = c2(xt) + x = xt + x + return x + + def remove_weight_norm(self): + for l in self.convs1: + remove_weight_norm(l) + for l in self.convs2: + remove_weight_norm(l) + + +class ResBlock2(torch.nn.Module): + def __init__( + self, channels: int, kernel_size: int = 3, dilation: Tuple[int, ...] = (1, 3) + ) -> None: + super().__init__() + self.convs = nn.ModuleList( + [ + weight_norm( + nn.Conv1d( + channels, + channels, + kernel_size, + 1, + dilation=dilation[0], + padding=get_padding(kernel_size, dilation[0]), + ) + ), + weight_norm( + nn.Conv1d( + channels, + channels, + kernel_size, + 1, + dilation=dilation[1], + padding=get_padding(kernel_size, dilation[1]), + ) + ), + ] + ) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + for c in self.convs: + xt = F.leaky_relu(x, LRELU_SLOPE) + xt = c(xt) + x = xt + x + return x + + def remove_weight_norm(self): + for l in self.convs: + remove_weight_norm(l) + + +def _hifigan( + name: str, + pretrained: bool = True, + progress: bool = True, + map_location=None, +) -> HifiganGenerator: + hifigan = HifiganGenerator() + if pretrained: + checkpoint = torch.hub.load_state_dict_from_url( + URLS[name], map_location=map_location, progress=progress + ) + consume_prefix_in_state_dict_if_present(checkpoint, "module.") + hifigan.load_state_dict(checkpoint) + hifigan.eval() + hifigan.remove_weight_norm() + return hifigan + + +def hifigan( + pretrained: bool = True, progress: bool = True, map_location=None +) -> HifiganGenerator: + return _hifigan("hifigan", pretrained, progress, map_location) + + +def hifigan_hubert_soft( + pretrained: bool = True, progress: bool = True, map_location=None +) -> HifiganGenerator: + return _hifigan("hifigan-hubert-soft", pretrained, progress, map_location=None) + + +def hifigan_hubert_discrete( + pretrained: bool = True, progress: bool = True, map_location=None +) -> HifiganGenerator: + return _hifigan("hifigan-hubert-discrete", pretrained, progress, map_location=None) diff --git a/voice_conversion/softvc/hifigan/hifigan/utils.py b/voice_conversion/softvc/hifigan/hifigan/utils.py new file mode 100644 index 0000000..3585663 --- /dev/null +++ b/voice_conversion/softvc/hifigan/hifigan/utils.py @@ -0,0 +1,84 @@ +import torch +import matplotlib + +matplotlib.use("Agg") +import matplotlib.pylab as plt + + +def get_padding(k, d): + return int((k * d - d) / 2) + + +def plot_spectrogram(spectrogram): + fig, ax = plt.subplots(figsize=(10, 2)) + im = ax.imshow(spectrogram, aspect="auto", origin="lower", interpolation="none") + plt.colorbar(im, ax=ax) + + fig.canvas.draw() + plt.close() + + return fig + + +def save_checkpoint( + checkpoint_dir, + generator, + discriminator, + optimizer_generator, + optimizer_discriminator, + scheduler_generator, + scheduler_discriminator, + step, + loss, + best, + logger, +): + state = { + "generator": { + "model": generator.state_dict(), + "optimizer": optimizer_generator.state_dict(), + "scheduler": scheduler_generator.state_dict(), + }, + "discriminator": { + "model": discriminator.state_dict(), + "optimizer": optimizer_discriminator.state_dict(), + "scheduler": scheduler_discriminator.state_dict(), + }, + "step": step, + "loss": loss, + } + checkpoint_dir.mkdir(exist_ok=True, parents=True) + checkpoint_path = checkpoint_dir / f"model-{step}.pt" + torch.save(state, checkpoint_path) + if best: + best_path = checkpoint_dir / "model-best.pt" + torch.save(state, best_path) + logger.info(f"Saved checkpoint: {checkpoint_path.stem}") + + +def load_checkpoint( + load_path, + generator, + discriminator, + optimizer_generator, + optimizer_discriminator, + scheduler_generator, + scheduler_discriminator, + rank, + logger, + finetune=False, +): + logger.info(f"Loading checkpoint from {load_path}") + checkpoint = torch.load(load_path, map_location={"cuda:0": f"cuda:{rank}"}) + generator.load_state_dict(checkpoint["generator"]["model"]) + discriminator.load_state_dict(checkpoint["discriminator"]["model"]) + if not finetune: + optimizer_generator.load_state_dict(checkpoint["generator"]["optimizer"]) + scheduler_generator.load_state_dict(checkpoint["generator"]["scheduler"]) + optimizer_discriminator.load_state_dict( + checkpoint["discriminator"]["optimizer"] + ) + scheduler_discriminator.load_state_dict( + checkpoint["discriminator"]["scheduler"] + ) + return checkpoint["step"], checkpoint["loss"] diff --git a/voice_conversion/softvc/hifigan/hubconf.py b/voice_conversion/softvc/hifigan/hubconf.py new file mode 100644 index 0000000..38dbd31 --- /dev/null +++ b/voice_conversion/softvc/hifigan/hubconf.py @@ -0,0 +1,3 @@ +dependencies = ["torch", "torchaudio"] + +from hifigan import hifigan, hifigan_hubert_discrete, hifigan_hubert_soft diff --git a/voice_conversion/softvc/hifigan/requirements.txt b/voice_conversion/softvc/hifigan/requirements.txt new file mode 100644 index 0000000..20dfcf8 --- /dev/null +++ b/voice_conversion/softvc/hifigan/requirements.txt @@ -0,0 +1,4 @@ +tensorboard==2.7.0 +torch==1.9.1 +torchaudio==0.9.1 +tqdm==4.62.3 \ No newline at end of file diff --git a/voice_conversion/softvc/hifigan/resample.py b/voice_conversion/softvc/hifigan/resample.py new file mode 100644 index 0000000..f6b46da --- /dev/null +++ b/voice_conversion/softvc/hifigan/resample.py @@ -0,0 +1,56 @@ +import argparse +from pathlib import Path +from concurrent.futures import ProcessPoolExecutor +from multiprocessing import cpu_count + +import torchaudio +from torchaudio.functional import resample + +from tqdm import tqdm + + +def process_wav(in_path, out_path, sample_rate): + wav, sr = torchaudio.load(in_path) + wav = resample(wav, sr, sample_rate) + torchaudio.save(out_path, wav, sample_rate) + return out_path, wav.size(-1) / sample_rate + + +def preprocess_dataset(args): + args.out_dir.mkdir(parents=True, exist_ok=True) + + futures = [] + executor = ProcessPoolExecutor(max_workers=cpu_count()) + print(f"Resampling audio in {args.in_dir}") + for in_path in args.in_dir.rglob("*.wav"): + relative_path = in_path.relative_to(args.in_dir) + out_path = args.out_dir / relative_path + out_path.parent.mkdir(parents=True, exist_ok=True) + futures.append( + executor.submit(process_wav, in_path, out_path, args.sample_rate) + ) + + results = [future.result() for future in tqdm(futures)] + + lengths = {path.stem: length for path, length in results} + seconds = sum(lengths.values()) + hours = seconds / 3600 + print(f"Wrote {len(lengths)} utterances ({hours:.2f} hours)") + + +if __name__ == "__main__": + parser = argparse.ArgumentParser(description="Resample an audio dataset.") + parser.add_argument( + "in_dir", metavar="in-dir", help="path to the dataset directory.", type=Path + ) + parser.add_argument( + "out_dir", metavar="out-dir", help="path to the output directory.", type=Path + ) + parser.add_argument( + "--sample-rate", + help="target sample rate (default 16kHz)", + type=int, + default=16000, + ) + args = parser.parse_args() + preprocess_dataset(args) diff --git a/voice_conversion/softvc/hifigan/train.py b/voice_conversion/softvc/hifigan/train.py new file mode 100644 index 0000000..0ff27ad --- /dev/null +++ b/voice_conversion/softvc/hifigan/train.py @@ -0,0 +1,347 @@ +import argparse +import logging +from pathlib import Path + +import torch +import torch.nn.functional as F +import torch.optim as optim +from torch.utils.data import DataLoader +from torch.utils.tensorboard import SummaryWriter +import torch.distributed as dist +from torch.utils.data.distributed import DistributedSampler +import torch.multiprocessing as mp +from torch.nn.parallel import DistributedDataParallel as DDP + +from hifigan.generator import HifiganGenerator +from hifigan.discriminator import ( + HifiganDiscriminator, + feature_loss, + discriminator_loss, + generator_loss, +) +from hifigan.dataset import MelDataset, LogMelSpectrogram +from hifigan.utils import load_checkpoint, save_checkpoint, plot_spectrogram + + +logging.basicConfig(level=logging.DEBUG) +logger = logging.getLogger(__name__) + + +BATCH_SIZE = 8 +SEGMENT_LENGTH = 8320 +HOP_LENGTH = 160 +SAMPLE_RATE = 16000 +BASE_LEARNING_RATE = 2e-4 +FINETUNE_LEARNING_RATE = 1e-4 +BETAS = (0.8, 0.99) +LEARNING_RATE_DECAY = 0.999 +WEIGHT_DECAY = 1e-5 +EPOCHS = 3100 +LOG_INTERVAL = 5 +VALIDATION_INTERVAL = 1000 +NUM_GENERATED_EXAMPLES = 10 +CHECKPOINT_INTERVAL = 5000 + + +def train_model(rank, world_size, args): + dist.init_process_group( + "nccl", + rank=rank, + world_size=world_size, + init_method="tcp://localhost:54321", + ) + + log_dir = args.checkpoint_dir / "logs" + log_dir.mkdir(exist_ok=True, parents=True) + + if rank == 0: + logger.setLevel(logging.DEBUG) + handler = logging.FileHandler(log_dir / f"{args.checkpoint_dir.stem}.log") + handler.setLevel(logging.DEBUG) + formatter = logging.Formatter( + "%(asctime)s [%(levelname)s] %(message)s", datefmt="%m/%d/%Y %I:%M:%S" + ) + handler.setFormatter(formatter) + logger.addHandler(handler) + else: + logger.setLevel(logging.ERROR) + + writer = SummaryWriter(log_dir) if rank == 0 else None + + generator = HifiganGenerator().to(rank) + discriminator = HifiganDiscriminator().to(rank) + + generator = DDP(generator, device_ids=[rank]) + discriminator = DDP(discriminator, device_ids=[rank]) + + optimizer_generator = optim.AdamW( + generator.parameters(), + lr=BASE_LEARNING_RATE if not args.finetune else FINETUNE_LEARNING_RATE, + betas=BETAS, + weight_decay=WEIGHT_DECAY, + ) + optimizer_discriminator = optim.AdamW( + discriminator.parameters(), + lr=BASE_LEARNING_RATE if not args.finetune else FINETUNE_LEARNING_RATE, + betas=BETAS, + weight_decay=WEIGHT_DECAY, + ) + + scheduler_generator = optim.lr_scheduler.ExponentialLR( + optimizer_generator, gamma=LEARNING_RATE_DECAY + ) + scheduler_discriminator = optim.lr_scheduler.ExponentialLR( + optimizer_discriminator, gamma=LEARNING_RATE_DECAY + ) + + train_dataset = MelDataset( + root=args.dataset_dir, + segment_length=SEGMENT_LENGTH, + sample_rate=SAMPLE_RATE, + hop_length=HOP_LENGTH, + train=True, + finetune=args.finetune, + ) + train_sampler = DistributedSampler(train_dataset, drop_last=True) + train_loader = DataLoader( + train_dataset, + batch_size=BATCH_SIZE, + sampler=train_sampler, + num_workers=8, + pin_memory=True, + shuffle=False, + drop_last=True, + ) + + validation_dataset = MelDataset( + root=args.dataset_dir, + segment_length=SEGMENT_LENGTH, + sample_rate=SAMPLE_RATE, + hop_length=HOP_LENGTH, + train=False, + finetune=args.finetune, + ) + validation_loader = DataLoader( + validation_dataset, + batch_size=1, + shuffle=False, + num_workers=8, + pin_memory=True, + ) + + melspectrogram = LogMelSpectrogram().to(rank) + + if args.resume is not None: + global_step, best_loss = load_checkpoint( + load_path=args.resume, + generator=generator, + discriminator=discriminator, + optimizer_generator=optimizer_generator, + optimizer_discriminator=optimizer_discriminator, + scheduler_generator=scheduler_generator, + scheduler_discriminator=scheduler_discriminator, + rank=rank, + logger=logger, + finetune=args.finetune, + ) + else: + global_step, best_loss = 0, float("inf") + + if args.finetune: + global_step, best_loss = 0, float("inf") + + n_epochs = EPOCHS + start_epoch = global_step // len(train_loader) + 1 + + logger.info("**" * 40) + logger.info(f"batch size: {BATCH_SIZE}") + logger.info(f"iterations per epoch: {len(train_loader)}") + logger.info(f"total of epochs: {n_epochs}") + logger.info(f"started at epoch: {start_epoch}") + logger.info("**" * 40 + "\n") + + for epoch in range(start_epoch, n_epochs + 1): + train_sampler.set_epoch(epoch) + + generator.train() + discriminator.train() + average_loss_mel = average_loss_discriminator = average_loss_generator = 0 + for i, (wavs, mels, tgts) in enumerate(train_loader, 1): + wavs, mels, tgts = wavs.to(rank), mels.to(rank), tgts.to(rank) + + # Discriminator + optimizer_discriminator.zero_grad() + + wavs_ = generator(mels.squeeze(1)) + mels_ = melspectrogram(wavs_) + + scores, _ = discriminator(wavs) + scores_, _ = discriminator(wavs_.detach()) + + loss_discriminator, _, _ = discriminator_loss(scores, scores_) + + loss_discriminator.backward() + optimizer_discriminator.step() + + # Generator + optimizer_generator.zero_grad() + + scores, features = discriminator(wavs) + scores_, features_ = discriminator(wavs_) + + loss_mel = F.l1_loss(mels_, tgts) + loss_features = feature_loss(features, features_) + loss_generator_adversarial, _ = generator_loss(scores_) + loss_generator = 45 * loss_mel + loss_features + loss_generator_adversarial + + loss_generator.backward() + optimizer_generator.step() + + global_step += 1 + + average_loss_mel += (loss_mel.item() - average_loss_mel) / i + average_loss_discriminator += ( + loss_discriminator.item() - average_loss_discriminator + ) / i + average_loss_generator += ( + loss_generator.item() - average_loss_generator + ) / i + + if rank == 0: + if global_step % LOG_INTERVAL == 0: + writer.add_scalar( + "train/loss_mel", + loss_mel.item(), + global_step, + ) + writer.add_scalar( + "train/loss_generator", + loss_generator.item(), + global_step, + ) + writer.add_scalar( + "train/loss_discriminator", + loss_discriminator.item(), + global_step, + ) + + if global_step % VALIDATION_INTERVAL == 0: + generator.eval() + + average_validation_loss = 0 + for j, (wavs, mels, tgts) in enumerate(validation_loader, 1): + wavs, mels, tgts = wavs.to(rank), mels.to(rank), tgts.to(rank) + + with torch.no_grad(): + wavs_ = generator(mels.squeeze(1)) + mels_ = melspectrogram(wavs_) + + length = min(mels_.size(-1), tgts.size(-1)) + + loss_mel = F.l1_loss(mels_[..., :length], tgts[..., :length]) + + average_validation_loss += ( + loss_mel.item() - average_validation_loss + ) / j + + if rank == 0: + if j <= NUM_GENERATED_EXAMPLES: + writer.add_audio( + f"generated/wav_{j}", + wavs_.squeeze(0), + global_step, + sample_rate=16000, + ) + writer.add_figure( + f"generated/mel_{j}", + plot_spectrogram(mels_.squeeze().cpu().numpy()), + global_step, + ) + + generator.train() + discriminator.train() + + if rank == 0: + writer.add_scalar( + "validation/mel_loss", average_validation_loss, global_step + ) + logger.info( + f"valid -- epoch: {epoch}, mel loss: {average_validation_loss:.4f}" + ) + + new_best = best_loss > average_validation_loss + if new_best or global_step % CHECKPOINT_INTERVAL == 0: + if new_best: + logger.info("-------- new best model found!") + best_loss = average_validation_loss + + if rank == 0: + save_checkpoint( + checkpoint_dir=args.checkpoint_dir, + generator=generator, + discriminator=discriminator, + optimizer_generator=optimizer_generator, + optimizer_discriminator=optimizer_discriminator, + scheduler_generator=scheduler_generator, + scheduler_discriminator=scheduler_discriminator, + step=global_step, + loss=average_validation_loss, + best=new_best, + logger=logger, + ) + + scheduler_discriminator.step() + scheduler_generator.step() + + logger.info( + f"train -- epoch: {epoch}, mel loss: {average_loss_mel:.4f}, generator loss: {average_loss_generator:.4f}, discriminator loss: {average_loss_discriminator:.4f}" + ) + + dist.destroy_process_group() + + +if __name__ == "__main__": + parser = argparse.ArgumentParser(description="Train or finetune HiFi-GAN.") + parser.add_argument( + "dataset_dir", + metavar="dataset-dir", + help="path to the preprocessed data directory", + type=Path, + ) + parser.add_argument( + "checkpoint_dir", + metavar="checkpoint-dir", + help="path to the checkpoint directory", + type=Path, + ) + parser.add_argument( + "--resume", + help="path to the checkpoint to resume from", + type=Path, + ) + parser.add_argument( + "--finetune", + help="whether to finetune (note that a resume path must be given)", + action="store_true", + ) + args = parser.parse_args() + + # display training setup info + logger.info(f"PyTorch version: {torch.__version__}") + logger.info(f"CUDA version: {torch.version.cuda}") + logger.info(f"CUDNN version: {torch.backends.cudnn.version()}") + logger.info(f"CUDNN enabled: {torch.backends.cudnn.enabled}") + logger.info(f"CUDNN deterministic: {torch.backends.cudnn.deterministic}") + logger.info(f"CUDNN benchmark: {torch.backends.cudnn.benchmark}") + logger.info(f"# of GPUS: {torch.cuda.device_count()}") + + # clear handlers + logger.handlers.clear() + + world_size = torch.cuda.device_count() + mp.spawn( + train_model, + args=(world_size, args), + nprocs=world_size, + join=True, + ) diff --git a/voice_conversion/softvc/hifigan/vocoder.png b/voice_conversion/softvc/hifigan/vocoder.png new file mode 100644 index 0000000..0ed2d54 Binary files /dev/null and b/voice_conversion/softvc/hifigan/vocoder.png differ diff --git a/voice_conversion/softvc/hubert/.gitignore b/voice_conversion/softvc/hubert/.gitignore new file mode 100644 index 0000000..0202868 --- /dev/null +++ b/voice_conversion/softvc/hubert/.gitignore @@ -0,0 +1,132 @@ +# Byte-compiled / optimized / DLL files +__pycache__/ +*.py[cod] +*$py.class + +# C extensions +*.so + +# Distribution / packaging +.Python +build/ +develop-eggs/ +dist/ +downloads/ +eggs/ +.eggs/ +lib/ +lib64/ +parts/ +sdist/ +var/ +wheels/ +pip-wheel-metadata/ +share/python-wheels/ +*.egg-info/ +.installed.cfg +*.egg +MANIFEST + +# PyInstaller +# Usually these files are written by a python script from a template +# before PyInstaller builds the exe, so as to inject date/other infos into it. +*.manifest +*.spec + +# Installer logs +pip-log.txt +pip-delete-this-directory.txt + +# Unit test / coverage reports +htmlcov/ +.tox/ +.nox/ +.coverage +.coverage.* +.cache +nosetests.xml +coverage.xml +*.cover +*.py,cover +.hypothesis/ +.pytest_cache/ + +# Translations +*.mo +*.pot + +# Django stuff: +*.log +local_settings.py +db.sqlite3 +db.sqlite3-journal + +# Flask stuff: +instance/ +.webassets-cache + +# Scrapy stuff: +.scrapy + +# Sphinx documentation +docs/_build/ + +# PyBuilder +target/ + +# Jupyter Notebook +.ipynb_checkpoints + +# IPython +profile_default/ +ipython_config.py + +# pyenv +.python-version + +# pipenv +# According to pypa/pipenv#598, it is recommended to include Pipfile.lock in version control. +# However, in case of collaboration, if having platform-specific dependencies or dependencies +# having no cross-platform support, pipenv may install dependencies that don't work, or not +# install all needed dependencies. +#Pipfile.lock + +# PEP 582; used by e.g. github.com/David-OConnor/pyflow +__pypackages__/ + +# Celery stuff +celerybeat-schedule +celerybeat.pid + +# SageMath parsed files +*.sage.py + +# Environments +.env +.venv +env/ +venv/ +ENV/ +env.bak/ +venv.bak/ + +# VSCode project settings +.vscode + +# Spyder project settings +.spyderproject +.spyproject + +# Rope project settings +.ropeproject + +# mkdocs documentation +/site + +# mypy +.mypy_cache/ +.dmypy.json +dmypy.json + +# Pyre type checker +.pyre/ diff --git a/voice_conversion/softvc/hubert/LICENSE b/voice_conversion/softvc/hubert/LICENSE new file mode 100644 index 0000000..6eb2af0 --- /dev/null +++ b/voice_conversion/softvc/hubert/LICENSE @@ -0,0 +1,21 @@ +MIT License + +Copyright (c) 2021 Benjamin van Niekerk + +Permission is hereby granted, free of charge, to any person obtaining a copy +of this software and associated documentation files (the "Software"), to deal +in the Software without restriction, including without limitation the rights +to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +copies of the Software, and to permit persons to whom the Software is +furnished to do so, subject to the following conditions: + +The above copyright notice and this permission notice shall be included in all +copies or substantial portions of the Software. + +THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE +SOFTWARE. diff --git a/voice_conversion/softvc/hubert/README.md b/voice_conversion/softvc/hubert/README.md new file mode 100644 index 0000000..ab45422 --- /dev/null +++ b/voice_conversion/softvc/hubert/README.md @@ -0,0 +1,163 @@ +

+ + Open In Colab + +

+ +# HuBERT + +Training and inference scripts for the HuBERT content encoders in [A Comparison of Discrete and Soft Speech Units for Improved Voice Conversion](https://ieeexplore.ieee.org/abstract/document/9746484). +For more details see [soft-vc](https://github.com/bshall/soft-vc). Audio samples can be found [here](https://bshall.github.io/soft-vc/). Colab demo can be found [here](https://colab.research.google.com/github/bshall/soft-vc/blob/main/soft-vc-demo.ipynb). + +
+ Soft-VC +
+
+ + Fig 1: Architecture of the voice conversion system. a) The discrete content encoder clusters audio features to produce a sequence of discrete speech units. b) The soft content encoder is trained to predict the discrete units. The acoustic model transforms the discrete/soft speech units into a target spectrogram. The vocoder converts the spectrogram into an audio waveform. + +
+ +## Example Usage + +### Programmatic Usage + +```python +import torch, torchaudio + +# Load checkpoint (either hubert_soft or hubert_discrete) +hubert = torch.hub.load("bshall/hubert:main", "hubert_soft").cuda() + +# Load audio +wav, sr = torchaudio.load("path/to/wav") +assert sr == 16000 +wav = wav.unsqueeze(0).cuda() + +# Extract speech units +units = hubert.units(x) +``` + +### Script-Based Usage + +``` +usage: encode.py [-h] [--extension EXTENSION] {soft,discrete} in-dir out-dir + +Encode an audio dataset. + +positional arguments: + {soft,discrete} available models (HuBERT-Soft or HuBERT-Discrete) + in-dir path to the dataset directory. + out-dir path to the output directory. + +optional arguments: + -h, --help show this help message and exit + --extension EXTENSION + extension of the audio files (defaults to .flac). +``` + +## Training + +### Step 1: Dataset Preparation + +Download and extract the [LibriSpeech](https://www.openslr.org/12) corpus. The training script expects the following tree structure for the dataset directory: + +``` +│ lengths.json +│ +└───wavs + ├───dev-* + │ ├───84 + │ ├───... + │ └───8842 + └───train-* + ├───19 + ├───... + └───8975 +``` + +The `train-*` and `dev-*` directories should contain the training and validation splits respectively. Note that there can be multiple `train` and `dev` folders e.g., `train-clean-100`, `train-other-500`, etc. Finally, the `lengths.json` file should contain key-value pairs with the file path and number of samples: + +```json +{ + "dev-clean/1272/128104/1272-128104-0000": 93680, + "dev-clean/1272/128104/1272-128104-0001": 77040, +} +``` + +### Step 2: Extract Discrete Speech Units + +Encode LibriSpeech using the HuBERT-Discrete model and `encode.py` script: + +``` +usage: encode.py [-h] [--extension EXTENSION] {soft,discrete} in-dir out-dir + +Encode an audio dataset. + +positional arguments: + {soft,discrete} available models (HuBERT-Soft or HuBERT-Discrete) + in-dir path to the dataset directory. + out-dir path to the output directory. + +optional arguments: + -h, --help show this help message and exit + --extension EXTENSION + extension of the audio files (defaults to .flac). +``` + +for example: + +``` +python encode.py discrete path/to/LibriSpeech/wavs path/to/LibriSpeech/units +``` + +At this point the directory tree should look like: + +``` +│ lengths.json +│ +├───discrete +│ ├───... +└───wavs + ├───... +``` + +### Step 3: Train the HuBERT-Soft Content Encoder + +``` +usage: train.py [-h] [--resume RESUME] [--warmstart] [--mask] [--alpha ALPHA] dataset-dir checkpoint-dir + +Train HuBERT soft content encoder. + +positional arguments: + dataset-dir path to the data directory. + checkpoint-dir path to the checkpoint directory. + +optional arguments: + -h, --help show this help message and exit + --resume RESUME path to the checkpoint to resume from. + --warmstart whether to initialize from the fairseq HuBERT checkpoint. + --mask whether to use input masking. + --alpha ALPHA weight for the masked loss. +``` + +## Links + +- [Soft-VC repo](https://github.com/bshall/soft-vc) +- [Soft-VC paper](https://ieeexplore.ieee.org/abstract/document/9746484) +- [Official HuBERT repo](https://github.com/pytorch/fairseq) +- [HuBERT paper](https://arxiv.org/abs/2106.07447) + +## Citation + +If you found this work helpful please consider citing our paper: + +``` +@inproceedings{ + soft-vc-2022, + author={van Niekerk, Benjamin and Carbonneau, Marc-André and Zaïdi, Julian and Baas, Matthew and Seuté, Hugo and Kamper, Herman}, + booktitle={ICASSP}, + title={A Comparison of Discrete and Soft Speech Units for Improved Voice Conversion}, + year={2022} +} +``` \ No newline at end of file diff --git a/voice_conversion/softvc/hubert/cluster.py b/voice_conversion/softvc/hubert/cluster.py new file mode 100644 index 0000000..18b754c --- /dev/null +++ b/voice_conversion/softvc/hubert/cluster.py @@ -0,0 +1,66 @@ +from pathlib import Path +import logging +import argparse + +import torch +import numpy as np +from sklearn.cluster import KMeans + +logging.basicConfig(level=logging.INFO) +logger = logging.getLogger(__name__) + + +def cluster(args): + with open(args.subset) as file: + subset = [line.strip() for line in file] + + logger.info(f"Loading features from {args.in_dir}") + features = [] + for path in subset: + in_path = args.in_dir / path + features.append(np.load(in_path.with_suffix(".npy"))) + features = np.concatenate(features, axis=0) + + logger.info(f"Clustering features of shape: {features.shape}") + kmeans = KMeans(n_clusters=args.n_clusters).fit(features) + + checkpoint_path = args.checkpoint_dir / f"kmeans_{args.n_clusters}.pt" + checkpoint_path.parent.mkdir(exist_ok=True, parents=True) + torch.save( + checkpoint_path, + { + "n_features_in_": kmeans.n_features_in_, + "_n_threads": kmeans._n_threads, + "cluster_centers_": kmeans.cluster_centers_, + }, + ) + + +if __name__ == "__main__": + parser = argparse.ArgumentParser(description="Cluster speech features features.") + parser.add_argument( + "in_dir", + metavar="in-dir", + help="path to the encoded dataset", + type=Path, + ) + parser.add_argument( + "subset", + matavar="subset", + help="path to the .txt file containing the list of files to cluster", + type=Path, + ) + parser.add_argument( + "checkpoint_dir", + metavar="checkpoint-dir", + help="path to the checkpoint directory", + type=Path, + ) + parser.add_argument( + "--n-clusters", + help="number of clusters", + type=int, + default=100, + ) + args = parser.parse_args() + cluster(args) diff --git a/voice_conversion/softvc/hubert/content-encoder.png b/voice_conversion/softvc/hubert/content-encoder.png new file mode 100644 index 0000000..fc59d53 Binary files /dev/null and b/voice_conversion/softvc/hubert/content-encoder.png differ diff --git a/voice_conversion/softvc/hubert/encode.py b/voice_conversion/softvc/hubert/encode.py new file mode 100644 index 0000000..42b300d --- /dev/null +++ b/voice_conversion/softvc/hubert/encode.py @@ -0,0 +1,56 @@ +import argparse +import logging +import numpy as np +from pathlib import Path +from tqdm import tqdm + +import torch +import torchaudio +from torchaudio.functional import resample + + +def encode_dataset(args): + print(f"Loading hubert checkpoint") + hubert = torch.hub.load("bshall/hubert:main", f"hubert_{args.model}").cuda() + + print(f"Encoding dataset at {args.in_dir}") + for in_path in tqdm(list(args.in_dir.rglob(f"*{args.extension}"))): + wav, sr = torchaudio.load(in_path) + wav = resample(wav, sr, 16000) + wav = wav.unsqueeze(0).cuda() + + with torch.inference_mode(): + units = hubert.units(wav) + + out_path = args.out_dir / in_path.relative_to(args.in_dir) + out_path.parent.mkdir(parents=True, exist_ok=True) + np.save(out_path.with_suffix(".npy"), units.squeeze().cpu().numpy()) + + +if __name__ == "__main__": + parser = argparse.ArgumentParser(description="Encode an audio dataset.") + parser.add_argument( + "model", + help="available models (HuBERT-Soft or HuBERT-Discrete)", + choices=["soft", "discrete"], + ) + parser.add_argument( + "in_dir", + metavar="in-dir", + help="path to the dataset directory.", + type=Path, + ) + parser.add_argument( + "out_dir", + metavar="out-dir", + help="path to the output directory.", + type=Path, + ) + parser.add_argument( + "--extension", + help="extension of the audio files (defaults to .flac).", + default=".flac", + type=str, + ) + args = parser.parse_args() + encode_dataset(args) diff --git a/voice_conversion/softvc/hubert/hubconf.py b/voice_conversion/softvc/hubert/hubconf.py new file mode 100644 index 0000000..14a62bc --- /dev/null +++ b/voice_conversion/softvc/hubert/hubconf.py @@ -0,0 +1,3 @@ +dependencies = ["torch", "torchaudio", "sklearn"] + +from hubert import hubert_discrete, hubert_soft, kmeans100 diff --git a/voice_conversion/softvc/hubert/hubert/__init__.py b/voice_conversion/softvc/hubert/hubert/__init__.py new file mode 100644 index 0000000..99727eb --- /dev/null +++ b/voice_conversion/softvc/hubert/hubert/__init__.py @@ -0,0 +1,8 @@ +from .model import ( + Hubert, + HubertDiscrete, + HubertSoft, + hubert_discrete, + hubert_soft, + kmeans100, +) diff --git a/voice_conversion/softvc/hubert/hubert/dataset.py b/voice_conversion/softvc/hubert/hubert/dataset.py new file mode 100644 index 0000000..ca7c5c4 --- /dev/null +++ b/voice_conversion/softvc/hubert/hubert/dataset.py @@ -0,0 +1,91 @@ +import random +from pathlib import Path +import numpy as np +import json + +import torch +import torch.nn.functional as F +from torch.utils.data import Dataset +import torchaudio + + +class AcousticUnitsDataset(Dataset): + def __init__( + self, + root: Path, + sample_rate: int = 16000, + label_rate: int = 50, + min_samples: int = 32000, + max_samples: int = 250000, + train: bool = True, + ): + self.wavs_dir = root / "wavs" + self.units_dir = root / "units" + + with open(root / "lengths.json") as file: + self.lenghts = json.load(file) + + pattern = "train-*/**/*.flac" if train else "dev-*/**/*.flac" + metadata = ( + (path, path.relative_to(self.wavs_dir).with_suffix("").as_posix()) + for path in self.wavs_dir.rglob(pattern) + ) + metadata = ((path, key) for path, key in metadata if key in self.lenghts) + self.metadata = [ + path for path, key in metadata if self.lenghts[key] > min_samples + ] + + self.sample_rate = sample_rate + self.label_rate = label_rate + self.min_samples = min_samples + self.max_samples = max_samples + self.train = train + + def __len__(self): + return len(self.metadata) + + def __getitem__(self, index): + wav_path = self.metadata[index] + units_path = self.units_dir / wav_path.relative_to(self.wavs_dir) + + wav, _ = torchaudio.load(wav_path) + wav = F.pad(wav, ((400 - 320) // 2, (400 - 320) // 2)) + codes = np.load(units_path.with_suffix(".npy")) + + return wav, torch.from_numpy(codes).long() + + def collate(self, batch): + wavs, codes = zip(*batch) + wavs, codes = list(wavs), list(codes) + + wav_lengths = [wav.size(-1) for wav in wavs] + code_lengths = [code.size(-1) for code in codes] + + wav_frames = min(self.max_samples, *wav_lengths) + + collated_wavs, wav_offsets = [], [] + for wav in wavs: + wav_diff = wav.size(-1) - wav_frames + wav_offset = random.randint(0, wav_diff) + wav = wav[:, wav_offset : wav_offset + wav_frames] + + collated_wavs.append(wav) + wav_offsets.append(wav_offset) + + rate = self.label_rate / self.sample_rate + code_offsets = [round(wav_offset * rate) for wav_offset in wav_offsets] + code_frames = round(wav_frames * rate) + remaining_code_frames = [ + length - offset for length, offset in zip(code_lengths, code_offsets) + ] + code_frames = min(code_frames, *remaining_code_frames) + + collated_codes = [] + for code, code_offset in zip(codes, code_offsets): + code = code[code_offset : code_offset + code_frames] + collated_codes.append(code) + + wavs = torch.stack(collated_wavs, dim=0) + codes = torch.stack(collated_codes, dim=0) + + return wavs, codes diff --git a/voice_conversion/softvc/hubert/hubert/model.py b/voice_conversion/softvc/hubert/hubert/model.py new file mode 100644 index 0000000..f275b1c --- /dev/null +++ b/voice_conversion/softvc/hubert/hubert/model.py @@ -0,0 +1,293 @@ +import copy +from typing import Optional, Tuple +import random + +from sklearn.cluster import KMeans + +import torch +import torch.nn as nn +import torch.nn.functional as F +from torch.nn.modules.utils import consume_prefix_in_state_dict_if_present + +URLS = { + "hubert-discrete": "https://github.com/bshall/hubert/releases/download/v0.1/hubert-discrete-e9416457.pt", + "hubert-soft": "https://github.com/bshall/hubert/releases/download/v0.1/hubert-soft-0d54a1f4.pt", + "kmeans100": "https://github.com/bshall/hubert/releases/download/v0.1/kmeans100-50f36a95.pt", +} + + +class Hubert(nn.Module): + def __init__(self, num_label_embeddings: int = 100, mask: bool = True): + super().__init__() + self._mask = mask + self.feature_extractor = FeatureExtractor() + self.feature_projection = FeatureProjection() + self.positional_embedding = PositionalConvEmbedding() + self.norm = nn.LayerNorm(768) + self.dropout = nn.Dropout(0.1) + self.encoder = TransformerEncoder( + nn.TransformerEncoderLayer( + 768, 12, 3072, activation="gelu", batch_first=True + ), + 12, + ) + self.proj = nn.Linear(768, 256) + + self.masked_spec_embed = nn.Parameter(torch.FloatTensor(768).uniform_()) + self.label_embedding = nn.Embedding(num_label_embeddings, 256) + + def mask(self, x: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]: + mask = None + if self.training and self._mask: + mask = _compute_mask((x.size(0), x.size(1)), 0.8, 10, x.device, 2) + x[mask] = self.masked_spec_embed.to(x.dtype) + return x, mask + + def encode( + self, x: torch.Tensor, layer: Optional[int] = None + ) -> Tuple[torch.Tensor, torch.Tensor]: + x = self.feature_extractor(x) + x = self.feature_projection(x.transpose(1, 2)) + x, mask = self.mask(x) + x = x + self.positional_embedding(x) + x = self.dropout(self.norm(x)) + x = self.encoder(x, output_layer=layer) + return x, mask + + def logits(self, x: torch.Tensor) -> torch.Tensor: + logits = torch.cosine_similarity( + x.unsqueeze(2), + self.label_embedding.weight.unsqueeze(0).unsqueeze(0), + dim=-1, + ) + return logits / 0.1 + + def forward(self, x: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]: + x, mask = self.encode(x) + x = self.proj(x) + logits = self.logits(x) + return logits, mask + + +class HubertSoft(Hubert): + def __init__(self): + super().__init__() + + @torch.inference_mode() + def units(self, wav: torch.Tensor) -> torch.Tensor: + wav = F.pad(wav, ((400 - 320) // 2, (400 - 320) // 2)) + x, _ = self.encode(wav) + return self.proj(x) + + +class HubertDiscrete(Hubert): + def __init__(self, kmeans): + super().__init__(504) + self.kmeans = kmeans + + @torch.inference_mode() + def units(self, wav: torch.Tensor) -> torch.LongTensor: + wav = F.pad(wav, ((400 - 320) // 2, (400 - 320) // 2)) + x, _ = self.encode(wav, layer=7) + x = self.kmeans.predict(x.squeeze().cpu().numpy()) + return torch.tensor(x, dtype=torch.long, device=wav.device) + + +class FeatureExtractor(nn.Module): + def __init__(self): + super().__init__() + self.conv0 = nn.Conv1d(1, 512, 10, 5, bias=False) + self.norm0 = nn.GroupNorm(512, 512) + self.conv1 = nn.Conv1d(512, 512, 3, 2, bias=False) + self.conv2 = nn.Conv1d(512, 512, 3, 2, bias=False) + self.conv3 = nn.Conv1d(512, 512, 3, 2, bias=False) + self.conv4 = nn.Conv1d(512, 512, 3, 2, bias=False) + self.conv5 = nn.Conv1d(512, 512, 2, 2, bias=False) + self.conv6 = nn.Conv1d(512, 512, 2, 2, bias=False) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + x = F.gelu(self.norm0(self.conv0(x))) + x = F.gelu(self.conv1(x)) + x = F.gelu(self.conv2(x)) + x = F.gelu(self.conv3(x)) + x = F.gelu(self.conv4(x)) + x = F.gelu(self.conv5(x)) + x = F.gelu(self.conv6(x)) + return x + + +class FeatureProjection(nn.Module): + def __init__(self): + super().__init__() + self.norm = nn.LayerNorm(512) + self.projection = nn.Linear(512, 768) + self.dropout = nn.Dropout(0.1) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + x = self.norm(x) + x = self.projection(x) + x = self.dropout(x) + return x + + +class PositionalConvEmbedding(nn.Module): + def __init__(self): + super().__init__() + self.conv = nn.Conv1d( + 768, + 768, + kernel_size=128, + padding=128 // 2, + groups=16, + ) + self.conv = nn.utils.weight_norm(self.conv, name="weight", dim=2) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + x = self.conv(x.transpose(1, 2)) + x = F.gelu(x[:, :, :-1]) + return x.transpose(1, 2) + + +class TransformerEncoder(nn.Module): + def __init__( + self, encoder_layer: nn.TransformerEncoderLayer, num_layers: int + ) -> None: + super(TransformerEncoder, self).__init__() + self.layers = nn.ModuleList( + [copy.deepcopy(encoder_layer) for _ in range(num_layers)] + ) + self.num_layers = num_layers + + def forward( + self, + src: torch.Tensor, + mask: torch.Tensor = None, + src_key_padding_mask: torch.Tensor = None, + output_layer: Optional[int] = None, + ) -> torch.Tensor: + output = src + for layer in self.layers[:output_layer]: + output = layer( + output, src_mask=mask, src_key_padding_mask=src_key_padding_mask + ) + return output + + +def _compute_mask( + shape: Tuple[int, int], + mask_prob: float, + mask_length: int, + device: torch.device, + min_masks: int = 0, +) -> torch.Tensor: + batch_size, sequence_length = shape + + if mask_length < 1: + raise ValueError("`mask_length` has to be bigger than 0.") + + if mask_length > sequence_length: + raise ValueError( + f"`mask_length` has to be smaller than `sequence_length`, but got `mask_length`: {mask_length} and `sequence_length`: {sequence_length}`" + ) + + # compute number of masked spans in batch + num_masked_spans = int(mask_prob * sequence_length / mask_length + random.random()) + num_masked_spans = max(num_masked_spans, min_masks) + + # make sure num masked indices <= sequence_length + if num_masked_spans * mask_length > sequence_length: + num_masked_spans = sequence_length // mask_length + + # SpecAugment mask to fill + mask = torch.zeros((batch_size, sequence_length), device=device, dtype=torch.bool) + + # uniform distribution to sample from, make sure that offset samples are < sequence_length + uniform_dist = torch.ones( + (batch_size, sequence_length - (mask_length - 1)), device=device + ) + + # get random indices to mask + mask_indices = torch.multinomial(uniform_dist, num_masked_spans) + + # expand masked indices to masked spans + mask_indices = ( + mask_indices.unsqueeze(dim=-1) + .expand((batch_size, num_masked_spans, mask_length)) + .reshape(batch_size, num_masked_spans * mask_length) + ) + offsets = ( + torch.arange(mask_length, device=device)[None, None, :] + .expand((batch_size, num_masked_spans, mask_length)) + .reshape(batch_size, num_masked_spans * mask_length) + ) + mask_idxs = mask_indices + offsets + + # scatter indices to mask + mask = mask.scatter(1, mask_idxs, True) + + return mask + + +def hubert_discrete( + pretrained: bool = True, + progress: bool = True, +) -> HubertDiscrete: + r"""HuBERT-Discrete from `"A Comparison of Discrete and Soft Speech Units for Improved Voice Conversion"`. + Args: + pretrained (bool): load pretrained weights into the model + progress (bool): show progress bar when downloading model + """ + kmeans = kmeans100(pretrained=pretrained, progress=progress) + hubert = HubertDiscrete(kmeans) + if pretrained: + checkpoint = torch.hub.load_state_dict_from_url( + URLS["hubert-discrete"], progress=progress + ) + consume_prefix_in_state_dict_if_present(checkpoint, "module.") + hubert.load_state_dict(checkpoint) + hubert.eval() + return hubert + + +def hubert_soft( + pretrained: bool = True, + progress: bool = True, +) -> HubertSoft: + r"""HuBERT-Soft from `"A Comparison of Discrete and Soft Speech Units for Improved Voice Conversion"`. + Args: + pretrained (bool): load pretrained weights into the model + progress (bool): show progress bar when downloading model + """ + hubert = HubertSoft() + if pretrained: + checkpoint = torch.hub.load_state_dict_from_url( + URLS["hubert-soft"], progress=progress + ) + consume_prefix_in_state_dict_if_present(checkpoint, "module.") + hubert.load_state_dict(checkpoint) + hubert.eval() + return hubert + + +def _kmeans( + num_clusters: int, pretrained: bool = True, progress: bool = True +) -> KMeans: + kmeans = KMeans(num_clusters) + if pretrained: + checkpoint = torch.hub.load_state_dict_from_url( + URLS[f"kmeans{num_clusters}"], progress=progress + ) + kmeans.__dict__["n_features_in_"] = checkpoint["n_features_in_"] + kmeans.__dict__["_n_threads"] = checkpoint["_n_threads"] + kmeans.__dict__["cluster_centers_"] = checkpoint["cluster_centers_"].numpy() + return kmeans + + +def kmeans100(pretrained: bool = True, progress: bool = True) -> KMeans: + r""" + k-means checkpoint for HuBERT-Discrete with 100 clusters. + Args: + pretrained (bool): load pretrained weights into the model + progress (bool): show progress bar when downloading model + """ + return _kmeans(100, pretrained, progress) diff --git a/voice_conversion/softvc/hubert/hubert/utils.py b/voice_conversion/softvc/hubert/hubert/utils.py new file mode 100644 index 0000000..8519e7a --- /dev/null +++ b/voice_conversion/softvc/hubert/hubert/utils.py @@ -0,0 +1,58 @@ +import torch + + +class Metric: + def __init__(self): + self.steps = 0 + self.value = 0 + + def update(self, value): + self.steps += 1 + self.value += (value - self.value) / self.steps + return self.value + + def reset(self): + self.steps = 0 + self.value = 0 + + +def save_checkpoint( + checkpoint_dir, + hubert, + optimizer, + scaler, + step, + loss, + best, + logger, +): + state = { + "hubert": hubert.state_dict(), + "optimizer": optimizer.state_dict(), + "scaler": scaler.state_dict(), + "step": step, + "loss": loss, + } + checkpoint_dir.mkdir(exist_ok=True, parents=True) + checkpoint_path = checkpoint_dir / f"model-{step}.pt" + torch.save(state, checkpoint_path) + if best: + best_path = checkpoint_dir / "model-best.pt" + torch.save(state, best_path) + logger.info(f"Saved checkpoint: {checkpoint_path.stem}") + + +def load_checkpoint( + load_path, + hubert, + optimizer, + scaler, + rank, + logger, +): + logger.info(f"Loading checkpoint from {load_path}") + checkpoint = torch.load(load_path, map_location={"cuda:0": f"cuda:{rank}"}) + hubert.load_state_dict(checkpoint["hubert"]) + scaler.load_state_dict(checkpoint["scaler"]) + optimizer.load_state_dict(checkpoint["optimizer"]) + return checkpoint["step"], checkpoint["loss"] diff --git a/voice_conversion/softvc/hubert/train.py b/voice_conversion/softvc/hubert/train.py new file mode 100644 index 0000000..fc01575 --- /dev/null +++ b/voice_conversion/softvc/hubert/train.py @@ -0,0 +1,453 @@ +import argparse +import logging +from pathlib import Path + +import torch +import torch.cuda.amp as amp +import torch.nn as nn +import torch.nn.functional as F +import torch.optim as optim +from torch.utils.data import DataLoader +from torch.utils.tensorboard import SummaryWriter +import torch.distributed as dist +from torch.utils.data.distributed import DistributedSampler +import torch.multiprocessing as mp +from torch.nn.parallel import DistributedDataParallel as DDP +from torch.nn.modules.utils import consume_prefix_in_state_dict_if_present + +from hubert.model import Hubert, URLS +from hubert.dataset import AcousticUnitsDataset +from hubert.utils import Metric, save_checkpoint, load_checkpoint + +logging.basicConfig(level=logging.INFO) +logger = logging.getLogger(__name__) + +######################################################################################## +# Define hyperparameters for training: +######################################################################################## + +BATCH_SIZE = 64 +LEARNING_RATE = 2e-5 +BETAS = (0.9, 0.98) +EPS = 1e-06 +WEIGHT_DECAY = 1e-2 +MAX_NORM = 10 +STEPS = 25000 +LOG_INTERVAL = 5 +VALIDATION_INTERVAL = 1000 +CHECKPOINT_INTERVAL = 1000 +BACKEND = "nccl" +INIT_METHOD = "tcp://localhost:54321" + + +def train(rank, world_size, args): + dist.init_process_group( + BACKEND, + rank=rank, + world_size=world_size, + init_method=INIT_METHOD, + ) + + #################################################################################### + # Setup logging utilities: + #################################################################################### + + log_dir = args.checkpoint_dir / "logs" + log_dir.mkdir(exist_ok=True, parents=True) + + if rank == 0: + logger.setLevel(logging.INFO) + handler = logging.FileHandler(log_dir / f"{args.checkpoint_dir.stem}.log") + handler.setLevel(logging.INFO) + formatter = logging.Formatter( + "%(asctime)s [%(levelname)s] %(message)s", datefmt="%m/%d/%Y %I:%M:%S" + ) + handler.setFormatter(formatter) + logger.addHandler(handler) + else: + logger.setLevel(logging.ERROR) + + writer = SummaryWriter(log_dir) if rank == 0 else None + + #################################################################################### + # Initialize models + #################################################################################### + + hubert = Hubert(mask=args.mask).to(rank) + + if args.warmstart: + checkpoint = torch.hub.load_state_dict_from_url( + URLS["hubert-discrete"], map_location={"cuda:0": f"cuda:{rank}"} + ) + consume_prefix_in_state_dict_if_present(checkpoint, "module.") + hubert.load_state_dict(checkpoint, strict=False) + + hubert = DDP(hubert, device_ids=[rank]) + + #################################################################################### + # Initialze optimizer and grad scaler + #################################################################################### + + optimizer = optim.AdamW( + hubert.parameters(), + lr=LEARNING_RATE, + betas=BETAS, + eps=EPS, + weight_decay=WEIGHT_DECAY, + ) + scaler = amp.GradScaler() + + #################################################################################### + # Initialize datasets and dataloaders + #################################################################################### + + train_dataset = AcousticUnitsDataset( + root=args.dataset_dir, + train=True, + ) + train_sampler = DistributedSampler(train_dataset, drop_last=True) + train_loader = DataLoader( + train_dataset, + collate_fn=train_dataset.collate, + batch_size=BATCH_SIZE, + sampler=train_sampler, + num_workers=8, + pin_memory=True, + shuffle=False, + drop_last=True, + ) + + validation_dataset = AcousticUnitsDataset( + root=args.dataset_dir, + train=False, + ) + validation_loader = DataLoader( + validation_dataset, + batch_size=1, + shuffle=False, + num_workers=8, + pin_memory=True, + ) + + #################################################################################### + # Load checkpoint if args.resume is set + #################################################################################### + + if args.resume is not None: + global_step, best_loss = load_checkpoint( + load_path=args.resume, + hubert=hubert, + optimizer=optimizer, + scaler=scaler, + rank=rank, + logger=logger, + ) + else: + global_step, best_loss = 0, float("inf") + + # =================================================================================# + # Start training loop + # =================================================================================# + + n_epochs = STEPS // len(train_loader) + 1 + start_epoch = global_step // len(train_loader) + 1 + + logger.info("**" * 40) + logger.info(f"PyTorch version: {torch.__version__}") + logger.info(f"CUDA version: {torch.version.cuda}") + logger.info(f"CUDNN version: {torch.backends.cudnn.version()}") + logger.info(f"CUDNN enabled: {torch.backends.cudnn.enabled}") + logger.info(f"CUDNN deterministic: {torch.backends.cudnn.deterministic}") + logger.info(f"CUDNN benchmark: {torch.backends.cudnn.benchmark}") + logger.info(f"# of GPUS: {torch.cuda.device_count()}") + logger.info(f"batch size: {BATCH_SIZE}") + logger.info(f"iterations per epoch: {len(train_loader)}") + logger.info(f"# of epochs: {n_epochs}") + logger.info(f"started at epoch: {start_epoch}") + logger.info("**" * 40 + "\n") + + if args.mask: + average_masked_loss = Metric() + average_unmasked_loss = Metric() + average_masked_accuracy = Metric() + average_unmasked_accuracy = Metric() + + epoch_masked_loss = Metric() + epoch_unmasked_loss = Metric() + epoch_masked_accuracy = Metric() + epoch_unmasked_accuracy = Metric() + else: + average_loss = Metric() + average_accuracy = Metric() + + epoch_loss = Metric() + epoch_accuracy = Metric() + + validation_loss = Metric() + validation_accuracy = Metric() + + for epoch in range(start_epoch, n_epochs + 1): + train_sampler.set_epoch(epoch) + + hubert.train() + if args.mask: + epoch_masked_loss.reset() + epoch_unmasked_loss.reset() + epoch_masked_accuracy.reset() + epoch_unmasked_accuracy.reset() + else: + epoch_loss.reset() + epoch_accuracy.reset() + + for wavs, codes in train_loader: + global_step += 1 + wavs, codes = wavs.to(rank), codes.to(rank) + + ############################################################################ + # Compute training loss + ############################################################################ + + optimizer.zero_grad() + + with amp.autocast(): + logits, mask = hubert(wavs) + length = min( + mask.size(-1) if args.mask else float("inf"), codes.size(-1) + ) + logits = logits[:, :length, :] + codes = codes[:, :length] + if args.mask: + mask = mask[:, :length] + + if args.mask: + masked_loss = F.cross_entropy(logits[mask], codes[mask]) + unmasked_loss = F.cross_entropy(logits[~mask], codes[~mask]) + loss = args.alpha * masked_loss + (1 - args.alpha) * unmasked_loss + else: + loss = F.cross_entropy(logits.transpose(1, 2), codes) + + scaler.scale(loss).backward() + scaler.unscale_(optimizer) + + nn.utils.clip_grad_norm_(hubert.parameters(), MAX_NORM) + + scaler.step(optimizer) + scaler.update() + + if args.mask: + masked_accuracy = logits[mask].argmax(dim=-1) == codes[mask] + masked_accuracy = torch.mean(masked_accuracy.float()) + + unmasked_accuracy = logits[~mask].argmax(dim=-1) == codes[~mask] + unmasked_accuracy = torch.mean(unmasked_accuracy.float()) + else: + accuracy = logits.argmax(dim=-1) == codes + accuracy = torch.mean(accuracy.float()) + + ############################################################################ + # Update and log training metrics + ############################################################################ + + if args.mask: + average_masked_loss.update(masked_loss.item()) + average_unmasked_loss.update(unmasked_loss.item()) + average_masked_accuracy.update(masked_accuracy.item()) + average_unmasked_accuracy.update(unmasked_accuracy.item()) + + epoch_masked_loss.update(masked_loss.item()) + epoch_unmasked_loss.update(unmasked_loss.item()) + epoch_masked_accuracy.update(masked_accuracy.item()) + epoch_unmasked_accuracy.update(unmasked_accuracy.item()) + else: + average_loss.update(loss.item()) + average_accuracy.update(accuracy.item()) + + epoch_loss.update(loss.item()) + epoch_accuracy.update(accuracy.item()) + + if rank == 0 and global_step % LOG_INTERVAL == 0: + if args.mask: + writer.add_scalar( + "train/masked_loss", + average_masked_loss.value, + global_step, + ) + writer.add_scalar( + "train/unmasked_loss", + average_unmasked_loss.value, + global_step, + ) + writer.add_scalar( + "train/masked_accuracy", + average_masked_accuracy.value * 100, + global_step, + ) + writer.add_scalar( + "train/unmasked_accuracy", + average_unmasked_accuracy.value * 100, + global_step, + ) + average_masked_loss.reset() + average_unmasked_loss.reset() + average_masked_accuracy.reset() + average_unmasked_accuracy.reset() + else: + writer.add_scalar( + "train/loss", + average_loss.value, + global_step, + ) + writer.add_scalar( + "train/accuracy", + average_accuracy.value, + global_step, + ) + average_loss.reset() + average_accuracy.reset() + + # --------------------------------------------------------------------------# + # Start validation loop + # --------------------------------------------------------------------------# + + if global_step % VALIDATION_INTERVAL == 0: + hubert.eval() + validation_loss.reset() + validation_accuracy.reset() + for wavs, codes in validation_loader: + wavs, codes = wavs.to(rank), codes.to(rank) + + with torch.no_grad(): + logits, _ = hubert(wavs) + logits = logits.transpose(1, 2) + + loss = F.cross_entropy(logits, codes) + + accuracy = logits.argmax(dim=1) == codes + accuracy = torch.mean(accuracy.float()) + + #################################################################### + # Update validation metrics + #################################################################### + + validation_loss.update(loss.item()) + validation_accuracy.update(accuracy.item()) + + hubert.train() + + ############################################################################ + # Log validation metrics + ############################################################################ + + if rank == 0: + writer.add_scalar( + "validation/unit_loss", + validation_loss.value, + global_step, + ) + writer.add_scalar( + "validation/unit_accuracy", + validation_accuracy.value * 100, + global_step, + ) + logger.info( + f"valid -- epoch: {epoch}, loss: {validation_loss.value:.4f}, accuracy: {validation_accuracy.value * 100:.2f}" + ) + + ############################################################################ + # Save model checkpoint + ############################################################################ + + new_best = best_loss > validation_loss.value + if new_best or global_step % CHECKPOINT_INTERVAL == 0: + if new_best: + logger.info("-------- new best model found!") + best_loss = validation_loss.value + + if rank == 0: + save_checkpoint( + checkpoint_dir=args.checkpoint_dir, + hubert=hubert, + optimizer=optimizer, + scaler=scaler, + step=global_step, + loss=validation_loss.value, + best=new_best, + logger=logger, + ) + + # -----------------------------------------------------------------------------# + # End validation loop + # -----------------------------------------------------------------------------# + + #################################################################################### + # Log training metrics + #################################################################################### + + logger.info( + f""" + train -- epoch: {epoch}, masked loss: {epoch_masked_loss.value:.4f}, unmasked loss: {epoch_unmasked_loss.value:.4f}, + masked accuracy: {epoch_masked_accuracy.value * 100:.2f}, umasked accuracy: {epoch_unmasked_accuracy.value * 100:.2f} + """ + ) + + # ==================================================================================# + # End training loop + # ==================================================================================# + + dist.destroy_process_group() + + +def train_hubert(args): + world_size = torch.cuda.device_count() + mp.spawn( + train, + args=(world_size, args), + nprocs=world_size, + join=True, + ) + + +if __name__ == "__main__": + parser = argparse.ArgumentParser(description="Train HuBERT soft content encoder.") + parser.add_argument( + "dataset_dir", + metavar="dataset-dir", + help="path to the data directory.", + type=Path, + ) + parser.add_argument( + "checkpoint_dir", + metavar="checkpoint-dir", + help="path to the checkpoint directory.", + type=Path, + ) + parser.add_argument( + "--resume", + help="path to the checkpoint to resume from.", + type=Path, + ) + parser.add_argument( + "--warmstart", + help="whether to initialize from the fairseq HuBERT checkpoint.", + action="store_true", + ) + parser.add_argument( + "--mask", + help="whether to use input masking.", + action="store_true", + ) + parser.add_argument( + "--alpha", + help="weight for the masked loss.", + default=1, + type=float, + ) + args = parser.parse_args() + + world_size = torch.cuda.device_count() + mp.spawn( + train, + args=(world_size, args), + nprocs=world_size, + join=True, + ) diff --git a/voice_conversion/softvc/softvctraceandonnx.py b/voice_conversion/softvc/softvctraceandonnx.py new file mode 100644 index 0000000..bea11ee --- /dev/null +++ b/voice_conversion/softvc/softvctraceandonnx.py @@ -0,0 +1,508 @@ +# -*- coding: utf-8 -*- +"""SoftVCTraceAndOnnx.ipynb + +Automatically generated by Colaboratory. + +Original file is located at + https://colab.research.google.com/drive/1g0YsQQY8fzXHmZUWAMzQO60c6OUlyKBl +""" + +!pip3 install torch torchvision torchaudio --extra-index-url https://download.pytorch.org/whl/cu116 +!pip3 install sklearn +!pip3 install soundfile +!pip3 install matplotlib + +!git clone https://github.com/bshall/hubert +!git clone https://github.com/bshall/acoustic-model +!git clone https://github.com/bshall/hifigan + +!nvidia-smi + +from google.colab import drive +drive.mount('/content/drive') + +# Torch Script Inference Hubert +import copy +from typing import Optional, Tuple +import random + +from sklearn.cluster import KMeans + +import torch +import torch.nn as nn +import torch.nn.functional as F +from torch.nn.modules.utils import consume_prefix_in_state_dict_if_present +import torch.nn.utils as utils + +from torch.nn.utils.weight_norm import WeightNorm +from typing import Dict, Any +from collections import OrderedDict + +def remove_weight_norm(module): + module_list = [mod for mod in module.children()] + if len(module_list) == 0: + for k, hook in module._forward_pre_hooks.items(): + if isinstance(hook, WeightNorm): + hook.remove(module) + del module._forward_pre_hooks[k] + else: + for mod in module_list: + remove_weight_norm(mod) + + +def modify_model(state_dict:Dict[str,Any]): + new_state_dict = OrderedDict() + for k in state_dict['acoustic-model'].keys(): + value = state_dict['acoustic-model'][k] + name = k[7:] # remove `module.` because they used ddp training to save + new_state_dict[name] = value + return new_state_dict + +# @torch.jit.export + # def generate(self, x: torch.Tensor) -> torch.Tensor: + # x = self.encoder(x) + # return self.decoder.generate(x) + class AcousticModelScripted(nn.Module): + def __init__(self, discrete: bool = False, upsample: bool = True): + super().__init__() + self.encoder = Encoder(discrete, upsample) + self.decoder = Decoder() + + def forward(self, x: torch.Tensor, mels: torch.Tensor) -> torch.Tensor: + x = self.encoder(x) + return self.decoder(x, mels) + + @torch.jit.export + def generate(self, x: torch.Tensor) -> torch.Tensor: + x = self.encoder(x) + return self.decoder.generate(x) + + + class Encoder(nn.Module): + def __init__(self, discrete: bool = False, upsample: bool = True): + super().__init__() + self.embedding = nn.Embedding(100 + 1, 256) if discrete else None + self.prenet = PreNet(256, 256, 256) + self.convs = nn.Sequential( + nn.Conv1d(256, 512, 5, 1, 2), + nn.ReLU(), + nn.InstanceNorm1d(512), + nn.ConvTranspose1d(512, 512, 4, 2, 1) if upsample else nn.Identity(), + nn.Conv1d(512, 512, 5, 1, 2), + nn.ReLU(), + nn.InstanceNorm1d(512), + nn.Conv1d(512, 512, 5, 1, 2), + nn.ReLU(), + nn.InstanceNorm1d(512), + ) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + if self.embedding is not None: + x = self.embedding(x) + x = self.prenet(x) + x = self.convs(x.transpose(1, 2)) + return x.transpose(1, 2) + + + class Decoder(nn.Module): + def __init__(self): + super().__init__() + self.prenet = PreNet(128, 256, 256) + self.lstm1 = nn.LSTM(512 + 256, 768, batch_first=True) + self.lstm2 = nn.LSTM(768, 768, batch_first=True) + self.lstm3 = nn.LSTM(768, 768, batch_first=True) + self.proj = nn.Linear(768, 128, bias=False) + + def forward(self, x: torch.Tensor, mels: torch.Tensor) -> torch.Tensor: + mels = self.prenet(mels) + x, _ = self.lstm1(torch.cat((x, mels), dim=-1)) + res = x + x, _ = self.lstm2(x) + x = res + x + res = x + x, _ = self.lstm3(x) + x = res + x + return self.proj(x) + + @torch.inference_mode() + def generate(self, xs: torch.Tensor) -> torch.Tensor: + m = torch.zeros(xs.size(0), 128, device=xs.device) + h1 = torch.zeros(1, xs.size(0), 768, device=xs.device) + c1 = torch.zeros(1, xs.size(0), 768, device=xs.device) + h2 = torch.zeros(1, xs.size(0), 768, device=xs.device) + c2 = torch.zeros(1, xs.size(0), 768, device=xs.device) + h3 = torch.zeros(1, xs.size(0), 768, device=xs.device) + c3 = torch.zeros(1, xs.size(0), 768, device=xs.device) + + mel = [] + for x in torch.unbind(xs, dim=1): + m = self.prenet(m) + x = torch.cat((x, m), dim=1).unsqueeze(1) + x1, (h1, c1) = self.lstm1(x, (h1, c1)) + x2, (h2, c2) = self.lstm2(x1, (h2, c2)) + x = x1 + x2 + x3, (h3, c3) = self.lstm3(x, (h3, c3)) + x = x + x3 + m = self.proj(x).squeeze(1) + mel.append(m) + return torch.stack(mel, dim=1) + + + class PreNet(nn.Module): + def __init__( + self, + input_size: int, + hidden_size: int, + output_size: int, + dropout: float = 0.5, + ): + super().__init__() + self.net = nn.Sequential( + nn.Linear(input_size, hidden_size), + nn.ReLU(), + nn.Dropout(dropout), + nn.Linear(hidden_size, output_size), + nn.ReLU(), + nn.Dropout(dropout), + ) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + return self.net(x) + +import IPython.display as display +import torch, torchaudio +import torch.nn as nn +import sys +import soundfile as sf + +from time import time +import torch.jit +sys.path.insert(0, './hubert/hubert') +from model import HubertSoft + +sys.path.insert(0, './acoustic-model/') +from acoustic.model import AcousticModel + +#@title select input wav +file_name = "/content/drive/MyDrive/VCModels/Source/brandon1.wav"#@param { type: "string"} + +source, sr = torchaudio.load(f"{file_name}") +source = torchaudio.functional.resample(source, sr, 16000) +source = source.unsqueeze(0).cuda() +print(f"Wav Shape and Type {source.shape} {source.dtype}") +if source.shape[1] == 2: # remove channel + source = torch.mean(source, dim=0).unsqueeze(0) + print(f"Removing Channel to mono Wav Shape and Type {source.shape} {source.dtype}") + +display.Audio(source.squeeze().cpu(), rate=16000) + +import datetime + +hubert_soft = torch.hub.load("bshall/hubert:main", "hubert_soft").cuda() +hubert_soft.eval() + +script_acoustic = True #@param {type: "boolean"} +print(script_acoustic) + +# Load the acoustic model (either hubert_soft or hubert_discrete) +input_acoustic_model_file = "/content/drive/MyDrive/VCModels/Padme/model-24000.pt" #@param {type: "string"} +acoustic_dict = torch.load(input_acoustic_model_file) # this used ddp ## REPLACE THIS! <-- ! +if script_acoustic == True: + acoustic = AcousticModelScripted().cuda() +else: + acoustic = AcousticModel().cuda() + +acoustic.eval() + +new_state_dict = modify_model(state_dict=acoustic_dict) +acoustic.load_state_dict(new_state_dict) + +import importlib.util +import sys + +spec = importlib.util.spec_from_file_location("hifigan.utils", "/content/hifigan/hifigan/utils.py") +u = importlib.util.module_from_spec(spec) +sys.modules["hifigan.utils"] = u +spec.loader.exec_module(u) + +spec = importlib.util.spec_from_file_location("hifigan.generator", "/content/hifigan/hifigan/generator.py") +g = importlib.util.module_from_spec(spec) +sys.modules["hifigan.generator"] = g +spec.loader.exec_module(g) + +new_state_dict = modify_model(state_dict=acoustic_dict) +acoustic.load_state_dict(new_state_dict) +input_hifigan_checkpoint_path = "/content/drive/MyDrive/VCModels/Hifigan/Padme.pt" #@param {type: "string"} +if input_hifigan_checkpoint_path: + checkpoint = torch.load(input_hifigan_checkpoint_path)["generator"]["model"] + hifigan = g.HifiganGenerator() + consume_prefix_in_state_dict_if_present(checkpoint, "module.") + hifigan.load_state_dict(checkpoint) + hifigan.remove_weight_norm() +else: + hifigan = torch.hub.load("bshall/hifigan:main", "hifigan_hubert_soft") + +hifigan.cuda() +hifigan.eval() + +#utils.remove_weight_norm(hubert_soft.positional_embedding.conv) +# Convert to the target speaker +units = None +mel = None + +unit_start_time = 0.0 +unit_end_time = 0.0 + +acoustic_start_time = 0.0 +acoustic_end_time = 0.0 + +hifi_start_time = 0.0 +hifi_end_time = 0.0 + +with torch.inference_mode(): + unit_start_time = datetime.datetime.now() + units = hubert_soft.units(source) + print(f"Units and Type {units.shape} {units.dtype}") + unit_end_time = datetime.datetime.now() + + hubert_soft_script = torch.jit.trace_module(hubert_soft,{"units":source}) + torch.jit.save(hubert_soft_script, 'hubert_soft.jit') + + print(f"Unit Generation Time {(unit_end_time - unit_start_time).total_seconds() * 1000}") + + print("Hubert Soft Exported") + + acoustic_start_time = datetime.datetime.now() + mel = acoustic.generate(units).transpose(1, 2) + print(f"Mel and Type {mel.shape} {mel.dtype}") + acoustic_end_time = datetime.datetime.now() + print(f"Acoustic Time {(acoustic_end_time - acoustic_start_time).total_seconds() * 1000}") + if script_acoustic == True: + output_acoustic_scripted = '/content/drive/MyDrive/VCModels/FinishedModels/Padme.recast' #@param {type: "string"} + print("Exporting Scripted") + acoustic_model_script = torch.jit.script(acoustic) + torch.jit.save(acoustic_model_script, output_acoustic_scripted) + else: + print("Exporting Traced") + acoustic_model_script = torch.jit.trace_module(acoustic,{"generate":units}) + output_acoustic_traced = '' #@param {type: "string"} + torch.jit.save(acoustic_model_script, output_acoustic_traced) + + print("Acoustic Model Exported") + hifi_start_time = datetime.datetime.now() + + target = hifigan(mel) + print(f"Wavform and Type {target.shape} {target.dtype}") + hifi_end_time = datetime.datetime.now() + print(f"MEL {target.shape}") + print(f"Hifi Gan Time {(hifi_end_time - hifi_start_time).total_seconds() * 1000}") + + hifi_gan_script = torch.jit.trace_module(hifigan,{"forward":mel}) + output_hifigan_jit = '/content/drive/MyDrive/VCModels/FinishedModels/Padme-hifigan-jit.pt' #@param {type: "string"} + torch.jit.save(hifi_gan_script, output_hifigan_jit) + print("Hifi Gan Model Exported") + +import datetime + +jit_unit_start_time = 0.0 +jit_unit_end_time = 0.0 +jit_unit_final = 0.0 + +jit_acoustic_start_time = 0.0 +jit_acoustic_end_time = 0.0 +jit_acoustic_final = 0.0 + +jit_hifi_start_time = 0.0 +jit_hifi_end_time = 0.0 +jit_hifi_final = 0.0 + +jit_model = torch.jit.load('hubert_soft.jit').cuda() +jit_model.eval() +with torch.no_grad(): + jit_unit_start_time = datetime.datetime.now() + jit_tensor = jit_model.units(source) + jit_unit_end_time = datetime.datetime.now() + jit_unit_final = jit_unit_end_time - jit_unit_start_time +print(jit_unit_final.total_seconds() * 1000) +print(jit_tensor.shape) +assert torch.allclose(jit_tensor,units) + +print(script_acoustic) +if script_acoustic == True: + acoustic_jit_model = torch.jit.load(output_acoustic_scripted).cuda() + print("Using Scripted Model") +else: + acoustic_jit_model = torch.jit.load(output_acoustic_traced).cuda() +acoustic_jit_model.eval() +with torch.no_grad(): + if script_acoustic == True: + print("Using Scripted Model") + jit_acoustic_start_time = datetime.datetime.now() + jit_mel = acoustic_jit_model.generate(jit_tensor) + else: + jit_acoustic_start_time = datetime.datetime.now() + jit_mel = acoustic_jit_model.generate(jit_tensor) + jit_acoustic_end_time = datetime.datetime.now() + jit_acoustic_final = jit_acoustic_end_time - jit_acoustic_start_time + print(jit_acoustic_final.total_seconds() * 1000) + +jit_mel = jit_mel.transpose(1,2) +print(jit_mel.shape) +assert torch.allclose(mel,jit_mel) + +hifi_gan = torch.jit.load(output_hifigan_jit).cuda() +hifi_gan.eval() +with torch.no_grad(): + jit_hifi_start_time = datetime.datetime.now() + jit_target = hifi_gan(jit_mel) + jit_hifi_end_time = datetime.datetime.now() + jit_hifi_final = jit_hifi_end_time - jit_hifi_start_time +print(jit_hifi_final.total_seconds() * 1000) +assert torch.allclose(target,jit_target) + +sf.write(f'./sample_jit.wav', jit_target.squeeze().detach().cpu().numpy(), samplerate=16000) +sf.write(f'./sample.wav', target.squeeze().cpu().detach().numpy(), samplerate=16000) + +#uploaded = files.upload() +import datetime + +jit_unit_start_time = 0.0 +jit_unit_end_time = 0.0 +jit_unit_final = 0.0 + +jit_acoustic_start_time = 0.0 +jit_acoustic_end_time = 0.0 +jit_acoustic_final = 0.0 + +jit_hifi_start_time = 0.0 +jit_hifi_end_time = 0.0 +jit_hifi_final = 0.0 + + +source, sr = torchaudio.load(f"{file_name}") +source = torchaudio.functional.resample(source, sr, 16000) +source = source.unsqueeze(0).cuda() +print(f"Wav Shape and Type {source.shape} {source.dtype}") +if source.shape[1] == 2: # remove channel + source = torch.mean(source, dim=1).unsqueeze(0) + print(f"Removing Channel to mono Wav Shape and Type {source.shape} {source.dtype}") + +display.Audio(source.squeeze().cpu(), rate=16000) + +jit_model = torch.jit.load('hubert_soft.jit').cuda() +jit_model.eval() + +if script_acoustic == True: + acoustic_jit_model = torch.jit.load(output_acoustic_scripted).cuda() + print("Loading Scripted Model") +else: + acoustic_jit_model = torch.jit.load(output_acoustic_traced).cuda() + print("Loading Tracing Model") + +acoustic_jit_model.eval() + +hifi_gan = torch.jit.load(output_hifigan_jit).cuda() +hifi_gan.eval() +print("Loaded Models") + +# Warm up the model to get the real time +for i in range(0,5): + with torch.no_grad(): + jit_unit_start_time = datetime.datetime.now() + jit_tensor = jit_model.units(source) + jit_unit_end_time = datetime.datetime.now() + + if script_acoustic == True: + print("Using Scripted Model TT2") + jit_acoustic_start_time = datetime.datetime.now() + jit_mel = acoustic_jit_model.generate(jit_tensor) + else: + print("Using Traced Model TT2") + jit_acoustic_start_time = datetime.datetime.now() + jit_mel = acoustic_jit_model.generate(jit_tensor) + jit_mel = jit_mel.transpose(1, 2) + + jit_acoustic_end_time = datetime.datetime.now() + + jit_hifi_start_time = datetime.datetime.now() + jit_target = hifi_gan(jit_mel) + jit_hifi_end_time = datetime.datetime.now() + + +jit_unit_final = jit_unit_end_time - jit_unit_start_time +jit_acoustic_final = jit_acoustic_end_time - jit_acoustic_start_time +jit_hifi_final = jit_hifi_end_time - jit_hifi_start_time + +print(f"16 Khz Wav Shape {source.shape}") + +print(f"Jit Tensor Shape for Unit Inputs: {jit_tensor.shape}") +print(f"Hubert Runtime : {jit_unit_final.total_seconds() * 1000}") + +print(f"Tactron 2 Runtime Mel Output Shape: {jit_acoustic_final.total_seconds() * 1000}") +print(f"Jit Tensor Shape for Mel output: {jit_mel.shape}") + +print(f"Hifi Runtime Mel Output Shape: {jit_hifi_final.total_seconds() * 1000}") +print(f"Jit Tensor Shape for Mel output: {jit_target.shape}") + +sf.write(f'./{file_name}_jit.wav', jit_target.squeeze().detach().cpu().numpy(), samplerate=16000) +display.Audio(jit_target.squeeze().cpu(), rate=16000) + +# silero vad +import datetime + + +source, sr = torchaudio.load(f"{file_name}") +source = torchaudio.functional.resample(source, sr, 16000) + +print(source.shape) + +input_vad_file = "/content/drive/MyDrive/tts/softvc-checkpoints/silero_vad.jit" #@param {type: "string"} + +class VADIterator: + # You will only need to verify the 100 ms chunks to see this checks a minute of audio + # sound_sample_size run vad in 30 ms, 60 ms or 100 ms + def __init__(self,threshold: float = 0.5, sampling_rate: int = 16000,sound_sample_size:float = 100): + self.silero_jit = torch.jit.load(input_vad_file).cuda() + self.silero_jit.eval() + self.sampling_rate = sampling_rate + self.sound_sample_size = sound_sample_size + + jit_silero_start_time = 0.0 + jit_silero_end_time = 0.0 + jit_silero_final = 0.0 + + def operating_frame_size(self,sampling_rate:int,ms_of_audio:float): + return int((ms_of_audio * sampling_rate) / 1000) # number of frames + + def check_voice_activity(self,source_sample) -> float: + probs_over_time = [] + # formula we know that we can get 1600 operating frames / 16000 sr * 1000 = 100 ms of audio from 1600 frames. + # this gives you the operating frames given the audio size you want. + op_frame_size = self.operating_frame_size(sampling_rate=self.sampling_rate,ms_of_audio=self.sound_sample_size) + print(op_frame_size) + iters = int(source_sample.shape[-1] / op_frame_size) # get the chunks + #print(iters) + for i in range(0,iters): # get sound_sample_size ms chunks + + #print(window_size) + chunk = source_sample[i * op_frame_size:(i + 1) * op_frame_size] + with torch.no_grad():# slide window across using VAD 100ms + jit_silero_start_time = datetime.datetime.now() + try: + print(f"Current {i}") + prob = self.silero_jit(chunk.cuda(),sr=self.sampling_rate).item() # input shape is [1*] just an array of samples. 1 = false + jit_silero_end_time = datetime.datetime.now() + print(f"Probability: {prob}") + jit_silero_final = jit_silero_end_time - jit_silero_start_time + #print(f"Silero VAD Time: {jit_silero_final.total_seconds() * 1000}") + probs_over_time.append(prob) + except Exception as e: + print(e) + + return probs_over_time + +iterator = VADIterator(sound_sample_size=100) +source_sample = source_sample=source.squeeze().squeeze() +print(source_sample.shape) +probs = iterator.check_voice_activity(source_sample=source_sample) +#print(probs) \ No newline at end of file