Add vits model and normalizing flow. Jupyter Notebook as example call

pvq_manipulation/Example_Notebook.ipynb

@@ -0,0 +1,331 @@
+{
+ "cells": [
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "f0e32cd2-4955-4140-8f48-9751a1a8c588",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import numpy as np \n",
+    "from pathlib import Path\n",
+    "import padertorch as pt\n",
+    "import paderbox as pb\n",
+    "import time\n",
+    "import torch\n",
+    "import torchaudio\n",
+    "import ipywidgets as widgets\n",
+    "from onnxruntime import InferenceSession\n",
+    "from pvq_manipulation.models.vits import Vits_NT\n",
+    "from pvq_manipulation.models.ffjord import FFJORD\n",
+    "from IPython.display import display, Audio, clear_output\n",
+    "from pvq_manipulation.models.hubert import HubertExtractor, SID_LARGE_LAYER\n",
+    "from paderbox.transform.module_resample import resample_sox\n",
+    "from pvq_manipulation.helper.vad import EnergyVAD\n",
+    "from train_tts_nt.helper.utils import rms_norm"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "d4df1db0-8439-4573-9dc2-5d578e8befa1",
+   "metadata": {},
+   "source": [
+    "# load TTS model"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "e6691176-6119-4bf0-9dcf-44d657c76074",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "storage_dir_tts = Path(\"./Saved_models/tts_model/\")\n",
+    "tts_model = Vits_NT.load_model(storage_dir_tts, checkpoint=\"checkpoint_390000.pth\")"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "c9c7541c-fab5-4d44-9b89-a26a34343e7c",
+   "metadata": {},
+   "source": [
+    "# load normalizing flow"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "e4a55082-c6c6-4283-96ed-217553f33bcd",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "storage_dir_normalizing_flow = Path(\"./Saved_models/norm_flow\")\n",
+    "config_norm_flow = pb.io.load_yaml(storage_dir_normalizing_flow / \"config.yaml\")\n",
+    "normalizing_flow = FFJORD.load_model(storage_dir_normalizing_flow, checkpoint=\"checkpoints/ckpt_best_loss.pth\")"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "deebed07-b28c-49de-b30f-d80b9e1c6899",
+   "metadata": {},
+   "source": [
+    "# load hubert features model"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "bc4627e1-bac7-4533-8cac-bbc296889855",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "hubert_model = HubertExtractor(\n",
+    "    layer=SID_LARGE_LAYER,\n",
+    "    model_name=\"HUBERT_LARGE\",\n",
+    "    backend=\"torchaudio\",\n",
+    "    device='cpu', \n",
+    "    storage_dir='/net/vol/rautenberg/storage/hubert'# target storage dir hubert model\n",
+    ")"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "c78fa11b-8617-4175-902c-8af0e4491201",
+   "metadata": {},
+   "source": [
+    "# Example Synthesis"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "4e8afa1b-b02e-4a40-982d-36aa78f37a57",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "speaker_id = 1034\n",
+    "example_id = \"1034_121119_000028_000001\"\n",
+    "\n",
+    "wav_1 = tts_model.synthesize_from_example({\n",
+    "    'text' : \"It took me quite a long time to develop a voice, and now that I have it I'm not going to be silent.\", \n",
+    "    'd_vector_storage_root': f\"./Saved_models/Dataset/Embeddings/{speaker_id}/{example_id}.pth\"\n",
+    "})\n",
+    "display(Audio(wav_1, rate=24_000, normalize=True))"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "feeb1d62-69f2-45c1-a172-16fcfbecd0da",
+   "metadata": {},
+   "source": [
+    "# Manipulation Block"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "625368d3-dd35-4da7-a358-7bbac448806c",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "def get_manipulation(\n",
+    "    example, \n",
+    "    d_vector,\n",
+    "    labels,\n",
+    "    flow, \n",
+    "    tts_model,\n",
+    "    manipulation_idx=0,\n",
+    "    manipulation_fkt=1,\n",
+    "):\n",
+    "    labels_manipulated = labels.clone()\n",
+    "    labels_manipulated[:,manipulation_idx] += manipulation_fkt\n",
+    "    \n",
+    "    output_forward = flow.forward((d_vector.float(), labels))[0]\n",
+    "    sampled_class_manipulated = flow.sample((output_forward, labels_manipulated))[0]\n",
+    "\n",
+    "    wav = tts_model.synthesize_from_example({\n",
+    "        'text': \"It took me quite a long time to develop a voice, and now that I have it I'm not going to be silent.\",\n",
+    "        'd_vector': d_vector.detach().numpy(),\n",
+    "        'd_vector_man': sampled_class_manipulated.detach().numpy(),\n",
+    "    })    \n",
+    "    return wav\n",
+    "\n",
+    "def extract_speaker_embedding(example):\n",
+    "    observation, sr = pb.io.load_audio(example['audio_path']['observation'], return_sample_rate=True)\n",
+    "    observation = resample_sox(observation, in_rate=sr, out_rate=16_000)\n",
+    "    \n",
+    "    vad = EnergyVAD(sample_rate=16_000)\n",
+    "    if observation.ndim == 1:\n",
+    "        observation = observation[None, :]\n",
+    "    \n",
+    "    observation = vad({'audio_data': observation})['audio_data']\n",
+    "    \n",
+    "    with torch.no_grad():\n",
+    "        example = tts_model.speaker_manager.prepare_example({'audio_data': {'observation': observation}, **example})\n",
+    "        example = pt.data.utils.collate_fn([example])\n",
+    "        example['features'] = torch.tensor(np.array(example['features']))\n",
+    "        d_vector = tts_model.speaker_manager.forward(example)[0]\n",
+    "    return d_vector"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "b722e503-a8f4-4702-acce-20bcdd828846",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "def load_speaker_labels(example, config_norm_flow, reg_stor_dir=Path('./Saved_models/pvq_extractor/')):\n",
+    "    audio, _ = torchaudio.load(example['audio_path']['observation'])\n",
+    "    num_samples = torch.tensor([audio.shape[-1]])\n",
+    "\n",
+    "    if torch.cuda.is_available():\n",
+    "        audio = audio.cuda()\n",
+    "        num_samples = num_samples.cuda()\n",
+    "    providers = [\"CPUExecutionProvider\"]\n",
+    "\n",
+    "    with torch.no_grad():\n",
+    "        features, seq_len = hubert_model(\n",
+    "            audio, \n",
+    "            24_000, \n",
+    "            sequence_lengths=num_samples,\n",
+    "        )\n",
+    "        features = np.mean(features.squeeze(0).detach().cpu().numpy(), axis=-1)\n",
+    "\n",
+    "        pvqd_predictions = {}\n",
+    "        for pvq in ['Breathiness', 'Loudness', 'Pitch', 'Resonance', 'Roughness', 'Strain', 'Weight']:\n",
+    "            with open(reg_stor_dir / f\"{pvq}.onnx\", \"rb\") as fid:\n",
+    "                onnx = fid.read()\n",
+    "            sess = InferenceSession(onnx, providers=providers)\n",
+    "            pred = sess.run(None, {\"X\": features[None]})[0].squeeze(1)\n",
+    "            pvqd_predictions[pvq] = pred.tolist()[0]\n",
+    "    labels = []\n",
+    "    for key in config_norm_flow['speaker_conditioning']:\n",
+    "        labels.append(pvqd_predictions[key]/100)\n",
+    "    return torch.tensor(labels)"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "008035ba-6054-4e6e-ab16-1aaaf68f584a",
+   "metadata": {},
+   "source": [
+    "# Get example manipulation"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "e921a3cd-1699-495c-b825-519fb706d89d",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "example = {\n",
+    "    'audio_path': {'observation': \"./Saved_models/Dataset/Audio_files/1034_121119_000028_000001.wav\"},\n",
+    "    'speaker_id': 1034,\n",
+    "    'example_id': \"1034_121119_000028_000001\",\n",
+    "}\n",
+    "\n",
+    "d_vector = extract_speaker_embedding(example)\n",
+    "labels = load_speaker_labels(example, config_norm_flow)\n",
+    "\n",
+    "wav_manipulated = get_manipulation(\n",
+    "    example=example, \n",
+    "    d_vector=d_vector, \n",
+    "    labels=labels[None, :], \n",
+    "    flow=normalizing_flow,\n",
+    "    tts_model=tts_model,\n",
+    "    manipulation_idx=0,\n",
+    "    manipulation_fkt=1,\n",
+    ")"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "09a04e5b-c2ab-43e5-b9df-171028100ab6",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "example = {\n",
+    "    'audio_path': {'observation': \"./Saved_models/Dataset/Audio_files/1034_121119_000028_000001.wav\"},\n",
+    "    'speaker_id': 1034,\n",
+    "    'example_id': \"1034_121119_000028_000001\",\n",
+    "}\n",
+    "\n",
+    "label_options = ['Weight', 'Resonance', 'Breathiness', 'Roughness', 'Loudness', 'Strain', 'Pitch']\n",
+    "\n",
+    "manipulation_idx_widget = widgets.Dropdown(\n",
+    "    options=[(label, i) for i, label in enumerate(label_options)],\n",
+    "    value=2,  # Standardwert: Breathiness\n",
+    "    description='Type:',\n",
+    "    style={'description_width': 'initial'}\n",
+    ")\n",
+    "\n",
+    "manipulation_fkt_widget = widgets.FloatSlider(\n",
+    "    value=1.0, min=-2.0, max=2.0, step=0.1,\n",
+    "    description='Strength:',\n",
+    "    style={'description_width': 'initial'}\n",
+    ")\n",
+    "\n",
+    "run_button = widgets.Button(description=\"Run Manipulation\")\n",
+    "\n",
+    "audio_output = widgets.Output()\n",
+    "\n",
+    "def update_manipulation(b):\n",
+    "    manipulation_idx = manipulation_idx_widget.value\n",
+    "    manipulation_fkt = manipulation_fkt_widget.value\n",
+    "    \n",
+    "    d_vector = extract_speaker_embedding(example)\n",
+    "    labels = load_speaker_labels(example, config_norm_flow)\n",
+    "\n",
+    "    with audio_output:\n",
+    "        clear_output(wait=True)\n",
+    "        display(widgets.Label(\"Processing...\"))\n",
+    "        \n",
+    "    time.sleep(1)  \n",
+    "    \n",
+    "    wav_manipulated = get_manipulation(\n",
+    "        example=example, \n",
+    "        d_vector=d_vector, \n",
+    "        labels=labels[None, :], \n",
+    "        flow=normalizing_flow,\n",
+    "        tts_model=tts_model,\n",
+    "        manipulation_idx=manipulation_idx,\n",
+    "        manipulation_fkt=manipulation_fkt,\n",
+    "    )\n",
+    "    \n",
+    "    with audio_output:\n",
+    "        clear_output(wait=True) \n",
+    "        display(Audio(wav_manipulated, rate=24_000, normalize=True))\n",
+    "        display(Audio(example['audio_path']['observation'], rate=24_000, normalize=True))\n",
+    "\n",
+    "    print(f\"Manipulated {label_options[manipulation_idx]} with strength {manipulation_fkt}\")\n",
+    "\n",
+    "run_button.on_click(update_manipulation)\n",
+    "display(manipulation_idx_widget, manipulation_fkt_widget, run_button, audio_output)"
+   ]
+  }
+ ],
+ "metadata": {
+  "kernelspec": {
+   "display_name": "voice editing",
+   "language": "python",
+   "name": "voice_editing"
+  },
+  "language_info": {
+   "codemirror_mode": {
+    "name": "ipython",
+    "version": 3
+   },
+   "file_extension": ".py",
+   "mimetype": "text/x-python",
+   "name": "python",
+   "nbconvert_exporter": "python",
+   "pygments_lexer": "ipython3",
+   "version": "3.11.9"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 5
+}

pvq_manipulation/helper/characters.yaml

@@ -0,0 +1,4 @@
+Yourtts:
+  "ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz\u00af\u00b7\u00df\u00e0\u00e1\u00e2\u00e3\u00e4\u00e6\u00e7\u00e8\u00e9\u00ea\u00eb\u00ec\u00ed\u00ee\u00ef\u00f1\u00f2\u00f3\u00f4\u00f5\u00f6\u00f9\u00fa\u00fb\u00fc\u00ff\u0101\u0105\u0107\u0113\u0119\u011b\u012b\u0131\u0142\u0144\u014d\u0151\u0153\u015b\u016b\u0171\u017a\u017c\u01ce\u01d0\u01d2\u01d4\u0430\u0431\u0432\u0433\u0434\u0435\u0436\u0437\u0438\u0439\u043a\u043b\u043c\u043d\u043e\u043f\u0440\u0441\u0442\u0443\u0444\u0445\u0446\u0447\u0448\u0449\u044a\u044b\u044c\u044d\u044e\u044f\u0451\u0454\u0456\u0457\u0491\u2013!'(),-.:;? "
+German:
+  "ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz!'(),-.:;?\u00af\u2013\u00fc\u00f6\u00e4\u00df\u201a\u2018\u2019"

pvq_manipulation/helper/moving_batch_norm.py

@@ -0,0 +1,140 @@
+"""
+This Code is adapted from https://github.com/RameenAbdal/StyleFlow/blob/master/module/normalization.py
+"""
+import torch
+import torch.nn as nn
+from torch.nn import Parameter
+
+
+class MovingBatchNormNd(nn.Module):
+    def __init__(self, num_features, eps=1e-4, decay=0.1, bn_lag=0., affine=True):
+        super(MovingBatchNormNd, self).__init__()
+        self.num_features = num_features
+        self.affine = affine
+        self.eps = eps
+        self.decay = decay
+        self.bn_lag = bn_lag
+        self.register_buffer('step', torch.zeros(1))
+        if self.affine:
+            self.weight = Parameter(torch.Tensor(num_features))
+            self.bias = Parameter(torch.Tensor(num_features))
+        else:
+            self.register_parameter('weight', None)
+            self.register_parameter('bias', None)
+        self.register_buffer('running_mean', torch.zeros(num_features))
+        self.register_buffer('running_var', torch.ones(num_features))
+        self.reset_parameters()
+
+    @property
+    def shape(self):
+        raise NotImplementedError
+
+    def reset_parameters(self):
+        self.running_mean.zero_()
+        self.running_var.fill_(1)
+        if self.affine:
+            self.weight.data.zero_()
+            self.bias.data.zero_()
+
+    def forward(self, x, c=None, logpx=None, reverse=False):
+        if reverse:
+            return self._reverse(x, logpx)
+        else:
+            return self._forward(x, logpx)
+
+    def _forward(self, x, logpx=None):
+        num_channels = x.size(-1)
+        used_mean = self.running_mean.clone().detach()
+        used_var = self.running_var.clone().detach()
+
+        if self.training:
+            # compute batch statistics
+            x_t = x.transpose(0, -1).reshape(num_channels, -1)
+            batch_mean = torch.mean(x_t, dim=1)
+
+            batch_var = torch.var(x_t, dim=1)
+
+            # moving average
+            if self.bn_lag > 0:
+                used_mean = batch_mean - (1 - self.bn_lag) * (batch_mean - used_mean.detach())
+                used_mean /= (1. - self.bn_lag**(self.step[0] + 1))
+                used_var = batch_var - (1 - self.bn_lag) * (batch_var - used_var.detach())
+                used_var /= (1. - self.bn_lag**(self.step[0] + 1))
+
+            # update running estimates
+            self.running_mean -= self.decay * (self.running_mean - batch_mean.data)
+            self.running_var -= self.decay * (self.running_var - batch_var.data)
+            self.step += 1
+
+        # perform normalization
+        used_mean = used_mean.view(*self.shape).expand_as(x)
+        used_var = used_var.view(*self.shape).expand_as(x)
+
+        y = (x - used_mean) * torch.exp(-0.5 * torch.log(used_var + self.eps))
+
+        if self.affine:
+            weight = self.weight.view(*self.shape).expand_as(x)
+            bias = self.bias.view(*self.shape).expand_as(x)
+            y = y * torch.exp(weight) + bias
+
+        if logpx is None:
+            return y
+        else:
+            #import ipdb
+            #ipdb.set_trace()
+            return y, logpx - self._logdetgrad(x, used_var).sum(-1, keepdim=True)
+
+    def _reverse(self, y, logpy=None):
+        used_mean = self.running_mean
+        used_var = self.running_var
+
+        if self.affine:
+            weight = self.weight.view(*self.shape).expand_as(y)
+            bias = self.bias.view(*self.shape).expand_as(y)
+            y = (y - bias) * torch.exp(-weight)
+
+        used_mean = used_mean.view(*self.shape).expand_as(y)
+        used_var = used_var.view(*self.shape).expand_as(y)
+        x = y * torch.exp(0.5 * torch.log(used_var + self.eps)) + used_mean
+
+        if logpy is None:
+            return x
+        else:
+            return x, logpy + self._logdetgrad(x, used_var).sum(-1, keepdim=True)
+
+    def _logdetgrad(self, x, used_var):
+        logdetgrad = -0.5 * torch.log(used_var + self.eps)
+        if self.affine:
+            weight = self.weight.view(*self.shape).expand(*x.size())
+            logdetgrad += weight
+        return logdetgrad
+
+    def __repr__(self):
+        return (
+            '{name}({num_features}, eps={eps}, decay={decay}, bn_lag={bn_lag},'
+            ' affine={affine})'.format(name=self.__class__.__name__, **self.__dict__)
+        )
+
+
+def stable_var(x, mean=None, dim=1):
+    if mean is None:
+        mean = x.mean(dim, keepdim=True)
+    mean = mean.view(-1, 1)
+    res = torch.pow(x - mean, 2)
+    max_sqr = torch.max(res, dim, keepdim=True)[0]
+    var = torch.mean(res / max_sqr, 1, keepdim=True) * max_sqr
+    var = var.view(-1)
+    # change nan to zero
+    var[var != var] = 0
+    return var
+
+
+class MovingBatchNorm1d(MovingBatchNormNd):
+    @property
+    def shape(self):
+        return [1, -1]
+
+    def forward(self, x, context=None, logpx=None, integration_times=None, reverse=False):
+        ret = super(MovingBatchNorm1d, self).forward(
+                x, context, logpx=logpx, reverse=reverse)
+        return ret

pvq_manipulation/helper/utils.py

@@ -0,0 +1,228 @@
+import paderbox as pb
+import torch
+
+from coqpit import Coqpit
+from dataclasses import dataclass, field
+from paderbox.transform.module_resample import resample_sox
+
+from typing import List
+
+from TTS.tts.models.vits import VitsAudioConfig, VitsArgs
+from TTS.tts.configs.shared_configs import BaseTTSConfig
+
+
+def load_audio(file_path, target_sr):
+    """Load the audio file normalized in [-1, 1]
+
+    Return Shapes:
+        - x: :math:`[1, T]`
+    """
+    if type(file_path) is dict:
+        if 'observation' in file_path:
+            file_path = file_path['observation']
+
+    x, sr = pb.io.load_audio(file_path, return_sample_rate=True)
+    if sr != target_sr:
+        x = resample_sox(x, in_rate=sr, out_rate=target_sr)
+    x = torch.tensor(x, dtype=torch.float32)[None, :]
+    x[x < -1] = -1
+    x[x > 1] = 1
+    assert (x > 1).sum() + (x < -1).sum() == 0
+    return x, target_sr
+
+
+@dataclass
+class VitsAudioConfig_NT(Coqpit):
+    fft_size: int = 1024
+    sample_rate: int = 16000
+    win_length: int = 1024
+    hop_length: int = 256
+    num_mels: int = 80
+    mel_fmin: int = 0
+    mel_fmax: int = None
+    fading: str = 'half'
+    window: str = 'hann'
+    pad: bool = True
+
+
+@dataclass
+class VitsConfig_NT(BaseTTSConfig):
+    """Defines parameters for VITS End2End TTS model.
+
+    Args:
+        model (str):
+            Model name. Do not change unless you know what you are doing.
+
+        model_args (VitsArgs):
+            Model architecture arguments. Defaults to `VitsArgs()`.
+
+        audio (VitsAudioConfig):
+            Audio processing configuration. Defaults to `VitsAudioConfig()`.
+
+        grad_clip (List):
+            Gradient clipping thresholds for each optimizer. Defaults to `[1000.0, 1000.0]`.
+
+        lr_gen (float):
+            Initial learning rate for the generator. Defaults to 0.0002.
+
+        lr_disc (float):
+            Initial learning rate for the discriminator. Defaults to 0.0002.
+
+        lr_scheduler_gen (str):
+            Name of the learning rate scheduler for the generator. One of the `torch.optim.lr_scheduler.*`. Defaults to
+            `ExponentialLR`.
+
+        lr_scheduler_gen_params (dict):
+            Parameters for the learning rate scheduler of the generator. Defaults to `{'gamma': 0.999875, "last_epoch":-1}`.
+
+        lr_scheduler_disc (str):
+            Name of the learning rate scheduler for the discriminator. One of the `torch.optim.lr_scheduler.*`. Defaults to
+            `ExponentialLR`.
+
+        lr_scheduler_disc_params (dict):
+            Parameters for the learning rate scheduler of the discriminator. Defaults to `{'gamma': 0.999875, "last_epoch":-1}`.
+
+        scheduler_after_epoch (bool):
+            If true, step the schedulers after each epoch else after each step. Defaults to `False`.
+
+        optimizer (str):
+            Name of the optimizer to use with both the generator and the discriminator networks. One of the
+            `torch.optim.*`. Defaults to `AdamW`.
+
+        kl_loss_alpha (float):
+            Loss weight for KL loss. Defaults to 1.0.
+
+        disc_loss_alpha (float):
+            Loss weight for the discriminator loss. Defaults to 1.0.
+
+        gen_loss_alpha (float):
+            Loss weight for the generator loss. Defaults to 1.0.
+
+        feat_loss_alpha (float):
+            Loss weight for the feature matching loss. Defaults to 1.0.
+
+        mel_loss_alpha (float):
+            Loss weight for the mel loss. Defaults to 45.0.
+
+        return_wav (bool):
+            If true, data loader returns the waveform as well as the other outputs. Do not change. Defaults to `True`.
+
+        compute_linear_spec (bool):
+            If true, the linear spectrogram is computed and returned alongside the mel output. Do not change. Defaults to `True`.
+
+        use_weighted_sampler (bool):
+            If true, use weighted sampler with bucketing for balancing samples between datasets used in training. Defaults to `False`.
+
+        weighted_sampler_attrs (dict):
+            Key retuned by the formatter to be used for weighted sampler. For example `{"root_path": 2.0, "speaker_name": 1.0}` sets sample probabilities
+            by overweighting `root_path` by 2.0. Defaults to `{}`.
+
+        weighted_sampler_multipliers (dict):
+            Weight each unique value of a key returned by the formatter for weighted sampling.
+            For example `{"root_path":{"/raid/datasets/libritts-clean-16khz-bwe-coqui_44khz/LibriTTS/train-clean-100/":1.0, "/raid/datasets/libritts-clean-16khz-bwe-coqui_44khz/LibriTTS/train-clean-360/": 0.5}`.
+            It will sample instances from `train-clean-100` 2 times more than `train-clean-360`. Defaults to `{}`.
+
+        r (int):
+            Number of spectrogram frames to be generated at a time. Do not change. Defaults to `1`.
+
+        add_blank (bool):
+            If true, a blank token is added in between every character. Defaults to `True`.
+
+        test_sentences (List[List]):
+            List of sentences with speaker and language information to be used for testing.
+
+        language_ids_file (str):
+            Path to the language ids file.
+
+        use_language_embedding (bool):
+            If true, language embedding is used. Defaults to `False`.
+
+    Note:
+        Check :class:`TTS.tts.configs.shared_configs.BaseTTSConfig` for the inherited parameters.
+
+    Example:
+
+        >>> from TTS.tts.configs.vits_config import VitsConfig
+        >>> config = VitsConfig()
+    """
+    model: str = "vits"
+    # model specific params
+    model_args: VitsArgs = field(default_factory=VitsArgs)
+    audio: VitsAudioConfig = field(default_factory=VitsAudioConfig)
+
+    # optimizer
+    grad_clip: List[float] = field(default_factory=lambda: [1000, 1000, 1000])
+    lr_gen: float = 0.0002
+    lr_disc: float = 0.0002
+    lr_scheduler_gen: str = "ExponentialLR"
+    lr_scheduler_gen_params: dict = field(default_factory=lambda: {"gamma": 0.999875, "last_epoch": -1})
+    lr_scheduler_disc: str = "ExponentialLR"
+    lr_scheduler_disc_params: dict = field(default_factory=lambda: {"gamma": 0.999875, "last_epoch": -1})
+    scheduler_after_epoch: bool = True
+    optimizer: str = "AdamW"
+    optimizer_params: dict = field(
+        default_factory=lambda: {"betas": [0.8, 0.99], "eps": 1e-9, "weight_decay": 0.01})
+
+    # loss params
+    kl_loss_alpha: float = 1.0
+    disc_loss_alpha: float = 1.0
+    gen_loss_alpha: float = 1.0
+    feat_loss_alpha: float = 1.0
+    mel_loss_alpha: float = 45.0
+    dur_loss_alpha: float = 1.0
+    speaker_encoder_loss_alpha: float = 1.0
+
+    # data loader params
+    return_wav: bool = True
+    compute_linear_spec: bool = True
+
+    # sampler params
+    use_weighted_sampler: bool = False  # TODO: move it to the base config
+    weighted_sampler_attrs: dict = field(default_factory=lambda: {})
+    weighted_sampler_multipliers: dict = field(default_factory=lambda: {})
+
+    # overrides
+    r: int = 1  # DO NOT CHANGE
+    add_blank: bool = True
+
+    # testing
+    test_sentences: List[List] = field(
+        default_factory=lambda: [
+            ["It took me quite a long time to develop a voice, and now that I have it I'm not going to be silent."],
+            ["Be a voice, not an echo."],
+            ["I'm sorry Dave. I'm afraid I can't do that."],
+            ["This cake is great. It's so delicious and moist."],
+            ["Prior to November 22, 1963."],
+        ]
+    )
+
+    # multi-speaker settings
+    # use speaker embedding layer
+    num_speakers: int = 0
+    use_speaker_embedding: bool = False
+    speakers_file: str = None
+    speaker_embedding_channels: int = 256
+    language_ids_file: str = None
+    use_language_embedding: bool = False
+
+    # use d-vectors
+    d_vectors_stor_file: bool = False
+    d_vector_model_file: str = None
+    d_vector_dim: int = None
+    d_vector_model: str = None
+    dataset_dict: dict = None
+    gan_speaker_conditioning: bool = True
+
+    sample_rate: int = 16_000
+    use_vad: bool = True
+    use_phone_labels: bool = False
+
+    CONFIG_SOLVER: str = ''
+
+    use_speaker_embedding_cond: bool = True
+
+    def __post_init__(self):
+        for key, val in self.model_args.items():
+            if hasattr(self, key):
+                self[key] = val
+

pvq_manipulation/helper/vad.py

@@ -0,0 +1,193 @@
+import numpy as np
+import paderbox as pb
+import padertorch as pt
+import typing
+
+from dataclasses import dataclass
+
+
+@pb.utils.functional.partial_decorator
+def conv_smoothing(signal, window_length=7, threshold=3):
+    """
+
+    Boundary effects are visible at beginning and end of signal.
+
+    Examples:
+        >>> conv_smoothing(np.array([False, True, True, True, False, False, False, True]), 3, 2)
+        array([False,  True,  True,  True, False, False, False, False])
+
+    Args:
+        signal:
+        window_length:
+        threshold:
+
+    Returns:
+
+    """
+    left_context = right_context = (window_length - 1) // 2
+    if window_length % 2 == 0:
+        right_context += 1
+    act_conv = np.sum(pb.array.segment_axis(
+        np.pad(signal, (left_context, right_context), mode='constant'),
+        length=window_length, shift=1, axis=0, end='cut'
+    ), axis=-1)
+    # act_conv = np.convolve(signal, np.ones(window_length), 's')
+    act = act_conv >= threshold
+    assert act.shape == signal.shape, (act.shape, signal.shape)
+    return act
+
+
+@dataclass
+class VAD(pt.Configurable):
+    smoothing: typing.Optional[typing.Callable] = None
+
+    def reset(self):
+        """Override for a stateful VAD"""
+        pass
+
+    def compute_vad(self, signal, time_resolution=True):
+        raise NotImplementedError()
+
+    def vad_to_time(self, vad, time_length):
+        raise NotImplementedError()
+
+    def __call__(self, signal, time_resolution=True, reset=True):
+        if reset:
+            self.reset()
+
+        vad = self.compute_vad(signal)
+
+        if self.smoothing is not None:
+            vad = pb.array.interval.ArrayInterval(self.smoothing(vad))
+
+        if time_resolution:
+            vad = self.vad_to_time(vad, time_length=signal.shape[-1])
+
+        return vad
+
+
+class EnergyVAD(VAD):
+    def __init__(self, sample_rate, threshold=0.3):
+        self.sample_rate = sample_rate
+        self.threshold = threshold
+
+    @staticmethod
+    def remove_silence(signal, vad_mask):
+        return signal[vad_mask == 1]
+
+    def __call__(self, example):
+        signal = example['audio_data']  # B T
+        vad_mask = self.get_vad_mask(signal)
+        signal = self.remove_silence(signal, vad_mask)
+        example['audio_data'] = signal
+        example['vad_mask'] = vad_mask
+        return example
+
+    def get_vad_mask(self, signal):
+        window_size = int(0.1 * self.sample_rate + 1)
+
+        half_context = (window_size - 1) // 2
+        std = np.std(signal, axis=-1, keepdims=True)
+        signal = signal - np.mean(signal, axis=-1, keepdims=True)
+        signal = np.abs(signal)
+        zeros = np.zeros(
+            [
+                signal.shape[0],
+                half_context,
+            ]
+        )
+        signal = np.concatenate([zeros, signal, zeros], axis=1)
+        sliding_max = np.max(pb.array.segment_axis(
+            signal,
+            length=window_size, shift=1, axis=1, end='cut'
+        ), axis=-1)
+        return sliding_max > self.threshold * std
+
+
+@dataclass
+class ThresholdVAD(VAD):
+    """
+    Energy-based VAD for almost clean files. Tested on WSJ clean data by Lukas
+    Drude.
+
+    Attributes:
+        threshold: Fraction of total signal standard deviation. Use 0.3 for
+            (almost) clean files (SNR >= 20dB, think LibriTTS) and 0.7 for less
+            clean files (think LibriSpeech).
+        window_size: Size of sliding max window.
+        sample_rate: Sampling rate of audio data.
+        smoothing: Optional callable that uses a sliding window over the raw
+            decision to return a smoothed VAD.
+    """
+    threshold: float = 0.3
+    window_size: typing.Optional[int] = None
+    sample_rate: int = 16_000
+    smoothing: typing.Optional[typing.Callable] = None
+
+    @classmethod
+    def finalize_dogmatic_config(cls, config):
+        rate = config['sample_rate']
+        config['smoothing'] = {
+            'partial': conv_smoothing,
+            'window_length': int(0.3 * rate),
+            'threshold': int(0.1 * rate),
+        }
+
+    def __post_init__(self):
+        if self.window_size is None:
+            self.window_size = int(0.1 * self.sample_rate + 1)
+
+        assert self.window_size % 2 == 1, self.window_size
+
+    def __call__(self, example):
+        if isinstance(example, dict):
+            signal = example['audio_data']
+            if signal.ndim == 2 and signal.shape[0] == 1:
+                signal = signal[0]
+            elif signal.ndim == 2 and signal.shape[0] != 1:
+                raise ValueError(
+                    'Only mono signals are supported but audio_data has shape '
+                    f'{signal.shape}'
+                )
+            vad = super().__call__(signal)
+            intervals = np.asarray(vad.intervals)
+            start, stop = zip(*intervals)
+            example['vad'] = vad
+            example['vad_start_samples'] = start
+            example['vad_stop_samples'] = stop
+        else:
+            example = super().__call__(example)
+        return example
+
+    def _detect_voice_activity(self, signal):
+        assert signal.ndim == 1, signal.shape
+
+        half_context = (self.window_size - 1) // 2
+        std = np.std(signal)
+        signal = signal - np.mean(signal)
+        assert np.min(signal) < 0
+        assert np.max(signal) > 0
+        signal = np.abs(signal)
+
+        sliding_max = np.max(pb.array.segment_axis(
+            np.pad(signal, (half_context, half_context), mode='constant'),
+            length=self.window_size, shift=1, axis=0, end='cut'
+        ), axis=-1)
+
+        assert sliding_max.shape == signal.shape, (
+            sliding_max.shape, signal.shape
+        )
+
+        unconstrained = sliding_max > self.threshold * std
+
+        return unconstrained
+
+    def compute_vad(self, signal, time_resolution=True):
+        assert time_resolution
+        return pb.array.interval.ArrayInterval(
+            self._detect_voice_activity(signal)
+        )
+
+    def vad_to_time(self, vad, time_length):
+        assert time_length == vad.shape[-1], (time_length, vad.shape[-1])
+        return vad

pvq_manipulation/models/ffjord.py

@@ -0,0 +1,247 @@
+import torch
+import paderbox as pb
+
+from padertorch.base import Model
+from torchdiffeq import odeint_adjoint as odeint
+from pvq_manipulation.helper.moving_batch_norm import MovingBatchNorm1d
+
+
+class ODEBlock(torch.nn.Module):
+    def __init__(
+        self,
+        ode_function,
+        train_flag=True,
+        reverse=False,
+    ):
+        super().__init__()
+        self.time_deriv_func = ode_function
+        self.noise = None
+        self.reverse = reverse
+        self.train_flag = train_flag
+
+    def forward(
+        self,
+        time: torch.Tensor,
+        states: tuple[torch.Tensor, torch.Tensor, torch.Tensor]
+    ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+        """
+        Helper function to use a neural network for dy(t)/dt = f_theta(t, y(t))
+
+        Hutchinson’s trace estimator, as proposed in the FFJORD Paper, was adapted from:
+        https://github.com/RameenAbdal/StyleFlow/blob/master/module/odefunc.py
+
+        Args:
+            time (torch.Tensor): Scalar tensor representing time
+            states (Tuple[torch.Tensor, torch.Tensor, torch.Tensor]):
+                - z (torch.Tensor): (batch_size, feature_dim) representing the input data.
+                - d_log_dz_dt (torch.Tensor): (batch_size, 1) representing the log derivative.
+                - labels (torch.Tensor): (batch_size, num_labeled_classes)
+
+        Returns:
+            Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+                - dz_dt (torch.Tensor): (batch_size, feature_dim) The derivative of z w.r.t. time
+                - d_log_dz_dt (torch.Tensor): (batch_size, 1) The negative log derivative
+                - labels (torch.Tensor): (batch_size, num_labeled_classes)
+        """
+
+        z, d_log_dz_dt, labels = states
+
+        if self.noise is None:
+            self.noise = self.sample_rademacher_like(z)
+
+        with torch.enable_grad():
+            z.requires_grad_(True)
+
+            dz_dt = self.time_deriv_func.forward(time, z, labels)
+            if self.train_flag:
+                d_log_dz_dt = self.divergence_approx(dz_dt, z, self.noise)
+            else:
+                d_log_dz_dt = torch.zeros_like(z[:, 0]).requires_grad_(True)
+
+        labels = torch.zeros_like(labels).requires_grad_(True)
+        return dz_dt, -d_log_dz_dt.view(z.shape[0], 1), labels
+
+    def divergence_approx(self, f, y, e=None):
+        e_dzdx = torch.autograd.grad(f, y, e, create_graph=True)[0]
+        e_dzdx_e = e_dzdx.mul(e)
+
+        cnt = 0
+        while not e_dzdx_e.requires_grad and cnt < 10:
+            e_dzdx = torch.autograd.grad(f, y, e, create_graph=True)[0]
+            e_dzdx_e = e_dzdx * e
+            cnt += 1
+
+        approx_tr_dzdx = e_dzdx_e.sum(dim=-1)
+        assert approx_tr_dzdx.requires_grad, \
+            "(failed to add node to graph) f=%s %s, y(rgrad)=%s, e_dzdx:%s, e:%s, e_dzdx_e:%s cnt:%s" \
+            % (
+                f.size(), f.requires_grad, y.requires_grad, e_dzdx.requires_grad, e.requires_grad,
+                e_dzdx_e.requires_grad, cnt)
+        return approx_tr_dzdx
+
+    def before_odeint(self, e=None):
+        self.noise = e
+
+    def sample_rademacher_like(self, z):
+        if not self.training:
+            torch.manual_seed(0)
+        return torch.randint(low=0, high=2, size=z.shape).to(z) * 2 - 1
+
+
+class FFJORD(Model):
+    """
+    This class is an implementation of the FFJORD model as proposed in
+    https://arxiv.org/pdf/1810.01367
+    """
+    def __init__(self, ode_function, normalize=True):
+        super().__init__()
+        device = 'cuda' if torch.cuda.is_available() else 'cpu'
+        self.input_dim = ode_function.input_dim
+        self.time_deriv_func = ODEBlock(ode_function=ode_function)
+        self.latent_dist = torch.distributions.MultivariateNormal(
+            torch.zeros(self.input_dim, device=device),
+            torch.eye(self.input_dim, device=device),
+        )
+        self.normalize = normalize
+        if self.normalize:
+            self.input_norm = MovingBatchNorm1d(self.input_dim, bn_lag=0)
+            self.output_norm = MovingBatchNorm1d(self.input_dim, bn_lag=0)
+
+    @staticmethod
+    def load_model(model_path, checkpoint):
+        model_dict = pb.io.load_yaml(model_path / "config.yaml")
+        model = Model.from_config(model_dict['model'])
+        cp = torch.load(
+            model_path / checkpoint,
+            map_location=torch.device('cuda' if torch.cuda.is_available() else 'cpu')
+        )
+        model_weights = cp.copy()
+        model.load_state_dict(model_weights['model'])
+        model.eval()
+        return model
+
+    def forward(
+        self,
+        state: tuple[torch.Tensor, torch.Tensor],
+        integration_times: torch.Tensor = torch.tensor([0.0, 1.0]
+        )
+    ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+        """
+        Integration from t_1 (data distribution) to t_0 (base distribution).
+        (training step)
+
+        Args:
+            state (Tuple[torch.Tensor, torch.Tensor]):
+                - z (torch.Tensor): (batch_size, feature_dim) representing the input data.
+                - labels (torch.Tensor): (batch_size, num_labeled_classes)
+
+            integration_times (torch.Tensor, optional): A tensor of shape (2,)
+            specifying the start and end times for integration. Defaults to torch.tensor([0.0, 1.0]).
+
+        Returns:
+            Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+                - dz_dt (torch.Tensor): A tensor of shape (batch_size, feature_dim) representing the derivative of z w.r.t. time.
+                - -d_log_dz_dt (torch.Tensor): (batch_size, 1) representing the negative log derivative.
+                - labels (torch.Tensor): (batch_size, num_labeled_classes)
+        """
+        z_1, labels = state
+
+        if z_1.dim() == 3:
+            z_1 = z_1.squeeze(1)
+
+        delta_logpz = torch.zeros(z_1.shape[0], 1).to(z_1.device)
+
+        if self.normalize:
+            z_1, delta_logpz = self.input_norm(z_1, context=labels, logpx=delta_logpz)
+
+        self.time_deriv_func.before_odeint()
+        state = odeint(
+            self.time_deriv_func,  # Calculates time derivatives.
+            (z_1, delta_logpz, labels),  # Values to update. init states
+            integration_times.to(z_1.device),  # When to evaluate.
+            method='dopri5',  # Runge-Kutta
+            atol=1e-5,  # Error tolerance
+            rtol=1e-5,  # Error tolerance
+        )
+        if self.normalize:
+            dz_dt, d_delta_log_dz_t = self.output_norm(state[0], context=state[2], logpx=state[1])
+        else:
+            dz_dt, d_delta_log_dz_t = state[0], state[1]
+
+        state = (dz_dt, d_delta_log_dz_t, labels)
+
+        if len(integration_times) == 2:
+            state = tuple(s[1] if s.shape[0] > 1 else s[0] for s in state)
+        return state
+
+    def sample(
+        self,
+        state: tuple[torch.Tensor, torch.Tensor],
+        integration_times: torch.Tensor = torch.tensor([1.0, 0.0])
+    ) -> tuple[torch.Tensor, torch.Tensor]:
+        """
+        This is the sampling step. Integration from t_0 (base distribution) to t_1 (data distribution).
+
+        Args:
+            state (Tuple[torch.Tensor, torch.Tensor]):
+                - z_0 (torch.Tensor): (batch_size, feature_dim) representing the initial state from the base distribution
+                - labels (torch.Tensor): (batch_size, num_labeled_classes)
+
+            integration_times (torch.Tensor, optional): A tensor of shape (2,) specifying the start (t_0) and end (t_1) times for integration.
+                Defaults to torch.tensor([1.0, 0.0])
+
+        Returns:
+            Tuple[torch.Tensor, torch.Tensor]:
+                - z_t1 (torch.Tensor): (batch_size, feature_dim) representing the sampled data at time t_1 (data distribution).
+                - labels (torch.Tensor): (batch_size, num_labeled_classes)
+        """
+        z_0, labels = state
+        delta_logpz = torch.zeros(z_0.shape[0], 1).to(z_0.device)
+        if self.normalize:
+            z_0, delta_logpz = self.output_norm(
+                z_0,
+                context=labels,
+                logpx=delta_logpz,
+                reverse=True
+            )
+
+        state = odeint(
+            self.time_deriv_func,  # Calculates time derivatives.
+            (z_0, delta_logpz, labels),  # Values to update. init states
+            integration_times.to(z_0.device),  # When to evaluate.
+            method='dopri5',  # Runge-Kutta
+            atol=1e-5,  # Error tolerance
+            rtol=1e-5,  # Error tolerance
+        )
+        if self.normalize:
+            dz_dt, d_delta_log_dz_t = self.input_norm(
+                state[0],
+                context=state[2],
+                logpx=state[1],
+                reverse=True
+            )
+        else:
+            dz_dt, d_delta_log_dz_t = state[0], state[1]
+        state = (dz_dt, d_delta_log_dz_t, labels)
+
+        if len(integration_times) == 2:
+            state = tuple(s[1] if s.shape[0] > 1 else s[0] for s in state)
+        return state
+
+    def example_to_device(self, examples, device):
+        observations = [example['observation'] for example in examples]
+        labels = [example['speaker_conditioning'].tolist() for example in examples if 'speaker_conditioning' in example]
+        observations_tensor = torch.tensor(observations, device=device, dtype=torch.float)
+        labels_tensor = torch.tensor(labels, device=device, dtype=torch.float) if labels else None
+        return observations_tensor, labels_tensor
+
+    def review(self, example, outputs):
+        z_t0, delta_logpz, _ = outputs
+        logpz_t1 = self.latent_dist.log_prob(z_t0) - delta_logpz
+        loss = -torch.mean(logpz_t1)
+        return dict(loss=loss, scalars=dict(loss=loss))
+
+    def modify_summary(self, summary):
+        summary = super().modify_summary(summary)
+        return summary
+

pvq_manipulation/models/hubert.py

@@ -0,0 +1,207 @@
+import os
+from pathlib import Path
+from contextlib import nullcontext
+import typing as tp
+from typing import List, Tuple, Optional
+
+import einops
+import numpy as np
+import torch
+import torch.nn.functional as F
+import torchaudio
+
+import padertorch as pt
+from padertorch.contrib.je.modules.conv_utils import (
+    compute_conv_output_sequence_lengths
+)
+from padertorch.utils import to_numpy
+from transformers.models.hubert.modeling_hubert import HubertModel
+
+# See https://ieeexplore.ieee.org/abstract/document/9814838, Fig. 2
+PR_BASE_LAYER = 11
+PR_LARGE_LAYER = 22
+SID_BASE_LAYER = 4
+SID_LARGE_LAYER = 6
+
+
+def tuple_to_int(sequence) -> list:
+    return list(map(lambda t: t[0], sequence))
+
+
+class HubertExtractor(pt.Module):
+    """Extract HuBERT features from raw waveform.
+
+    Args:
+        model_name (str): Name of the pretrained HuBERT model on huggingface.co.
+            Defaults to "facebook/hubert-large-ll60k".
+        layer (int): Index of the layer to extract features from. Defaults to
+            22.
+        freeze (bool): If True, freeze the weights of the encoder
+            (i.e., no finetuning of Transformer layers). Defaults to True.
+    """
+
+    def __init__(
+            self,
+            model_name: str = "facebook/hubert-large-ll60k",
+            layer: tp.Union[int, str] = PR_LARGE_LAYER,
+            freeze: bool = True,
+            detach: bool = False,
+            device: str = "cpu",
+            backend: str = "torchaudio",
+            storage_dir: str = None,
+    ):
+        super().__init__()
+
+        if not freeze and detach:
+            raise ValueError(
+                'detach=True only supported if freeze=True\n'
+                f'Got: freeze={freeze}, detach={detach}'
+            )
+        if backend == "torchaudio":
+            bundle = getattr(torchaudio.pipelines, model_name)
+            self.model = bundle.get_model(dl_kwargs={'model_dir': storage_dir}).eval().to(device)
+            self.sampling_rate = bundle.sample_rate
+        else:
+            raise ValueError(f'Unknown backend: {backend}')
+        self.backend = backend
+
+        if freeze:
+            for param in self.model.parameters():
+                param.requires_grad = False
+        else:
+            # Always freeze feature extractor and feature projection layers
+            for param in self.model.feature_extractor.parameters():
+                param.requires_grad = False
+            for param in self.model.feature_projection.parameters():
+                param.requires_grad = False
+
+        self.layer = layer
+        self.freeze = freeze
+        self.detach = detach
+
+        if self.layer == 'all':
+            self.weights = torch.nn.Parameter(
+                torch.ones(24), requires_grad=True
+            )
+
+    @property
+    def cache_dir(self):
+        return Path(os.environ['STORAGE_ROOT']) / 'huggingface' / 'hub'
+
+    @property
+    def context(self):
+        if self.detach:
+            return torch.no_grad()
+        else:
+            return nullcontext()
+
+    def compute_output_lengths(
+            self, input_lengths: Optional[List[int]]
+    ) -> Optional[List[int]]:
+        """Compute the number of time frames for each batch entry.
+
+        Args:
+            input_lengths: List with number of samples per batch entry.
+
+        Returns:
+            List with number of time frames per batch entry.
+        """
+        if input_lengths is None:
+            return input_lengths
+        output_lengths = np.asarray(input_lengths) + self.window_size - 1
+        for kernel_size, dilation, stride in zip(
+                self.kernel_sizes, self.dilations, self.strides,
+        ):
+            output_lengths = compute_conv_output_sequence_lengths(
+                output_lengths,
+                kernel_size=kernel_size,
+                dilation=dilation,
+                pad_type=None,
+                stride=stride,
+            )
+        return output_lengths.tolist()
+
+    def forward(
+        self,
+        time_signal: torch.Tensor,
+        sampling_rate: int,
+        sequence_lengths: Optional[List[int]] = None,
+        extract_features: bool = False,
+        other_inputs: Optional[dict] = None,
+    ) -> Tuple[torch.Tensor, Optional[List[int]]]:
+        """Extract HuBERT features from raw waveform.
+
+        Args:
+            time_signal: Time signal of shape (batch, 1, time) or (batch, time)
+                sampled at 16 kHz.
+            sequence_lengths: List with number of samples per batch entry.
+
+        Returns:
+            x (torch.Tensor): HuBERT features of shape
+                (batch, D, time frames).
+            sequence_lengths (List[int]): List with number of time frames per
+                batch entry.
+        """
+        del other_inputs
+
+        if time_signal.ndim == 3:
+            time_signal = einops.rearrange(time_signal, 'b c t -> (b c) t')
+
+        time_signal = torchaudio.functional.resample(
+            time_signal, sampling_rate, self.sampling_rate
+        )
+        if sequence_lengths is not None:
+            if isinstance(sequence_lengths, (list, tuple)):
+                sequence_lengths = torch.tensor(sequence_lengths).long() \
+                    .to(time_signal.device)
+            sequence_lengths = (
+                    sequence_lengths / sampling_rate * self.sampling_rate
+            ).long()
+
+        if self.freeze or self.detach:
+            self.model.eval()
+        with self.context:
+            if self.backend == "torchaudio":
+                self.model: torchaudio.models.Wav2Vec2Model
+                x, sequence_lengths = self.model.extract_features(
+                    time_signal, lengths=sequence_lengths,
+                    num_layers=self.layer,
+                )
+                if isinstance(self.layer, int):
+                    x = x[-1].transpose(1, 2)
+                else:
+                    raise NotImplementedError(self.layer)
+                return x, sequence_lengths
+
+            self.model: HubertModel
+            n_pad = self.window_size - 1
+            time_signal = F.pad(time_signal, (0, n_pad), value=0)
+            if extract_features:
+                features = self.model.feature_extractor(time_signal.float()) \
+                    .transpose(1, 2)
+                x = self.model.feature_projection(features).transpose(1, 2)
+            else:
+                outputs = self.model(
+                    time_signal.float(), output_hidden_states=True
+                )
+                if isinstance(self.layer, int):
+                    x = outputs.hidden_states[self.layer].transpose(1, 2)
+                    if self.detach:
+                        x = x.detach()
+                elif self.layer == 'all':
+                    hidden_states = []
+                    for _, hidden_state in enumerate(outputs.hidden_states):
+                        x = hidden_state.transpose(1, 2)
+                        if self.detach:
+                            x = x.detach()
+                        hidden_states.append(x)
+                    hidden_states = torch.stack(hidden_states)  # (L, B, D, T)
+                    x = (hidden_states * self.weights[:, None, None, None]) \
+                        .sum(dim=0)
+                else:
+                    raise ValueError(f'Unknown layer: {self.layer}')
+
+        sequence_lengths = to_numpy(sequence_lengths)
+        sequence_lengths = self.compute_output_lengths(sequence_lengths)
+
+        return x.unsqueeze(1), sequence_lengths

pvq_manipulation/models/ode_functions.py

@@ -0,0 +1,96 @@
+"""
+Implementation of Δz = f(t, z, labels)
+f() is a neural network with the architecture defined in StyleFlow
+StyleFlow: Attribute-conditioned Exploration of StyleGAN-Generated Images using Conditional Continuous Normalizing Flows
+"""
+import torch
+
+
+class CNFNN(torch.nn.Module):
+    def __init__(
+            self,
+            input_dim,
+            condition_dim,
+            hidden_channels,
+    ):
+        super().__init__()
+        self.layers = torch.nn.ModuleList()
+        hidden_dims = hidden_channels + [input_dim]
+        self.input_dim = input_dim
+
+        for idx, hidden_dim in enumerate(hidden_dims):
+            self.layers.append(CNFBlock(
+                input_dim=input_dim,
+                condition_dim=condition_dim,
+                output_dim=hidden_dim,
+                output_layer=False if idx < len(hidden_dims) - 1 else True,
+            ))
+            input_dim = hidden_dim
+
+    def forward(self, t, z, labels):
+        """
+        This function computes: Δz = f(t, z, labels)
+
+        Args:
+            t (torch.Tensor): () Time step of the ODE
+            z (torch.Tensor): (Batch_size, Input_dim) Intermediate value
+            labels (torch.Tensor): (Batch_size, condition_dim) Speaker attributes 
+
+        Returns:
+            Δz (torch.Tensor): : (Batch_size, Input_dim) Computed delta
+        """
+        for layer in self.layers:
+            z = layer(t, z, labels)
+        return z
+
+
+class CNFBlock(torch.nn.Module):
+    def __init__(
+            self,
+            input_dim,
+            output_dim,
+            condition_dim,
+            output_layer,
+    ):
+        super().__init__()
+        self._layer = torch.nn.Linear(input_dim, output_dim)
+        self._hyper_bias = torch.nn.Linear(
+            1 + condition_dim,
+            output_dim,
+            bias=False
+        )
+        self._hyper_gate = torch.nn.Linear(
+            1 + condition_dim,
+            output_dim
+        )
+        self.output_layer = output_layer
+
+    def forward(self, t, z, labels):
+        """
+        Args:
+            t (torch.Tensor): () Time step of the ODE
+            z (torch.Tensor): (Batch_size, Input_dim) Intermediate value
+            labels (torch.Tensor): (Batch_size, condition_dim) Speaker attributes 
+
+        Returns:
+            z (torch.Tensor): : (Batch_size, Output_dim) Intermediate value
+        """
+        if labels.dim() == 1:
+            labels = labels[:, None]
+        elif labels.dim() == 3:
+            labels = labels.squeeze(1)
+
+        tz_cat = torch.cat((t.expand(z.shape[0], 1), labels), dim=1)
+
+        gate = torch.sigmoid(self._hyper_gate(tz_cat))
+        bias = self._hyper_bias(tz_cat)
+
+        if z.dim() == 3:
+            gate = gate.unsqueeze(1)
+            bias = bias.unsqueeze(1)
+
+        z = self._layer(z) * gate + bias
+
+        if not self.output_layer:
+            z = torch.tanh(z)
+        return z

pvq_manipulation/models/vits.py

@@ -0,0 +1,742 @@
+"""
+This is a wrapper for the TTS VITS model.
+TTS.tts.models.vits
+https://github.com/coqui-ai/TTS/blob/dev/TTS/tts/models/vits.py
+"""
+import os
+import numpy as np
+import paderbox as pb
+import padertorch as pt
+import torch
+
+from coqpit import Coqpit
+from padertorch.ops._stft import STFT
+from pathlib import Path
+from pvq_manipulation.helper.utils import VitsAudioConfig_NT, VitsConfig_NT, load_audio
+
+from torch.utils.data import DataLoader
+from torch.cuda.amp.autocast_mode import autocast
+from TTS.tts.configs.shared_configs import CharactersConfig
+from TTS.tts.layers.vits.networks import PosteriorEncoder, ResidualCouplingBlocks, TextEncoder
+from TTS.tts.models.vits import Vits, VitsArgs, VitsDataset, spec_to_mel, wav_to_spec
+from TTS.tts.utils.languages import LanguageManager
+from TTS.tts.utils.speakers import SpeakerManager
+from TTS.tts.utils.synthesis import embedding_to_torch, numpy_to_torch
+from TTS.tts.utils.text.tokenizer import TTSTokenizer
+from TTS.tts.utils.helpers import generate_path, rand_segments, segment, sequence_mask
+from TTS.utils.audio import AudioProcessor
+from TTS.vocoder.models.hifigan_generator import HifiganGenerator
+from trainer.trainer import to_cuda
+from typing import Dict, List, Union
+
+
+STORAGE_ROOT = Path(os.getenv('STORAGE_ROOT')).expanduser()
+
+
+class Vits_NT(Vits):
+    def __init__(
+        self,
+        config: Coqpit,
+        ap: "AudioProcessor" = None,
+        tokenizer: "TTSTokenizer" = None,
+        speaker_manager: SpeakerManager = None,
+        language_manager: LanguageManager = None,
+        sample_rate: int = None,
+    ):
+        super().__init__(
+            config,
+            ap,
+            tokenizer,
+            speaker_manager,
+            language_manager
+        )
+        self.sample_rate = sample_rate
+        self.embedded_speaker_dim = self.args.d_vector_dim
+        self.posterior_encoder = PosteriorEncoder(
+            self.args.out_channels,
+            self.args.hidden_channels,
+            self.args.hidden_channels,
+            kernel_size=self.args.kernel_size_posterior_encoder,
+            dilation_rate=self.args.dilation_rate_posterior_encoder,
+            num_layers=self.args.num_layers_posterior_encoder,
+            cond_channels=self.embedded_speaker_dim,
+        )
+
+        self.flow = ResidualCouplingBlocks(
+            self.args.hidden_channels,
+            self.args.hidden_channels,
+            kernel_size=self.args.kernel_size_flow,
+            dilation_rate=self.args.dilation_rate_flow,
+            num_layers=self.args.num_layers_flow,
+            cond_channels=self.embedded_speaker_dim,
+        )
+
+        self.text_encoder = TextEncoder(
+            self.args.num_chars,
+            self.args.hidden_channels,
+            self.args.hidden_channels,
+            self.args.hidden_channels_ffn_text_encoder,
+            self.args.num_heads_text_encoder,
+            self.args.num_layers_text_encoder,
+            self.args.kernel_size_text_encoder,
+            self.args.dropout_p_text_encoder,
+            language_emb_dim=self.embedded_language_dim,
+        )
+        self.waveform_decoder = HifiganGenerator(
+            self.args.hidden_channels,
+            1,
+            self.args.resblock_type_decoder,
+            self.args.resblock_dilation_sizes_decoder,
+            self.args.resblock_kernel_sizes_decoder,
+            self.args.upsample_kernel_sizes_decoder,
+            self.args.upsample_initial_channel_decoder,
+            self.args.upsample_rates_decoder,
+            inference_padding=0,
+            cond_channels=self.embedded_speaker_dim if self.config.gan_speaker_conditioning else 0,
+            conv_pre_weight_norm=False,
+            conv_post_weight_norm=False,
+            conv_post_bias=False,
+        )
+        self.speaker_manager = self.speaker_manager
+        self.speaker_encoder = self.speaker_manager
+
+        self.speaker_manager.eval()
+
+        self.epoch = 0
+        self.num_epochs = config['epochs']
+        self.lr_lambda = 0
+        self.config_solver = config['CONFIG_SOLVER']
+        self.config = config
+
+        self.stft = STFT(
+            size=self.config.audio.win_length,
+            shift=self.config.audio.hop_length,
+            window_length=self.config.audio.win_length,
+            fading=self.config.audio.fading,
+            window=self.config.audio.window,
+            pad=self.config.audio.pad
+        )
+
+    def get_spectogram_nt(self, wav):
+        """
+        Extracts spectrogram from audio
+        Args:
+            wav (torch.Tensor): (Batch_size, Num_samples)
+        Returns:
+            spectrogram (torch.Tensor): (Batch_size, Frequency_bins, Time) spectrogram
+        """
+        wav = wav.squeeze(1)
+        stft_signal = self.stft(wav)
+        stft_signal = torch.einsum('btf-> bft', stft_signal)
+        spectrogram = stft_signal.real ** 2 + stft_signal.imag ** 2
+        spectrogram = torch.sqrt(spectrogram + 1e-6)
+        return spectrogram
+
+    def get_aux_input_from_test_sentences(self, sentence_info):
+        """
+        Get aux input for the inference step from test sentences
+        Args:
+            sentence_info (dict): Expected keys:
+                - "d_vector_storage_root" (str)
+                - "d_vector" (torch.Tensor)
+                - "d_vector_man" (torch.Tensor) (optional)
+        Returns:
+            aux_input (dict): aux input for the inference step
+        """
+        if 'd_vector' not in sentence_info.keys():
+            d_vector_file = sentence_info['d_vector_storage_root']
+            d_vector = torch.load(d_vector_file)
+            return {"d_vector": d_vector, **sentence_info}
+        else:
+            return sentence_info
+
+    @staticmethod
+    def init_from_config(
+            config: "VitsConfig",
+            samples= None,
+            verbose=True
+    ):
+        """
+        Initiate model from config
+        Args:
+            config (VitsConfig): Model config.
+            samples (Union[List[List], List[Dict]]): Training samples to parse speaker ids for training.
+                Defaults to None.
+        Returns:
+            model (Vits): Initialized model.
+        """
+        upsample_rate = torch.prod(torch.as_tensor(config.model_args.upsample_rates_decoder)).item()
+        assert (upsample_rate == config.audio.hop_length), f" [!] Product of upsample rates must be equal to the hop length - {upsample_rate} vs {config.audio.hop_length}"
+        ap = AudioProcessor.init_from_config(config, verbose=verbose)
+        tokenizer, new_config = TTSTokenizer.init_from_config(config)
+        language_manager = LanguageManager.init_from_config(config)
+        speaker_manager = pt.Module.from_storage_dir(
+            config['d_vector_model_file'],
+            checkpoint_name='ckpt_latest.pth',
+            consider_mpi=False,
+            config_name='config.json',
+        )
+        speaker_manager.num_speakers = config['num_speakers']
+        for param in speaker_manager.parameters():
+            param.requires_grad = False
+
+        return Vits_NT(
+            new_config,
+            ap,
+            tokenizer,
+            speaker_manager=speaker_manager,
+            language_manager=language_manager,
+            sample_rate=config['sample_rate'],
+        )
+
+    @torch.no_grad()
+    def inference(self, x, aux_input=None):
+        """
+        Note:
+            To run in batch mode, provide `x_lengths` else model assumes that the batch size is 1.
+
+        Args:
+            x (torch.Tensor): (batch_size, T_seq) or (T_seq) Input character sequence IDs
+            aux_input (dict): Expected keys:
+                - d_vector (torch.Tensor): (batch_size, Feature_dim) speaker_embedding
+                - x_lengths: (torch.Tensor): (batch_size) length of each text token
+
+        Returns:
+            - model_outputs (torch.Tensor): (batch_size, T_wav) Synthesized waveform
+        """
+        speaker_embedding = aux_input['d_vector'].detach()[:, :, None]
+        if aux_input['d_vector_man'] is not None:
+            speaker_embedding_man = aux_input['d_vector_man'].detach()[:, :, None]
+        else:
+            speaker_embedding_man = speaker_embedding
+        aux_input['tokens'] = x.clone()
+        x_lengths = self._set_x_lengths(x, aux_input)
+        x, m_p, logs_p, x_mask = self.text_encoder(
+            x,
+            x_lengths,
+            lang_emb=None
+        )
+        logw = self.duration_predictor(
+            x,
+            x_mask,
+            g=speaker_embedding,
+            lang_emb=None,
+        )
+
+        w = torch.exp(logw) * x_mask * self.length_scale
+        w_ceil = torch.ceil(w)
+        y_lengths = torch.clamp_min(torch.sum(w_ceil, [1, 2]), 1).long()
+        y_mask = sequence_mask(y_lengths, None).to(x_mask.dtype).unsqueeze(1)  # [B, 1, T_dec]
+
+        attn_mask = x_mask * y_mask.transpose(1, 2)
+        attn = generate_path(w_ceil.squeeze(1), attn_mask.squeeze(1).transpose(1, 2))
+        m_p = torch.matmul(attn.transpose(1, 2), m_p.transpose(1, 2)).transpose(1, 2)
+        logs_p = torch.matmul(attn.transpose(1, 2), logs_p.transpose(1, 2)).transpose(1, 2)
+
+        z_p = m_p + torch.randn_like(m_p) * torch.exp(logs_p) * self.inference_noise_scale
+
+        z = self.flow(z_p, y_mask, g=speaker_embedding_man, reverse=True)
+        z, _, _, y_mask = self.upsampling_z(
+            z,
+            y_lengths=y_lengths,
+            y_mask=y_mask
+        )
+        o = self.waveform_decoder(
+            (z * y_mask)[:, :, : self.max_inference_len],
+            g=speaker_embedding_man if self.config.gan_speaker_conditioning else None
+        )
+        return o
+
+    def forward(self, x, x_lengths, y, y_lengths, aux_input, inference=False):
+        """
+        Forward pass of the model.
+
+        Args:
+            x (torch.tensor): (Batch, T_seq) Input character sequence IDs
+            x_lengths (torch.tensor): (Batch) Input character sequence lengths.
+            y (torch.tensor): (Batch_size, Frequency_bins, Time) Input spectrograms.
+            y_lengths (torch.tensor): (Batch) Input spectrogram lengths.
+            aux_input (dict, optional): Expected keys:
+                - d_vector (torch.Tensor): (batch_size, Feature_dim) speaker_embedding
+                - waveform: (torch.Tensor): (Batch_size, Num_samples) Target waveform
+        Returns:
+            Dict: model outputs keyed by the output name.
+        """
+        outputs = {}
+        speaker_embedding = aux_input['d_vector'].detach()[:, :, None]
+        x, m_p, logs_p, x_mask = self.text_encoder(
+            x,
+            x_lengths,
+            lang_emb=None
+        )
+        z, m_q, logs_q, y_mask = self.posterior_encoder(
+            y,
+            y_lengths,
+            g=speaker_embedding,
+        )
+        z_p = self.flow(z, y_mask, g=speaker_embedding)
+        outputs, attn = self.forward_mas(
+            outputs,
+            z_p,
+            m_p,
+            logs_p,
+            x,
+            x_mask,
+            y_mask,
+            g=speaker_embedding,
+            lang_emb=None,
+        )
+        m_p = torch.einsum("klmn, kjm -> kjn", [attn, m_p])
+        logs_p = torch.einsum("klmn, kjm -> kjn", [attn, logs_p])
+
+        z_slice, slice_ids = rand_segments(
+            z,
+            y_lengths,
+            self.spec_segment_size,
+            let_short_samples=True,
+            pad_short=True
+        )
+        z_slice, spec_segment_size, slice_ids, _ = self.upsampling_z(
+            z_slice,
+            slice_ids=slice_ids,
+        )
+
+        wav_seg = segment(
+            aux_input['waveform'],
+            slice_ids * self.config.audio.hop_length,
+            spec_segment_size * self.config.audio.hop_length,
+            pad_short=True,
+        )
+        o = self.waveform_decoder(
+            z_slice,
+            g=speaker_embedding if self.config.gan_speaker_conditioning else None
+        )
+
+        if self.args.use_speaker_encoder_as_loss and self.speaker_manager.encoder is not None:
+            wavs_batch = torch.cat((wav_seg, o), dim=0)
+            if self.audio_transform is not None:
+                wavs_batch = self.audio_transform(wavs_batch)
+            with torch.no_grad():
+                pred_embs = self.speaker_manager.encoder.forward(wavs_batch, l2_norm=True)
+            gt_spk_emb, syn_spk_emb = torch.chunk(pred_embs, 2, dim=0)
+        else:
+            gt_spk_emb, syn_spk_emb = None, None
+
+        outputs.update(
+            {
+                "model_outputs": o,
+                "alignments": attn.squeeze(1),
+                "m_p": m_p,
+                "logs_p": logs_p,
+                "z": z,
+                "z_p": z_p,
+                "m_q": m_q,
+                "logs_q": logs_q,
+                "waveform_seg": wav_seg,
+                "gt_spk_emb": gt_spk_emb,
+                "syn_spk_emb": syn_spk_emb,
+                "slice_ids": slice_ids,
+                "z_slice": z_slice,
+                "speaker_embedding": speaker_embedding,
+            }
+        )
+        return outputs
+
+    @staticmethod
+    def load_model(model_path, checkpoint):
+        """
+        Load model from checkpoint
+
+        Args:
+            model_path (str): model path
+            checkpoint (str): checkpoint name
+
+        Returns:
+            model (pvq_manipulation.models.vits.Vits_NT): model
+        """
+        config = pb.io.load_json(model_path / "config.json")
+        model_args = VitsArgs(**config['model_args'])
+        audio_config = VitsAudioConfig_NT(**config['audio'])
+        characters_config = CharactersConfig(**config['characters'])
+        del config['audio']
+        del config['characters']
+        del config['model_args']
+
+        config = VitsConfig_NT(
+            model_args=model_args,
+            audio=audio_config,
+            characters=characters_config,
+            **config,
+        )
+        model = Vits_NT.init_from_config(config)
+        cp = torch.load(
+            model_path / checkpoint,
+            map_location=torch.device('cpu')
+        )
+        model_weights = cp['model'].copy()
+        model.load_state_dict(model_weights, strict=False)
+        model.eval()
+        return model
+
+    def synthesize_from_example(self, s_info):
+        """
+        Synthesize voice from example
+
+        Args:
+            s_info (dict): Expected keys:
+                - "speaker_id" (str),
+                - "example_id" (str),
+                - "audio_path" (str),
+                - "d_vector_storage_root" (str),
+                - "text" (str) specifying the text to synthesize
+        Returns:
+            wav (torch.Tensor): synthesized waveform
+        """
+        aux_inputs = self.get_aux_input_from_test_sentences(s_info)
+        use_cuda = "cuda" in str(next(self.parameters()).device)
+
+        device = next(self.parameters()).device
+        if use_cuda:
+            device = "cuda"
+
+        text_inputs = np.asarray(
+            self.tokenizer.text_to_ids(aux_inputs["text"], language=None),
+            dtype=np.int32,
+        )
+        d_vector = embedding_to_torch(aux_inputs["d_vector"], device=device)
+
+        if "d_vector_man" in aux_inputs.keys():
+            d_vector_man = embedding_to_torch(aux_inputs["d_vector_man"], device=device)
+
+        text_inputs = numpy_to_torch(text_inputs, torch.long, device=device)
+        text_inputs = text_inputs.unsqueeze(0)
+
+        wav = self.inference(
+            text_inputs,
+            aux_input={
+                "x_lengths": torch.tensor(
+                    text_inputs.shape[1:2]
+                ).to(text_inputs.device),
+                "d_vector": d_vector,
+                "d_vector_man": d_vector_man if "d_vector_man" in aux_inputs.keys() else None
+            }
+        )[0].data.cpu().numpy().squeeze()
+        return wav
+
+    def format_batch_on_device(self, batch):
+        """Format batch on device"""
+        ac = self.config.audio
+
+        batch['waveform'] = to_cuda(batch['waveform'])
+        wav = batch["waveform"]
+
+        batch['spec'] = self.get_spectogram_nt(wav)
+
+        if self.args.encoder_sample_rate:
+            spec_mel = wav_to_spec(batch["waveform"], ac.fft_size, ac.hop_length, ac.win_length, center=False)
+            if spec_mel.size(2) > int(batch["spec"].size(2) * self.interpolate_factor):
+                spec_mel = spec_mel[:, :, : int(batch["spec"].size(2) * self.interpolate_factor)]
+            else:
+                batch["spec"] = batch["spec"][:, :, : int(spec_mel.size(2) / self.interpolate_factor)]
+        else:
+            spec_mel = batch["spec"]
+
+        batch["mel"] = spec_to_mel(
+            spec=spec_mel,
+            n_fft=ac.fft_size,
+            num_mels=ac.num_mels,
+            sample_rate=ac.sample_rate,
+            fmin=ac.mel_fmin,
+            fmax=ac.mel_fmax,
+        )
+
+        if self.args.encoder_sample_rate:
+            assert batch["spec"].shape[2] == int(
+                batch["mel"].shape[2] / self.interpolate_factor
+            ), f"{batch['spec'].shape[2]}, {batch['mel'].shape[2]}"
+        else:
+            assert batch["spec"].shape[2] == batch["mel"].shape[2], f"{batch['spec'].shape[2]}, {batch['mel'].shape[2]}"
+
+        batch["spec_lens"] = (batch["spec"].shape[2] * batch["waveform_rel_lens"]).int()
+        batch["mel_lens"] = (batch["mel"].shape[2] * batch["waveform_rel_lens"]).int()
+
+        if self.args.encoder_sample_rate:
+            assert (batch["spec_lens"] - (batch["mel_lens"] / self.interpolate_factor).int()).sum() == 0
+        else:
+            assert (batch["spec_lens"] - batch["mel_lens"]).sum() == 0
+
+        batch["spec"] = batch["spec"] * sequence_mask(batch["spec_lens"]).unsqueeze(1)
+        batch["mel"] = batch["mel"] * sequence_mask(batch["mel_lens"]).unsqueeze(1)
+        return batch
+
+    def train_step(
+        self,
+        batch: dict,
+        criterion: torch.nn.Module,
+        optimizer_idx: int,
+    ):
+        """
+        Perform a single training step. Run the model forward pass and compute losses.
+
+        Args:
+            batch (Dict): Input tensors.
+            criterion (nn.Module): Loss layer designed for the model.
+            optimizer_idx (int): Index of optimizer to use. 0 for the generator and 1 for the discriminator networks.
+
+        Returns:
+            Tuple[Dict, Dict]: Model ouputs and computed losses.
+        """
+        if optimizer_idx == 0:
+            # generator pass
+            outputs = self.forward(
+                batch["tokens"],
+                batch["token_lens"],
+                batch["spec"],
+                batch["spec_lens"],
+                aux_input={
+                    **batch,
+                },
+            )
+            # cache tensors for the generator pass
+            self.model_outputs_cache = outputs  # pylint: disable=attribute-defined-outside-init
+            scores_disc_fake, _, scores_disc_real, _ = self.disc(
+                outputs["model_outputs"].detach(),
+                outputs["waveform_seg"]
+            )
+            # compute loss
+            with autocast(enabled=False):  # use float32 for the criterion
+                loss_dict = criterion[optimizer_idx](
+                    scores_disc_real,
+                    scores_disc_fake,
+                )
+            return outputs, loss_dict
+
+        if optimizer_idx == 1:
+            # compute melspec segment
+            with autocast(enabled=False):
+                if self.args.encoder_sample_rate:
+                    spec_segment_size = self.spec_segment_size * int(self.interpolate_factor)
+                else:
+                    spec_segment_size = self.spec_segment_size
+                mel_slice = segment(
+                    batch["mel"].float(),
+                    self.model_outputs_cache["slice_ids"],
+                    spec_segment_size,
+                    pad_short=True
+                )
+
+                spec = self.get_spectogram_nt(
+                    self.model_outputs_cache["model_outputs"].float(),
+                )
+                mel_slice_hat = spec_to_mel(
+                    spec=spec,
+                    n_fft=self.config.audio.fft_size,
+                    num_mels=self.config.audio.num_mels,
+                    sample_rate=self.config.audio.sample_rate,
+                    fmin=self.config.audio.mel_fmin,
+                    fmax=self.config.audio.mel_fmax,
+                )
+
+            # compute discriminator scores and features
+            scores_disc_fake, feats_disc_fake, _, feats_disc_real = self.disc(
+                self.model_outputs_cache["model_outputs"],
+                self.model_outputs_cache["waveform_seg"],
+            )
+
+            # compute losses
+            with autocast(enabled=False):  # use float32 for the criterion
+                loss_dict = criterion[optimizer_idx](
+                    mel_slice_hat=mel_slice.float(),
+                    mel_slice=mel_slice_hat.float(),
+                    z_p=self.model_outputs_cache["z_p"].float(),
+                    logs_q=self.model_outputs_cache["logs_q"].float(),
+                    m_p=self.model_outputs_cache["m_p"].float(),
+                    logs_p=self.model_outputs_cache["logs_p"].float(),
+                    z_len=batch["spec_lens"],
+                    scores_disc_fake=scores_disc_fake,
+                    feats_disc_fake=feats_disc_fake,
+                    feats_disc_real=feats_disc_real,
+                    loss_duration=self.model_outputs_cache["loss_duration"],
+                    use_speaker_encoder_as_loss=self.args.use_speaker_encoder_as_loss,
+                    gt_spk_emb=self.model_outputs_cache["gt_spk_emb"],
+                    syn_spk_emb=self.model_outputs_cache["syn_spk_emb"],
+                )
+            return self.model_outputs_cache, loss_dict
+        raise ValueError(" [!] Unexpected `optimizer_idx`.")
+
+    @torch.no_grad()
+    def test_run(self, assets):
+        """Generic test run for `tts` models used by `Trainer`.
+
+        You can override this for a different behaviour.
+
+        Returns:
+            Tuple[Dict, Dict]: Test figures and audios to be projected to Tensorboard.
+        """
+        print(" | > Synthesizing test sentences.")
+        test_audios = {}
+        test_figures = {}
+        test_sentences = self.config.test_sentences
+        for idx, s_info in enumerate(test_sentences):
+            wav = self.synthesize_from_example(s_info)
+            test_audios["{}-audio".format(idx)] = wav
+        return {"figures": test_figures, "audios": test_audios}
+
+    def get_data_loader(
+        self,
+        config: Coqpit,
+        assets: Dict,
+        is_eval: bool,
+        samples: Union[List[Dict], List[List]],
+        verbose: bool,
+        num_gpus: int,
+        rank: int = None,
+    ) -> "DataLoader":
+        dataset = VitsDataset_NT(
+            model_args=self.args,
+            speaker_manager=self.speaker_manager,
+            config=self.config,
+            use_phone_labels=config.use_phone_labels,
+            sample_rate=self.sample_rate,
+            samples=samples,
+            batch_group_size=0 if is_eval else config.batch_group_size * config.batch_size,
+            min_text_len=config.min_text_len,
+            max_text_len=config.max_text_len,
+            min_audio_len=config.min_audio_len,
+            max_audio_len=config.max_audio_len,
+            phoneme_cache_path=config.phoneme_cache_path,
+            precompute_num_workers=config.precompute_num_workers,
+            verbose=verbose,
+            tokenizer=self.tokenizer,
+            start_by_longest=config.start_by_longest,
+        )
+
+        # sort input sequences from short to long
+        dataset.preprocess_samples()
+
+        # get samplers
+        sampler = self.get_sampler(config, dataset, num_gpus)
+        loader = DataLoader(
+            dataset,
+            batch_sampler=sampler,
+            collate_fn=dataset.collate_fn,
+            num_workers=config.num_eval_loader_workers if is_eval else config.num_loader_workers,
+            pin_memory=False,
+        )
+        return loader
+
+
+class VitsDataset_NT(VitsDataset):
+    def __init__(
+        self,
+        model_args,
+        speaker_manager,
+        sample_rate,
+        config,
+        use_phone_labels,
+        *args,
+        **kwargs
+    ):
+        super().__init__(model_args, *args, **kwargs)
+        self.speaker_manager = speaker_manager
+        self.sample_rate = sample_rate
+        self.config = config
+        self.use_phone_labels = use_phone_labels
+
+    def __getitem__(self, idx):
+        example = self.samples[idx]
+        token_ids = self.get_token_ids(idx, example["text"])
+
+        wav, _ = load_audio(example["audio_file"], target_sr=self.sample_rate)
+
+        speaker_id = example['speaker_name']
+        example_id = example['example_id']
+        d_vector = None
+        for dataset_dict_sub in self.config.dataset_dict['datasets'].values():
+            d_vector_file = dataset_dict_sub['d_vector_storage_root']
+            if (Path(d_vector_file) / f'{speaker_id}/{example_id}.pth').is_file():
+                d_vector = torch.load(Path(d_vector_file) / f'{speaker_id}/{example_id}.pth')
+                break
+        if d_vector is None:
+            raise ValueError(f'Could not find d_vector for example {example_id}')
+
+        if d_vector.dim() == 1:
+            d_vector = d_vector[None, :]
+        return {
+            "raw_text": example['text'],
+            "token_ids": token_ids,
+            "token_len": len(token_ids),
+            "wav": wav,
+            "d_vector": d_vector,
+            "speaker_name": example["speaker_name"]
+        }
+
+    def collate_fn(self, batch):
+        """
+        Collate a list of samples from a Dataset into a batch for VITS.
+
+        Args:
+            batch (dict): Expeted keys:
+                - wav (list): list of tensors
+                - token_ids (list):
+                - token_len (list):
+                - speaker_name (list):
+                - language_name (list):
+                - audiofile_path (list):
+                - raw_text (list):
+                - wav_d_vector (list):
+        Returns:
+            - tokens (torch.Tensor): (B, T)
+            - token_lens (torch.Tensor): (B)
+            - token_rel_lens (torch.Tensor): (B)
+            - wav (torch.Tensor): (B, 1, T)
+            - wav_lens (torch.Tensor): (B)
+            - wav_rel_lens (torch.Tensor): (B)
+            - speaker_names (torch.Tensor): (B)
+            - language_names (torch.Tensor): (B)
+            - audiofile_paths (torch.Tensor): (B)
+            - raw_texts (torch.Tensor): (B)
+            - audio_unique_names (torch.Tensor): (B)
+        """
+        B = len(batch)
+        batch = {k: [dic[k] for dic in batch] for k in batch[0]}
+
+        _, ids_sorted_decreasing = torch.sort(
+            torch.LongTensor(
+                [
+                    x.size(1) for x in batch["wav"]]
+            ),
+            dim=0,
+            descending=True
+        )
+
+        max_text_len = max([len(x) for x in batch["token_ids"]])
+        token_lens = torch.LongTensor(batch["token_len"])
+        token_rel_lens = token_lens / token_lens.max()
+
+        wav_lens = [w.shape[1] for w in batch["wav"]]
+        wav_lens = torch.LongTensor(wav_lens)
+        wav_lens_max = torch.max(wav_lens)
+        wav_rel_lens = wav_lens / wav_lens_max
+
+        token_padded = torch.LongTensor(B, max_text_len)
+        wav_padded = torch.FloatTensor(B, 1, wav_lens_max)
+        token_padded = token_padded.zero_() + self.pad_id
+        wav_padded = wav_padded.zero_() + self.pad_id
+        for i in range(len(ids_sorted_decreasing)):
+            token_ids = batch["token_ids"][i]
+            token_padded[i, : batch["token_len"][i]] = torch.LongTensor(token_ids)
+            wav = batch["wav"][i]
+            wav_padded[i, :, : wav.size(1)] = torch.FloatTensor(wav)
+
+        return {
+            "tokens": token_padded,
+            "token_lens": token_lens,
+            "token_rel_lens": token_rel_lens,
+            "waveform": wav_padded,
+            "waveform_lens": wav_lens,
+            "waveform_rel_lens": wav_rel_lens,
+            "speaker_names": batch["speaker_name"],
+            "raw_text": batch["raw_text"],
+            "d_vector": torch.concatenate(batch["d_vector"]) if 'd_vector' in batch.keys() else None,
+        }

setup.py

@@ -0,0 +1,13 @@
+from distutils.core import setup
+
+setup(
+    name='pvq_manipulation',
+    version='0.0.0',
+    author='Department of Communications Engineering, Paderborn University',
+    author_email='sek@nt.upb.de',
+    license='MIT',
+    keywords='audio speech',
+    install_requires=[
+        'torchdiffeq',
+    ],
+)