48khz softvc (#16)

* 48KHz softvc * add some notes * almost forgot the inference script * fix readme link * requirements.txt
2026-10-16 00:13:07 +00:00 · 2023-05-24 21:43:03 -07:00
parent c18b3f878a
commit f86b44d8fc
39 changed files with 3840 additions and 0 deletions
@@ -0,0 +1,15 @@
+# 48KHz SoftVC
+
+Herein lies 48KHz softvc, a failed experiment to produce a high samplerate version of https://github.com/bshall/soft-vc
+The goal was to leverage the strong character quality that softvc has, ie it's ability to work well with weird voices
+such as Evanescence, GladOS, Darth Vader, and Yoda, while providing a higher samplerate than the 16KHz of the original
+version. We used the same technique as so-vits-svc, to upscale the 16KHz HuBERT "speech units" to 48KHz equivalents,
+and trained on 48KHz and 44.1 KHz source data. Unfortunately, the results have a robotic, metallic sound to them,
+and sound basically the same as if it was 16KHz (with weird robotic noises added for flavour).
+
+
+Additionally, warm-starting didn't seem to work, maybe there's some trick to remove the speaker embedding of the base
+model, but it's not part of the code provided by bshall.
+
+
+To hear the results for yourself, visit https://drive.google.com/drive/folders/1byrr6R70qr4kUVj39KXi9NCXmbJVJgNU?usp=sharing
@@ -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.
@@ -0,0 +1,180 @@
+<p align="center">
+    <a target="_blank" href="https://colab.research.google.com/github/bshall/soft-vc/blob/main/soft-vc-demo.ipynb">
+        <img src="https://colab.research.google.com/assets/colab-badge.svg" alt="Open In Colab"/>
+    </a>
+</p>
+
+# 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).
+
+<div align="center">
+    <img width="100%" alt="Soft-VC"
+      src="https://raw.githubusercontent.com/bshall/acoustic-model/main/acoustic-model.png">
+</div>
+<div>
+  <sup>
+    <strong>Fig 1:</strong> Architecture of the voice conversion system. a) The <strong>discrete</strong> content encoder clusters audio features to produce a sequence of discrete speech units. b) The <strong>soft</strong> 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.
+  </sup>
+</div>
+
+## 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}
+}
+```
@@ -0,0 +1 @@
+from .model import AcousticModel, hubert_discrete, hubert_soft
@@ -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).T
+        units = np.load(units_path)
+
+        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
@@ -0,0 +1,169 @@
+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 = False,
+    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 = False,
+    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 = False,
+    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,
+    )
+
+
@@ -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=48000,
+            n_fft=1024,
+            win_length=1024,
+            hop_length=480,
+            center=False,
+            power=1.0,
+            norm="slaney",
+            onesided=True,
+            n_mels=128,
+            mel_scale="slaney",
+        )
+
+    def forward(self, wav):
+        padding = (1024 - 480) // 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
@@ -0,0 +1,75 @@
+import argparse
+
+from pathlib import Path
+import numpy as np
+from tqdm import tqdm
+
+import torch
+
+from acoustic.model import AcousticModel
+from torch.nn.modules.utils import consume_prefix_in_state_dict_if_present
+
+
+def repeat_expand_2d(content: torch.Tensor, target_len: int) -> torch.Tensor:
+    # content : [h, t]
+    src_len = content.shape[-1]
+    if target_len < src_len:
+        return content[:, :target_len]
+    else:
+        return torch.nn.functional.interpolate(
+            content.unsqueeze(0), size=target_len, mode="nearest"
+        ).squeeze(0)
+
+
+def generate(args):
+    print("Loading acoustic model checkpoint")
+    #acoustic = torch.hub.load("bshall/acoustic-model:main", f"hubert_{args.model}").cuda()
+    acoustic = AcousticModel().cuda().eval()
+    checkpoint = torch.load("exp/rue/model-2000.pt", map_location="cuda")["acoustic-model"]
+    consume_prefix_in_state_dict_if_present(checkpoint, "module.")
+    acoustic.load_state_dict(checkpoint)
+
+    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()
+        units = units.T
+        print(units.shape)
+        target_len = units.shape[1]*3
+        units = units.squeeze(0)
+        units = repeat_expand_2d(units, target_len).T.unsqueeze(0)
+
+        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)
+
@@ -0,0 +1,3 @@
+dependencies = ["torch", "torchaudio"]
+
+from acoustic import hubert_discrete, hubert_soft
@@ -0,0 +1,65 @@
+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, 48000)
+
+    logmel = melspectrogram(wav.unsqueeze(0))
+
+    np.save(str(out_path).replace(".wav", ".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
+        out_path.parent.mkdir(parents=True, exist_ok=True)
+        process_wav(in_path, out_path)
+        #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 = 441 / 44100
+    #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)
@@ -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 48kHz)",
+        type=int,
+        default=48000,
+    )
+    args = parser.parse_args()
+    preprocess_dataset(args)
@@ -0,0 +1,319 @@
+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 = 64
+LEARNING_RATE = 4e-4
+BETAS = (0.8, 0.99)
+WEIGHT_DECAY = 1e-5
+STEPS = 9999999999999999999
+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()
+
+
+                if global_step % CHECKPOINT_INTERVAL==0:
+                    if rank == 0:
+                        save_checkpoint(
+                            checkpoint_dir=args.checkpoint_dir,
+                            acoustic=acoustic,
+                            optimizer=optimizer,
+                            step=global_step,
+                            loss=validation_loss.value,
+                            best=9999,
+                            logger=logger,
+                        )
+
+
+                ############################################################################
+                # 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:
+                    logger.info("-------- new best model found!")
+                    best_loss = validation_loss.value
+
+
+            # -----------------------------------------------------------------------------#
+            # 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,
+    )
@@ -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.
@@ -0,0 +1,139 @@
+<p align="center">
+    <a target="_blank" href="https://colab.research.google.com/github/bshall/soft-vc/blob/main/soft-vc-demo.ipynb">
+        <img src="https://colab.research.google.com/assets/colab-badge.svg" alt="Open In Colab"/>
+    </a>
+</p>
+
+# 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).
+
+<div align="center">
+    <img width="100%" alt="Soft-VC"
+      src="https://raw.githubusercontent.com/bshall/hifigan/main/vocoder.png">
+</div>
+<div>
+  <sup>
+    <strong>Fig 1:</strong> Architecture of the voice conversion system. a) The <strong>discrete</strong> content encoder clusters audio features to produce a sequence of discrete speech units. b) The <strong>soft</strong> 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.
+  </sup>
+</div>
+
+## 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).
@@ -0,0 +1,60 @@
+from pathlib import Path
+import numpy as np
+import argparse
+import torch
+import torchaudio
+from tqdm import tqdm
+from torch.nn.modules.utils import consume_prefix_in_state_dict_if_present
+
+from hifigan.generator import HifiganGenerator
+
+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()
+    checkpoint = torch.load("exp/vctk/model-66000.pt")["generator"]["model"]
+
+    consume_prefix_in_state_dict_if_present(checkpoint, "module.")
+    hifigan = HifiganGenerator()
+    hifigan.load_state_dict(checkpoint)
+    hifigan.remove_weight_norm()
+    hifigan.eval()
+    hifigan.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, 48000)
+
+
+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)
@@ -0,0 +1 @@
+from .generator import hifigan, hifigan_hubert_discrete, hifigan_hubert_soft
@@ -0,0 +1,127 @@
+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=48000,
+            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/**/*.wav" if train else "dev/**/*.wav}"
+
+        self.metadata = [
+            path.relative_to(self.data_dir)
+            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)
+        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.replace(".wav", ".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,
+            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()
+            wav = flip * gain * wav / max([wav.abs().max(), 0.001])
+
+        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
@@ -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
@@ -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 = 48000,
+    ) -> 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 = False,
+    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 = False, progress: bool = True, map_location=None
+) -> HifiganGenerator:
+    return _hifigan("hifigan", pretrained, progress, map_location)
+
+
+def hifigan_hubert_soft(
+    pretrained: bool = False, progress: bool = True, map_location=None
+) -> HifiganGenerator:
+    return _hifigan("hifigan-hubert-soft", pretrained, progress, map_location=None)
+
+
+def hifigan_hubert_discrete(
+    pretrained: bool = False, progress: bool = True, map_location=None
+) -> HifiganGenerator:
+    return _hifigan("hifigan-hubert-discrete", pretrained, progress, map_location=None)
@@ -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"]
@@ -0,0 +1,3 @@
+dependencies = ["torch", "torchaudio"]
+
+from hifigan import hifigan, hifigan_hubert_discrete, hifigan_hubert_soft
@@ -0,0 +1,4 @@
+tensorboard==2.7.0
+torch==1.9.1
+torchaudio==0.9.1
+tqdm==4.62.3
@@ -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 48kHz)",
+        type=int,
+        default=48000,
+    )
+    args = parser.parse_args()
+    preprocess_dataset(args)
@@ -0,0 +1,351 @@
+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 = 32
+SEGMENT_LENGTH = 8320
+HOP_LENGTH = 160
+SAMPLE_RATE = 48000
+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 = 1000
+
+
+def train_model(rank, world_size, args):
+    dist.init_process_group(
+        "nccl",
+        rank=rank,
+        world_size=world_size,
+        init_method="tcp://localhost:12345",
+    )
+
+    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)
+            #length = min(mels_.size(-1), tgts.size(-1))
+            #loss_mel = F.l1_loss(mels_[..., :length], tgts[..., :length])
+
+            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=48000,
+                            )
+                            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 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=9999,
+                            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()
+    print(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,
+    )
@@ -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.
@@ -0,0 +1,163 @@
+<p align="center">
+    <a target="_blank" href="https://colab.research.google.com/github/bshall/soft-vc/blob/main/soft-vc-demo.ipynb">
+        <img src="https://colab.research.google.com/assets/colab-badge.svg" alt="Open In Colab"/>
+    </a>
+</p>
+
+# 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).
+
+<div align="center">
+    <img width="100%" alt="Soft-VC"
+      src="https://raw.githubusercontent.com/bshall/hubert/main/content-encoder.png">
+</div>
+<div>
+  <sup>
+    <strong>Fig 1:</strong> Architecture of the voice conversion system. a) The <strong>discrete</strong> content encoder clusters audio features to produce a sequence of discrete speech units. b) The <strong>soft</strong> 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.
+  </sup>
+</div>
+
+## 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}
+}
+```
@@ -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)
@@ -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)
@@ -0,0 +1,3 @@
+dependencies = ["torch", "torchaudio", "sklearn"]
+
+from hubert import hubert_discrete, hubert_soft, kmeans100
@@ -0,0 +1,8 @@
+from .model import (
+    Hubert,
+    HubertDiscrete,
+    HubertSoft,
+    hubert_discrete,
+    hubert_soft,
+    kmeans100,
+)
@@ -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
@@ -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)
@@ -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"]
@@ -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,
+    )
@@ -0,0 +1,121 @@
+#@title Load the models 
+#@markdown leave hifigan blank to use defaults
+import sys
+import torch, torchaudio
+from typing import Dict, Any
+from collections import OrderedDict
+from torch.nn.modules.utils import consume_prefix_in_state_dict_if_present
+
+# Inputs 
+check_point_path = './acoustic-model/exp/glados/model-35000.pt' #@param {type:"string"}
+hifigan_checkpoint_path = "./hifigan/exp/vctk/model-156000.pt" #@param {type: "string"}
+file_name = "/mnt/8terra/tts/voice-conversion-sources/ieatmeat.wav" #@param {type: "string"}
+
+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
+
+def repeat_expand_2d(content: torch.Tensor, target_len: int) -> torch.Tensor:
+    # content : [h, t]
+    src_len = content.shape[1]
+    if target_len < src_len:
+        return content[:, target_len]
+    else:
+        return torch.nn.functional.interpolate(
+            content.unsqueeze(0), size=target_len, mode="nearest"
+        )
+
+
+sys.path.append("./acoustic-model/")
+from acoustic import AcousticModel
+
+hubert = torch.hub.load("bshall/hubert:main", "hubert_soft").cuda()
+
+
+# Load the acoustic model (either hubert_soft or hubert_discrete) 
+acoustic_dict = torch.load(check_point_path) # this used ddp
+acoustic = AcousticModel().cuda()
+acoustic.eval()
+
+import importlib.util
+import sys
+
+spec = importlib.util.spec_from_file_location("hifigan.utils", "./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", "./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)
+if hifigan_checkpoint_path:
+  checkpoint = torch.load(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()
+  hifigan.eval()
+else:
+  hifigan = torch.hub.load("bshall/hifigan:main", "hifigan_hubert_soft")
+hifigan.cuda()
+
+#@title This code selects the source file
+#import IPython.display as display
+sys.path.insert(0, './acoustic-model/acoustic/')
+from model import AcousticModel
+import soundfile as sf
+#from google.colab import files
+from time import time
+sys.path.append("./hifigan/")
+from hifigan.generator import HifiganGenerator
+from hifigan.generator import consume_prefix_in_state_dict_if_present
+#uploaded = files.upload()
+
+#for name in uploaded.keys():
+    #file_name = name    
+source, sr = torchaudio.load(f"{file_name}")
+source = torchaudio.functional.resample(source, sr, 16000)
+source = torch.mean(source, dim=0).unsqueeze(0)
+source = source.unsqueeze(0).cuda()
+#display.Audio(source.squeeze().cpu(), rate=16000)
+
+#@title Run the inference and display the result
+# Convert to the target speaker
+with torch.inference_mode():
+    # Extract speech units
+    start = time()
+    units = hubert.units(source)
+    print(units.shape)
+    #units = repeat_expand_2d(units, target_len)
+    units = torch.nn.functional.interpolate(
+            units.transpose(1,2), size=units.shape[1]*3, mode="nearest"
+    ).transpose(1,2)
+
+    print(units.shape)
+    #units = units.unsqueeze(0)
+    #units = units.transpose(1,2)
+
+    end = time()
+    # Generate target spectrogram
+    print(f"time for feature embedding:{end-start}")
+
+    start = time()
+    mel = acoustic.generate(units).transpose(1,2)
+    end = time()
+    print(f"time for inference:{end-start}")
+    # Generate audio waveform
+    start = time()
+    target = hifigan(mel)
+    end = time()
+    print(f"time for wavform {end-start}")
+    print(target.squeeze().cpu().unsqueeze(0).shape)
+
+torchaudio.save("output.wav", target.squeeze().cpu().unsqueeze(0), 48000)
@@ -0,0 +1,59 @@
+absl-py==1.4.0
+cachetools==5.3.0
+certifi==2023.5.7
+charset-normalizer==3.1.0
+cmake==3.26.3
+contourpy==1.0.7
+cycler==0.11.0
+filelock==3.12.0
+fonttools==4.39.4
+google-auth==2.18.1
+google-auth-oauthlib==1.0.0
+grpcio==1.54.2
+idna==3.4
+Jinja2==3.1.2
+joblib==1.2.0
+kiwisolver==1.4.4
+lit==16.0.5
+Markdown==3.4.3
+MarkupSafe==2.1.2
+matplotlib==3.7.1
+mpmath==1.3.0
+networkx==3.1
+numpy==1.24.3
+nvidia-cublas-cu11==11.10.3.66
+nvidia-cuda-cupti-cu11==11.7.101
+nvidia-cuda-nvrtc-cu11==11.7.99
+nvidia-cuda-runtime-cu11==11.7.99
+nvidia-cudnn-cu11==8.5.0.96
+nvidia-cufft-cu11==10.9.0.58
+nvidia-curand-cu11==10.2.10.91
+nvidia-cusolver-cu11==11.4.0.1
+nvidia-cusparse-cu11==11.7.4.91
+nvidia-nccl-cu11==2.14.3
+nvidia-nvtx-cu11==11.7.91
+oauthlib==3.2.2
+packaging==23.1
+Pillow==9.5.0
+protobuf==4.23.1
+pyasn1==0.5.0
+pyasn1-modules==0.3.0
+pyparsing==3.0.9
+python-dateutil==2.8.2
+requests==2.30.0
+requests-oauthlib==1.3.1
+rsa==4.9
+scikit-learn==1.2.2
+scipy==1.10.1
+six==1.16.0
+sympy==1.12
+tensorboard==2.13.0
+tensorboard-data-server==0.7.0
+threadpoolctl==3.1.0
+torch==2.0.1
+torchaudio==2.0.2
+tqdm==4.65.0
+triton==2.0.0
+typing_extensions==4.5.0
+urllib3==1.26.15
+Werkzeug==2.3.4
				`@@ -0,0 +1 @@`
				`from .model import AcousticModel, hubert_discrete, hubert_soft`
				`@@ -0,0 +1 @@`
				`from .generator import hifigan, hifigan_hubert_discrete, hifigan_hubert_soft`