Unofficial implementation of NANSY++ in Pytorch Lightning
import os
from datetime import datetime
from typing import Any, Dict, Optional, Tuple

import fleep
import librosa
import streamlit as st
import torch
import torchaudio.transforms as T


def select_checkpoint(log_dir: str) -> Tuple[Any, Any, Any, Any]:
    cfg_options = sorted(os.listdir(log_dir))
    cfg_name = st.sidebar.selectbox("Select config", options=cfg_options, index=0)
    exp_options = sorted(
        os.listdir(os.path.join(log_dir, cfg_name)) if cfg_name is not None else [],
        key=lambda date_string: datetime.strptime(
            "-".join(date_string.split("-")[:5]), "%Y-%m-%d_%H-%M-%S"
        ),
    )
    exp_date = st.sidebar.selectbox("Select experience", options=exp_options, index=0)
    exp_dir = os.path.join(log_dir, cfg_name, exp_date)
    ckpt_options = sorted(
        [
            ckpt_name
            for ckpt_name in os.listdir(os.path.join(exp_dir, "checkpoints"))
            if "last" not in ckpt_name
        ]
        if os.path.isdir(os.path.join(exp_dir, "checkpoints"))
        else [],
        key=lambda ckpt_name: int(os.path.splitext(ckpt_name)[0].split("=")[-1]),
    )
    ckpt_name = st.sidebar.selectbox("Select checkpoint", options=ckpt_options, index=0)
    return cfg_name, exp_date, ckpt_name, exp_dir


def load_audio(
    audio_name: Optional[str],
    sample_rate: int,
    tmp_dir: str,
    sample_paths: Dict[str, str],
    key: int,
) -> Tuple[Optional[str], Optional[str], Optional[torch.Tensor]]:
    audio_selection_method = st.selectbox(
        "Audio selection method",
        options=["Upload an audio file", "Choose from VCTK dataset"],
        index=1,
        key=100 * key,
    )
    audio_path, audio = None, None
    if audio_selection_method == "Upload an audio file":
        uploaded_file = st.file_uploader(
            "", type=["wav", "mp3", "m4a", "flac", "ogg"], key=100 * key + 1
        )
        if uploaded_file is not None:
            ext = fleep.get(uploaded_file.read(128)).extension[0]
            audio_path = f"{tmp_dir}/{audio_name}.{ext}"
            with open(audio_path, "wb") as file:
                file.write(uploaded_file.getbuffer())
            st.audio(audio_path)
    elif audio_selection_method == "Choose from VCTK dataset":
        audio_name = st.selectbox(
            "List of audios",
            options=list(sample_paths.keys()),
            key=100 * key + 2,
            index=key,
        )
        if audio_name is not None:
            audio_path = sample_paths[audio_name]
            st.audio(audio_path)
    if audio_path is not None:
        audio = torch.from_numpy(
            librosa.load(audio_path, sr=sample_rate, mono=True)[0]
        ).unsqueeze(0)

    return audio_name, audio_path, audio

app/components.py
MWM-io/nansypp
# pylint: disable-msg=C0103
# pylint: disable-msg=C0116
import os
import socket
from datetime import datetime

import pyrootutils

root = pyrootutils.setup_root(__file__, dotenv=True, pythonpath=True, cwd=False)

# autoroot should be imported here for PYTHON_PATH-setup, do not remove it!
import soundfile as sf
import streamlit as st
import torch
import torchaudio
import torchaudio.transforms as T

from app.components import load_audio, select_checkpoint
from src.dataclasses.backbone import RunDirBase
from src.inference.backbone import BackboneInferencer

# Title stack
st.title("🎙️ Nansy ++")
st.write("")

machine = socket.gethostname()
task = st.sidebar.selectbox("Select task", options=["Reconstruct", "Convert"], index=1)

log_dir = os.path.join(root, "static/runs", RunDirBase.RUNS_BACKBONE)
sample_dir = os.path.join(root, "static/samples/vctk")
tmp_dir = os.path.join(root, "static/tmp")
os.makedirs(tmp_dir, exist_ok=True)

sample_paths = {
    "p225 [female]": f"{sample_dir}/p225_001.wav",
    "p226 [male]": f"{sample_dir}/p226_002.wav",
    "p227 [male]": f"{sample_dir}/p227_003.wav",
    "p228 [female]": f"{sample_dir}/p228_004.wav",
    "p229 [female]": f"{sample_dir}/p229_005.wav",
}


@st.cache_resource
def load_backbone(
    exp_dir: str, checkpoint_name: str, device: str
) -> BackboneInferencer:
    return BackboneInferencer(exp_dir, checkpoint_name, device)


# Select device
cuda_device_ordinal = st.sidebar.selectbox(
    "Select GPU", options=list(range(torch.cuda.device_count())), index=0
)
device = f"cuda:{cuda_device_ordinal}"

# Load generator
cfg_name, exp_date, ckpt_name, exp_dir = select_checkpoint(log_dir)


if ckpt_name is not None and cfg_name is not None:
    inferencer = load_backbone(exp_dir, ckpt_name, device)
    cfg, gen = inferencer.config, inferencer.generator
    input_sr = inferencer.input_sr
    output_sr = inferencer.output_sr

    source_audio, target_audio = None, None
    if task == "Reconstruct":
        st.info("Original")
        source_name, source_path, source_audio = load_audio(
            "uploaded_source", input_sr, tmp_dir, sample_paths, 0
        )
        target_name, target_audio = source_name, source_audio
    elif task == "Convert":
        st.info("Source")
        source_name, source_path, source_audio = load_audio(
            "uploaded_source", input_sr, tmp_dir, sample_paths, 0
        )
        st.info("Target")
        target_name, target_path, target_audio = load_audio(
            "uploaded_target", input_sr, tmp_dir, sample_paths, 1
        )

    if source_audio is not None and target_audio is not None and task is not None:
        # Inference stack
        st.info(task + "ed 🦜")
        raw_path = f"{tmp_dir}/{source_name}-to-{target_name}-{cfg_name}-{exp_date}-{ckpt_name}.wav"
        raw_converted, _ = inferencer.voice_conversion(
            source_audio=source_audio, target_audio=target_audio
        )
        sf.write(
            raw_path,
            raw_converted,
            samplerate=output_sr,
        )
        st.audio(raw_path)

    torch.cuda.empty_cache()

app/reconstruction_and_vc.py
MWM-io/nansypp
import os

import pyrootutils
import soundfile as sf
import streamlit as st
import torch

root = pyrootutils.setup_root(__file__, dotenv=True, pythonpath=True, cwd=False)

from app.components import load_audio, select_checkpoint
from src.dataclasses.backbone import RunDirBase
from src.inference.tts import TextToSpeechInferencer

st.title("Nansy-TTS")
st.write("")


log_dir = os.path.join(root, "static/runs", RunDirBase.RUNS_TTS)
sample_dir = os.path.join(root, "static/samples/vctk")
tmp_dir = os.path.join(root, "static/tmp")
os.makedirs(tmp_dir, exist_ok=True)

sample_paths = {
    "p238 [female]": f"{sample_dir}/p238_001.wav",
    "p248 [female]": f"{sample_dir}/p248_002.wav",
    "p261 [female]": f"{sample_dir}/p261_003.wav",
    "p326 [male]": f"{sample_dir}/p326_004.wav",
    "p347 [male]": f"{sample_dir}/p347_005.wav",
}


# Select device
cuda_device_ordinal = st.sidebar.selectbox(
    "Select GPU", options=list(range(torch.cuda.device_count())), index=0
)
device = f"cuda:{cuda_device_ordinal}"


@st.cache_resource()
def load_nansytts(
    exp_dir: str,
    ckpt_name: str,
) -> TextToSpeechInferencer:
    inferencer = TextToSpeechInferencer(device)
    inferencer.instantiate_tts(exp_dir)
    inferencer.load_tts_weights(ckpt_name)
    inferencer.instantiate_backbone_from_tts()
    return inferencer


_, exp_date, ckpt_name, exp_dir = select_checkpoint(log_dir)

if exp_dir is not None and ckpt_name is not None:
    inferencer = load_nansytts(
        exp_dir,
        ckpt_name,
    )
    input_sr = inferencer.generator.input_sample_rate
    output_sr = inferencer.generator.output_sample_rate
    st.info("Style audio")
    style_name, style_path, style_audio = load_audio(
        "uploaded_style", input_sr, tmp_dir, sample_paths, 0
    )
    style_audio = style_audio.to(device)
    st.info("Text input")
    text = st.text_input("Write text", "To be or not to be that is the question")
    st.info("Generated audio")
    path = os.path.join(tmp_dir, "tts_output.wav")
    signal = inferencer(text, style_audio)
    sf.write(path, signal.squeeze(0).detach().cpu().numpy(), samplerate=output_sr)
    st.audio(path)
    torch.cuda.empty_cache()

app/text_to_speech.py
MWM-io/nansypp
from typing import List, Optional

import pytorch_lightning as pl
import torch
from torch.utils.data import ConcatDataset, DataLoader, Dataset, random_split
from typing_extensions import (  # must import TypedDict from typing_extensions for use with NotRequired
    NotRequired,
    TypedDict,
)

from src.data.dataset.backbone import BackboneDataset
from src.networks.misc.perturbator import InformationPerturbator
from src.networks.misc.transform import ConstantQTransform, MelSpectrogram


class AudioData(TypedDict):
    src_audios: torch.Tensor
    tgt_audios: torch.Tensor
    src_audios: torch.Tensor
    tgt_audios: torch.Tensor
    perturbed_audio_1: NotRequired[torch.Tensor]
    perturbed_audio_2: NotRequired[torch.Tensor]
    cqt: NotRequired[torch.Tensor]
    mel_spec: NotRequired[torch.Tensor]


class BackboneDataModule(pl.LightningDataModule):
    train_dataset: Dataset[AudioData]
    val_dataset: Dataset[AudioData]

    def __init__(
        self,
        batch_size: int,
        pin_memory: bool,
        num_workers: int,
        train_dataset: BackboneDataset,
        val_dataset: Optional[BackboneDataset],
        val_split: Optional[float],
        val_seed: Optional[int],
        information_perturbator: Optional[InformationPerturbator],
        cqt: Optional[ConstantQTransform],
        mel_spec: Optional[MelSpectrogram],
    ) -> None:
        """
        Args:
            batch_size: batch size.
            pin_memory: pin memory.
            num_workers: number of workers.
            train_dataset: training dataset.
            val_dataset: validation dataset. Optional.
            val_split: fraction of training dataset to be used as validation dataset if validation not specified. Optional.
            val_seed: seed to be used when sampling fraction in case validation not specified. Optional.
            information_perturbator: information perturbator.
            cqt: constant Q-transform computer.
            mel_spec: log-mel scale spectrogram computer.
        """
        super().__init__()
        assert (
            val_dataset is not None or val_split is not None
        ), "either val_dataset or val_split should be specified."

        self.train_dataset = train_dataset
        self.val_dataset = val_dataset
        self.segment_length_s = train_dataset.segment_length_s
        self.val_split = val_split
        self.val_seed = val_seed

        self.batch_size = batch_size
        self.pin_memory = pin_memory
        self.num_workers = num_workers

        self.information_perturbator = (
            information_perturbator.cpu()
            if information_perturbator is not None
            else None
        )
        self.cqt = cqt
        self.mel_spec = mel_spec

    def collate_fn(self, items):
        # FIXME: should provide information_perturbator=None and cqt=None to validation
        src_audios = torch.vstack([item["src_audio"] for item in items])
        tgt_audios = torch.vstack([item["tgt_audio"] for item in items])
        batch: AudioData = AudioData(src_audios=src_audios, tgt_audios=tgt_audios)
        if self.information_perturbator is not None:
            batch["perturbed_audio_1"] = self.information_perturbator(src_audios)
            batch["perturbed_audio_2"] = self.information_perturbator(src_audios)
        if self.cqt is not None:
            batch["cqt"] = self.cqt(src_audios)
        if self.mel_spec is not None:
            batch["mel_spec"] = self.mel_spec(src_audios)
        return batch

    def setup(self, stage: Optional[str] = None) -> None:
        if self.val_dataset is None:
            self.train_dataset, self.val_dataset = random_split(
                self.train_dataset,
                lengths=[1 - self.val_split, self.val_split],
                generator=(
                    torch.Generator().manual_seed(self.val_seed)
                    if self.val_seed is not None
                    else None
                ),
            )

        print(
            "Training dataset length:",
            round(len(self.train_dataset) * self.segment_length_s / 3600, 1),
            f"hours of audio.",
        )
        print(
            "Evaluation dataset length:",
            round(len(self.val_dataset) * self.segment_length_s / 3600, 1),
            "hours of audio.",
        )

    def train_dataloader(self) -> DataLoader[AudioData]:
        return DataLoader(
            self.train_dataset,
            batch_size=self.batch_size,
            shuffle=True,
            pin_memory=self.pin_memory,
            num_workers=self.num_workers,
            collate_fn=self.collate_fn,
            drop_last=True,
        )

    def val_dataloader(self) -> DataLoader[AudioData]:
        return DataLoader(
            self.val_dataset,
            batch_size=1,
            shuffle=False,
            pin_memory=self.pin_memory,
            num_workers=self.num_workers,
            collate_fn=self.collate_fn,
        )

src/data/datamodule/backbone.py
MWM-io/nansypp
import os
from typing import List, Optional

import pytorch_lightning as pl
import torch
from torch.nn.utils.rnn import pad_sequence
from torch.utils.data import DataLoader, Dataset
from typing_extensions import TypedDict

from src.networks.backbone.generator import AnalysisFeatures

FIELDS_TO_PAD = ["phoneme_sequence", "phoneme_duration_f_sequence"]
FIELD_WITH_DICT = "backbone_analysis_features"

FEATURES_LENGTH_PARENT = {
    "cqt": "cqt",
    "pitch": "cqt",
    "p_amp": "cqt",
    "ap_amp": "cqt",
    "linguistic": "linguistic",
}
FIX_LEN_FEATURES = ["timbre_global", "timbre_bank"]


class TextToSpeechData(TypedDict):
    phoneme_sequence: torch.Tensor
    phoneme_duration_f_sequence: torch.Tensor
    style_audio: torch.Tensor
    backbone_analysis_features: AnalysisFeatures


class PitchStats:
    def __init__(self, stats_dir: str, file: str) -> None:
        self.stats = torch.load(os.path.join(stats_dir, file))

    def __getitem__(self, attr):
        return self.stats[attr]


class TextToSpeechDataModule(pl.LightningDataModule):
    train_dataset: Dataset[TextToSpeechData]
    val_dataset: Dataset[TextToSpeechData]

    def __init__(
        self,
        train_dataset: Dataset,
        val_dataset: Dataset,
        batch_size: int,
        pin_memory: bool,
        num_workers: int,
        num_val_workers: int,
        shuffle: bool,
    ) -> None:
        super().__init__()
        self.train_dataset = train_dataset
        self.val_dataset = val_dataset

        self.batch_size = batch_size
        self.pin_memory = pin_memory
        self.num_workers = num_workers
        self.num_val_workers = num_val_workers
        self.shuffle = shuffle

    def setup(self, stage: Optional[str] = None) -> None:
        pass

    @staticmethod
    def collate_tts(items: List) -> TextToSpeechData:
        if isinstance(items, list) and isinstance(items[0], list):
            items = [i for item in items for i in item]
        other_fields = set(items[0].keys()) - set(FIELDS_TO_PAD + [FIELD_WITH_DICT])
        batch = {
            field: torch.vstack([i[field] for i in items]) for field in other_fields
        }
        for field_to_pad in FIELDS_TO_PAD:
            batch[field_to_pad] = (
                pad_sequence(
                    [i[field_to_pad] for i in items],
                    batch_first=True,
                    padding_value=0,
                )
                .clone()
                .detach()
            )
        batch[FIELD_WITH_DICT] = {
            key: torch.vstack([i[FIELD_WITH_DICT][key] for i in items]).clone().detach()
            for key in items[0][FIELD_WITH_DICT].keys()
        }
        return batch

    def train_dataloader(self) -> DataLoader[TextToSpeechData]:
        return DataLoader(
            self.train_dataset,
            batch_size=self.batch_size,
            collate_fn=self.collate_tts,
            num_workers=self.num_workers,
            pin_memory=True,
            drop_last=True,
            shuffle=self.shuffle,
        )

    def val_dataloader(self) -> DataLoader[TextToSpeechData]:
        return DataLoader(
            self.val_dataset,
            batch_size=1,
            collate_fn=self.collate_tts,
            num_workers=self.num_val_workers,
            pin_memory=True,
            shuffle=False,
        )

src/data/datamodule/tts.py
MWM-io/nansypp
import glob
import itertools
import os
from typing import List, Optional, Tuple

import numpy as np
import torch
from torch.utils.data import Dataset
from tqdm import tqdm

from src.data.utils import decoded_audio_duration, fast_load


class BackboneDataset(Dataset[torch.Tensor]):
    """Audio dataset returning randomly drawn chunks of audio from audio files."""

    def __init__(
        self,
        segment_length_s: float,
        input_data_dirs: Optional[List[str]],
        input_file_list: Optional[str],
        input_sample_rate: int,
        output_data_dirs: Optional[List[str]],
        output_file_list: Optional[str],
        output_sample_rate: int,
        verbose: bool = True,
    ):
        """
        Args:
            segment_length_s: duration in seconds of sliced segments returned by the dataset.
            input_data_dirs: directory containing decoded audio that will be used as input
            input_file_list: path to .npy file containing a list of files to be used as input. Alternative to input_data_dir.
            input_sample_rate: sample rate of decoded files used as input (specified through input_data_dir or input_file_list)
            output_data_dirs: directory containing decoded audio that will be used as output. Optional.
            output_file_list: path to .npy file containing a list of files to be used as output. Alternative to output_data_dir. Optional.
            output_sample_rate: sample rate of decoded files used as output (specified through output_data_dir or output_file_list)
        """
        self.segment_length_s = segment_length_s
        self.input_sample_rate = input_sample_rate
        self.output_sample_rate = (
            output_sample_rate if output_sample_rate is not None else input_sample_rate
        )

        self.input_num_frames = int(self.input_sample_rate * self.segment_length_s)
        self.output_num_frames = int(self.output_sample_rate * self.segment_length_s)

        assert (
            input_data_dirs is not None or input_file_list is not None
        ), "either input_data_dir or input_file_list must be defined"

        input_segments = self._get_segments(
            input_data_dirs, None, input_sample_rate, verbose=verbose
        )

        # Intersect segments if input and output sample rates are different
        if self.output_sample_rate != self.input_sample_rate:
            assert (
                output_data_dirs is not None or output_file_list is not None
            ), "either output_data_dir or output_file_list must be defined if output_sample_rate is not None"

            output_segments = self._get_segments(
                output_data_dirs,
                output_file_list,
                output_sample_rate,
                verbose=verbose,
            )

            input_segments_dict = {
                (os.path.basename(path), offset): (path, offset)
                for path, offset in input_segments
            }
            output_segments_dict = {
                (os.path.basename(path), offset): (path, offset)
                for path, offset in output_segments
            }
            self.segments_dict = {
                (filename, offset): {
                    "input_path": input_segments_dict[filename, offset][0],
                    "output_path": output_segments_dict[filename, offset][0],
                }
                for (filename, offset) in list(
                    set(input_segments_dict.keys()).intersection(
                        set(output_segments_dict.keys())
                    )
                )
            }
        else:
            self.segments_dict = {
                (os.path.basename(path), offset): {
                    "input_path": path,
                    "output_path": path,
                }
                for path, offset in input_segments
            }
        self.segments_keys = list(self.segments_dict.keys())

    def __len__(self) -> int:
        """Number of audio segments in the dataset."""
        return len(self.segments_dict)

    def _get_segments(
        self,
        data_dirs: Optional[List[str]],
        file_list: Optional[str],
        sample_rate: int,
        verbose: bool,
    ) -> List[Tuple[str, float]]:
        if file_list is not None:
            filepaths = np.load(file_list)
        else:
            filepaths = []
            for data_dir in data_dirs:
                filepaths += glob.glob(f"{data_dir}/**/*.npy", recursive=True)
        return list(
            itertools.chain.from_iterable(
                [
                    [
                        (filepath, num_seconds_offset)
                        for num_seconds_offset in range(
                            int(
                                decoded_audio_duration(filepath, sample_rate)
                                - self.segment_length_s
                            )
                            + 1
                        )
                    ]
                    for filepath in (tqdm(filepaths) if verbose else filepaths)
                    if decoded_audio_duration(filepath, sample_rate)
                    >= self.segment_length_s
                ]
            )
        )

    def __getitem__(self, index: int) -> torch.Tensor:
        audio_filename, time_offset = self.segments_keys[index]
        audio_dict = self.segments_dict[(audio_filename, time_offset)]
        src_audio: torch.Tensor = fast_load(
            path=audio_dict["input_path"],
            num_frames=self.input_num_frames,
            frame_offset=int(time_offset * self.input_sample_rate),
        )
        tgt_audio: torch.Tensor = fast_load(
            path=audio_dict["output_path"],
            num_frames=self.output_num_frames,
            frame_offset=int(time_offset * self.output_sample_rate),
        )
        return {"src_audio": src_audio, "tgt_audio": tgt_audio}

src/data/dataset/backbone.py
MWM-io/nansypp
import os

import pandas as pd
import torch
from torch.utils.data import Dataset

from src.data.datamodule.tts import TextToSpeechData
from src.data.text.text_processor import TextProcessor
from src.data.utils import fast_load
from src.networks.backbone.generator import AnalysisFeatures


class TTSBaseDataset(Dataset[TextToSpeechData]):
    """
    TextToSpeechDataset.
    Base dataset for TTS. As TTS expects ground truth segment + style segment,
    the simpliest dataset expects a csv file
    where each line precies a pair of ground truth and style segments.
    """

    def __init__(
        self,
        dataset_dir: str,
        descriptor_file: str,
        alignment_file: str,
        sample_rate: int,
        text_processor: TextProcessor,
    ):
        self.descriptor_file = os.path.join(dataset_dir, descriptor_file)
        self.segments = pd.read_csv(self.descriptor_file)
        alignments = pd.read_csv(alignment_file)
        self.segments = self.segments.merge(alignments, on="audio")

        self.sample_rate = sample_rate
        self.text_processor = text_processor

    def __len__(self) -> int:
        return len(self.segments)

    def __getitem__(self, index: int) -> TextToSpeechData:
        sample = self.segments.iloc[index]
        duration = sample["duration"]
        style_start = sample.get("style_start", 0)
        style_end = sample.get("style_end", duration)

        style_segment: torch.Tensor = fast_load(
            path=sample["decoded"],
            num_frames=int((style_end - style_start) * self.sample_rate),
            frame_offset=int(style_start * self.sample_rate),
        )
        alignment = self.text_processor.load_alignment(sample["align"])
        (
            phoneme_sequence,
            phoneme_duration_f_sequence,
        ) = self.text_processor.get_alignment_segment(alignment, 0, duration)
        batch = torch.load(sample["batch"], map_location="cpu")
        items = AnalysisFeatures(
            {key: value[sample["batch_i"]].unsqueeze(0) for key, value in batch.items()}
        )
        return {
            "phoneme_sequence": phoneme_sequence,
            "phoneme_duration_f_sequence": phoneme_duration_f_sequence,
            "style_audio": style_segment,
            "backbone_analysis_features": items,
        }

src/data/dataset/tts_base.py
MWM-io/nansypp
import json
import os
from typing import Dict, List

import autoroot
import numpy as np
import torch

from src.data.datamodule.tts import (
    FEATURES_LENGTH_PARENT,
    FIX_LEN_FEATURES,
    TextToSpeechData,
)
from src.data.dataset.tts_base import TTSBaseDataset
from src.data.utils import fast_load
from src.networks.backbone.generator import AnalysisFeatures


class TTSTrainDataset(TTSBaseDataset):
    """
    Text to Speech dataset that returns a list of items of same length for each index.
    """

    def __init__(
        self,
        max_seconds_per_batch: float,
        features_lengths_file: str,
        *args,
        **kwargs,
    ):
        super().__init__(*args, **kwargs)
        if not os.path.isabs(features_lengths_file):
            features_lengths_file = os.path.join(autoroot.root, features_lengths_file)
        with open(features_lengths_file, "r") as file:
            self.features_lengths = json.load(file)[str(self.sample_rate)]

        self.max_seconds_per_batch = max_seconds_per_batch
        min_duration = self.segments["duration"].min() * 10 // 10
        max_diration = self.segments["duration"].max() * 10 // 10 + 1
        durations = np.arange(min_duration, max_diration, 0.05)
        self.segments_per_duration = {}
        self.num_samples_per_duration = {}
        for segment_min_duration, segment_max_duration in zip(
            durations[:-1], durations[1:]
        ):
            segment_min_duration = round(segment_min_duration, 3)
            curr_segments = self.segments[
                (segment_min_duration < self.segments["duration"])
                & (self.segments["duration"] < segment_max_duration)
            ].reset_index(drop=True)
            if len(curr_segments) > 0:
                self.segments_per_duration[segment_min_duration] = curr_segments
                self.num_samples_per_duration[segment_min_duration] = len(curr_segments)
        self.durations = np.array(list(self.segments_per_duration.keys()))
        self.batches_per_duration = {
            key: int(np.ceil(key * value / self.max_seconds_per_batch))
            for key, value in self.num_samples_per_duration.items()
        }
        self.samples_per_batch_per_duration = {
            key: int(value / self.batches_per_duration[key])
            for key, value in self.num_samples_per_duration.items()
        }
        self.length = sum(self.batches_per_duration.values())
        self.order = []
        for key, value in self.batches_per_duration.items():
            for batch_i in range(value):
                self.order.append((key, batch_i))
        self.order = np.array(self.order)
        self.precompute_batch_indeces()

    def precompute_batch_indeces(self):
        self.batches_indices = {}
        for duration in self.durations:
            positions = np.repeat(
                np.arange(self.batches_per_duration[duration]),
                self.samples_per_batch_per_duration[duration],
            )[: self.num_samples_per_duration[duration]]
            np.random.shuffle(positions)
            self.batches_indices[duration] = positions
        np.random.shuffle(self.order)

    def __len__(self) -> int:
        return self.length

    def get_batch(self, order_i):
        duration, batch_i = order_i
        return self.segments_per_duration[duration].iloc[
            np.where(self.batches_indices[duration] == batch_i)[0]
        ]

    def _get_features_segment(
        self, complete_features: AnalysisFeatures, start: float, end: float
    ) -> AnalysisFeatures:
        items = {feature: complete_features[feature] for feature in FIX_LEN_FEATURES}
        for feature, feature_len in FEATURES_LENGTH_PARENT.items():
            start_f = self.features_lengths[str(float(start))][feature_len]
            delta_f = self.features_lengths[str(float(end - start))][feature_len]
            items[feature] = (
                complete_features[feature][
                    ...,
                    start_f : start_f + delta_f,
                ]
                .clone()
                .detach()
                .to("cpu")
            )
        return AnalysisFeatures(items)

    def get_train_pair(self, sample: Dict, duration: float) -> TextToSpeechData:
        style_segment: torch.Tensor = fast_load(
            path=sample["decoded"],
            num_frames=int(duration * self.sample_rate),
            frame_offset=0,
        )
        alignment = self.text_processor.load_alignment(sample["align"])
        (
            phoneme_sequence,
            phoneme_duration_f_sequence,
        ) = self.text_processor.get_alignment_segment(
            alignment,
            0,
            duration,
        )
        gt_audio_features = torch.load(sample["batch"], map_location="cpu")
        items = self._get_features_segment(
            gt_audio_features,
            0,
            duration,
        )
        return {
            "phoneme_sequence": phoneme_sequence,
            "phoneme_duration_f_sequence": phoneme_duration_f_sequence,
            "style_audio": style_segment,
            "backbone_analysis_features": items,
        }

    def __getitem__(self, index: int) -> List[TextToSpeechData]:
        batch_data = self.get_batch(self.order[index])
        duration = self.order[index][0]
        if index == self.length - 1:
            self.precompute_batch_indeces()
        return [
            self.get_train_pair(sample.to_dict(), duration)
            for _, sample in batch_data.iterrows()
        ]

src/data/dataset/tts_train.py
MWM-io/nansypp
"""
https://jonathanbgn.com/2021/08/30/audio-augmentation.html
"""

import math
import os
import pathlib
import random
from abc import ABC, abstractmethod

import numpy as np
import torch
import torchaudio
from tqdm import tqdm

from src.data.utils import decoded_audio_duration, fast_load


class Augmentation(ABC):
    @abstractmethod
    def __init__(self):
        super().__init__()
        self.augmentation_number = 0

    @abstractmethod
    def __call__(self, audio_data):
        pass

    @abstractmethod
    def sample(self, size, audio_length):
        pass


class RandomSpeedChange(Augmentation):
    def __init__(self, sample_rate: int):
        self.augmentation_number = 1
        self.sample_rate = sample_rate

    def __call__(self, audio_data):
        speed_factor = random.choice([0.9, 1.0, 1.1])
        if speed_factor == 1.0:  # no change
            return audio_data

        # change speed and resample to original rate:
        sox_effects = [
            ["speed", str(speed_factor)],
            ["rate", str(self.sample_rate)],
        ]
        transformed_audio, _ = torchaudio.sox_effects.apply_effects_tensor(
            audio_data, self.sample_rate, sox_effects
        )
        return transformed_audio

    def sample(self, size, audio_length):
        pass


class RandomBackgroundNoise(Augmentation):
    def __init__(
        self,
        sample_rate: int,
        noise_dir: str,
        min_snr_db: int,  # 0
        max_snr_db: int,  # 15
        noise_scale: float,
        augmentation_number: int,
        length_s: float,
    ):
        """
        Args:
            sample_rate: sample rate.
            noise_dir: directory containing noise audio samples
            min_snr_db: minimum source-to-noise-ratio in decibel used to generate signal scaling audio data
            max_snr_db: maximum source-to-noise-ratio in decibel used to generate signal scaling audio data
            noise_scale: noise signal weight
            augmentation_number: augmentation index used when composing
            length_s: segment length of audio data from dataset in seconds
        """
        self.augmentation_number = augmentation_number
        self.sample_rate = sample_rate
        self.min_snr_db = min_snr_db
        self.max_snr_db = max_snr_db
        self.noise_scale = noise_scale
        self.length_s = length_s
        self.length_f = int(sample_rate * length_s)

        if not os.path.exists(noise_dir):
            raise IOError(f"Noise directory `{noise_dir}` does not exist")
        # find all NPY files including in sub-folders:
        self.noise_files_list = list(pathlib.Path(noise_dir).glob("*.npy"))
        if len(self.noise_files_list) == 0:
            raise IOError(
                f"No decoded .npy file found in the noise directory `{noise_dir}`"
            )
        self.noise_files_dict = {
            path: int(decoded_audio_duration(path, sample_rate) * sample_rate)
            for path in tqdm(self.noise_files_list)
        }

    def __call__(self, audio_data, noises=None):
        """Add random noise to the audio_data.
        Args:
            audio_data: [torch.float32; [B, T]], input tensor.
        Returns:
            [torch.float32; [B, T]], generated augmentation.
        """
        shape = audio_data.shape
        if len(audio_data.shape) == 1:
            audio_data = audio_data.reshape((1, -1))
        N, audio_length = audio_data.shape
        if noises is None:
            noises = self.sample(N, audio_length)
        noises_to_add = noises[:N, :audio_length].to(audio_data.device)
        snr_db = random.randint(self.min_snr_db, self.max_snr_db)
        snr = math.exp(snr_db / 10)
        audio_power = audio_data.norm(p=2, dim=1)
        noise_power = noises_to_add.norm(p=2, dim=1)
        scale = (snr * noise_power / (audio_power + 1e-6)).reshape(-1, 1)
        result = (scale * audio_data + self.noise_scale * noises_to_add) / 2
        result = result.reshape(shape)
        return result

    def sample(self, size, audio_length):
        file_indices = np.random.choice(len(self.noise_files_list), size, replace=False)
        return torch.vstack(
            [
                fast_load(
                    self.noise_files_list[file_idx],
                    audio_length,
                    np.random.randint(
                        0,
                        self.noise_files_dict[self.noise_files_list[file_idx]]
                        - audio_length,
                    ),
                )
                for file_idx in file_indices
            ]
        )


class ComposeTransform(Augmentation):
    def __init__(self, transforms):
        self.augmentation_number = len(transforms)
        self.transforms = transforms

    def __call__(self, audio_data):
        for t in self.transforms:
            audio_data = t(audio_data)
        return audio_data

    def sample(self, size, audio_length):
        pass

src/data/preprocessing/augmentation.py
MWM-io/nansypp
import json

import numpy as np
import torch
import typer
from tqdm import tqdm

from src.inference.backbone import BackboneInferencer

app = typer.Typer()


def main(
    exp_dir: str,
    ckpt_name: str,
    save_path: str = "static/tts/features_lengths.json",
    device="cuda:0",
):
    generator = BackboneInferencer(str(exp_dir), ckpt_name, device).generator

    keys = ["cqt", "linguistic"]
    sample_rates = [16000, 44100]
    all_results = {}

    for sample_rate in tqdm(sample_rates):
        results = {}
        results[0] = {"cqt": 0, "linguistic": 0}
        for audio_length in tqdm(np.arange(0.25, 25.1, 0.05)):
            results[audio_length] = {}
            len_f = int(audio_length * sample_rate)
            sample = torch.zeros((1, len_f)).to("cuda")
            features = generator.analyze(sample, enable_information_perturbator=False)
            for key in keys:
                results[audio_length][key] = features[key].shape[-1]
            assert features["pitch"].shape[-1] == features["cqt"].shape[-1]
            assert features["p_amp"].shape[-1] == features["cqt"].shape[-1]
            assert features["ap_amp"].shape[-1] == features["cqt"].shape[-1]

            assert features["timbre_global"].shape[-1] == 192
            assert features["timbre_bank"].shape[-1] == 50
        all_results[sample_rate] = results

    with open(save_path, "w") as file:
        json.dump(all_results, file)


if __name__ == "__main__":
    typer.run(main)

src/data/preprocessing/compute_features_lengths.py
MWM-io/nansypp
# https://drive.google.com/file/d/1JiCMBVTG4BMREK8cT3MYck1MgYvwASL0/view
GLOW_TTS_CKPT_PATH = "path/to/glow-tts/pretrained.pth"

import os
import sys

import autoroot

GLOW_TTS_REPO_PATH = os.path.join(autoroot.root, "glow-tts/")

sys.path.append(GLOW_TTS_REPO_PATH)


import warnings

import pandas as pd
import torch
import typer
from data_utils import TextMelCollate, TextMelLoader
from text import _id_to_symbol, symbols
from torch.utils.data import DataLoader
from tqdm import tqdm
from typing_extensions import Annotated

import models
import utils

warnings.simplefilter(action="ignore", category=FutureWarning)

app = typer.Typer()


def repr_seg(label, start, end, score=1.0):
    return f"{label}\t({score:4.2f}): [{start:5d}, {end:5d})"


def save_from_attn(attn, x, x_lengths, txt_path):
    index_phonemes = [_id_to_symbol[i] for i in x[:x_lengths].tolist()]
    limits = [0] + attn.sum(axis=1)[:x_lengths].cumsum(dim=0).int().tolist()

    os.makedirs(os.path.dirname(txt_path), exist_ok=True)
    with open(txt_path, "w") as file:
        for i, start, end in zip(index_phonemes, limits[:-1], limits[1:]):
            print(repr_seg("{" + i + "}", start, end), file=file)


def main(
    config: Annotated[str, typer.Argument(help="GlowTTS config file path")],
    input_dir: Annotated[
        str, typer.Argument(help="path to input directory containing raw audio")
    ],
    output_dir: Annotated[
        str,
        typer.Argument(help="path to output directory that will contain alignments"),
    ],
    device: Annotated[str, typer.Argument()] = "cuda:0",
):
    audio_path_to_txt_path = (
        lambda i: os.path.join(output_dir, i.replace(input_dir, "")).rsplit(".", 1)[0]
        + ".txt"
    )

    hps = utils.get_hparams_from_file(config)
    torch.manual_seed(hps.train.seed)

    model = models.FlowGenerator(
        len(symbols) + getattr(hps.data, "add_blank", False),
        out_channels=hps.data.n_mel_channels,
        **hps.model,
    ).to(device)
    utils.load_checkpoint(GLOW_TTS_CKPT_PATH, model)
    model.decoder.store_inverse()  # do not calcuate jacobians for fast decoding
    _ = model.to(device)
    _ = model.eval()

    train_dataset = TextMelLoader(hps.data.training_files, hps.data)
    train_sampler = torch.utils.data.distributed.DistributedSampler(
        train_dataset, num_replicas=1, rank=0, shuffle=False
    )
    train_loader = DataLoader(
        train_dataset,
        num_workers=16,
        shuffle=False,
        batch_size=hps.train.batch_size,
        pin_memory=True,
        drop_last=False,
        collate_fn=TextMelCollate(1),
        sampler=train_sampler,
    )

    paths_to_align_paths = {}
    for batch in tqdm(train_loader, dynamic_ncols=True):
        x, x_lengths, y, y_lengths, paths = batch
        (
            (_, _, _, _, _),
            (_, _, _),
            (attn, _, _),
        ) = model(
            x.to(device),
            x_lengths.to(device),
            y.to(device),
            y_lengths.to(device),
            gen=False,
        )
        attn = attn.squeeze(1)
        size = len(paths)
        for index in range(size):
            txt_path = audio_path_to_txt_path(paths[index][0])
            save_from_attn(
                attn[index],
                x[index],
                x_lengths[index],
                txt_path,
            )
            paths_to_align_paths[paths[index][0]] = txt_path
        del x, x_lengths, y, y_lengths
        torch.cuda.empty_cache()
    dataframe = pd.DataFrame.from_dict(
        paths_to_align_paths, orient="index", columns=["align"]
    ).reset_index(names="audio")
    dataframe.to_csv(os.path.join(output_dir, "dataset.csv"), index=False)


if __name__ == "__main__":
    typer.run(main)

src/data/preprocessing/compute_glowtts_alignments.py
MWM-io/nansypp
import glob
import os
from multiprocessing import Process
from pathlib import Path
from typing import Dict, List

import librosa
import numpy as np
import pandas as pd
import typer
from tqdm import tqdm
from typing_extensions import Annotated

app = typer.Typer()


def decode_audio(audio_path: str, sample_rate: int) -> np.ndarray:
    """
    Decode audio at given sampling rate.

    Args:
        audio_path: path to audio file.
        sample_rate: sampling rate to use for decoding.
    """
    audio, _ = librosa.load(audio_path, sr=sample_rate, mono=True)
    return audio


def process_subset(
    audio_dict: Dict[str, str], list_idx: List[int], sample_rate: int
) -> None:
    """
    Decode list of audio.
    Args:
        audio_list: list of audio paths.
        list_idx: list of indices to process in loop.
        output_dir: path to output directory.
        sample_rate: sampling rate to use for decoding.
    """
    audio_list = list(audio_dict.keys())
    for idx in tqdm(list_idx):
        output_path = audio_list[idx]
        audio_path = audio_dict[output_path]
        if not os.path.isfile(output_path):
            audio = decode_audio(audio_path, sample_rate)
            os.makedirs(os.path.dirname(output_path), exist_ok=True)
            np.save(output_path, audio)


def main(
    input_dir: Annotated[
        str, typer.Option("-i", help="path to input directory containing raw audio")
    ],
    output_dir: Annotated[
        str,
        typer.Option(
            "-o", help="path to output directory that will contain decoded audio"
        ),
    ],
    sample_rate: Annotated[
        int,
        typer.Option("-sr", help="sampling rate"),
    ] = 44100,
    n_processes: Annotated[
        int,
        typer.Option("-np", help="number of processes"),
    ] = 6,
):
    audio_path_to_decode_path = (
        lambda i: os.path.join(output_dir, i.replace(input_dir, "")).rsplit(".", 1)[0]
        + ".npy"
    )

    os.makedirs(output_dir, exist_ok=True)
    audio_list = (
        glob.glob(os.path.join(input_dir, "**/*.wav"), recursive=True)
        + glob.glob(os.path.join(input_dir, "**/*.flac"), recursive=True)
        + glob.glob(os.path.join(input_dir, "**/*.mp4"), recursive=True)
        + glob.glob(os.path.join(input_dir, "**/*.mp3"), recursive=True)
        + glob.glob(os.path.join(input_dir, "**/*.sph"), recursive=True)
        + glob.glob(os.path.join(input_dir, "**/*.ogg"), recursive=True)
    )
    audio_dict = {audio_path_to_decode_path(path): path for path in audio_list}
    dataframe = pd.DataFrame.from_dict(
        audio_dict, orient="index", columns=["audio"]
    ).reset_index(names="decoded")
    dataframe.to_csv(os.path.join(output_dir, "dataset.csv"), index=False)
    # Divide dataset equally into sublists for each process
    sublists = np.array_split(np.arange(len(audio_dict)), n_processes)

    procs = []
    # instantiating processes
    for sublist in sublists:
        proc = Process(
            target=process_subset,
            args=(audio_dict, sublist, sample_rate),
        )
        procs.append(proc)
        proc.start()

    # complete the processes
    for proc in procs:
        proc.join()

    print("\nAll process finished")


if __name__ == "__main__":
    typer.run(main)

src/data/preprocessing/decode.py
MWM-io/nansypp
import json
from pathlib import Path
from typing import Optional

import pandas as pd
import torch
import torchaudio
import torchaudio.transforms as T
import typer
from torch.utils.data import DataLoader, Dataset
from tqdm import tqdm

from src.data.utils import decoded_audio_duration, fast_load
from src.inference.backbone import BackboneInferencer

app = typer.Typer()


class ExtractStatsDataset(Dataset):
    def __init__(
        self,
        sample_rate: int,
        file: Optional[str] = None,
        segments: Optional[pd.DataFrame] = None,
    ) -> None:
        super().__init__()
        if file is not None:
            self.segments = pd.read_csv(file)
        elif segments is not None:
            self.segments = segments.copy()
        else:
            raise AttributeError("Either file or segments should be provided.")
        self.sample_rate = sample_rate

    def __len__(self) -> int:
        return len(self.segments)

    def __getitem__(self, index: int) -> torch.Tensor:
        sample = self.segments.iloc[index]
        sample_start = sample.get("gt_start", 0)
        sample_end = sample.get("gt_end", 100)
        if "decoded" in sample:
            num_frames = int(
                min(
                    sample_end - sample_start,
                    decoded_audio_duration(sample.decoded, self.sample_rate),
                )
                * self.sample_rate
            )
            audio_sample = fast_load(
                path=sample.decoded,
                num_frames=num_frames,
                frame_offset=int(sample_start) * self.sample_rate,
            )
        elif "audio" in sample:
            audio_sample, sample_rate = torchaudio.load(sample.audio)
            if sample_rate != self.sample_rate:
                audio_sample = T.Resample(sample_rate, self.sample_rate)(audio_sample)
            audio_sample = audio_sample[0]
        else:
            raise AttributeError
        return audio_sample


def main(
    data_file: Path,
    sample_rate: int,
    path_to_save: Path,
    exp_dir: Path,
    ckpt_name: str,
    batch_size: int = 1,
    device: str = "cuda:0",
):
    """
    Script to compute TTS targets for backbone model.
    Args:
        - data_file: path to csv file describing TTS training segments
        - path_to_save: path to dir to save resulting data description file
            and computed statistics
        - exp_dir: experiment dir of backbone
        - ckpt_name: checkpoint name to use
        - config_name: backbone config name from experiment dir
        - batch_size: for statistics computation
        - device
    Example:
    python -m src.data.preprocessing.precompute_tts_targets \
        path/to/dataset.csv \
        path/to/save/ \
        path/to/backbone/experiment/ \
        <checkpoint_name>.ckpt
    """
    backbone_info = {
        "exp_dir": str(exp_dir),
        "ckpt_name": ckpt_name,
    }

    inferencer = BackboneInferencer(str(exp_dir), ckpt_name, device)
    dataset = ExtractStatsDataset(sample_rate=sample_rate, file=data_file)
    dataloader = DataLoader(dataset, batch_size=batch_size, shuffle=False)

    batch_path = path_to_save.joinpath("batches/")
    batch_path.mkdir(parents=True, exist_ok=True)
    batch_to_path = lambda i: batch_path.joinpath(f"{i}.pt")

    with open(path_to_save.joinpath("backbone_info.json"), "w") as backbone_file:
        json.dump(backbone_info, backbone_file)

    dataframe = pd.read_csv(data_file)
    dataframe["duration"] = dataframe["decoded"].apply(
        lambda i: decoded_audio_duration(i, sample_rate)
    )
    dataframe["batch"] = (dataframe.index // batch_size).map(batch_to_path)
    dataframe["batch_i"] = dataframe.index % batch_size
    dataframe.to_csv(path_to_save.joinpath(data_file.name), index=False)

    pitch_stats = {
        "p_amp": {
            "min": 1e8,
            "max": -1e8,
        },
        "ap_amp": {
            "min": 1e8,
            "max": -1e8,
        },
        "pitch": {
            "min": 1e8,
            "max": -1e8,
        },
    }
    for i, batch in tqdm(enumerate(dataloader), total=len(dataloader)):
        stats = inferencer.generator.analyze(
            batch.to(device), enable_information_perturbator=False
        )
        for key, value in pitch_stats.items():
            pitch_stats[key]["min"] = min(value["min"], stats[key].min().item())
            pitch_stats[key]["max"] = max(value["max"], stats[key].max().item())
        torch.save(stats, batch_to_path(i))
    torch.save(pitch_stats, path_to_save.joinpath("pitch_stats.pt"))


if __name__ == "__main__":
    typer.run(main)

src/data/preprocessing/precompute_tts_targets.py
MWM-io/nansypp
""" from https://github.com/keithito/tacotron """
import re

from src.data.text import cleaners
from src.data.text.symbols import punctuation_to_parse, symbols

# Mappings from symbol to numeric ID and vice versa:
_symbol_to_id = {s: i for i, s in enumerate(symbols)}
_id_to_symbol = {i: s for i, s in enumerate(symbols)}

# Regular expression matching text enclosed in curly braces:
_curly_re = re.compile(r"(.*?)\{(.+?)\}(.*)")


def get_arpabet(word, dictionary):
    word_arpabet = dictionary.lookup(word)
    if word_arpabet is not None:
        return "{" + word_arpabet[0] + "}"
    else:
        return word


def text_to_sequence(text, cleaner_names, dictionary=None):
    """Converts a string of text to a sequence of IDs corresponding to the symbols in the text.

    The text can optionally have ARPAbet sequences enclosed in curly braces embedded
    in it. For example, "Turn left on {HH AW1 S S T AH0 N} Street."

    Args:
      text: string to convert to a sequence
      cleaner_names: names of the cleaner functions to run the text through
      dictionary: arpabet class with arpabet dictionary

    Returns:
      List of integers corresponding to the symbols in the text
    """
    sequence = []

    space = _symbols_to_sequence(" ")
    # Check for curly braces and treat their contents as ARPAbet:
    while len(text):
        m = _curly_re.match(text)
        if not m:
            clean_text = _clean_text(text, cleaner_names)
            clean_text = clean_text.translate(
                str.maketrans({key: f" {key}" for key in punctuation_to_parse})
            ).lower()
            if dictionary is not None:
                clean_text = [get_arpabet(w, dictionary) for w in clean_text.split(" ")]
                for i in range(len(clean_text)):
                    t = clean_text[i]
                    if t.startswith("{"):
                        sequence += _arpabet_to_sequence(t[1:-1])
                    else:
                        sequence += _symbols_to_sequence(t)
                    sequence += space
            else:
                sequence += _symbols_to_sequence(clean_text)
            break
        sequence += _symbols_to_sequence(_clean_text(m.group(1), cleaner_names))
        sequence += _arpabet_to_sequence(m.group(2))
        text = m.group(3)

    # remove trailing space
    if dictionary is not None:
        sequence = sequence[:-1] if sequence[-1] == space[0] else sequence
    return sequence


def sequence_to_text(sequence):
    """Converts a sequence of IDs back to a string"""
    result = ""
    for symbol_id in sequence:
        if symbol_id in _id_to_symbol:
            s = _id_to_symbol[symbol_id]
            # Enclose ARPAbet back in curly braces:
            if len(s) > 1 and s[0] == "@":
                s = "{%s}" % s[1:]
            result += s
    return result.replace("}{", " ")


def _clean_text(text, cleaner_names):
    for name in cleaner_names:
        cleaner = getattr(cleaners, name)
        if not cleaner:
            raise Exception("Unknown cleaner: %s" % name)
        text = cleaner(text)
    return text


def _symbols_to_sequence(symbols):
    return [_symbol_to_id[s] for s in symbols if _should_keep_symbol(s)]


def _arpabet_to_sequence(text):
    return _symbols_to_sequence(["@" + s for s in text.split()])


def _should_keep_symbol(s):
    return s in _symbol_to_id and s is not "_" and s is not "~"

src/data/text/__init__.py
MWM-io/nansypp
""" from https://github.com/keithito/tacotron """

"""
Cleaners are transformations that run over the input text at both training and eval time.

Cleaners can be selected by passing a comma-delimited list of cleaner names as the "cleaners"
hyperparameter. Some cleaners are English-specific. You'll typically want to use:
  1. "english_cleaners" for English text
  2. "transliteration_cleaners" for non-English text that can be transliterated to ASCII using
     the Unidecode library (https://pypi.python.org/pypi/Unidecode)
  3. "basic_cleaners" if you do not want to transliterate (in this case, you should also update
     the symbols in symbols.py to match your data).
"""

import re

from unidecode import unidecode

from src.data.text.numbers import normalize_numbers

# Regular expression matching whitespace:
_whitespace_re = re.compile(r"\s+")

# List of (regular expression, replacement) pairs for abbreviations:
_abbreviations = [
    (re.compile("\\b%s\\." % x[0], re.IGNORECASE), x[1])
    for x in [
        ("mrs", "misess"),
        ("mr", "mister"),
        ("dr", "doctor"),
        ("st", "saint"),
        ("co", "company"),
        ("jr", "junior"),
        ("maj", "major"),
        ("gen", "general"),
        ("drs", "doctors"),
        ("rev", "reverend"),
        ("lt", "lieutenant"),
        ("hon", "honorable"),
        ("sgt", "sergeant"),
        ("capt", "captain"),
        ("esq", "esquire"),
        ("ltd", "limited"),
        ("col", "colonel"),
        ("ft", "fort"),
    ]
]


def expand_abbreviations(text):
    for regex, replacement in _abbreviations:
        text = re.sub(regex, replacement, text)
    return text


def expand_numbers(text):
    return normalize_numbers(text)


def lowercase(text):
    return text.lower()


def collapse_whitespace(text):
    return re.sub(_whitespace_re, " ", text)


def convert_to_ascii(text):
    return unidecode(text)


def basic_cleaners(text):
    """Basic pipeline that lowercases and collapses whitespace without transliteration."""
    text = lowercase(text)
    text = collapse_whitespace(text)
    return text


def transliteration_cleaners(text):
    """Pipeline for non-English text that transliterates to ASCII."""
    text = convert_to_ascii(text)
    text = lowercase(text)
    text = collapse_whitespace(text)
    return text


def english_cleaners(text):
    """Pipeline for English text, including number and abbreviation expansion."""
    text = convert_to_ascii(text)
    text = lowercase(text)
    text = expand_numbers(text)
    text = expand_abbreviations(text)
    text = collapse_whitespace(text)
    return text

src/data/text/cleaners.py
MWM-io/nansypp
""" from https://github.com/keithito/tacotron """

"""
Defines the set of symbols used in text input to the model.

The default is a set of ASCII characters that works well for English or text that has been run through Unidecode. For other data, you can modify _characters. See TRAINING_DATA.md for details. """
from src.data.text import cmudict

_pad = "_"
_punctuation = "!'(),.:;? "
_special = "-"
# _letters = "ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz"
_letters = "abcdefghijklmnopqrstuvwxyz"

# Prepend "@" to ARPAbet symbols to ensure uniqueness (some are the same as uppercase letters):
_arpabet = ["@" + s for s in cmudict.valid_symbols]

# Export all symbols:
symbols = [_pad] + list(_special) + list(_punctuation) + list(_letters) + _arpabet

punctuation_to_parse = _punctuation.replace("'", "").replace(" ", "")

src/data/text/symbols.py
MWM-io/nansypp
import os
from typing import List, Tuple

import autoroot
import torch

from src.data.text import _symbol_to_id, cmudict, text_to_sequence
from src.data.text.symbols import _pad, symbols

LINGUISTIC_FRAMES_PER_S = 50
PITCH_FRAMES_PER_S = {
    16000: 63,
    44100: 173,
}


class TextProcessor:
    def __init__(self, sample_rate: int):
        self.sample_rate = sample_rate
        self.labels = symbols
        self.alignment_whitelist = "1234567890. "
        self.dictionary = {c: i for i, c in enumerate(self.labels)}
        self.align_frame_to_wave_frame = 86
        self.linguistic_features_length_ratio = (
            LINGUISTIC_FRAMES_PER_S / self.align_frame_to_wave_frame
        )
        self.pitch_features_length_ratio = (
            PITCH_FRAMES_PER_S[self.sample_rate] / self.align_frame_to_wave_frame
        )
        self.cmudict = cmudict.CMUDict(
            os.path.join(autoroot.root, "static/tts/cmu_dictionary")
        )
        self.text_cleaners = ["english_cleaners"]
        self.pad_frame = _pad

    def text_to_encoded_phoneme_sequence(self, text: str) -> torch.Tensor:
        return torch.IntTensor(
            text_to_sequence(text.strip(), self.text_cleaners, self.cmudict)
        )

    def encode_phoneme_sequence(self, sequence: List[str]) -> torch.Tensor:
        return torch.tensor([_symbol_to_id[char] for char in sequence])

    def load_alignment(
        self,
        alignment_path: str,
    ) -> List[List]:
        with open(alignment_path, "r") as file:
            lines = file.readlines()
        alignment = []
        for line in lines:
            phoneme = line.split("\t")[0]
            if "{" in phoneme and "}" in phoneme:
                phoneme = phoneme.replace("{", "").replace("}", "")
            start, stop = list(
                map(
                    int,
                    "".join(
                        filter(self.alignment_whitelist.__contains__, line)
                    ).split()[1:],
                )
            )
            alignment.append([phoneme, start, stop])
        return alignment

    def get_alignment_segment(
        self,
        alignment: List[List],
        start_s: float,
        end_s: float,
    ) -> Tuple[torch.Tensor, torch.Tensor]:
        alignment_frame_start = int(start_s * self.align_frame_to_wave_frame)
        alignment_frame_end = int(end_s * self.align_frame_to_wave_frame)
        alignment_segment = list(
            filter(
                lambda i: i[1] < alignment_frame_end and alignment_frame_start < i[2],
                alignment,
            )
        )
        alignment_segment[0][1] = max(alignment_frame_start, alignment_segment[0][1])
        alignment_segment[-1][2] = min(alignment_frame_end, alignment_segment[-1][2])

        start_segment, end_segment = [], []
        if alignment_segment[0][1] > alignment_frame_start:
            start_segment.append(
                [self.pad_frame, alignment_frame_start, alignment_segment[0][1]]
            )
        if alignment_segment[-1][2] < alignment_frame_end:
            end_segment.append(
                [self.pad_frame, alignment_segment[-1][2], alignment_frame_end]
            )
        alignment_segment = start_segment + alignment_segment + end_segment

        phoneme_sequence = self.encode_phoneme_sequence(
            [char for (char, _, _) in alignment_segment]
        )
        phoneme_duration_f_sequence = torch.tensor(
            [int((stop - start)) for (_, start, stop) in alignment_segment]
        )
        return phoneme_sequence, phoneme_duration_f_sequence

src/data/text/text_processor.py
MWM-io/nansypp
import os

import numpy as np
import torch


def fast_load(
    path: str, num_frames: int, frame_offset: int, dtype: str = "float32"
) -> torch.Tensor:
    """
    Fast loading of decoded audio .npy file.
    """
    if dtype == "float32":
        with open(path, "rb") as file:
            file.seek(128 + 4 * frame_offset)
            content = file.read(4 * num_frames)
            array = np.frombuffer(content, dtype="float32")
    elif dtype == "float64":
        with open(path, "rb") as file:
            file.seek(128 + 8 * frame_offset)
            content = file.read(8 * num_frames)
            array = np.float32(np.frombuffer(content, dtype="float64"))
    else:
        raise ValueError(f"Unexpected dtype {dtype}")

    return torch.from_numpy(array.copy())


def decoded_audio_duration(
    path: str, sample_rate: int, dtype: str = "float32"
) -> float:
    """
    Assess decoded audio .npy file duration from file size.
    """
    if dtype == "float32":
        return (os.path.getsize(path) - 128) / (4 * sample_rate)
    if dtype == "float64":
        return (os.path.getsize(path) - 128) / (8 * sample_rate)
    raise ValueError(f"Unexpected dtype {dtype}")

src/data/utils.py
MWM-io/nansypp
"""Typed hydra configurations for the whole project

Notes and gotchas:
    - Make sure to decorate each config class with @dataclass, even classes that are
    themselves derived from @dataclass-decorated configuration classes,
    otherwise Hydra fails at composing the configurations.

    - In the case of hierarchies of structured configurations such as:
    ```
    @dataclass
    class SubConfigSchema:
        foo: str = MISSING

    @dataclass
    class SubConfigA(SubConfigSchema):
        foo: str = 'A'
        
    @dataclass
    class SubConfigB(SubConfigSchema):
        foo: str = 'B'

    @dataclass
    class Config
        sub_config: SubConfig  # uninitialized
    ```
    
    Do not pass a default value of SubConfigA() or SubConfigB() to config.sub_config via a default class attribute,
    since this would prevent Hydra from overriding those:
    DON'T
    ```
    @dataclass
    class Config:
        sub_config: SubConfig = SubConfigA()
    ```
    Because trying to override config.sub_config with SubConfigB would yield:
    ```Merge error: SubConfigB is not a subclass of SubConfigA.```

    Default values should be provided via the Defaults list, either directly from the structured config,
    or through a yaml configuration file.
"""

from dataclasses import dataclass, field
from enum import Enum
from typing import Dict, List, Optional

from hydra.core.config_store import ConfigStore
from omegaconf import II, MISSING, OmegaConf


class StringEnum(str, Enum):
    def __str__(self) -> str:
        return self.value


cs = ConfigStore.instance()


# -------- DATA -------- #
@dataclass
class AdditiveBackgroundNoiseConfig:
    _target_: str = "src.data.preprocessing.augmentation.RandomBackgroundNoise"
    sample_rate: int = II("train_dataset.input_sample_rate")
    noise_dir: str = "/path/to/demand/decoded/audio/directory"
    min_snr_db: int = 0
    max_snr_db: int = 15
    noise_scale: float = MISSING
    augmentation_number: int = 1
    length_s: float = II("train_dataset.segment_length_s")


@dataclass
class SmallScaleAdditiveBackgroundNoiseConfig(AdditiveBackgroundNoiseConfig):
    noise_scale: float = 0.8


cs.store(
    group="datamodule/augmentation",
    name="nansypp_random_background_noise",
    node=AdditiveBackgroundNoiseConfig,
)

cs.store(
    group="datamodule/augmentation",
    name="small_scale_datamodule_random_background_noise",
    node=SmallScaleAdditiveBackgroundNoiseConfig,
)


@dataclass
class DatasetConfig:
    _target_: str = "src.data.dataset.backbone.BackboneDataset"
    segment_length_s: float = 1.0
    input_data_dirs: Optional[List[str]] = None
    input_file_list: Optional[str] = None
    input_sample_rate: int = MISSING
    output_data_dirs: Optional[List[str]] = None
    output_file_list: Optional[str] = None
    output_sample_rate: int = MISSING


@dataclass
class HifiTTSDatasetConfig(DatasetConfig):
    input_data_dirs: List[str] = field(
        default_factory=lambda: [
            "/path/to/hifitts/decoded/audio/directory",
        ]
    )
    input_sample_rate: int = 44100
    output_sample_rate: int = 44100


cs.store(group="train_dataset", name="base", node=DatasetConfig)
cs.store(
    group="train_dataset",
    name="hifitts",
    node=HifiTTSDatasetConfig,
)


@dataclass
class DataModuleConfig:
    _target_: str = "src.data.datamodule.backbone.BackboneDataModule"
    batch_size: int = 4
    num_workers: int = 2
    pin_memory: bool = True

    train_dataset: DatasetConfig = II("train_dataset")
    val_dataset: Optional[DatasetConfig] = II("val_dataset")
    val_split: Optional[float] = None
    val_seed: Optional[int] = None


cs.store(group="datamodule", name="base", node=DataModuleConfig)


# -------- LINGUISTIC-INFORMATION-PRESERVING PERTURBATOR -------- #
@dataclass
class PraatAugmentConfig:
    _target_: str = "src.networks.misc.praat.PraatAugment"
    sample_rate: int = II("train_dataset.input_sample_rate")
    pitch_steps: float = 0.01
    pitch_floor: float = 75
    pitch_ceil: float = 600


@dataclass
class ParametricEqualizerAugmentConfig:
    _target_: str = "src.networks.misc.peq.ParametricEqualizer"
    sample_rate: int = II("train_dataset.input_sample_rate")
    window_length: int = II("generator.mel_spectrogram_transform.window_length")


@dataclass
class InformationPerturbatorConfig:
    _target_: str = "src.networks.misc.perturbator.InformationPerturbator"
    formant_shift: float = 1.4
    pitch_shift: float = 2.0
    pitch_range: float = 1.5
    cutoff_lowpass: float = 60.0
    cutoff_highpass: float = 10000.0
    q_min: float = 2.0
    q_max: float = 5.0
    num_peaks: int = 8
    gain_range: float = 12.0

    stft_window_length: int = II("generator.mel_spectrogram_transform.window_length")
    stft_hop_length: int = II("generator.mel_spectrogram_transform.hop_length")

    praat_augment: PraatAugmentConfig = PraatAugmentConfig()
    parametric_equalizer: ParametricEqualizerAugmentConfig = (
        ParametricEqualizerAugmentConfig()
    )

    additive_noise: Optional[
        AdditiveBackgroundNoiseConfig
    ] = SmallScaleAdditiveBackgroundNoiseConfig()


@dataclass
class BypassedInformationPerturbatorConfig(InformationPerturbatorConfig):
    _target_: str = "torch.nn.Identity"


# -------- GENERATOR -------- #
@dataclass
class PretrainedWav2Vec2Config:
    name: str = MISSING
    sample_rate: int = MISSING
    """The sample rate at which the pre-trained wav2vec2 model was trained and operates.

    Make sure to provide the appropriate value here!

    Can be different from the general NANSY++ sample-rate, in which case
    inputs to the wav2vec2 encoder will be resampled.
    """


@dataclass
class PretrainedWav2Vec2XLSR53Config(PretrainedWav2Vec2Config):
    name: str = "facebook/wav2vec2-large-xlsr-53"
    sample_rate: int = 16000


cs.store(
    group="generator/linguistic_encoder/wav2vec2/pretrained_model",
    # group="pretrained_wav2vec2_model",
    name="large-xlsr-53",
    node=PretrainedWav2Vec2XLSR53Config,
)


@dataclass
class Wav2Vec2Config:
    """Operates at a sampling-frequency independent of the rest of the NANSY++ network"""

    _target_: str = "src.networks.backbone.encoders.wav2vec2.Wav2Vec2Wrapper"
    global_sample_rate: int = II("train_dataset.input_sample_rate")

    pretrained_model: PretrainedWav2Vec2Config = PretrainedWav2Vec2Config()

    layer_for_linguistic_features: int = MISSING
    """The Wav2Vec2 Transformer layer at which hidden linguistic features are extracted from.

    TODO: check the below comment from @revsic, this is not very clear
    NOTE(@revsic, legacy repository): One-base indexing required here because the
    hidden state at dimension 0 corresponds to the position-informed convolution features
    """
    trim_unused_layers: bool = True


@dataclass
class NansyppWav2Vec2Config(Wav2Vec2Config):
    layer_for_linguistic_features: int = 15  # cf. Figure 8.c in NANSY++ [choi2023]
    pretrained_model: PretrainedWav2Vec2Config = PretrainedWav2Vec2XLSR53Config()


cs.store(
    group="generator/linguistic_encoder/wav2vec2",
    name="nansypp-wav2vec2",
    node=NansyppWav2Vec2Config,
)


# all original STFT hyperparameters are UNKNOWN
@dataclass
class MelSTFTConfig:
    _target_: str = "src.networks.misc.transform.MelSpectrogram"

    hop_length: int = 256
    """Hop length, the number of frames between adjacent windows."""

    window_length: int = 1024
    """The length of the windows."""

    mel: int = 80
    """The size of the mel filterbanks."""

    fmin: float = 0
    """The minimum frequency."""

    fmax: Optional[float] = 8000
    """The maximum frequency. If None, uses half of the sample rate as default."""

    sample_rate: int = II("train_dataset.input_sample_rate")
    """The sample rate."""


@dataclass
class CQTConfig:
    _target_: str = "src.networks.misc.transform.ConstantQTransform"

    hop_length: int = 256
    """The number of samples between adjacent frame."""

    fmin: float = 32.7
    """The minimum frequency."""

    bins: int = 191
    """The number of output bins."""

    bins_per_octave: int = 24
    """The number of frequency bins per octave."""

    sample_rate: int = II("train_dataset.input_sample_rate")
    """The sampling rate."""


@dataclass
class FrameLevelSynthesizerConfig:
    _target_: str = (
        "src.networks.backbone.synthesizers.framelevel.FrameLevelSynthesizer"
    )

    in_channels: int = II("generator.linguistic_encoder.hidden_channels")
    """Size of the input channels."""

    kernels: int = 3
    """Size of the convolutional kernels."""

    dilations: List[int] = field(default_factory=lambda: [1, 3, 9, 27, 1, 3, 9, 27])
    """Dilation rates."""

    blocks: int = 2
    """Number of the 1x1 ConvGLU blocks after dilated ConvGLU."""

    leak: float = II("generator.leak")
    """Negative slope of the leaky ReLUs."""

    dropout_rate: float = II("generator.dropout_rate")
    """Dropout rates."""

    timbre_embedding_channels: int = II("generator.timbre_encoder.out_channels")
    """Size of the time-varying timbre embeddings."""


@dataclass
class LinguisticEncoderConfig:
    _target_: str = "src.networks.backbone.encoders.linguistic.LinguisticEncoder"
    wav2vec2: Wav2Vec2Config = Wav2Vec2Config()

    hidden_channels: int = 128
    """Size of the hidden channels."""

    preconv_blocks: int = 2
    """Number of pre-convolution blocks."""

    convglu_kernel_sizes: List[int] = field(
        default_factory=lambda: [3, 3, 3, 3, 3, 3, 3, 3, 1, 1]
    )  # 3 * 8 + [1] * 2
    """Size of the ConvGLU kernels."""

    leak: float = II("generator.leak")
    """Negative slope of leaky ReLUs."""

    dropout_rate: float = II("generator.dropout_rate")
    """Dropout rate."""


@dataclass
class PitchEncoderConfig:
    _target_: str = "src.networks.backbone.encoders.pitch.PitchEncoder"

    freq: int = 160
    """Number of frequency bins."""

    min_pitch: float = 50
    """The minimum predicted pitch."""

    max_pitch: float = 1000
    """The maximum predicted pitch."""

    prekernels: int = 7
    """Size of the first convolutional kernels."""

    kernels: int = 3
    """Size of the frequency-convolution kernels."""

    channels: int = 128
    """Size of the channels."""

    blocks: int = 2
    """Number of residual blocks."""

    gru_dim: int = 256  # NANSY++ value UNKNOWN
    """Size of the GRU hidden states."""

    hidden_channels: int = 256  # NANSY++ value UNKNOWN
    """Size of the hidden channels."""

    f0_bins: int = 64
    """Size of the output f0-bins."""

    f0_activation: str = "sigmoid"
    """F0 activation function."""


@dataclass
class SynthesizerConfig:
    _target_: str = "src.networks.backbone.synthesizers.synthesizer.Synthesizer"

    channels: int = 64
    kernels: int = 3
    dilation_rate: int = 2
    layers: int = 10
    cycles: int = 3
    input_sample_rate: int = II("train_dataset.input_sample_rate")
    output_sample_rate: int = II("train_dataset.output_sample_rate")
    scale: int = II("generator.cqt.hop_length")
    aux: int = II("generator.linguistic_encoder.hidden_channels")


@dataclass
class AttentiveStatisticsPoolingConfig:
    _target_: str = "src.networks.backbone.encoders.timbre.AttentiveStatisticsPooling"

    channels: int = II("generator.timbre_encoder.hidden_channels")
    bottleneck: int = II("generator.timbre_encoder.bottleneck")


@dataclass
class TimbreEncoderConfig:
    _target_: str = "src.networks.backbone.encoders.timbre.TimbreEncoder"

    in_channels: int = II("generator.mel_spectrogram_transform.mel")
    """Size of the input channels."""

    out_channels: int = 192
    """Size of the output embeddings."""

    channels: int = 512
    """Size of the major states."""

    prekernels: int = 5
    """Size of the convolutional kernels before feed to SERes2Block."""

    scale: int = 8
    """The number of the resolutions, for SERes2Block."""

    kernels: int = 3
    """Size of the convolutional kernels, for SERes2Block."""

    dilations: List[int] = field(default_factory=lambda: [2, 3, 4])
    """Dilation factors."""

    bottleneck: int = 128
    """Size of the bottleneck layers, both SERes2Block and AttentiveStatisticsPooling."""

    # NANSY++: 3072
    hidden_channels: int = 1536
    """Size of the hidden channels for attentive statistics pooling."""

    latent: int = 512
    """Size of the timbre latent query vectors."""

    timbre: int = 128
    """Size of the timbre tokens."""

    tokens: int = 50
    """The number of timbre tokens."""

    # UNKNOWN
    heads: int = 8
    """The number of attention heads, for timbre token block."""

    linguistic_encoder_hidden_channels: int = II(
        "generator.linguistic_encoder.hidden_channels"
    )
    """The size of the features returned by the linguistic encoder."""

    # UNKNOWN
    slerp: float = 0.5
    """Weight value for spherical interpolation."""

    mel_spectrogram_transform: MelSTFTConfig = II("generator.mel_spectrogram_transform")

    attentive_statistics_pooling: AttentiveStatisticsPoolingConfig = (
        AttentiveStatisticsPoolingConfig()
    )


def compute_cqt_center(cqt_bins: int, pitch_encoder_freq: int) -> int:
    return (cqt_bins - pitch_encoder_freq) // 2


OmegaConf.register_new_resolver("cqt_center", compute_cqt_center)


# -------- GENERATOR -------- #
@dataclass
class GeneratorConfig:
    _target_: str = "src.networks.backbone.generator.Generator"

    input_sample_rate: int = II("train_dataset.input_sample_rate")
    output_sample_rate: int = II("train_dataset.output_sample_rate")
    """Input and output sample-rate for the generator"""

    # UNKNOWN, default negative-slope of nn.LeakyReLU units
    leak: float = 0.01

    # UNKNOWN, default dropout rate of nn.Transformer
    dropout_rate: float = 0.1

    cqt_center: int = II(
        "cqt_center:${generator.cqt.bins},${generator.pitch_encoder.freq}"
    )

    cqt: CQTConfig = CQTConfig()
    mel_spectrogram_transform: MelSTFTConfig = MelSTFTConfig()
    frame_level_synthesizer: FrameLevelSynthesizerConfig = FrameLevelSynthesizerConfig()
    linguistic_encoder: LinguisticEncoderConfig = LinguisticEncoderConfig()
    pitch_encoder: PitchEncoderConfig = PitchEncoderConfig()
    synthesizer: SynthesizerConfig = SynthesizerConfig()
    timbre_encoder: TimbreEncoderConfig = TimbreEncoderConfig()
    information_perturbator: InformationPerturbatorConfig = MISSING


# -------- DISCRIMINATOR -------- #
@dataclass
class DiscriminatorConfig:
    _target_: str = MISSING


@dataclass
class MultiPeriodDiscriminatorConfig(DiscriminatorConfig):
    _target_: str = "src.networks.misc.mpd.MultiPeriodDiscriminator"

    periods: List[int] = field(default_factory=lambda: [2, 3, 5, 7, 11])
    channels: List[int] = field(default_factory=lambda: [32, 128, 512, 1024])
    kernels: int = 5
    stride: int = 3
    postkernels: int = 3
    leak: float = 0.1


# -------- MODEL -------- #
@dataclass
class ModelConfig:
    _target_: str = "src.models.backbone.Backbone"
    input_sample_rate: float = II("train_dataset.input_sample_rate")
    output_sample_rate: float = II("train_dataset.output_sample_rate")
    accumulate_grad_batches: int = 1
    # gradient norm logging
    grad_norm_logging_interval_batches: int = 100
    # gradient clipping
    gradient_clip_val: Optional[float] = None
    gradient_clip_algorithm: Optional[str] = None


# -------- OPTIMIZERS -------- #
@dataclass
class OptimizerConfig:
    _target_: str = "torch.optim.Adam"

    lr: float = 1e-4
    betas: List[float] = field(default_factory=lambda: [0.9, 0.999])


@dataclass
class GeneratorOptimizerConfig(OptimizerConfig):
    pass


@dataclass
class DiscriminatorOptimizerConfig(OptimizerConfig):
    lr: float = 2e-4


# -------- GENERATOR LOSS -------- #
@dataclass
class ContrastiveLossConfig:
    _target_: str = "src.losses.backbone.generator.ContrastiveLoss"

    negative_samples_minimum_distance_to_positive: int = 5
    # UNKNOWN
    num_candidates: int = 15
    temperature: float = 0.1


@dataclass
class ReconstructionLossConfig:
    _target_: str = MISSING


@dataclass
class MelL1LossConfig(ReconstructionLossConfig):
    _target_: str = "src.losses.backbone.generator.MelL1Loss"
    mel_spectrogram_transform: MelSTFTConfig = II("generator.mel_spectrogram_transform")


@dataclass
class MultiScaleFFTLossConfig(ReconstructionLossConfig):
    _target_: str = "src.losses.backbone.generator.MultiScaleFFTLoss"

    # objective
    # 16khz sample-rate,
    # default window_sizes=[1920, 320, 80], hop_sizes=[640, 80, 40] in NSF
    window_sizes: List[int] = field(default_factory=lambda: [2048, 512, 128])
    hop_sizes: List[int] = field(default_factory=lambda: [512, 128, 32])


@dataclass
class PitchLossConfig:
    pass


@dataclass
class RelativePitchDifferenceLossConfig(PitchLossConfig):
    _target_: str = "src.losses.backbone.generator.RelativePitchDifferenceLoss"

    # pitch consistency
    cqt_center: int = II("generator.cqt_center")
    cqt_shift_min: int = -12
    cqt_shift_max: int = 12

    # huber norm
    # UNKNOWN
    sigma: float = 0.5

    pitch_freq: int = II("generator.pitch_encoder.freq")

    # delta: 1.0
    # (0.25 * sigma) as stated in Spice paper and issue: https://github.com/revsic/torch-nansypp/issues/1
    delta: float = 0.125  # 0.25 * 0.5
    pitch_encoder = II("generator.pitch_encoder")


cs.store(
    group="generator_loss/pitch_prediction_loss",
    name="relative_pitch_difference_loss",
    node=RelativePitchDifferenceLossConfig,
)


@dataclass
class PitchSupervisionLossConfig(PitchLossConfig):
    _target_: str = "src.losses.backbone.generator.PitchSupervisionLoss"
    sample_rate: int = II("train_dataset.input_sample_rate")
    cqt_hop_length: int = II("generator.cqt.hop_length")


cs.store(
    group="generator_loss/pitch_prediction_loss",
    name="pitch_supervision_loss",
    node=PitchSupervisionLossConfig,
)


@dataclass
class GeneratorLossConfig:
    _target_: str = "src.losses.backbone.generator.GeneratorLoss"

    input_sample_rate: int = II("train_dataset.input_sample_rate")
    output_sample_rate: int = II("train_dataset.output_sample_rate")

    cqt_bins_per_octave: int = II("generator.cqt.bins_per_octave")

    segment_length_s: float = II("train_dataset.segment_length_s")

    # content loss warmup
    linguistic_loss_start_weight: float = 1e-5
    linguistic_loss_end_weight: int = 10

    pitch_prediction_loss: PitchLossConfig = MISSING
    contrastive_loss: ContrastiveLossConfig = ContrastiveLossConfig()
    reconstruction_losses: List[ReconstructionLossConfig] = field(
        default_factory=lambda: [MelL1LossConfig(), MultiScaleFFTLossConfig()]
    )

    non_scalar_logging_steps: int = 5000
    vc_source_audio_paths: Dict[str, str] = field(
        default_factory=lambda: {"p225": "static/samples/vctk/p225_001.wav"}
    )
    vc_target_audio_paths: Dict[str, str] = field(
        default_factory=lambda: {
            "p226": "static/samples/vctk/p226_002.wav",
            "p227": "static/samples/vctk/p227_003.wav",
            "p228": "static/samples/vctk/p228_004.wav",
        }
    )

    # required for logging plots
    cqt_hop_length: int = II("generator.cqt.hop_length")


# -------- LOSS DISCRIMINATOR -------- #
@dataclass
class DiscriminatorLossConfig:
    _target_: str = "src.losses.backbone.discriminator.LeastSquaresDiscriminatorLoss"


# -------- LOGGING -------- #
@dataclass
class LoggerConfig:
    _target_: str = "pytorch_lightning.loggers.tensorboard.TensorBoardLogger"
    name: str = "default"
    save_dir: str = "log/"
    default_hp_metric: bool = False


class LightiningTrainerPrecision(StringEnum):
    """
    Valid pytorch-lightining training precisions as of April 2023.

    See: https://github.com/Lightning-AI/lightning/pull/16783#issue-1587848352
    """

    mixed_16 = "16-mixed"
    mixed_bf16 = "bf16-mixed"
    full_32 = "32-true"
    double_64 = "64-true"


# -------- TRAIN -------- #
@dataclass
class TrainerConfig:
    _target_: str = "pytorch_lightning.Trainer"
    default_root_dir: str = "../static"

    # maximum optimizer steps - put 2*N to aim for N training steps as there are 2 optimizers
    max_steps: int = 2 * 1000000
    val_check_interval: Optional[int] = None

    limit_train_batches: Optional[int] = None
    limit_val_batches: Optional[int] = None

    benchmark: Optional[bool] = False
    precision: LightiningTrainerPrecision = LightiningTrainerPrecision.full_32


@dataclass
class CriterionConfig:
    pass


@dataclass
class MetricConfig:
    pass


@dataclass
class CallbackConfig:
    _target_: str = MISSING
    logging_interval: str = MISSING


@dataclass
class LRMonitorCallbackConfig(CallbackConfig):
    _target_: str = "pytorch_lightning.callbacks.LearningRateMonitor"
    logging_interval: str = "epoch"


@dataclass
class CallbacksConfig:
    lr_monitor: CallbackConfig = field(default_factory=LRMonitorCallbackConfig)


@dataclass
class ModelCheckpointConfig:
    _target_: str = "pytorch_lightning.callbacks.ModelCheckpoint"
    dirpath: str = "checkpoints/"
    # step is the optimizer step - divide by 2 to access training step
    filename: str = "opt-steps={step}"
    every_n_train_steps: int = 50000
    save_top_k: int = -1
    verbose: bool = False


@dataclass
class ResumeConfig:
    checkpoint_path: str = MISSING


@dataclass
class LightningProfilerConfig:
    _target_: str = MISSING
    filename: str = MISSING


@dataclass
class PassThroughLightningProfilerConfig(LightningProfilerConfig):
    _target_: str = "pytorch_lightning.profilers.PassThroughProfiler"
    filename: str = "no_profiler"


@dataclass
class SimpleLightningProfilerConfig(LightningProfilerConfig):
    _target_: str = "pytorch_lightning.profilers.SimpleProfiler"
    filename: str = "simple_profiler"


@dataclass
class AdvancedLightningProfilerConfig(LightningProfilerConfig):
    _target_: str = "pytorch_lightning.profilers.AdvancedProfiler"
    filename: str = "advanced_profiler"


@dataclass
class PytorchLightningProfilerConfig(LightningProfilerConfig):
    _target_: str = "pytorch_lightning.profilers.PyTorchProfiler"
    filename: str = "pytorch_profiler"
    with_stack = True


# Hydra-related configuration


class RunDirBase(StringEnum):
    RUNS_BACKBONE = "runs_backbone"
    RUNS_TTS = "runs_tts"
    DEBUG = "debug"


@dataclass
class Config:
    """
    NANSY++ Hydra baseline config.
    """

    model_name: str = "baseline"
    experiment: str = (
        f'{II("model_name")}-{II("dataset.name")}-{II("generator.sample_rate")}hz'
    )

    train_dataset: DatasetConfig = MISSING
    val_dataset: Optional[DatasetConfig] = None
    datamodule: DataModuleConfig = DataModuleConfig()
    perturbator: InformationPerturbatorConfig = InformationPerturbatorConfig()
    generator: GeneratorConfig = GeneratorConfig()
    discriminator: DiscriminatorConfig = MultiPeriodDiscriminatorConfig()
    generator_optimizer: GeneratorOptimizerConfig = GeneratorOptimizerConfig()
    discriminator_optimizer: DiscriminatorOptimizerConfig = (
        DiscriminatorOptimizerConfig()
    )
    generator_loss: GeneratorLossConfig = GeneratorLossConfig()
    discriminator_loss: DiscriminatorLossConfig = DiscriminatorLossConfig()
    model: ModelConfig = ModelConfig()
    model_checkpoint: ModelCheckpointConfig = ModelCheckpointConfig()
    logger: LoggerConfig = LoggerConfig()
    trainer: TrainerConfig = TrainerConfig()

    profiler: LightningProfilerConfig = MISSING

    callbacks: Dict[str, CallbackConfig] = field(
        default_factory=lambda: {"lr_monitor": LRMonitorCallbackConfig()}
    )
    criteria: Dict[str, CriterionConfig] = field(default_factory=lambda: {})
    metrics: Dict[str, MetricConfig] = field(default_factory=lambda: {})

    seed: Optional[int] = None
    resume: Optional[ResumeConfig] = None

    verbose: bool = False

    # hydra-configurations
    run_dir_base: RunDirBase = RunDirBase.RUNS_BACKBONE


cs.store(name="nansypp_config", node=Config)

cs.store(group="datamodule", name="base", node=DataModuleConfig)

cs.store(name="nansypp_generator", node=GeneratorConfig)
cs.store(
    group="generator/information_perturbator",
    name="nansypp_information_perturbator",
    node=InformationPerturbatorConfig,
)
cs.store(
    group="generator/information_perturbator",
    name="bypassed_information_perturbator",
    node=BypassedInformationPerturbatorConfig,
)
cs.store(group="generator/cqt", name="nansypp_cqt", node=CQTConfig)
cs.store(
    group="generator/frame_level_synthesizer",
    name="nansypp_frame_level_synthesizer",
    node=FrameLevelSynthesizerConfig,
)
cs.store(
    group="generator/linguistic_encoder",
    name="nansypp_linguistic_encoder",
    node=LinguisticEncoderConfig,
)
cs.store(
    group="generator/loss",
    name="nansypp_generator_loss",
    node=GeneratorLossConfig,
)
cs.store(
    group="generator/mel_spectrogram_transform",
    name="nansypp_mel_spectrogram_transform",
    node=MelSTFTConfig,
)
cs.store(
    group="generator/optimizer",
    name="nansypp_generator_optimizer",
    node=GeneratorOptimizerConfig,
)
cs.store(
    group="generator/pitch_encoder",
    name="nansypp_pitch_encoder",
    node=PitchEncoderConfig,
)
cs.store(
    group="generator/timbre_encoder",
    name="nansypp_timbre_encoder",
    node=TimbreEncoderConfig,
)
cs.store(
    group="generator/wav2vec2_wrapper",
    name="nansypp_wav2vec2_wrapper",
    node=Wav2Vec2Config,
)

# discriminator
cs.store(
    name="nansypp_discriminator",
    node=DiscriminatorConfig,
)
cs.store(
    group="discriminator/loss",
    name="nansypp_discriminator_loss",
    node=DiscriminatorLossConfig,
)
cs.store(
    group="discriminator/optimizer",
    name="nansypp_discriminator_optimizer",
    node=DiscriminatorOptimizerConfig,
)

cs.store(group="logger", name="base_logger", node=LoggerConfig)

cs.store(group="profiler", name="disabled", node=PassThroughLightningProfilerConfig)
cs.store(group="profiler", name="simple", node=SimpleLightningProfilerConfig)
cs.store(group="profiler", name="advanced", node=AdvancedLightningProfilerConfig)
cs.store(group="profiler", name="pytorch", node=PytorchLightningProfilerConfig)

src/dataclasses/backbone.py
MWM-io/nansypp
from dataclasses import dataclass, field
from typing import Dict, List, Optional

from hydra.core.config_store import ConfigStore
from omegaconf import II, MISSING

from src.dataclasses.backbone import (
    CallbackConfig,
    LightiningTrainerPrecision,
    LoggerConfig,
    LRMonitorCallbackConfig,
    PretrainedWav2Vec2Config,
    PretrainedWav2Vec2XLSR53Config,
    ResumeConfig,
    RunDirBase,
    Wav2Vec2Config,
)

cs = ConfigStore.instance()


# -------- DATA -------- #


@dataclass
class CmuTextProcessorConfig:
    _target_: str = "src.data.text.text_processor.TextProcessor"
    sample_rate: int = II("datamodule.train_dataset.sample_rate")


@dataclass
class ValDatasetConfig:
    _target_: str = "src.data.dataset.tts_base.TTSBaseDataset"
    dataset_dir: str = MISSING
    descriptor_file: str = MISSING
    alignment_file: str = MISSING
    sample_rate: int = MISSING
    text_processor: CmuTextProcessorConfig = CmuTextProcessorConfig()


@dataclass
class TrainDatasetConfig(ValDatasetConfig):
    _target_: str = "src.data.dataset.tts_train.TTSTrainDataset"
    dataset_dir: str = II("data_path")
    text_processor: CmuTextProcessorConfig = CmuTextProcessorConfig()
    max_seconds_per_batch: int = 128
    features_lengths_file: str = "static/tts/features_lengths.json"


@dataclass
class DataModuleConfig:
    _target_: str = "src.data.datamodule.tts.TextToSpeechDataModule"
    train_dataset: TrainDatasetConfig = TrainDatasetConfig()
    val_dataset: Optional[ValDatasetConfig] = ValDatasetConfig()
    batch_size: int = 1
    pin_memory: bool = True
    num_workers: int = 16
    num_val_workers: int = 0
    shuffle: bool = False


cs.store(group="datamodule", name="base", node=DataModuleConfig)


# -------- NETWORK -------- #
@dataclass
class PhonemeEncoderConfig:
    _target_: str = "src.networks.tts.encoders.phoneme.PhonemeEncoder"
    conv_in_channels: int = 1024
    conv_out_channels: int = 1024
    conv_kernel_size: int = 5
    conv_stride: int = 1
    dropout: float = II("dropout")
    n_convs: int = 3
    conditionning_channels: int = II("network.style_encoder.out_linear")
    nhead: int = 2
    num_encoder_layers: int = 2
    num_decoder_layers: int = 2
    dim_feedforward: int = 1024
    n_transformers: int = 3
    negative_slope: float = 0.01
    out_linear: int = 128
    num_labels: int = 122
    conv_middle_channels: int = 1024


@dataclass
class NansyppTTSWav2Vec2Config(Wav2Vec2Config):
    layer_for_linguistic_features: int = 3  # cf. Figure 11.c in NANSY++ [choi2023]
    pretrained_model: PretrainedWav2Vec2Config = PretrainedWav2Vec2XLSR53Config()


cs.store(
    group="generator/linguistic_encoder/wav2vec2",
    name="nansypp-wav2vec2-tts",
    node=NansyppTTSWav2Vec2Config,
)


@dataclass
class StyleEncoderConfig:
    _target_: str = "src.networks.tts.encoders.style.StyleEncoder"
    conv_in_channels: int = 1024
    conv_out_channels: int = 1024
    conv_kernel_size: int = 5
    conv_stride: int = 2
    dropout: float = II("dropout")
    n_convs: int = 3
    bottleneck: int = 128
    wav2vec2: Wav2Vec2Config = NansyppTTSWav2Vec2Config(
        global_sample_rate=II("datamodule.train_dataset.sample_rate")
    )
    out_linear: int = 128
    conv_middle_channels: int = 1024


@dataclass
class DurationPredictorConfig:
    _target_: str = "src.networks.tts.decoders.duration.DurationPredictor"
    conv_in_channels: int = 128
    conv_out_channels: int = 128
    conv_kernel_size: int = 5
    conv_stride: int = 1
    style_dim: int = II("network.style_encoder.out_linear")
    dropout: float = II("dropout")
    n_convs: int = 2
    conv_middle_channels: int = 128


@dataclass
class LinguisticDecoderConfig:
    _target_: str = "src.networks.tts.decoders.linguistic.LinguisticDecoder"
    first_conv_in_channels: int = 128
    first_conv_out_channels: int = 128
    first_conv_kernel_size: int = 5
    first_conv_stride: int = 1
    dropout: float = II("dropout")
    n_transformer_blocks: int = 2
    transformer_nhead: int = 2
    transformer_num_encoder_layers: int = 2
    transformer_num_decoder_layers: int = 2
    transformer_dim_feedforward: int = 1024
    cln_convs_middle_channels: int = 128
    cln_convs_out_channels: int = 128
    cln_convs_kernel_size: int = 1
    cln_convs_stride: int = 1
    style_dim: int = II("network.style_encoder.out_linear")
    n_cln_convs: int = 3


@dataclass
class AmplitudeDecoderConfig:
    _target_: str = "src.networks.tts.decoders.amplitude.AmplitudeDecoder"
    conv_in_channels: int = 128
    conv_out_channels: int = 128
    conv_kernel_size: int = 7
    conv_stride: int = 1
    style_dim: int = II("network.style_encoder.out_linear")
    dropout: float = II("dropout")
    n_convs: int = 3
    conv_middle_channels: int = 128


@dataclass
class F0DecoderConfig:
    _target_: str = "src.networks.tts.decoders.f0.F0Decoder"
    conv_in_channels: int = 128
    conv_out_channels: int = 128
    conv_kernel_size: int = 7
    conv_stride: int = 1
    style_dim: int = II("network.style_encoder.out_linear")
    dropout: float = II("dropout")
    n_convs: int = 2
    gru_hidden_size: int = 64
    gru_num_layers: int = 2
    conv_middle_channels: int = 128


@dataclass
class PitchStatsConfig:
    _target_: str = "src.data.datamodule.tts.PitchStats"
    stats_dir: str = II("data_path")
    file: str = "pitch_stats.pt"


# -------- NETWORK -------- #
@dataclass
class NetworkConfig:
    _target_: str = "src.networks.tts.tts.TextToSpeechNetwork"
    phoneme_encoder: PhonemeEncoderConfig = PhonemeEncoderConfig()
    style_encoder: StyleEncoderConfig = StyleEncoderConfig()
    duration_predictor: DurationPredictorConfig = DurationPredictorConfig()
    linguistic_decoder: LinguisticDecoderConfig = LinguisticDecoderConfig()
    amplitude_decoder: AmplitudeDecoderConfig = AmplitudeDecoderConfig()
    f0_decoder: F0DecoderConfig = F0DecoderConfig()
    pitch_stats: PitchStatsConfig = PitchStatsConfig()
    backbone_info_file: str = field(
        default_factory=lambda: II("data_path") + "backbone_info.json"
    )
    sample_rate: int = II("datamodule.train_dataset.sample_rate")
    text_processor: CmuTextProcessorConfig = CmuTextProcessorConfig()


# -------- OPTIMIZERS -------- #
@dataclass
class OptimizerConfig:
    _target_: str = "torch.optim.Adam"

    lr: float = 1e-4
    betas: List[float] = field(default_factory=lambda: [0.9, 0.999])


@dataclass
class SchedulerConfig:
    _target_: str = "torch.optim.lr_scheduler.ExponentialLR"
    gamma: float = (1e-2) ** (1 / 20)


# -------- LOSS -------- #
@dataclass
class LossConfig:
    _target_: str = "src.losses.tts.TextToSpeechLoss"
    pitch_weight: float = 1
    duration_weight: float = 1
    linguistic_weight: float = 1
    linguistic_multiplicator: float = 100


# -------- MODEL -------- #
@dataclass
class ModelConfig:
    _target_: str = "src.models.tts.TextToSpeechModel"
    grad_norm_logging_interval_batches: int = 100
    cer_texts_file: Optional[str] = "static/tts/generation_test.json"


# -------- TRAIN -------- #
@dataclass
class ModelCheckpointConfig:
    _target_: str = "pytorch_lightning.callbacks.ModelCheckpoint"
    dirpath: str = "checkpoints/"
    filename: str = "steps={step}"
    every_n_train_steps: int = 1000
    save_top_k: int = -1
    verbose: bool = False


@dataclass
class TrainerConfig:
    _target_: str = "pytorch_lightning.Trainer"
    default_root_dir: str = "../static"

    val_check_interval: int = 1000
    max_steps: int = 200000

    limit_train_batches: Optional[int] = None
    limit_val_batches: Optional[int] = None

    benchmark: Optional[bool] = False
    precision: LightiningTrainerPrecision = LightiningTrainerPrecision.full_32

    # gradient clipping
    gradient_clip_val: Optional[float] = None
    gradient_clip_algorithm: Optional[str] = None


@dataclass
class Config:
    data_path: str
    model_name: str = "tts_baseline"
    experiment: str = "test"

    seed: Optional[int] = None
    datamodule: DataModuleConfig = DataModuleConfig()

    network: NetworkConfig = NetworkConfig()
    loss: LossConfig = LossConfig()
    optimizer: OptimizerConfig = OptimizerConfig()
    scheduler: Optional[SchedulerConfig] = SchedulerConfig()

    model: ModelConfig = ModelConfig()
    model_checkpoint: ModelCheckpointConfig = ModelCheckpointConfig()
    logger: LoggerConfig = LoggerConfig()
    callbacks: Dict[str, CallbackConfig] = field(
        default_factory=lambda: {"lr_monitor": LRMonitorCallbackConfig()}
    )
    trainer: TrainerConfig = TrainerConfig()
    resume: Optional[ResumeConfig] = None
    run_dir_base: RunDirBase = RunDirBase.RUNS_TTS
    dropout: float = 0.0


cs.store(name="nansy_tts_config", node=Config)

src/dataclasses/tts.py
MWM-io/nansypp
import os
import warnings
from typing import Tuple

import hydra
import hydra.utils as hu
import numpy as np
import pyrootutils
import soundfile as sf
import torch
import torchaudio
import torchaudio.transforms as T
import typer
from omegaconf import DictConfig
from typing_extensions import Annotated

root = pyrootutils.setup_root(__file__, dotenv=True, pythonpath=True, cwd=False)

import src.dataclasses.backbone
from src.networks.backbone.generator import Generator

app = typer.Typer()


warnings.filterwarnings("ignore")


class BackboneInferencer:
    generator: Generator
    config: DictConfig
    input_sr: int
    output_sr: int

    def __init__(
        self,
        exp_dir: str,
        checkpoint_name: str,
        device: str,
    ):
        """
        Args:
            exp_dir: experiment directory. Should be an absolute path (Hydra constraint).
            checkpoint_name: checkpoint file name.
            device: device on which inferencer should run. Example: "cpu", "cuda", "cuda:0".
        """
        self.device = device
        self.instantiate_generator(exp_dir)

        checkpoint_path = os.path.join(exp_dir, "checkpoints", checkpoint_name)
        self.load_generator_weights(checkpoint_path)

    def instantiate_generator(self, exp_dir: str) -> None:
        config_dir = os.path.join(exp_dir, ".hydra")
        hydra.core.global_hydra.GlobalHydra.get_state().clear()
        hydra.core.global_hydra.GlobalHydra.instance().clear()
        with hydra.initialize_config_dir(
            config_dir=config_dir
            if os.path.isabs(config_dir)
            else os.path.join(root, config_dir)
        ):
            config = hydra.compose(config_name="config")
        self.input_sr = config.train_dataset.input_sample_rate
        self.output_sr = config.train_dataset.output_sample_rate
        if "additive_noise" in config.generator.information_perturbator:
            config.generator.information_perturbator.additive_noise = None
        self.generator: Generator = hu.instantiate(config.generator)
        self.config = config

    def load_generator_weights(self, checkpoint_path: str) -> None:
        # load weights
        state_dict = torch.load(checkpoint_path, map_location="cpu")["state_dict"]
        generator_state_dict = {
            k.replace("generator.", ""): v
            for k, v in state_dict.items()
            if "generator." in k and "generator_loss" not in k
        }
        inconsistent_keys = self.generator.load_state_dict(
            generator_state_dict, strict=False
        )
        assert (
            len(
                [
                    key
                    for key in inconsistent_keys.missing_keys
                    if ("wav2vec2" not in key) and ("cqt" not in key)
                ]
            )
            == 0
        )
        _ = self.generator.to(self.device)
        _ = self.generator.eval()

    def voice_conversion(
        self, source_audio: torch.Tensor, target_audio: torch.Tensor
    ) -> Tuple[np.ndarray, int]:
        """Apply target audio voice timbre to source audio content.
        Args:
            source_audio: [torch.Tensor; [1, Ns]] tensor with mono audio content at sampling rate defined by model config.
            target_audio: [torch.Tensor; [1, Nt]] tensor with mono audio content at sampling rate defined by model config.
        """
        return (
            self.generator.voice_conversion(
                source_audio=source_audio,
                target_audio=target_audio,
                device=torch.device(self.device),
                f0_statistics=True,
                enable_information_perturbator=False,
            )
            .cpu()
            .squeeze(0)
            .numpy(),
            self.output_sr,
        )

    def reconstruct(self, audio: torch.Tensor) -> Tuple[np.ndarray, int]:
        """Reconstruct audio signal through features extraction followed by audio synthesis.
        Args:
            audio: [torch.Tensor; [1, Ns]] tensor with mono audio content at sampling rate defined by model config.
        """
        return self.voice_conversion(audio, audio)


def main(
    exp_dir: Annotated[str, typer.Argument(help="path to experiment directory")],
    checkpoint_name: Annotated[str, typer.Argument(help="checkpoint name")],
    device: Annotated[
        str, typer.Argument(help="device name: cpu, cuda:0, cuda:1, etc.")
    ] = "cuda",
    source_audio_path: Annotated[
        str, typer.Argument(help="path to source raw audio")
    ] = "static/samples/vctk/p225_001.wav",
    target_audio_path: Annotated[
        str, typer.Argument(help="path to target raw audio")
    ] = "static/samples/vctk/p226_002.wav",
    output_dir: Annotated[
        str, typer.Argument(help="path to directory where to save generated audio")
    ] = "static/tmp",
):
    os.makedirs(output_dir, exist_ok=True)
    inferencer = BackboneInferencer(
        exp_dir if os.path.isabs(exp_dir) else os.path.join(root, exp_dir),
        checkpoint_name,
        device,
    )
    source_audio, source_sr = torchaudio.load(source_audio_path)
    source_audio = T.Resample(source_sr, inferencer.input_sr)(source_audio)
    target_audio, target_sr = torchaudio.load(target_audio_path)
    target_audio = T.Resample(target_sr, inferencer.output_sr)(target_audio)
    converted, sr = inferencer.voice_conversion(source_audio, target_audio)
    output_path = os.path.join(
        output_dir,
        f"{os.path.splitext(os.path.basename(source_audio_path))[0]}-to-{os.path.splitext(os.path.basename(target_audio_path))[0]}.wav",
    )
    print(f"Synthesized audio saved at: {output_path}")
    sf.write(output_path, converted, sr)


if __name__ == "__main__":
    typer.run(main)

src/inference/backbone.py
MWM-io/nansypp
import os
import warnings

import hydra
import hydra.utils as hu
import pyrootutils
import torch
import torchaudio
import torchaudio.transforms as T
import typer
from typing_extensions import Annotated

root = pyrootutils.setup_root(__file__, dotenv=True, pythonpath=True, cwd=False)

from src.inference.backbone import BackboneInferencer
from src.networks.tts.tts import TextToSpeechNetwork

app = typer.Typer()

warnings.filterwarnings("ignore")


class TextToSpeechInferencer:
    def __init__(
        self,
        device="cuda:0",
    ):
        super().__init__()
        self.device = device
        self.backbone_exp_dir = None
        self.backbone_checkpoint_name = None
        self.tts_exp_dir = None
        self.tts_checkpoint_path = None
        self.generator = None
        self.tts_network = None

    def instantiate_backbone(self, exp_dir: str, checkpoint_name: str) -> None:
        if (exp_dir, checkpoint_name) != (
            self.backbone_exp_dir,
            self.backbone_checkpoint_name,
        ):
            self.backbone_exp_dir = exp_dir
            self.backbone_checkpoint_name = checkpoint_name
            self.generator = BackboneInferencer(
                exp_dir,
                checkpoint_name,
                self.device,
            ).generator

    def instantiate_backbone_from_tts(self) -> None:
        self.instantiate_backbone(
            self.tts_network.backbone_info["exp_dir"],
            self.tts_network.backbone_info["ckpt_name"],
        )

    def instantiate_tts(self, exp_dir: str) -> None:
        if not os.path.isabs(exp_dir):
            exp_dir = os.path.join(root, exp_dir)
        if exp_dir == self.tts_exp_dir:
            return
        self.tts_exp_dir = exp_dir
        hydra.core.global_hydra.GlobalHydra.get_state().clear()
        hydra.core.global_hydra.GlobalHydra.instance().clear()
        with hydra.initialize_config_dir(config_dir=os.path.join(exp_dir, ".hydra")):
            config = hydra.compose(config_name="config")
        self.tts_network: TextToSpeechNetwork = hu.instantiate(config.network)

    def load_tts_weights(self, checkpoint_name: str) -> None:
        tts_checkpoint_path = os.path.join(
            self.tts_exp_dir, "checkpoints", checkpoint_name
        )
        if self.tts_checkpoint_path == tts_checkpoint_path:
            return
        self.tts_checkpoint_path = tts_checkpoint_path
        state_dict = torch.load(self.tts_checkpoint_path, map_location="cpu")[
            "state_dict"
        ]
        tts_state_dict = {k.replace("network.", ""): v for k, v in state_dict.items()}

        inconsistent_keys = self.tts_network.load_state_dict(
            tts_state_dict, strict=False
        )
        assert (
            len(
                [
                    key
                    for key in inconsistent_keys.missing_keys
                    if "wav2vec2" not in key and "w2v_3rd" not in key
                ]
            )
            == 0
        )
        _ = self.tts_network.to(self.device)
        _ = self.tts_network.eval()

    def __call__(
        self,
        text,
        style_audio,
        phoneme_duration_f_sequence=None,
    ):
        if phoneme_duration_f_sequence is not None:
            phoneme_duration_f_sequence = phoneme_duration_f_sequence.reshape(
                (1, -1)
            ).to(self.device)

        phoneme_sequence = (
            self.tts_network.text_processor.text_to_encoded_phoneme_sequence(text).to(
                self.device
            )
        )

        _, tts_decoding_features = self.tts_network(
            phoneme_sequence.reshape((1, -1)),
            style_audio.reshape((1, -1)).to(self.device),
            phoneme_duration_f_sequence,
        )
        timbre_global, timbre_bank = self.generator.analyze_timbre(style_audio)
        _, signal = self.generator.synthesize(
            tts_decoding_features["pitch"].squeeze(1),
            tts_decoding_features["p_amp"].squeeze(1),
            tts_decoding_features["ap_amp"].squeeze(1),
            tts_decoding_features["linguistic"],
            timbre_global,
            timbre_bank,
        )
        return signal


def main(
    exp_dir: Annotated[str, typer.Argument(help="path TTS training directory")],
    checkpoint_name: Annotated[str, typer.Argument(help="name of TTS checkpoint")],
    audio_path: Annotated[
        str, typer.Argument(help="path to audio sample of voice to generate")
    ],
    text: Annotated[str, typer.Argument(help="text to generate")],
    out_path: Annotated[str, typer.Argument(help="path to save generated audio")],
    device: Annotated[str, typer.Option("-d", help="text to generate")] = "cuda:0",
):
    out_dir = os.path.dirname(out_path)
    os.makedirs(out_dir, exist_ok=True)
    inferencer = TextToSpeechInferencer(device=device)
    inferencer.instantiate_tts(exp_dir)
    inferencer.load_tts_weights(checkpoint_name)
    inferencer.instantiate_backbone_from_tts()
    audio, sample_rate = torchaudio.load(audio_path)
    if sample_rate != inferencer.generator.input_sample_rate:
        audio = T.Resample(sample_rate, inferencer.generator.input_sample_rate)(audio)
    signal = inferencer(text, audio.to(device))
    torchaudio.save(out_path, signal.cpu(), inferencer.generator.output_sample_rate)
    print(f"Synthesized audio saved at: {out_path}")


if __name__ == "__main__":
    typer.run(main)

src/inference/tts.py
MWM-io/nansypp
from typing import Dict, List, Tuple, Union

import torch
from torch import nn

from src.utilities.profiling import DO_PROFILING, cuda_synchronized_timer
from src.utilities.types import copy_docstring_and_signature

DiscriminatorLossMetrics = Dict[str, Union[float, torch.Tensor]]


class LeastSquaresDiscriminatorLoss(nn.Module):
    """LS-GAN's L2-based GAN discriminator-loss.

    Introduced in Mao et al. ICCV'2017, _Least Squares Generative Adversarial Networks_.
    """

    @cuda_synchronized_timer(DO_PROFILING, prefix="L2DiscriminatorLoss")
    def forward(
        self, logits_fake: List[torch.Tensor], logits_real: List[torch.Tensor]
    ) -> Tuple[torch.Tensor, DiscriminatorLossMetrics]:
        if len(logits_fake) > 0 and len(logits_real) > 0:
            d_fake = logits_fake[0].new_zeros(1)
            d_real = logits_real[0].new_zeros(1)
        else:
            d_fake = torch.zeros(1)
            d_real = torch.zeros(1)

        for logit_fake, logit_real in zip(logits_fake, logits_real):
            d_fake = d_fake + logit_fake.square().mean()
            d_real = d_real + (1 - logit_real).square().mean()

        loss: torch.Tensor = d_fake + d_real
        metrics: DiscriminatorLossMetrics = {
            "disc/loss": loss.detach(),
            "disc/d-real": d_real.mean().detach(),
            "disc/d-fake": d_fake.mean().detach(),
        }
        return loss, metrics

    @copy_docstring_and_signature(forward)
    def __call__(self, *args, **kwargs):  # type: ignore[no-untyped-def]
        return super().__call__(*args, **kwargs)

src/losses/backbone/discriminator.py
MWM-io/nansypp
import itertools
import os
from abc import abstractmethod
from typing import Dict, Iterable, List, Optional, Tuple, TypedDict, Union

import librosa
import librosa.display
import matplotlib.pyplot as plt
import numpy as np
import pyrootutils
import torch
import torch.nn.functional as F
import torchaudio
import torchaudio.transforms as T
from matplotlib.figure import Figure
from torch import nn
from torch.optim import Optimizer

from src.networks.backbone.encoders.pitch import PitchEncoder
from src.networks.backbone.generator import Generator
from src.networks.misc.mpd import MultiPeriodDiscriminator
from src.networks.misc.transform import MelSpectrogram
from src.utilities.crepe import crepe_pitch_estimation
from src.utilities.profiling import DO_PROFILING, cuda_synchronized_timer
from src.utilities.types import copy_docstring_and_signature

LOG_IDX = 0
LOG_MAXLEN = 1.5
LOG_AUDIO = 3
LOG_PLOTS = 2

GeneratorLossMetrics = Dict[str, Union[float, torch.Tensor]]
GeneratorLossImages = Dict[str, np.ndarray]
GeneratorLossAudios = Dict[str, np.ndarray]
GeneratorLossPlots = Dict[str, Figure]

root = pyrootutils.setup_root(__file__, dotenv=True, pythonpath=True, cwd=False)


class ContrastiveMetrics(TypedDict):
    cont_loss: torch.Tensor
    metric_pos: torch.Tensor
    metric_neg: torch.Tensor


class AdversarialDiscriminativeLosses(TypedDict):
    d_fake: torch.Tensor
    features_matching_loss: torch.Tensor


class ContrastiveLoss(nn.Module):
    """Contrastive loss.

    Introduced in Kaizhi Qian et al., _Contentvec: An improved self-supervised speech representation by disentangling speakers_
    """

    def __init__(
        self,
        num_candidates: int,
        negative_samples_minimum_distance_to_positive: int,
        temperature: float,
    ) -> None:
        super().__init__()

        self.negative_samples_minimum_distance_to_positive = (
            negative_samples_minimum_distance_to_positive
        )
        self.temperature = temperature
        self.num_candidates = num_candidates

    def make_negative_sampling_mask(
        self, num_items: int, device: torch.device
    ) -> torch.Tensor:
        upper_triu_mask = torch.triu(
            torch.ones(num_items, num_items),
            self.negative_samples_minimum_distance_to_positive + 1,
        )
        all_mask = upper_triu_mask.T + upper_triu_mask
        random_values = all_mask * torch.rand(num_items, num_items)
        k_th_quant = torch.topk(random_values, min(self.num_candidates, num_items))[0][
            :, -1:
        ]
        random_mask = (random_values >= k_th_quant) * all_mask + torch.eye(num_items)
        return random_mask.to(device)

    @cuda_synchronized_timer(DO_PROFILING, prefix="ContrastiveLoss")
    def forward(
        self, x: torch.Tensor, y: torch.Tensor
    ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
        # [2, B, linguistic_hidden_channels, S], normalize for cosine similarity.
        cosine_similarity = F.normalize(
            torch.stack(
                [x, y],
                dim=0,
            ),
            p=2,
            dim=2,
        )

        # N
        num_tokens = cosine_similarity.shape[-1]
        # [B, N]
        positive = cosine_similarity.prod(dim=0).sum(dim=1) / self.temperature
        # [2, B, N, N]
        confusion_matrix = (
            torch.matmul(cosine_similarity.transpose(2, 3), cosine_similarity)
            / self.temperature
        )

        # [N, N]
        negative_sampling_mask = self.make_negative_sampling_mask(num_tokens, x.device)

        # [2, B, N, N(sum = candidates)], negative case
        masked_confusion_matrix = confusion_matrix.masked_fill(
            ~negative_sampling_mask.to(torch.bool), -np.inf
        )
        # [2, B, N], negative case
        negative = masked_confusion_matrix.exp().sum(dim=-1)
        # []
        contrastive_loss = (
            -torch.sum(positive / negative, dim=2).log().sum(dim=0).mean()
        )
        mean_positive = positive.mean() * self.temperature
        mean_negative = (
            (confusion_matrix * negative_sampling_mask).sum(dim=-1)
            / self.num_candidates
        ).mean() * self.temperature

        return contrastive_loss, mean_positive, mean_negative

    @copy_docstring_and_signature(forward)
    def __call__(self, *args, **kwargs):  # type: ignore[no-untyped-def]
        return super().__call__(*args, **kwargs)


Features = Dict[str, torch.Tensor]


class ReconstructionLossWithFeatures(nn.Module):
    @abstractmethod
    def forward(
        self, input: torch.Tensor, target: torch.Tensor
    ) -> Tuple[torch.Tensor, Features]:
        ...

    @copy_docstring_and_signature(forward)
    def __call__(self, *args, **kwargs):  # type: ignore[no-untyped-def]
        return super().__call__(*args, **kwargs)


class MelL1Loss(ReconstructionLossWithFeatures):
    """Mel spectrogram loss.

    As described in Kong et al., 2020., _Neural source-filter waveform models for sta- tistical parametric speech synthesis._
    """

    def __init__(
        self,
        mel_spectrogram_transform: MelSpectrogram,
    ) -> None:
        super().__init__()
        self.mel_spectrogram_transform = mel_spectrogram_transform

    def forward(
        self, input: torch.Tensor, target: torch.Tensor
    ) -> Tuple[torch.Tensor, Features]:
        mel_reconstructed = self.mel_spectrogram_transform(input)
        mel_original = self.mel_spectrogram_transform(target)

        mel_features = {
            "mel_reconstructed": mel_reconstructed.detach(),
            "mel_original": mel_original.detach(),
        }
        return F.l1_loss(mel_reconstructed, mel_original), mel_features


class MultiScaleFFTLoss(ReconstructionLossWithFeatures):
    """Multi-scale spectrogram loss (MSS).

    As described in Wang et al., 2019, _Hifi-gan: Generative adversarial networks for efficient and high fidelity speech synthesis._
    """

    windows: Iterable[nn.Parameter]

    def __init__(
        self,
        hop_sizes: List[int],
        window_sizes: List[int],
    ) -> None:
        super().__init__()
        self.hop_sizes = hop_sizes
        self.window_sizes = window_sizes
        self.windows = nn.ParameterList(
            [
                nn.Parameter(torch.hann_window(window_size), requires_grad=False)
                for window_size in self.window_sizes
            ]
        )

    def forward(
        self, input: torch.Tensor, target: torch.Tensor
    ) -> Tuple[torch.Tensor, Features]:
        # [B x 2, T]
        mss_loss = input.new_zeros(1)

        # batch-process input and target together
        input_and_target = torch.cat([input, target], dim=0)
        for hop_size, window_size, window in zip(
            self.hop_sizes, self.window_sizes, self.windows
        ):
            # [B x 2, win // 2 + 1, T / hop]
            stft = torch.stft(
                input_and_target,
                window_size,
                hop_size,
                window=window,
                return_complex=True,
            )
            # [B, win // 2 + 1, T / hop]
            mag_stft_input, mag_stft_target = stft.abs().chunk(2, dim=0)
            # []
            mss_loss = mss_loss + F.mse_loss(mag_stft_input, mag_stft_target)
        return mss_loss, {}


class AggregatedReconstructionLossWithFeatures(ReconstructionLossWithFeatures):
    def __init__(self, losses: List[ReconstructionLossWithFeatures]) -> None:
        super().__init__()
        self.losses = nn.ModuleList(losses)

    @cuda_synchronized_timer(DO_PROFILING, prefix="AggregatedReconstructionLoss")
    def forward(
        self, input: torch.Tensor, target: torch.Tensor
    ) -> Tuple[torch.Tensor, Features]:
        aggregated_loss = input.new_zeros(1)
        aggregated_features: Features = {}
        for loss_function in self.losses:
            loss, features = loss_function(input, target)
            aggregated_loss += loss
            aggregated_features.update(features)
        return aggregated_loss, aggregated_features


class RelativePitchDifferenceLoss(nn.Module):
    """Relative pitch difference loss.

    Introduced in Gfeller et al., 2020, _Spice: Self-supervised pitch estimation_
    """

    def __init__(
        self,
        cqt_center: int,
        cqt_shift_min: int,
        cqt_shift_max: int,
        pitch_freq: int,
        delta: float,
        sigma: float,
    ) -> None:
        super().__init__()
        self.cqt_center = cqt_center
        self.cqt_shift_min = cqt_shift_min
        self.cqt_shift_max = cqt_shift_max
        self.pitch_freq = pitch_freq
        self.delta = delta
        self.sigma = sigma

    @cuda_synchronized_timer(DO_PROFILING, prefix="RelativePitchDifferenceLoss")
    def forward(
        self,
        cqt: torch.Tensor,
        pitch: torch.Tensor,
        pitch_encoder: PitchEncoder,
        bsize: int,
    ) -> torch.Tensor:
        random_transposition = torch.randint(
            self.cqt_shift_min,
            self.cqt_shift_max + 1,  # for inclusive range
            (bsize,),
            device=cqt.device,
        )
        # real start index
        start = random_transposition + self.cqt_center
        # sampled
        transposed_cqt = torch.stack(
            [cqt_[i : i + self.pitch_freq] for cqt_, i in zip(cqt, start)]
        )
        # [B, N]
        predicted_transposed_pitch, _, _, _ = pitch_encoder(transposed_cqt)

        # pitch consistency
        return F.huber_loss(
            predicted_transposed_pitch.log2()
            + self.sigma * random_transposition[:, None],
            pitch.log2(),
            delta=self.delta,
        )

    @copy_docstring_and_signature(forward)
    def __call__(self, *args, **kwargs):  # type: ignore[no-untyped-def]
        return super().__call__(*args, **kwargs)


class PitchSupervisionLoss(nn.Module):
    """Mean-square error loss between predicted pitches and ground truth pitches."""

    def __init__(self, sample_rate: int, cqt_hop_length: int):
        super().__init__()
        self.sample_rate = sample_rate
        self.cqt_hop_length = cqt_hop_length

    def forward(
        self, audios: torch.Tensor, cqts: torch.Tensor, pred_pitches: torch.Tensor
    ) -> torch.Tensor:
        device = pred_pitches.device
        gt_pitches = []
        remove_pred_indices = []
        for idx, audio in enumerate(audios):
            gt_pitch = crepe_pitch_estimation(
                audio.unsqueeze(0),
                self.sample_rate,
                self.cqt_hop_length,
                device,
            )
            # handle crepe potential crash
            if gt_pitch is not None:
                gt_pitches.append(
                    crepe_pitch_estimation(
                        audio.unsqueeze(0),
                        self.sample_rate,
                        self.cqt_hop_length,
                        device,
                    )
                )
            else:
                remove_pred_indices.append(idx)
        gt_pitches = torch.cat(gt_pitches, dim=0).to(pred_pitches.device)[
            ..., : cqts.shape[-1]
        ]
        # delete in reverse order to not throw off subsequent indexes in list
        for idx in sorted(remove_pred_indices, reverse=True):
            del pred_pitches[indx]
        return nn.functional.mse_loss(pred_pitches, gt_pitches)

    @copy_docstring_and_signature(forward)
    def __call__(self, *args, **kwargs):  # type: ignore[no-untyped-def]
        return super().__call__(*args, **kwargs)


class GeneratorLoss(nn.Module):
    """Class to compute losses of generator modules."""

    def __init__(
        self,
        input_sample_rate: int,
        output_sample_rate: int,
        segment_length_s: float,
        cqt_hop_length: int,
        cqt_bins_per_octave: int,
        linguistic_loss_start_weight: float,
        linguistic_loss_end_weight: int,
        non_scalar_logging_steps: int,
        vc_source_audio_paths: Dict[str, str],
        vc_target_audio_paths: Dict[str, str],
        contrastive_loss: ContrastiveLoss,
        reconstruction_losses: List[ReconstructionLossWithFeatures],
        pitch_prediction_loss: RelativePitchDifferenceLoss,
        generator: Generator,
        discriminator: MultiPeriodDiscriminator,
        generator_optimizer: Optimizer,
        discriminator_optimizer: Optimizer,
        initial_training_step: int,
        max_training_step: int,
    ):
        super().__init__()

        self.input_sample_rate = input_sample_rate
        self.output_sample_rate = output_sample_rate
        self.cqt_hop_length = cqt_hop_length
        self.cqt_bins_per_octave = cqt_bins_per_octave
        self.segment_length_s = segment_length_s
        self.linguistic_loss_start_weight = linguistic_loss_start_weight
        self.linguistic_loss_end_weight = linguistic_loss_end_weight
        self.non_scalar_logging_steps = non_scalar_logging_steps
        self.vc_source_audio_paths = vc_source_audio_paths
        self.vc_target_audio_paths = vc_target_audio_paths
        self.max_training_step = max_training_step
        self.initial_training_step = initial_training_step

        self.content_weight_increment = (
            linguistic_loss_end_weight - linguistic_loss_start_weight
        ) / max_training_step
        self.content_weight = (
            linguistic_loss_start_weight
            + initial_training_step * self.content_weight_increment
        )

        self.contrastive_loss = contrastive_loss
        self.reconstruction_losses = AggregatedReconstructionLossWithFeatures(
            reconstruction_losses
        )
        self.pitch_prediction_loss = pitch_prediction_loss
        self.generator = generator
        self.discriminator = discriminator
        self.generator_optimizer = generator_optimizer
        self.discriminator_optimizer = discriminator_optimizer

        self.cmap = np.array(plt.get_cmap("viridis").colors)
        self.source_dict = {}
        for src_audio_name, src_audio_path in vc_source_audio_paths.items():
            source_audio, src_sr = torchaudio.load(os.path.join(root, src_audio_path))
            self.source_dict[src_audio_name] = T.Resample(
                src_sr, self.input_sample_rate
            )(source_audio).float()
        self.target_dict = {}
        for tgt_audio_name, tgt_audio_path in vc_target_audio_paths.items():
            if tgt_audio_path is not None:
                target_audio, tgt_sr = torchaudio.load(
                    os.path.join(root, tgt_audio_path)
                )
                self.target_dict[tgt_audio_name] = T.Resample(
                    tgt_sr, self.output_sample_rate
                )(target_audio).float()

    def update_warmup(self) -> None:
        """Update the content loss weights."""
        self.content_weight += self.content_weight_increment

    @cuda_synchronized_timer(DO_PROFILING, prefix="GeneratorLoss")
    def linguistic_loss(
        self,
        audio: torch.Tensor,
        linguistic_features: torch.Tensor,
        perturbed_audio_secondary: Optional[torch.Tensor] = None,
    ) -> Tuple[torch.Tensor, ContrastiveMetrics]:
        # [B, linguistic_hidden_channels, S]
        linguistic_features_alternate_perturbations = self.generator.analyze_linguistic(
            audio,
            enable_information_perturbator=True,
            perturbed_inputs=perturbed_audio_secondary,
        )
        contrastive_loss, positive_loss, negative_loss = self.contrastive_loss(
            linguistic_features, linguistic_features_alternate_perturbations
        )
        return contrastive_loss, {
            "cont_loss": contrastive_loss.detach(),
            "metric_pos": positive_loss.detach(),
            "metric_neg": negative_loss.detach(),
        }

    @cuda_synchronized_timer(DO_PROFILING, prefix="GeneratorLoss")
    def discriminative_loss(
        self, audio: torch.Tensor, synth: torch.Tensor
    ) -> AdversarialDiscriminativeLosses:
        (
            logits_per_period_reconstructed,
            feature_maps_per_period_reconstructed,
        ) = self.discriminator(synth)
        _, feature_maps_per_period_original = self.discriminator(audio)

        discriminator_fake_loss = audio.new_zeros(1)
        features_matching_loss = audio.new_zeros(1)

        # discriminator loss
        for logits_reconstructed in logits_per_period_reconstructed:
            discriminator_fake_loss = (
                discriminator_fake_loss + (1 - logits_reconstructed).square().mean()
            )

        # feature-matching loss
        for fmap_f, fmap_r in zip(
            feature_maps_per_period_reconstructed, feature_maps_per_period_original
        ):
            for ff, fr in zip(fmap_f, fmap_r):
                features_matching_loss = features_matching_loss + (ff - fr).abs().mean()

        return {
            "d_fake": discriminator_fake_loss,
            "features_matching_loss": features_matching_loss,
        }

    @cuda_synchronized_timer(DO_PROFILING, prefix="GeneratorLoss")
    def forward(
        self,
        audio: torch.Tensor,
        synth: torch.Tensor,
        pitch: torch.Tensor,
        cqt: torch.Tensor,
        linguistic_features: torch.Tensor,
        it: int,
        perturbed_audio_secondary: Optional[torch.Tensor] = None,
    ) -> Tuple[
        torch.Tensor,
        GeneratorLossMetrics,
        GeneratorLossImages,
        GeneratorLossAudios,
        GeneratorLossPlots,
    ]:
        device = audio.device
        bsize, _ = audio.shape

        (
            aggregated_reconstruction_loss,
            reconstruction_features,
        ) = self.reconstruction_losses(audio, synth)

        if isinstance(self.pitch_prediction_loss, RelativePitchDifferenceLoss):
            pitch_loss = self.pitch_prediction_loss(
                cqt=cqt,
                pitch=pitch,
                pitch_encoder=self.generator.pitch_encoder,
                bsize=bsize,
            )
        elif isinstance(self.pitch_prediction_loss, PitchSupervisionLoss):
            pitch_loss = self.pitch_prediction_loss(
                audios=audio, cqts=cqt, pred_pitches=pitch
            )

        (
            linguistic_contrastive_loss,
            linguistic_contrastive_metrics,
        ) = self.linguistic_loss(
            audio,
            linguistic_features=linguistic_features,
            perturbed_audio_secondary=perturbed_audio_secondary,
        )

        discriminative_losses = self.discriminative_loss(audio, synth)

        # NOTE: this rough scaling was introduced in the legacy repository (@revsic) and appears
        # to be taken from the paper GANSpeech by Yang et al., published at INTERSPEECH'21
        # Paper available at: https://arxiv.org/abs/2106.15153
        # An open-source implementation of GANSpeech is not available, so this implementation
        # with a .detach() call is just hypothetical.
        weight = (
            aggregated_reconstruction_loss
            / discriminative_losses["features_matching_loss"]
        ).detach()

        self.update_warmup()

        loss = (
            discriminative_losses["d_fake"]
            + weight * discriminative_losses["features_matching_loss"]
            + aggregated_reconstruction_loss
            + self.content_weight * linguistic_contrastive_loss
        )
        if pitch_loss.requires_grad:
            loss += pitch_loss

        metrics: GeneratorLossMetrics = {
            "gen/loss": loss.detach().item(),
            "gen/d-fake": discriminative_losses["d_fake"].detach().item(),
            "gen/fmap": discriminative_losses["features_matching_loss"].detach().item(),
            "gen/rctor": aggregated_reconstruction_loss.detach().item(),
            "gen/pitch": pitch_loss.detach().item(),
            "gen/cont": linguistic_contrastive_loss.detach().item(),
            "metric/cont-pos": linguistic_contrastive_metrics["metric_pos"]
            .detach()
            .item(),
            "metric/cont-neg": linguistic_contrastive_metrics["metric_neg"]
            .detach()
            .item(),
            "common/warmup": self.content_weight,
            "common/weight": weight.item(),
            "common/learning-rate-g": self.generator_optimizer.param_groups[0]["lr"],
            "common/learning-rate-d": self.discriminator_optimizer.param_groups[0][
                "lr"
            ],
        }

        images: GeneratorLossImages = {}
        audios: GeneratorLossAudios = {}
        plots: GeneratorLossPlots = {}
        mode = "train" if self.generator.training else "validation"
        if it % self.non_scalar_logging_steps == 0:
            images = {
                f"mel-gt/{mode}": self.mel_img(
                    reconstruction_features["mel_original"]
                    .detach()
                    .cpu()
                    .numpy()[LOG_IDX]
                ),
                f"mel-synth/{mode}": self.mel_img(
                    reconstruction_features["mel_reconstructed"]
                    .detach()
                    .cpu()
                    .numpy()[LOG_IDX]
                ),
                f"log-cqt/{mode}": self.mel_img(
                    cqt.clamp_min(1e-5).log().cpu().detach().numpy()[LOG_IDX]
                ),
            }
            audios = {
                f"speech/{mode}": audio.detach().cpu().numpy()[LOG_IDX],
                f"synth/{mode}": synth.detach().cpu().numpy()[LOG_IDX],
            }
            if mode == "validation":
                for (src_audio_name, src_audio), (
                    tgt_audio_name,
                    tgt_audio,
                ) in itertools.product(
                    self.source_dict.items(), self.target_dict.items()
                ):
                    audios.update(
                        {
                            f"conversion/vctk-{src_audio_name}-to-{tgt_audio_name}": self.generator.voice_conversion(
                                src_audio, tgt_audio, device
                            )
                            .squeeze(0)
                            .cpu()
                            .numpy()
                        }
                    )
            plots = {
                f"pitch/{mode}_{i}": self.plot_pitch(cqt_item, pitch_item)
                for i, (cqt_item, pitch_item) in enumerate(
                    zip(
                        cqt.cpu().numpy()[:LOG_PLOTS, ...],
                        pitch.cpu().detach().numpy()[:LOG_PLOTS, ...],
                    )
                )
            }
        return loss, metrics, images, audios, plots

    @cuda_synchronized_timer(DO_PROFILING, prefix="GeneratorLoss")
    def mel_img(self, mel: np.ndarray) -> np.ndarray:
        """Generate mel-spectrogram images.
        Args:
            mel: [np.float32; [mel, T]], mel-spectrogram.
        Returns:
            [np.float32; [3, mel, T]], mel-spectrogram in viridis color map.
        """
        # minmax norm in range(0, 1)
        mel = (mel - mel.min()) / (mel.max() - mel.min() + 1e-7)
        # in range(0, 255)
        mel = (mel * 255).astype(np.uint8)
        # [mel, T, 3]
        mel = self.cmap[mel]
        # [3, mel, T], make origin lower
        mel = np.flip(mel, axis=0).transpose(2, 0, 1)
        return mel

    @cuda_synchronized_timer(DO_PROFILING, prefix="GeneratorLoss")
    def plot_pitch(self, cqt: np.ndarray, pitch: np.ndarray) -> Figure:
        fig, ax = plt.subplots()
        img = librosa.display.specshow(
            librosa.amplitude_to_db(cqt, ref=np.max),
            sr=self.input_sample_rate,
            x_axis="time",
            y_axis="cqt_hz",
            hop_length=self.cqt_hop_length,
            fmin=32.7,
            bins_per_octave=self.cqt_bins_per_octave,
            ax=ax,
        )
        plt.plot(
            np.arange(0, self.segment_length_s, self.segment_length_s / len(pitch)),
            pitch,
            color="green",
            label="f0",
        )
        ax.legend()
        fig.colorbar(img, ax=ax, format="%+2.0f dB")
        return fig

    @copy_docstring_and_signature(forward)
    def __call__(self, *args, **kwargs):  # type: ignore[no-untyped-def]
        return super().__call__(*args, **kwargs)

src/losses/backbone/generator.py
MWM-io/nansypp
from typing import Dict, Tuple, Union

import torch
from torch import nn

TextToSpeechLossMetrics = Dict[str, Union[float, torch.Tensor]]


class TextToSpeechLoss(nn.Module):
    """Text-to-speech loss module."""

    def __init__(
        self,
        pitch_weight: float,
        linguistic_weight: float,
        duration_weight: float,
        linguistic_multiplicator: float = 100,
    ):
        """Initializer."""
        super().__init__()

        self.pitch_weight = pitch_weight
        self.linguistic_weight = linguistic_weight
        self.duration_weight = duration_weight
        self.linguistic_multiplicator = linguistic_multiplicator
        self.mae = nn.L1Loss()
        self.mse = nn.MSELoss()

    @staticmethod
    def min_max_norm(value: torch.Tensor, min_max: Dict) -> torch.Tensor:
        return (value - min_max["min"]) / (min_max["max"] - min_max["min"])

    def pitch_loss(
        self,
        prediction: torch.Tensor,
        target: torch.Tensor,
        mask: torch.Tensor,
        stats: Dict,
    ) -> torch.Tensor:
        return self.mae(
            self.min_max_norm(prediction, stats) * mask,
            self.min_max_norm(target, stats) * mask,
        )

    def forward(
        self,
        pred_duration: torch.Tensor,
        decoded_linguistic_features: torch.Tensor,
        decoded_p_amp: torch.Tensor,
        decoded_ap_amp: torch.Tensor,
        decoded_pitch: torch.Tensor,
        ref_duration: torch.Tensor,
        ref_linguistic_features: torch.Tensor,
        ref_p_amp: torch.Tensor,
        ref_ap_amp: torch.Tensor,
        ref_pitch: torch.Tensor,
        pitch_stats: Dict,
        mode: str,
    ) -> Tuple[torch.Tensor, TextToSpeechLossMetrics]:
        linguistic_mask = ref_linguistic_features != 0
        pitch_mask = ref_pitch != 0
        duration_mask = ref_duration != 0

        linguistic_loss = self.mse(
            self.linguistic_multiplicator
            * decoded_linguistic_features
            * linguistic_mask,
            self.linguistic_multiplicator * ref_linguistic_features * linguistic_mask,
        )
        p_amp_loss = self.pitch_loss(
            decoded_p_amp, ref_p_amp, pitch_mask, pitch_stats["p_amp"]
        )
        ap_amp_loss = self.pitch_loss(
            decoded_ap_amp, ref_ap_amp, pitch_mask, pitch_stats["ap_amp"]
        )
        pitch_loss = self.pitch_loss(
            decoded_pitch, ref_pitch, pitch_mask, pitch_stats["pitch"]
        )
        duration_loss = self.mse(
            pred_duration * duration_mask,
            ref_duration.to(torch.float32) * duration_mask,
        )

        loss = (
            self.linguistic_weight * linguistic_loss
            + self.pitch_weight * (p_amp_loss + ap_amp_loss + pitch_loss)
            + self.duration_weight * duration_loss
        )
        metrics: TextToSpeechLossMetrics = {
            f"tts/{mode}/loss": loss.detach().item(),
            f"tts/{mode}/linguistic": linguistic_loss.detach().item(),
            f"tts/{mode}/p_amp": p_amp_loss.detach().item(),
            f"tts/{mode}/ap_amp": ap_amp_loss.detach().item(),
            f"tts/{mode}/pitch": pitch_loss.detach().item(),
            f"tts/{mode}/duration": duration_loss.detach().item(),
        }

        return loss, metrics

src/losses/tts.py
MWM-io/nansypp
from typing import Any, Dict, List, Mapping, Optional, Tuple, Union

import numpy as np
import pytorch_lightning as pl
import torch
from pytorch_lightning.loggers.tensorboard import TensorBoardLogger
from pytorch_lightning.utilities.grads import grad_norm as lightning_grad_norm
from torch import nn
from torch.utils.tensorboard.writer import SummaryWriter

from src.data.datamodule.backbone import AudioData
from src.losses.backbone.discriminator import LeastSquaresDiscriminatorLoss
from src.losses.backbone.generator import (
    GeneratorLoss,
    GeneratorLossAudios,
    GeneratorLossImages,
    GeneratorLossMetrics,
    GeneratorLossPlots,
)
from src.networks.backbone.generator import Generator
from src.networks.misc.mpd import MultiPeriodDiscriminator
from src.utilities.profiling import DO_PROFILING, cuda_synchronized_timer


class Backbone(pl.LightningModule):
    def __init__(
        self,
        input_sample_rate: int,
        output_sample_rate: int,
        generator: Generator,
        discriminator: MultiPeriodDiscriminator,
        generator_optimizer: torch.optim.Optimizer,
        discriminator_optimizer: torch.optim.Optimizer,
        generator_loss: GeneratorLoss,
        discriminator_loss: LeastSquaresDiscriminatorLoss,
        grad_norm_logging_interval_batches: int,
        gradient_clip_val: Optional[float],
        gradient_clip_algorithm: Optional[str],
        accumulate_grad_batches: int,
    ):
        super().__init__()
        self.input_sample_rate = input_sample_rate
        self.output_sample_rate = output_sample_rate
        self.generator = generator
        self.discriminator = discriminator
        self.generator_optimizer = generator_optimizer
        self.discriminator_optimizer = discriminator_optimizer
        self.generator_loss = generator_loss
        self.discriminator_loss = discriminator_loss
        self.configure_optimizers()
        self.automatic_optimization = False
        self.grad_norm_logging_interval_batches = grad_norm_logging_interval_batches
        self.gradient_clip_val = gradient_clip_val
        self.gradient_clip_algorithm = gradient_clip_algorithm
        self.accumulate_grad_batches = accumulate_grad_batches

    def configure_optimizers(
        self,
    ) -> Tuple[List[torch.optim.Optimizer], List[torch.optim.lr_scheduler.LRScheduler]]:
        return [self.generator_optimizer, self.discriminator_optimizer], []

    def training_step(
        self,
        batch: AudioData,
        batch_idx: int,
    ) -> None:
        src_audios = batch["src_audios"]
        tgt_audios = batch["tgt_audios"]
        perturbed_audio_1 = batch.get("perturbed_audio_1", None)
        perturbed_audio_2 = batch.get("perturbed_audio_2", None)
        precomputed_cqt = batch.get("cqt", None)
        precomputed_mel = batch.get("mel_spec", None)

        generator_optimizer, discriminator_optimizer = self.optimizers()

        # Synthesize "fake" data
        synth, analysis_features, _ = self.generator(
            src_audios,
            perturbed_inputs=perturbed_audio_1,
            enable_information_perturbator=True,
            precomputed_cqt=precomputed_cqt,
            precomputed_mel=precomputed_mel,
        )

        # Train discriminator
        self.toggle_optimizer(discriminator_optimizer)
        discriminator_optimizer.zero_grad()

        logits_real, _ = self.discriminator(tgt_audios)
        logits_fake, _ = self.discriminator(synth.detach())

        discriminator_loss, discriminator_metrics = self.discriminator_loss(
            logits_fake, logits_real
        )

        self.manual_backward(discriminator_loss)

        # gradient clipping needs to be performed manually in case of manual optimization
        if (
            self.gradient_clip_val is not None
            and self.gradient_clip_algorithm is not None
        ):
            self.clip_gradients(
                discriminator_optimizer,
                gradient_clip_val=self.gradient_clip_val,
                gradient_clip_algorithm=self.gradient_clip_algorithm,
            )

        self._track_grad_norm(
            self.discriminator,
            discriminator_metrics,
            batch_index=batch_idx,
            prefix="grad-norm/disc/",
        )
        self.log_dict(discriminator_metrics, sync_dist=False, rank_zero_only=True)

        if (batch_idx + 1) % self.accumulate_grad_batches == 0:
            discriminator_optimizer.step()
            discriminator_optimizer.zero_grad()
        self.untoggle_optimizer(discriminator_optimizer)

        # Train generator
        self.toggle_optimizer(generator_optimizer)
        generator_optimizer.zero_grad()

        # We have only updated the discriminator in the previous steps,
        # so the fake data remains up-to-date
        (
            aggregated_generator_loss,
            generator_metrics,
            generator_images,
            generator_audios,
            generator_plots,
        ) = self.generator_loss(
            audio=tgt_audios,
            synth=synth,
            pitch=analysis_features["pitch"],
            cqt=analysis_features["cqt"],
            linguistic_features=analysis_features["linguistic"],
            it=batch_idx,
            perturbed_audio_secondary=perturbed_audio_2,
        )

        self.manual_backward(aggregated_generator_loss)

        # gradient clipping needs to be performed manually in case of manual optimization
        if (
            self.gradient_clip_val is not None
            and self.gradient_clip_algorithm is not None
        ):
            self.clip_gradients(
                generator_optimizer,
                gradient_clip_val=self.gradient_clip_val,
                gradient_clip_algorithm=self.gradient_clip_algorithm,
            )

        self._track_grad_norm(
            self.generator,
            generator_metrics,
            batch_index=batch_idx,
            prefix="grad-norm/gen/",
        )

        self.log_all(
            generator_metrics, generator_images, generator_audios, generator_plots
        )

        if (batch_idx + 1) % self.accumulate_grad_batches == 0:
            generator_optimizer.step()
            generator_optimizer.zero_grad()
        self.untoggle_optimizer(generator_optimizer)

    def validation_step(
        self,
        batch: AudioData,
        batch_idx: int,
    ) -> None:
        src_audios = batch["src_audios"]
        tgt_audios = batch["tgt_audios"]
        perturbed_audio_2 = batch.get("perturbed_audio_2", None)

        # Evaluate generator
        synth, analysis_features, _ = self.generator(
            src_audios,
            enable_information_perturbator=False,
        )
        _, generator_metrics, images, audios, plots = self.generator_loss(
            audio=tgt_audios,
            synth=synth,
            pitch=analysis_features["pitch"],
            cqt=analysis_features["cqt"],
            linguistic_features=analysis_features["linguistic"],
            it=batch_idx,
            perturbed_audio_secondary=perturbed_audio_2,
        )
        self.log_all(generator_metrics, images, audios, plots)

        # Evaluate discriminator
        logits_real, _ = self.discriminator(tgt_audios)
        logits_fake, _ = self.discriminator(synth)
        _, discriminator_metrics = self.discriminator_loss(logits_fake, logits_real)
        self.log_dict(discriminator_metrics, sync_dist=True)

    @cuda_synchronized_timer(DO_PROFILING, prefix="Backbone")
    def log_all(
        self,
        losses: GeneratorLossMetrics,
        images: GeneratorLossImages,
        audios: GeneratorLossAudios,
        plots: GeneratorLossPlots,
    ) -> None:
        assert self.logger is not None, "you didn't initialize a logger"
        assert isinstance(
            self.logger, TensorBoardLogger
        ), "expects self.logger to be a TensorBoardLogger"

        tensorboard: SummaryWriter = self.logger.experiment
        self.log_dict(losses, sync_dist=False, rank_zero_only=True)
        for name, image in images.items():
            tensorboard.add_image(name, image, int(self.global_step / 2))
        for name, audio in audios.items():
            tensorboard.add_audio(
                name,
                audio,
                int(self.global_step / 2),
                sample_rate=self.output_sample_rate,
            )
        for name, plot in plots.items():
            tensorboard.add_figure(name, plot, int(self.global_step / 2))

    def load_state_dict(self, state_dict: Mapping[str, Any], strict: bool = True):
        if strict:
            missing_keys = super().load_state_dict(state_dict=state_dict, strict=False)
            not_wav2vec2_keys = [
                key for key in missing_keys[0] if "wav2vec2" not in key
            ]
            if len(not_wav2vec2_keys) == 0:
                raise Exception(f"Missing state dict keys: {not_wav2vec2_keys}")
        else:
            return super().load_state_dict(state_dict=state_dict, strict=strict)

    def _track_grad_norm(
        self,
        model: nn.Module,
        metrics_dict: Dict[str, Union[float, torch.Tensor]],
        batch_index: int,
        prefix: str = "",
    ) -> None:
        """
        Gotchas:
        - Must be careful to only retrieve norm of the gradients AFTER the backward pass!
        - Retrieving this is pretty slow, so do not compute this everytime!
        """
        if batch_index % self.grad_norm_logging_interval_batches != 0:
            return
        (
            grad_norms,
            mean_param_norm,
        ) = Backbone._mean_grad_and_param_norms(model)
        if prefix != "":
            grad_norms = {prefix + key: value for key, value in grad_norms.items()}
        metrics_dict.update(grad_norms)
        metrics_dict[f"{prefix}param-norm"] = mean_param_norm

    @staticmethod
    @cuda_synchronized_timer(DO_PROFILING, prefix="Backbone")
    @torch.no_grad()
    def _mean_grad_and_param_norms(
        model: nn.Module,
    ) -> Tuple[Dict[str, float], float]:
        grad_norms = lightning_grad_norm(model, 2)

        parameters = model.parameters()
        param_norm = np.mean(
            [torch.norm(p).item() for p in parameters if p.dtype == torch.float32]
        ).item()
        return grad_norms, param_norm

src/models/backbone.py
MWM-io/nansypp
from typing import Any, Dict, Mapping, Optional, Tuple, Union

import numpy as np
import pytorch_lightning as pl
import torch
from pytorch_lightning.loggers.tensorboard import TensorBoardLogger
from pytorch_lightning.utilities.grads import grad_norm as lightning_grad_norm
from torch import nn

from src.data.datamodule.tts import TextToSpeechData
from src.inference.tts import TextToSpeechInferencer
from src.losses.tts import TextToSpeechLoss
from src.networks.tts.tts import TextToSpeechNetwork
from src.utilities.profiling import DO_PROFILING, cuda_synchronized_timer
from src.utilities.tts_evaluation import TTSEvaluationModule


class TextToSpeechModel(pl.LightningModule):
    """NANSY-TTS model."""

    def __init__(
        self,
        network: TextToSpeechNetwork,
        optimizer: torch.optim.Optimizer,
        loss: TextToSpeechLoss,
        scheduler: Optional[torch.optim.lr_scheduler.LRScheduler],
        grad_norm_logging_interval_batches: int,
        cer_texts_file: Optional[str],
    ):
        """Initializer."""
        super().__init__()

        self.network = network
        self.optimizer = optimizer
        self.scheduler = scheduler
        self.loss = loss

        self.grad_norm_logging_interval_batches = grad_norm_logging_interval_batches

        self.inferencer = None
        self.tensorboard = None
        self.cer_texts_file = cer_texts_file
        self.tts_cer_module = None

    def training_step(self, batch: TextToSpeechData, batch_idx: int) -> torch.Tensor:
        (
            phoneme_sequence,
            audio_phoneme_alignment,
            backbone_analysis_features,
            randomly_sliced_audio,
        ) = (
            batch["phoneme_sequence"],
            batch["phoneme_duration_f_sequence"],
            batch["backbone_analysis_features"],
            batch["style_audio"],
        )
        upsample_lens = {
            "cqt_len": (backbone_analysis_features["pitch"] != 0).sum(axis=1),
            "linguistic_len": (
                backbone_analysis_features["linguistic"][:, 0, :] != 0
            ).sum(axis=1),
        }
        encoding_features, decoding_features = self.network(
            phoneme_sequence,
            randomly_sliced_audio,
            audio_phoneme_alignment,
            upsample_lens,
        )
        loss, metrics = self.loss(
            encoding_features["duration"].squeeze(1),
            decoding_features["linguistic"],
            decoding_features["p_amp"].squeeze(1),
            decoding_features["ap_amp"].squeeze(1),
            decoding_features["pitch"].squeeze(1),
            audio_phoneme_alignment,
            backbone_analysis_features["linguistic"],
            backbone_analysis_features["p_amp"],
            backbone_analysis_features["ap_amp"],
            backbone_analysis_features["pitch"],
            self.network.pitch_stats,
            "training",
        )
        self._track_grad_norm(
            self.network,
            metrics,
            batch_index=batch_idx,
            prefix="grad-norm/",
        )
        self.log_dict(metrics, sync_dist=False, rank_zero_only=True)

        return loss

    def validation_step(self, batch: TextToSpeechData, batch_idx: int) -> torch.Tensor:
        (
            phoneme_sequence,
            audio_phoneme_alignment,
            backbone_analysis_features,
            randomly_sliced_audio,
        ) = (
            batch["phoneme_sequence"],
            batch["phoneme_duration_f_sequence"],
            batch["backbone_analysis_features"],
            batch["style_audio"],
        )
        upsample_lens = {
            "cqt_len": (backbone_analysis_features["pitch"] != 0).sum(axis=1),
            "linguistic_len": (
                backbone_analysis_features["linguistic"][:, 0, :] != 0
            ).sum(axis=1),
        }
        encoding_features, decoding_features = self.network(
            phoneme_sequence,
            randomly_sliced_audio,
            audio_phoneme_alignment,
            upsample_lens,
        )
        loss, metrics = self.loss(
            encoding_features["duration"].squeeze(1),
            decoding_features["linguistic"],
            decoding_features["p_amp"].squeeze(1),
            decoding_features["ap_amp"].squeeze(1),
            decoding_features["pitch"].squeeze(1),
            audio_phoneme_alignment,
            backbone_analysis_features["linguistic"],
            backbone_analysis_features["p_amp"],
            backbone_analysis_features["ap_amp"],
            backbone_analysis_features["pitch"],
            self.network.pitch_stats,
            "validation",
        )
        self.log_dict(metrics, sync_dist=False, rank_zero_only=True)
        return loss

    def configure_optimizers(self) -> Any:
        if self.scheduler:
            return [self.optimizer], [
                {"scheduler": self.scheduler, "interval": "epoch"}
            ]
        return self.optimizer

    def _track_grad_norm(
        self,
        model: nn.Module,
        metrics_dict: Dict[str, Union[float, torch.Tensor]],
        batch_index: int,
        prefix: str = "",
    ) -> None:
        """
        Gotchas:
        - Must be careful to only retrieve norm of the gradients AFTER the backward pass!
        - Retrieving this is pretty slow, so do not compute this everytime!
        """
        if batch_index % self.grad_norm_logging_interval_batches != 0:
            return
        (
            grad_norms,
            mean_param_norm,
        ) = TextToSpeechModel._mean_grad_and_param_norms(model)
        if prefix != "":
            grad_norms = {prefix + key: value for key, value in grad_norms.items()}
        metrics_dict.update(grad_norms)
        metrics_dict[f"{prefix}param-norm"] = mean_param_norm

    @staticmethod
    @cuda_synchronized_timer(DO_PROFILING, prefix="TextToSpeechModel")
    @torch.no_grad()
    def _mean_grad_and_param_norms(
        model: nn.Module,
    ) -> Tuple[Dict[str, float], float]:
        grad_norms = lightning_grad_norm(model, 2)

        parameters = model.parameters()
        param_norm = np.mean(
            [torch.norm(p).item() for p in parameters if p.dtype == torch.float32]
        ).item()
        return grad_norms, param_norm

    def load_state_dict(self, state_dict: Mapping[str, Any], strict: bool = True):
        if strict:
            missing_keys = super().load_state_dict(state_dict=state_dict, strict=False)
            not_wav2vec2_keys = [key for key in missing_keys[0] if "w2v_3rd" not in key]
            assert (
                len(not_wav2vec2_keys) == 0
            ), f"Following keys are missing in the checkpoint: {not_wav2vec2_keys}"
            return self
        return super().load_state_dict(state_dict=state_dict, strict=strict)

    def on_validation_epoch_end(self):
        if not self.cer_texts_file:
            return

        assert self.logger is not None, "you didn't initialize a logger"
        assert isinstance(
            self.logger, TensorBoardLogger
        ), "expects self.logger to be a TensorBoardLogger"

        if self.tensorboard is None:
            self.tensorboard = self.logger.experiment
        if self.tts_cer_module is None:
            self.tts_cer_module = TTSEvaluationModule(
                evaluation_data=self.cer_texts_file
            )
        if self.inferencer is None:
            self.inferencer = TextToSpeechInferencer(device=self.device)
            self.inferencer.tts_network = self.network
            self.inferencer.instantiate_backbone_from_tts()

        self.network.eval()
        self.inferencer.generator.eval()
        with torch.no_grad():
            generated_audios = self.tts_cer_module.generate_audios(self.inferencer)
        for (_, text), audio in generated_audios.items():
            self.tensorboard.add_audio(
                text,
                audio,
                self.global_step,
                sample_rate=self.network.sample_rate,
            )
        transcriptions = self.tts_cer_module.generate_transcriptions(generated_audios)
        cer_score = self.tts_cer_module.compute_score(transcriptions)
        self.tensorboard.add_scalar("CER/validation", cer_score, self.global_step)

src/models/tts.py
MWM-io/nansypp
from typing import List

import torch
import torch.nn.functional as F
from torch import nn

from src.networks.backbone.encoders.wav2vec2 import Wav2Vec2Wrapper


class ConvGLU(nn.Module):
    """Dropout - Conv1d - GLU - residual connection."""

    def __init__(self, in_channels: int, kernels: int, dropout: float):
        """Initializer.
        Args:
            in_channels: size of the input channels.
            kernels: size of the convolutional kernels.
            dropout: dropout rate.
        """
        super().__init__()
        self.conv = nn.Sequential(
            nn.Dropout(dropout),
            nn.Conv1d(in_channels, in_channels * 2, kernels, padding=kernels // 2),
            nn.GLU(dim=1),
        )

    def forward(self, inputs: torch.Tensor) -> torch.Tensor:
        """Transform the inputs.
        Args:
            inputs: [torch.float32; [B, C, N]], input 1D tensor.
        Returns:
            [torch.float32; [B, C, N]], transformed.
        """
        return inputs + self.conv(inputs)


class LinguisticEncoder(nn.Module):
    """Additional linguistic information encoder."""

    in_channels: int
    """The number of channels in the input and intermediate layers."""

    hidden_channels: int
    """The number of hidden channels."""

    preconv_blocks: int
    """The number of pre-convolution blocks."""

    convglu_kernel_sizes: List[int]
    """The sizes of the ConvGLU kernels."""

    leak: float
    """The negative slope of leaky ReLUs."""

    dropout_rate: float
    """The dropout rate."""

    wav2vec2: Wav2Vec2Wrapper
    """A wav2vec2 model for extracting linguistic features."""

    def __init__(
        self,
        wav2vec2: Wav2Vec2Wrapper,
        hidden_channels: int,
        preconv_blocks: int,
        convglu_kernel_sizes: List[int],
        leak: float,
        dropout_rate: float,
    ):
        """Initializer.
        Args:
            wav2vec2: a Wav2Vec2 model for extracting linguistic features.
            hidden_channels: The number of hidden channels.
            preconv_blocks: The number of pre-convolution blocks.
            convglu_kernel_sizes: The sizes of the ConvGLU kernels.
            leak: The negative slope of leaky ReLUs.
            dropout_rate: The dropout rate.
        """
        super().__init__()

        self.hidden_channels = hidden_channels
        self.preconv_blocks = preconv_blocks
        self.convglu_kernel_sizes = convglu_kernel_sizes
        self.leak = leak
        self.dropout_rate = dropout_rate

        self.wav2vec2 = wav2vec2
        self.in_channels = wav2vec2.channels

        # in_channels=1024, hidden_channels=128, preconv=2
        # unknown `leak`, `dropout`
        self.preconv = nn.Sequential(
            nn.Conv1d(self.in_channels, hidden_channels, 1),
            nn.LeakyReLU(leak),
            nn.Dropout(dropout_rate),
            *[
                nn.Sequential(
                    nn.Conv1d(hidden_channels, hidden_channels, 1),
                    nn.LeakyReLU(leak),
                    nn.Dropout(dropout_rate),
                )
                for _ in range(preconv_blocks - 1)
            ]
        )
        # kernels=[3] * 8 + [1] * 2
        self.convglu = nn.Sequential(
            *[
                ConvGLU(hidden_channels, kernel_size, dropout_rate)
                for kernel_size in convglu_kernel_sizes
            ]
        )

        self.proj = nn.Conv1d(hidden_channels, hidden_channels, 1)

    def forward(self, inputs: torch.Tensor) -> torch.Tensor:
        """Filter the linguistic information.
        Args:
            inputs: [torch.float32; [B, in_channels, N]], input features.
        Returns:
            [torch.float32; [B, hidden_channels, N]], linguistic information.
        """
        # [B, S, w2v2_channels]
        w2v2 = self.wav2vec2(inputs)
        # [B, hidden_channels, N]
        x = self.preconv(w2v2.transpose(1, 2))
        # [B, hidden_channels, N]
        x = self.convglu(x)
        # [B, hidden_channels, N]
        return F.normalize(self.proj(x), dim=1)

src/networks/backbone/encoders/linguistic.py
MWM-io/nansypp
from typing import Tuple

import numpy as np
import torch
import torch.nn.functional as F
from torch import nn


class ResBlock(nn.Module):
    """Residual block,"""

    def __init__(self, in_channels: int, out_channels: int, kernels: int):
        """Initializer.
        Args:
            in_channels: size of the input channels.
            out_channels: size of the output channels.
            kernels: size of the convolutional kernels.
        """
        super().__init__()
        self.branch = nn.Sequential(
            nn.BatchNorm2d(in_channels),
            nn.GELU(),
            nn.Conv2d(
                in_channels, out_channels, (kernels, 1), padding=(kernels // 2, 0)
            ),
            nn.BatchNorm2d(out_channels),
            nn.GELU(),
            nn.Conv2d(
                out_channels, out_channels, (kernels, 1), padding=(kernels // 2, 0)
            ),
        )

        self.shortcut = nn.Conv2d(in_channels, out_channels, 1)

    def forward(self, inputs: torch.Tensor) -> torch.Tensor:
        """Transform the inputs.
        Args:
            inputs: [torch.float32; [B, in_channels, F, N]], input channels.
        Returns:
            [torch.float32; [B, out_channels, F // 2, N]], output channels.
        """
        # [B, out_channels, F, N]
        outputs = self.branch(inputs)
        # [B, out_channels, F, N]
        shortcut = self.shortcut(inputs)
        # [B, out_channels, F // 2, N]
        return F.avg_pool2d(outputs + shortcut, (2, 1))


def exponential_sigmoid(x: torch.Tensor) -> torch.Tensor:
    """Exponential sigmoid.
    Args:
        x: [torch.float32; [...]], input tensors.
    Returns:
        sigmoid outputs.
    """
    return 2.0 * torch.sigmoid(x) ** np.log(10) + 1e-7


class PitchEncoder(nn.Module):
    """Pitch-encoder."""

    freq: int
    """Number of frequency bins."""
    min_pitch: float
    """The minimum predicted pitch."""
    max_pitch: float
    """The maximum predicted pitch."""
    prekernels: int
    """Size of the first convolutional kernels."""
    kernels: int
    """Size of the frequency-convolution kernels."""
    channels: int
    """Size of the channels."""
    blocks: int
    """Number of residual blocks."""
    gru_dim: int
    """Size of the GRU hidden states."""
    hidden_channels: int
    """Size of the hidden channels."""
    f0_bins: int
    """Size of the output f0-bins."""
    f0_activation: str
    """F0 activation function."""

    def __init__(
        self,
        freq: int,
        min_pitch: float,
        max_pitch: float,
        prekernels: int,
        kernels: int,
        channels: int,
        blocks: int,
        gru_dim: int,
        hidden_channels: int,
        f0_bins: int,
        f0_activation: str,
    ):
        """Initializer.
        Args:
            freq: Number of frequency bins.
            min_pitch: The minimum predicted pitch.
            max_pitch: The maximum predicted pitch.
            prekernels: Size of the first convolutional kernels.
            kernels: Size of the frequency-convolution kernels.
            channels: Size of the channels.
            blocks: Number of residual blocks.
            gru_dim: Size of the GRU hidden states.
            hidden_channels: Size of the hidden channels.
            f0_bins: Size of the output f0-bins.
            f0_activation: f0 activation function.
        """
        super().__init__()

        self.freq = freq
        self.min_pitch = min_pitch
        self.max_pitch = max_pitch
        self.prekernels = prekernels
        self.kernels = kernels
        self.channels = channels
        self.blocks = blocks
        self.gru_dim = gru_dim
        self.hidden_channels = hidden_channels
        self.f0_bins = f0_bins
        self.f0_activation = f0_activation

        assert f0_activation in ["softmax", "sigmoid", "exp_sigmoid"]

        # prekernels=7
        self.preconv = nn.Conv2d(
            1, channels, (prekernels, 1), padding=(prekernels // 2, 0)
        )
        # channels=128, kernels=3, blocks=2
        self.resblock = nn.Sequential(
            *[ResBlock(channels, channels, kernels) for _ in range(blocks)]
        )
        # unknown `gru`
        self.gru = nn.GRU(
            freq * channels // (2 * blocks),
            gru_dim,
            batch_first=True,
            bidirectional=True,
        )
        # unknown `hidden_channels`
        # f0_bins=64
        self.proj = nn.Sequential(
            nn.Linear(gru_dim * 2, hidden_channels * 2),
            nn.ReLU(),
            nn.Linear(hidden_channels * 2, f0_bins + 2),
        )

        self.register_buffer(
            "pitch_bins",
            # linear space in log-scale
            torch.linspace(
                np.log(self.min_pitch),
                np.log(self.max_pitch),
                self.f0_bins,
            ).exp(),
        )

    def forward(
        self, cqt_slice: torch.Tensor
    ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
        """Compute the pitch from inputs.
        Args:
            cqt_slice: [torch.float32; [B, F, N]]
                The input tensor. A frequency-axis slice of a Constant-Q Transform.
        Returns:
            pitch: [torch.float32; [B, N]], predicted pitch, based on frequency bins.
            f0_logits: [torch.float32; [B, N, f0_bins]], predicted f0 activation weights,
                based on the pitch bins.
            p_amp, ap_amp: [torch.float32; [B, N]], amplitude values.
        """
        # B, _, N
        batch_size, _, timesteps = cqt_slice.shape
        # [B, C, F, N]
        x = self.preconv(cqt_slice[:, None])
        # [B, C F // 4, N]
        x = self.resblock(x)
        # [B, N, C x F // 4]
        x = x.permute(0, 3, 1, 2).reshape(batch_size, timesteps, -1)
        # [B, N, G x 2]
        x, _ = self.gru(x)
        # [B, N, f0_bins], [B, N, 1], [B, N, 1]
        f0_weights, p_amp, ap_amp = torch.split(
            self.proj(x), [self.f0_bins, 1, 1], dim=-1
        )

        # potentially apply activation function
        if self.f0_activation == "softmax":
            f0_weights = torch.softmax(f0_weights, dim=-1)
        elif self.f0_activation == "sigmoid":
            f0_weights = torch.sigmoid(f0_weights) / torch.sigmoid(f0_weights).sum(
                dim=-1, keepdim=True
            )
        elif self.f0_activation == "exp_sigmoid":
            f0_weights = exponential_sigmoid(f0_weights) / exponential_sigmoid(
                f0_weights
            ).sum(dim=-1, keepdim=True)

        # [B, N]
        pitch = (f0_weights * self.pitch_bins).sum(dim=-1)

        return (
            pitch,
            f0_weights,
            exponential_sigmoid(p_amp).squeeze(dim=-1),
            exponential_sigmoid(ap_amp).squeeze(dim=-1),
        )

src/networks/backbone/encoders/pitch.py
MWM-io/nansypp
from typing import List, Optional, Tuple

import numpy as np
import torch
import torch.nn.functional as F
from torch import nn

from src.networks.misc.transform import MelSpectrogram
from src.utilities.profiling import DO_PROFILING, cuda_synchronized_timer
from src.utilities.types import copy_docstring_and_signature


class Res2Block(nn.Module):
    """Multi-scale residual blocks."""

    def __init__(self, channels: int, scale: int, kernels: int, dilation: int):
        """Initializer.
        Args:
            channels: size of the input channels.
            scale: the number of the blocks.
            kenels: size of the convolutional kernels.
            dilation: dilation factors.
        """
        super().__init__()
        assert channels % scale == 0, (
            f"size of the input channels(={channels})"
            f" should be factorized by scale(={scale})"
        )
        width = channels // scale
        self.scale = scale
        self.convs = nn.ModuleList(
            [
                nn.Sequential(
                    nn.Conv1d(
                        width,
                        width,
                        kernels,
                        padding=(kernels - 1) * dilation // 2,
                        dilation=dilation,
                    ),
                    nn.ReLU(),
                    nn.BatchNorm1d(width),
                )
                for _ in range(scale - 1)
            ]
        )

    def forward(self, inputs: torch.Tensor) -> torch.Tensor:
        """Transform the inputs.
        Args:
            inputs: [torch.float32; [B, C, T]], input 1D tensor,
                where C = `channels`.
        Returns:
            [torch.float32; [B, C, T]], transformed.
        """
        # [B, W, T], (S - 1) x [B, W, T] where W = C // S
        straight, *xs = inputs.chunk(self.scale, dim=1)
        # [B, W, T]
        base = torch.zeros_like(straight)
        # S x [B, W, T]
        outs = [straight]
        for x, conv in zip(xs, self.convs):
            # [B, W, T], increasing receptive field progressively
            base = conv(x + base)
            outs.append(base)
        # [B, C, T]
        return torch.cat(outs, dim=1)


class SERes2Block(nn.Module):
    """Multiscale residual block with Squeeze-Excitation modules."""

    def __init__(
        self, channels: int, scale: int, kernels: int, dilation: int, bottleneck: int
    ):
        """Initializer.
        Args:
            channels: size of the input channels.
            scale: the number of the resolutions, for res2block.
            kernels: size of the convolutional kernels.
            dilation: dilation factor.
            bottleneck: size of the bottleneck layers for squeeze and excitation.
        """
        super().__init__()
        self.preblock = nn.Sequential(
            nn.Conv1d(channels, channels, 1), nn.ReLU(), nn.BatchNorm1d(channels)
        )

        self.res2block = Res2Block(channels, scale, kernels, dilation)

        self.postblock = nn.Sequential(
            nn.Conv1d(channels, channels, 1), nn.ReLU(), nn.BatchNorm1d(channels)
        )

        self.excitation = nn.Sequential(
            nn.Linear(channels, bottleneck),
            nn.ReLU(),
            nn.Linear(bottleneck, channels),
            nn.Sigmoid(),
        )

    def forward(self, inputs: torch.Tensor) -> torch.Tensor:
        """Transform the inputs.
        Args:
            inputs: [torch.float32; [B, C, T]], input tensors,
                where C = `channels`.
        Returns:
            [torch.float32; [B, C, T]], transformed.
        """
        # [B, C, T]
        x = self.preblock(inputs)
        # [B, C, T], res2net, multi-scale architecture
        x = self.res2block(x)
        # [B, C, T]
        x = self.postblock(x)
        # [B, C], squeeze and excitation
        scale = self.excitation(x.mean(dim=-1))
        # [B, C, T]
        x = x * scale[..., None]
        # residual connection
        return x + inputs


class AttentiveStatisticsPooling(nn.Module):
    """Attentive statistics pooling."""

    def __init__(self, channels: int, bottleneck: int):
        """Initializer.
        Args:
            channels: size of the input channels.
            bottleneck: size of the bottleneck.
        """
        super().__init__()
        # nonlinear=Tanh
        # ref: https://github.com/KrishnaDN/Attentive-Statistics-Pooling-for-Deep-Speaker-Embedding
        # ref: https://github.com/TaoRuijie/ECAPA-TDNN
        self.attention = nn.Sequential(
            nn.Conv1d(channels, bottleneck, 1),
            nn.Tanh(),
            nn.Conv1d(bottleneck, channels, 1),
            nn.Softmax(dim=-1),
        )

    def forward(self, inputs: torch.Tensor) -> torch.Tensor:
        """Pooling with weighted statistics.
        Args:
            inputs: [torch.float32; [B, C, T]], input tensors,
                where C = `channels`.
        Returns:
            [torch.float32; [B, C x 2]], weighted statistics.
        """
        # [B, C, T]
        weights = self.attention(inputs)
        # [B, C]
        mean = torch.sum(weights * inputs, dim=-1)
        # var = torch.sum(weights * inputs**2, dim=-1) - mean**2
        var = torch.sum(weights * inputs**2, dim=-1) - mean**2
        # [B, C x 2], for numerical stability of square root
        return torch.cat([mean, (var + 1e-7).sqrt()], dim=-1)

class TTSAttentiveStatisticsPooling(AttentiveStatisticsPooling):
    """Fixed Attentive statistics pooling."""

    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)

    def forward(self, inputs: torch.Tensor) -> torch.Tensor:
        """Pooling with weighted statistics.
        Args:
            inputs: [torch.float32; [B, C, T]], input tensors,
                where C = `channels`.
        Returns:
            [torch.float32; [B, C x 2]], weighted statistics.
        """
        # [B, C, T]
        weights = self.attention(inputs)
        # [B, C]
        mean = torch.sum(weights * inputs, dim=-1)
        # var = torch.sum(weights * inputs**2, dim=-1) - mean**2
        var = torch.sqrt((weights * (inputs - mean.unsqueeze(-1))).pow(2).sum(dim=-1))
        # [B, C x 2], for numerical stability of square root
        return torch.cat([mean, var], dim=-1)


class MultiheadAttention(nn.Module):
    """Multi-head scaled dot-product attention."""

    def __init__(
        self,
        keys: int,
        values: int,
        queries: int,
        out_channels: int,
        hidden_channels: int,
        heads: int,
    ):
        """Initializer.
        Args:
            keys, valeus, queries: size of the input channels.
            out_channels: size of the output channels.
            hidden_channels: size of the hidden channels.
            heads: the number of the attention heads.
        """
        super().__init__()
        assert (
            hidden_channels % heads == 0
        ), f"size of hidden_channels channels(={hidden_channels}) should be factorized by heads(={heads})"
        self.channels, self.heads = hidden_channels // heads, heads
        self.proj_key = nn.Conv1d(keys, hidden_channels, 1)
        self.proj_value = nn.Conv1d(values, hidden_channels, 1)
        self.proj_query = nn.Conv1d(queries, hidden_channels, 1)
        self.proj_out = nn.Conv1d(hidden_channels, out_channels, 1)

    def forward(
        self,
        keys: torch.Tensor,
        values: torch.Tensor,
        queries: torch.Tensor,
        mask: Optional[torch.Tensor] = None,
    ) -> torch.Tensor:
        """Transform the inputs.
        Args:
            keys: [torch.float32; [B, keys, S]], attention key.
            values: [torch.float32; [B, values, S]], attention value.
            queries: [torch.float32; [B, queries, T]], attention query.
            mask: [torch.float32; [B, S, T]], attention mask, 0 for paddings.
        Returns:
            [torch.float32; [B, out_channels, T]], transformed outputs.
        """
        # B, T
        bsize, _, querylen = queries.shape
        # S
        keylen = keys.shape[-1]
        assert keylen == values.shape[-1], "lengths of key and value are not matched"
        # [B, H, hidden_channels // H, S]
        keys = self.proj_key(keys).view(bsize, self.heads, -1, keylen)
        values = self.proj_value(values).view(bsize, self.heads, -1, keylen)
        # [B, H, hidden_channels // H, T]
        queries = self.proj_query(queries).view(bsize, self.heads, -1, querylen)
        # [B, H, S, T]
        score = torch.matmul(keys.transpose(2, 3), queries) * (self.channels**-0.5)
        if mask is not None:
            score.masked_fill_(~mask[:, None, :, :1].to(torch.bool), -np.inf)
        # [B, H, S, T]
        weights = torch.softmax(score, dim=2)
        # [B, out_channels, T]
        out = self.proj_out(torch.matmul(values, weights).view(bsize, -1, querylen))
        if mask is not None:
            out = out * mask[:, :1]
        return out


class TimbreEncoder(nn.Module):
    """ECAPA-TDNN: Emphasized Channel Attention,
    [1] Propagation and Aggregation in TDNN Based Speaker Verification,
        Desplanques et al., 2020, arXiv:2005.07143.
    [2] Res2Net: A New Multi-scale Backbone architecture,
        Gao et al., 2019, arXiv:1904.01169.
    [3] Squeeze-and-Excitation Networks, Hu et al., 2017, arXiv:1709.01507.
    [4] Attentive Statistics Pooling for Deep Speaker Embedding,
        Okabe et al., 2018, arXiv:1803.10963.
    """

    in_channels: int
    """Size of the input channels."""
    out_channels: int
    """Size of the output embeddings."""
    channels: int
    """Size of the major states."""
    prekernels: int
    """Size of the convolutional kernels before feed to SERes2Block."""
    scale: int
    """The number of the resolutions, for SERes2Block."""
    kernels: int
    """Size of the convolutional kernels, for SERes2Block."""
    dilations: List[int]
    """Dilation factors."""
    bottleneck: int
    """Size of the bottleneck layers, both SERes2Block and AttentiveStatisticsPooling."""
    # NANSY++: 3072
    hidden_channels: int
    """Size of the hidden channels for attentive statistics pooling."""
    latent: int
    """Size of the timbre latent query."""
    timbre: int
    """Size of the timbre tokens."""
    tokens: int
    """The number of the timbre tokens."""
    # unknown
    heads: int
    """The number of the attention heads, for timbre token block."""
    linguistic_encoder_hidden_channels: int
    """The size of the features returned by the linguistic encoder."""
    # unknown
    slerp: float
    """Weight value for spherical interpolation."""
    mel_spectrogram_transform: MelSpectrogram
    """Utility class for Mel Spectrogram transform."""

    def __init__(
        self,
        in_channels: int,
        out_channels: int,
        channels: int,
        prekernels: int,
        scale: int,
        kernels: int,
        dilations: List[int],
        bottleneck: int,
        hidden_channels: int,
        latent: int,
        timbre: int,
        tokens: int,
        heads: int,
        linguistic_encoder_hidden_channels: int,
        slerp: float,
        mel_spectrogram_transform: MelSpectrogram,
        attentive_statistics_pooling: AttentiveStatisticsPooling,
    ):
        """Initializer.
        Args:
            in_channels: size of the input channels.
            out_channels: size of the output embeddings.
            channels: size of the major states.
            prekernels: size of the convolutional kernels before feed to SERes2Block.
            scale: the number of the resolutions, for SERes2Block.
            kernels: size of the convolutional kernels, for SERes2Block.
            dilations: dilation factors.
            bottleneck: size of the bottleneck layers,
                both SERes2Block and AttentiveStatisticsPooling.
            hidden_channels: size of the hidden channels for attentive statistics pooling.
            latent: size of the timbre latent query.
            timbre: size of the timbre tokens.
            tokens: the number of the timbre tokens.
            heads: the number of the attention heads, for timbre token block.
            linguistic_encoder_hidden_channels:
                The size of the features returned by the linguistic encoder.
            slerp: weight value for spherical interpolation.
        """
        super().__init__()

        self.in_channels = in_channels
        self.out_channels = out_channels
        self.channels = channels
        self.prekernels = prekernels
        self.scale = scale
        self.kernels = kernels
        self.dilations = dilations
        self.bottleneck = bottleneck
        self.hidden_channels = hidden_channels
        self.latent = latent
        self.timbre = timbre
        self.tokens = tokens
        self.heads = heads
        self.linguistic_encoder_hidden_channels = linguistic_encoder_hidden_channels
        # unknown `slerp`
        assert 0 <= slerp <= 1, f"value slerp(={slerp:.2f}) should be in range [0, 1]"
        self.slerp = slerp

        self.mel_spectrogram_transform = mel_spectrogram_transform

        contents = self.contents = linguistic_encoder_hidden_channels + out_channels + 3

        # channels=512, prekernels=5
        # ref:[1], Figure2 and Page3, "architecture with either 512 or 1024 channels"
        self.preblock = nn.Sequential(
            nn.Conv1d(in_channels, channels, prekernels, padding=prekernels // 2),
            nn.ReLU(),
            nn.BatchNorm1d(channels),
        )
        # scale=8, kernels=3, dilations=[2, 3, 4], bottleneck=128
        self.blocks = nn.ModuleList(
            [
                SERes2Block(channels, scale, kernels, dilation, bottleneck)
                for dilation in dilations
            ]
        )
        # hidden_channels=1536
        # TODO(@revsic): hidden_channels=3072 on NANSY++
        self.conv1x1 = nn.Sequential(
            nn.Conv1d(len(dilations) * channels, hidden_channels, 1), nn.ReLU()
        )
        # multi-head attention for time-varying timbre
        # NANSY++, latent=512, tokens=50
        self.timbre_query = nn.Parameter(torch.randn(1, latent, tokens))
        # NANSY++, timbre=128
        # unknown `heads`
        self.pre_mha = MultiheadAttention(
            keys=hidden_channels,
            values=hidden_channels,
            queries=latent,
            out_channels=latent,
            hidden_channels=latent,
            heads=heads,
        )
        self.post_mha = MultiheadAttention(
            keys=hidden_channels,
            values=hidden_channels,
            queries=latent,
            out_channels=timbre,
            hidden_channels=latent,
            heads=heads,
        )
        # attentive pooling and additional projector
        # out_channels=192
        self.pool = nn.Sequential(
            attentive_statistics_pooling,
            nn.BatchNorm1d(hidden_channels * 2),
            nn.Linear(hidden_channels * 2, out_channels),
            nn.BatchNorm1d(out_channels),
        )

        # time-varying timbre encoder
        self.timbre_key = nn.Parameter(torch.randn(1, timbre, tokens))
        self.sampler = MultiheadAttention(
            keys=timbre,
            values=timbre,
            queries=contents,
            out_channels=timbre,
            hidden_channels=latent,
            heads=heads,
        )
        self.proj = nn.Conv1d(timbre, out_channels, 1)

    @cuda_synchronized_timer(DO_PROFILING, prefix="TimbreEncoder")
    def forward(
        self, inputs: torch.Tensor, precomputed_mel: Optional[torch.Tensor] = None
    ) -> Tuple[torch.Tensor, torch.Tensor]:
        """Evaluate the x-vectors from the input sequence.
        Args:
            inputs: [torch.float32; [B, in_channels, T]], input sequences,
        Returns:
            [torch.float32; [B, out_channels]], global x-vectors,
            [torch.float32; [B, timbre, tokens]], timbre token bank.
        """
        if precomputed_mel is None:
            mel_spectrogram = self.mel_spectrogram_transform(inputs)
        else:
            mel_spectrogram = precomputed_mel

        # [B, C, T]
        x = self.preblock(mel_spectrogram)
        # N x [B, C, T]
        xs = []
        for block in self.blocks:
            # [B, C, T]
            x = block(x)
            xs.append(x)
        # [B, H, T]
        mfa = self.conv1x1(torch.cat(xs, dim=1))
        # [B, O]
        global_ = F.normalize(self.pool(mfa), p=2, dim=-1)
        # B
        bsize, _ = global_.shape
        # [B, latent, tokens]
        query = self.timbre_query.repeat(bsize, 1, 1)
        # [B, latent, tokens]
        query = self.pre_mha(mfa, mfa, query) + query
        # [B, timbre, tokens]
        local = self.post_mha(mfa, mfa, query)
        # [B, out_channels], [B, timbre, tokens]
        return global_, local

    def sample_timbre(
        self,
        contents: torch.Tensor,
        global_: torch.Tensor,
        tokens: torch.Tensor,
        eps: float = 1e-5,
    ) -> torch.Tensor:
        """Sample the timbre tokens w.r.t. the contents.
        Args:
            contents: [torch.float32; [B, contents, T]], content queries.
            global_: [torch.float32; [B, out_channels]], global x-vectors, L2-normalized.
            tokens: [torch.float32; [B, timbre, tokens]], timbre token bank.
            eps: small value for preventing train instability of arccos in slerp.
        Returns:
            [torch.float32; [B, out_channels, T]], time-varying timbre embeddings.
        """
        # [B, timbre, tokens]
        key = self.timbre_key.repeat(contents.shape[0], 1, 1)
        # [B, timbre, T]
        sampled = self.sampler(key, tokens, contents)
        # [B, out_channels, T]
        sampled = F.normalize(self.proj(sampled), p=2, dim=1)
        # [B, 1, T]
        theta = torch.matmul(global_[:, None], sampled).clamp(-1 + eps, 1 - eps).acos()
        # [B, 1, T], slerp
        # clamp the theta is not necessary since cos(theta) is already clampped
        return (
            torch.sin(self.slerp * theta) * sampled
            + torch.sin((1 - self.slerp) * theta) * global_[..., None]
        ) / theta.sin()

    @copy_docstring_and_signature(forward)
    def __call__(self, *args, **kwargs):  # type: ignore[no-untyped-def]
        return super().__call__(*args, **kwargs)

src/networks/backbone/encoders/timbre.py
MWM-io/nansypp
from typing import Any, Dict, Optional

import numpy as np
import torch
import torchaudio
from torch import nn
from transformers import Wav2Vec2Model
from transformers.modeling_outputs import Wav2Vec2BaseModelOutput

from src.dataclasses.backbone import PretrainedWav2Vec2Config


class Wav2Vec2Wrapper(nn.Module):
    """Wrapping huggingface wav2vec2.0."""

    wav2vec2: Wav2Vec2Model
    resample: torchaudio.transforms.Resample

    channels: int
    layer_for_linguistic_features: int
    sample_rate: int

    def __init__(
        self,
        global_sample_rate: int,
        pretrained_model: PretrainedWav2Vec2Config,
        layer_for_linguistic_features: int,
        trim_unused_layers: bool,
    ):
        """Load the wav2vec2.0 pretrained model and instantiate an nn.Resampler for adapting sample-rates.

        Args:
            global_sample_rate:
                The general sample rate of the NANSY++ model.
            pretrained_model_name_and_sample_rate (default uses Facebook's XLSR-53):
                Name and sample-rate of the pretrained model, .
            layer_for_linguistic_features:
                Hidden layer from the wav2vec2 model used for extracting linguistic features.
            trim_unused_layers:
                If True, removes all layers in the Wav2Vec2Encoder after `layer_for_linguistic_features`.
        """
        super().__init__()

        self.name = pretrained_model.name
        self.wav2vec2_sample_rate = pretrained_model.sample_rate
        self.trim_unused_layers = trim_unused_layers

        self.wav2vec2: Wav2Vec2Model = Wav2Vec2Model.from_pretrained(
            self.name, local_files_only=True
        )
        assert isinstance(
            self.wav2vec2, Wav2Vec2Model
        ), "Wav2Vec2Model initialization failed"

        # aliases
        self.channels = self.wav2vec2.config.output_hidden_size

        self.strides = np.prod(self.wav2vec2.config.conv_stride).item()
        self.layer_for_linguistic_features = (
            layer_for_linguistic_features
            or Wav2Vec2Wrapper.LAYER_FOR_LINGUISTIC_FEATURES
        )
        if self.trim_unused_layers:
            self._trim_wav2vec2_encoder_at_layer_(self.layer_for_linguistic_features)

        # resampler
        self.global_sample_rate = global_sample_rate
        self.resample = torchaudio.transforms.Resample(
            self.global_sample_rate, self.wav2vec2_sample_rate
        )
        self.eval()

    @torch.no_grad()
    def forward(
        self, audio: torch.Tensor, audio_lengths: Optional[torch.Tensor] = None
    ) -> torch.Tensor:
        """Extract the features from audio.
        Args:
            audio: [torch.float32; [B, T']], audio, [-1, 1]-ranged.
            audio_lengths: [torch.long; [B]], individual length of the audio samples in the batch,
                masking the inputs if provided.
        Returns:
            linguistic: [torch.float32; [B, S, C]], linguistic encodings,
                where C = `channels`
                      S = T // `conv_stride`, T = ceil(T' / `sr` x `sr_w2v2`)
        """
        # [B, T]
        audio = self.resample(audio)

        # B, T
        batch_size, duration = audio.shape
        if audio_lengths is None:
            audio_lengths = torch.full(
                (batch_size,), duration, dtype=torch.long, device=audio.device
            )
        else:
            # rearange to 16khz audio frames
            audio_lengths = torch.ceil(
                audio_lengths / self.global_sample_rate * self.wav2vec2_sample_rate
            ).to(torch.long)

        # [B, T]
        mask = (
            torch.arange(duration, device=audio_lengths.device)[None]
            < audio_lengths[:, None]
        ).to(torch.float32)

        ## normalize the inputs before feeding them to wav2vec2
        ## reference: Wav2VecFeatureExtractor
        # [B]
        mean = (audio * mask).sum(dim=-1) / audio_lengths.to(torch.float32)
        # [B]
        var = ((audio - mean[:, None]) * mask).square().sum(dim=-1) / audio_lengths.to(
            torch.float32
        )
        # small additive value for numerical stability of square root
        eps = 1e-7
        # [B, T]
        normed = (audio - mean[:, None]) / (var[:, None] + eps).sqrt() * mask

        # run predict
        output: Wav2Vec2BaseModelOutput = self.wav2vec2(
            normed, attention_mask=mask.to(torch.long), output_hidden_states=True
        )
        assert output.hidden_states is not None  # helps the type-checker

        # [B, S, C(=1024)]
        linguistic_features = output.hidden_states[self.layer_for_linguistic_features]
        return linguistic_features

    def _trim_wav2vec2_encoder_at_layer_(self, layer: int) -> None:
        del self.wav2vec2.encoder.layers[layer + 1 :]

    def train(self, mode: bool = True) -> "Wav2Vec2Wrapper":
        """Supports only evaluation."""
        if mode:
            import warnings

            warnings.warn(
                "Wav2Vec2Wrapper does not support training mode, using frozen model."
            )
        return super().train(False)

    def state_dict(self, *args, **kwargs) -> Dict[str, Any]:
        """Do not return the state dict."""
        return {}

    def _load_from_state_dict(self, *args, **kwargs):
        """Do not load state dict."""
        pass

src/networks/backbone/encoders/wav2vec2.py
MWM-io/nansypp
from typing import Optional, Tuple, TypedDict, Union

import torch
import torch.nn.functional as F
from torch import nn

from src.networks.backbone.encoders.linguistic import LinguisticEncoder
from src.networks.backbone.encoders.pitch import PitchEncoder
from src.networks.backbone.encoders.timbre import TimbreEncoder
from src.networks.backbone.synthesizers.framelevel import FrameLevelSynthesizer
from src.networks.backbone.synthesizers.synthesizer import Synthesizer
from src.networks.misc.perturbator import InformationPerturbator
from src.networks.misc.transform import ConstantQTransform, MelSpectrogram
from src.utilities.profiling import DO_PROFILING, cuda_synchronized_timer
from src.utilities.types import copy_docstring_and_signature


class AnalysisFeatures(TypedDict):
    cqt: torch.Tensor

    pitch: torch.Tensor
    ap_amp: torch.Tensor
    p_amp: torch.Tensor

    linguistic: torch.Tensor

    timbre_global: Optional[torch.Tensor]
    timbre_bank: Optional[torch.Tensor]


class SynthesisFeatures(TypedDict):
    excitation: torch.Tensor


class Generator(nn.Module):
    input_sample_rate: int
    output_sample_rate: int
    leak: float
    dropout_rate: float
    pitch_start: float
    pitch_end: float

    def __init__(
        self,
        input_sample_rate: int,
        output_sample_rate: int,
        leak: float,
        dropout_rate: float,
        cqt_center: int,
        linguistic_encoder: LinguisticEncoder,
        cqt: ConstantQTransform,
        pitch_encoder: PitchEncoder,
        mel_spectrogram_transform: MelSpectrogram,
        synthesizer: Synthesizer,
        frame_level_synthesizer: FrameLevelSynthesizer,
        information_perturbator: InformationPerturbator,
        timbre_encoder: Optional[TimbreEncoder],
    ):
        super().__init__()
        self.input_sample_rate = input_sample_rate
        self.output_sample_rate = output_sample_rate
        self.leak = leak
        self.dropout_rate = dropout_rate

        # dependency-injection all the way
        self.linguistic_encoder = linguistic_encoder
        self.cqt = cqt
        self.pitch_encoder = pitch_encoder
        self.mel_spectrogram_transform = mel_spectrogram_transform
        self.timbre_encoder = timbre_encoder
        self.frame_level_synthesizer = frame_level_synthesizer
        self.information_perturbator = information_perturbator
        self.synthesizer = synthesizer

        self.cqt_center = cqt_center

    @cuda_synchronized_timer(DO_PROFILING, prefix="Generator")
    def analyze_pitch(
        self,
        inputs: torch.Tensor,
        index: Optional[int] = None,
        precomputed_cqt: Optional[torch.Tensor] = None,
    ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
        """Estimate the pitch and periodical, aperiodical amplitudes.
        Args:
            inputs: [torch.float32; [B, T]], input speech signal.
            index: CQT start index. If None, use `cqt_center`.
        Returns:
            [torch.float32; [B, cqt_bins, N]], CQT features.
            [torch.float2; [B, N]], frame-level pitch and amplitude sequence.
        """
        # TODO(@revsic): use log-scaled CQT or not?
        if precomputed_cqt is None:
            # [B, cqt_bins, N(=T / cqt_hop)]
            cqt = self.cqt(inputs)
        else:
            cqt = precomputed_cqt

        # alias
        freq = self.pitch_encoder.freq
        if index is None:
            index = self.cqt_center

        # [B, N], [B, N, f0_bins], [B, N], [B, N]
        pitch, _, p_amp, ap_amp = self.pitch_encoder(cqt[:, index : index + freq])

        # [B, cqt_bins, N], [B, N]
        return cqt, pitch, p_amp, ap_amp

    @cuda_synchronized_timer(DO_PROFILING, prefix="Generator")
    def analyze_linguistic(
        self,
        inputs: torch.Tensor,
        enable_information_perturbator: bool = False,
        perturbed_inputs: Optional[torch.Tensor] = None,
    ) -> torch.Tensor:
        """Analyze the linguistic informations from inputs.
        Args:
            inputs: [torch.float32; [B, T]], input speech signal.
            enable_information_perturbator: bool
                Whether to apply linguistic-features-preserving data augmentations.
                Used during training only, do not enable at inference time.
        Returns:
            [torch.float32; [B, linguistic_hidden_channels, S]], linguistic informations.
        """
        if enable_information_perturbator and perturbed_inputs is None:
            inputs = self.information_perturbator(inputs)
        elif perturbed_inputs is not None:
            inputs = perturbed_inputs
        # [B, linguistic_hidden_channels, S]
        return self.linguistic_encoder(inputs)

    @cuda_synchronized_timer(DO_PROFILING, prefix="Generator")
    def analyze_timbre(
        self, inputs: torch.Tensor, precomputed_mel: Optional[torch.Tensor] = None
    ) -> Union[Tuple[torch.Tensor, torch.Tensor], Tuple[None, None]]:
        """Analyze the timbre informations from inputs.
        Args:
            inputs: [torch.float32; [B, T]], input speech signal.
        Returns:
            [torch.float32; [B, timb_global]], global timbre emebddings.
            [torch.float32; [B, timb_timbre, timb_tokens]], timbre token bank.
        """
        if self.timbre_encoder is not None:
            # [B, timb_global], [B, timb_timbre, timb_tokens]
            return self.timbre_encoder(inputs, precomputed_mel=precomputed_mel)
        else:
            return None, None

    @cuda_synchronized_timer(DO_PROFILING, prefix="Generator")
    def analyze(
        self,
        inputs: torch.Tensor,
        enable_information_perturbator: bool,
        perturbed_inputs: Optional[torch.Tensor] = None,
        precomputed_cqt: Optional[torch.Tensor] = None,
        precomputed_mel: Optional[torch.Tensor] = None,
    ) -> AnalysisFeatures:
        """Analyze the input signal.
        Args:
            inputs: [torch.float32; [B, T]], input speech signal.
            enable_information_perturbator: bool
                Whether to apply linguistic-features-preserving data augmentations.
                Used during training only, do not enable at inference time.
        Returns;
            Analysis features:
                cqt: [torch.float32; []], CQT features.
                pitch, p_amp, ap_amp: [torch.float2; [B, N]],
                    frame-level pitch and amplitude sequence.
                linguistic: [torch.float32; [B, linguistic_hidden_channels, S]], linguistic informations.
                timbre_global: [torch.float32; [B, timb_global]], global timbre emebddings.
                timbre_bank: [torch.float32; [B, timb_timbre, timb_tokens]], timbre token bank.
        """
        # [], [B, N]
        cqt, pitch, p_amp, ap_amp = self.analyze_pitch(
            inputs, precomputed_cqt=precomputed_cqt
        )
        # [B, linguistic_hidden_channels, S]
        linguistic_features = self.analyze_linguistic(
            inputs, enable_information_perturbator, perturbed_inputs
        )
        # [B, timb_global], [B, timb_timbre, timb_tokens]
        timbre_global, timbre_bank = self.analyze_timbre(
            inputs, precomputed_mel=precomputed_mel
        )

        return {
            "cqt": cqt,
            "pitch": pitch,
            "p_amp": p_amp,
            "ap_amp": ap_amp,
            "linguistic": linguistic_features,
            "timbre_global": timbre_global,
            "timbre_bank": timbre_bank,
        }

    @cuda_synchronized_timer(DO_PROFILING, prefix="Generator")
    def synthesize(
        self,
        pitch: torch.Tensor,
        p_amp: torch.Tensor,
        ap_amp: torch.Tensor,
        linguistic: torch.Tensor,
        timbre_global: Optional[torch.Tensor],
        timbre_bank: Optional[torch.Tensor],
        noise: Optional[torch.Tensor] = None,
    ) -> Tuple[torch.Tensor, torch.Tensor]:
        """Synthesize the signal.
        Args:
            pitch, p_amp, ap_amp: [torch.float32; [B, N]], frame-level pitch, amplitude sequence.
            linguistic: [torch.float32; [B, linguistic_hidden_channels, S]], linguistic features.
            timbre_global: [torch.float32; [B, timb_global]], global timbre.
            timbre_bank: [torch.float32; [B, timb_timbre, timb_tokens]], timbre token bank.
            noise: [torch.float32; [B, T]], predefined noise for excitation signal, if provided.
        Returns:
            [torch.float32; [B, T]], excitation and synthesized speech signal.
        """
        # S
        ling_len = linguistic.shape[-1]
        # [B, 3, S]
        pitch_rel = F.interpolate(
            torch.stack([pitch, p_amp, ap_amp], dim=1), size=ling_len
        )
        if self.timbre_encoder is not None:
            assert (
                timbre_global is not None and timbre_bank is not None
            ), "timbre_global and timbre_bank should be provided if a non-null timbre_encoder is used"

            # [B, 3 + linguistic_hidden_channels + timb_global, S]
            contents = torch.cat(
                [
                    pitch_rel,
                    linguistic,
                    timbre_global[..., None].repeat(1, 1, ling_len),
                ],
                dim=1,
            )
            # [B, timbre_global, S]
            timbre_sampled = self.timbre_encoder.sample_timbre(
                contents, timbre_global, timbre_bank
            )
        else:
            timbre_sampled = None

        # [B, linguistic_hidden_channels, S]
        frame = self.frame_level_synthesizer(linguistic, timbre_sampled)
        # [B, T], [B, T]
        return self.synthesizer(pitch, p_amp, ap_amp, frame, noise)

    @cuda_synchronized_timer(DO_PROFILING, prefix="Generator")
    def forward(
        self,
        inputs: torch.Tensor,
        noise: Optional[torch.Tensor] = None,
        enable_information_perturbator: bool = True,
        perturbed_inputs: Optional[torch.Tensor] = None,
        precomputed_cqt: Optional[torch.Tensor] = None,
        precomputed_mel: Optional[torch.Tensor] = None,
    ) -> Tuple[torch.Tensor, AnalysisFeatures, SynthesisFeatures]:
        """Reconstruct input audio.
        Args:
            inputs: [torch.float32; [B, T]], input signal.
            noise: [torch.float32; [B, T]], predefined noise for excitation, if provided.
            enable_information_perturbator: bool
                Whether to apply linguistic-features-preserving data augmentations.
                Used during training only, do not enable at inference time.
        Returns:
            [torch.float32; [B, T]], reconstructed.
            auxiliary outputs, reference `Nansypp.analyze`.
        """
        try:
            _, timesteps = inputs.shape
        except AttributeError:
            inputs, perturbed_inputs, _ = inputs
            _, timesteps = inputs.shape

        analysis_features = self.analyze(
            inputs,
            enable_information_perturbator,
            perturbed_inputs=perturbed_inputs,
            precomputed_cqt=precomputed_cqt,
        )

        # [B, T]
        excitation, synth = self.synthesize(
            analysis_features["pitch"],
            analysis_features["p_amp"],
            analysis_features["ap_amp"],
            analysis_features["linguistic"],
            analysis_features["timbre_global"],
            analysis_features["timbre_bank"],
            noise=noise,
        )
        # truncate
        synth = synth[
            :, : int(timesteps / self.input_sample_rate * self.output_sample_rate)
        ]
        # update
        synthesis_features = SynthesisFeatures(excitation=excitation)
        return synth, analysis_features, synthesis_features

    @torch.no_grad()
    def voice_conversion(
        self,
        source_audio: torch.Tensor,
        target_audio: torch.Tensor,
        device: torch.device,
        f0_statistics: bool = True,
        enable_information_perturbator: bool = False,
    ) -> torch.Tensor:
        """Apply target audio voice timbre to source audio content.
        Args:
            source_audio: [torch.Tensor; [1, Ns]] tensor with mono audio content at sampling rate defined by model config.
            target_audio: [torch.Tensor; [1, Nt]] tensor with mono audio content at sampling rate defined by model config.
            device: [torch.device] device, CPU or GPU, on which computation are expected to run.
            f0_statistics: [boolean], whether to replace statistics of estimated f0 of the source (mean, std) with those of the target.
            enable_information_perturbator: bool
                Whether to apply linguistic-features-preserving data augmentations.
                Used during training only, do not enable at inference time.
        Return:
            f0_corrected_synth: [torch.Tensor; [1, Ns]] tensor with mono audio content at sampling rate defined by model config
        """
        source_features = self.analyze(
            source_audio.to(device),
            enable_information_perturbator=enable_information_perturbator,
        )
        target_features = self.analyze(
            target_audio.to(device),
            enable_information_perturbator=enable_information_perturbator,
        )

        # Transform the F0 statistics of the source speaker into the target speakers'
        source_pitch = source_features["pitch"]
        if f0_statistics:
            standardized_source_pitch = (
                source_pitch - source_pitch.mean()
            ) / source_pitch.std()
            source_pitch = (
                standardized_source_pitch * target_features["pitch"].std()
                + target_features["pitch"].mean()
            )

        noise = None
        _, f0_corrected_synth = self.synthesize(
            source_pitch,
            source_features["p_amp"],
            source_features["ap_amp"],
            source_features["linguistic"],
            target_features.get("timbre_global", None),
            target_features.get("timbre_bank", None),
            noise=noise,
        )
        return f0_corrected_synth

    @copy_docstring_and_signature(forward)
    def __call__(self, *args, **kwargs):  # type: ignore[no-untyped-def]
        return super().__call__(*args, **kwargs)

src/networks/backbone/generator.py
MWM-io/nansypp
from typing import List, Optional, Sequence

import torch
import torch.nn.functional as F
from torch import nn

from src.utilities.profiling import DO_PROFILING, cuda_synchronized_timer


class ConvGLU(nn.Module):
    """Dropout - Conv1d - GLU and conditional layer normalization."""

    def __init__(
        self,
        channels: int,
        kernels: int,
        dilations: int,
        dropout: float,
        conditionning_channels: Optional[int] = None,
    ):
        """Initializer.
        Args:
            channels: size of the input channels.
            kernels: size of the convolutional kernels.
            dilations: dilation rate of the convolution.
            dropout: dropout rate.
            conditionning_channels: size of the condition channels, if provided.
        """
        super().__init__()
        self.conv = nn.Sequential(
            nn.Dropout(dropout),
            nn.Conv1d(
                channels,
                channels * 2,
                kernels,
                dilation=dilations,
                padding=(kernels - 1) * dilations // 2,
            ),
            nn.GLU(dim=1),
        )

        self.conditionning_channels = conditionning_channels
        if self.conditionning_channels is not None:
            self.cond = nn.Conv1d(self.conditionning_channels, channels * 2, 1)

    def forward(
        self, inputs: torch.Tensor, cond: Optional[torch.Tensor] = None
    ) -> torch.Tensor:
        """Transform the inputs with given conditions.
        Args:
            inputs: [torch.float32; [B, channels, T]], input channels.
            cond: [torch.float32; [B, cond, T]], if provided.
        Returns:
            [torch.float32; [B, channels, T]], transformed.
        """
        # [B, channels, T]
        x = inputs + self.conv(inputs)
        if cond is not None:
            assert self.cond is not None, "condition module does not exists"
            # [B, channels, T]
            x = F.instance_norm(x, use_input_stats=True)
            # [B, channels, T]
            weight, bias = self.cond(cond).chunk(2, dim=1)
            # [B, channels, T]
            x = x * weight + bias
        return x


class CondSequential(nn.Module):
    """Sequential pass with conditional inputs."""

    def __init__(self, modules: Sequence[nn.Module]):
        """Initializer.
        Args:
            modules: list of torch modules.
        """
        super().__init__()
        self.lists = nn.ModuleList(modules)

    # def forward(self, inputs: torch.Tensor, *args, **kwargs) -> torch.Tensor:
    def forward(
        self, inputs: torch.Tensor, cond: Optional[torch.Tensor] = None
    ) -> torch.Tensor:
        """Pass the inputs to modules.
        Args:
            inputs: arbitary input tensors.
            args, kwargs: positional, keyword arguments.
        Returns:
            output tensor.
        """
        x = inputs
        for module in self.lists:
            # x = module.forward(x, *args, **kwargs)
            x = module.forward(x, cond)
        return x


class FrameLevelSynthesizer(nn.Module):
    """Frame-level synthesizer."""

    def __init__(
        self,
        in_channels: int,
        kernels: int,
        dilations: List[int],
        blocks: int,
        leak: float,
        dropout_rate: float,
        timbre_embedding_channels: Optional[int],
    ):
        """Initializer.
        Args:
            in_channels: The size of the input channels.
            kernels: The size of the convolutional kernels.
            dilations: The dilation rates.
            blocks: The number of the ConvGLU blocks after dilated ConvGLU.
            leak: The negative slope of the leaky relu.
            dropout_rate: The dropout rate.
            timbre_embedding_channels: The size of the time-varying timbre embeddings.
        """
        super().__init__()

        self.in_channels = in_channels
        self.timbre_embedding_channels = timbre_embedding_channels
        self.kernels = kernels
        self.dilations = dilations
        self.blocks = blocks
        self.leak = leak
        self.dropout_rate = dropout_rate

        # channels=1024
        # unknown `leak`, `dropout`
        self.preconv = nn.Sequential(
            nn.Conv1d(in_channels, in_channels, 1),
            nn.LeakyReLU(leak),
            nn.Dropout(dropout_rate),
        )
        # kernels=3, dilations=[1, 3, 9, 27, 1, 3, 9, 27], blocks=2
        self.convglus = CondSequential(
            [
                ConvGLU(
                    in_channels,
                    kernels,
                    dilation,
                    dropout_rate,
                    conditionning_channels=timbre_embedding_channels,
                )
                for dilation in dilations
            ]
            + [
                ConvGLU(
                    in_channels,
                    1,
                    1,
                    dropout_rate,
                    conditionning_channels=timbre_embedding_channels,
                )
                for _ in range(blocks)
            ]
        )

        self.proj = nn.Conv1d(in_channels, in_channels, 1)

    @cuda_synchronized_timer(DO_PROFILING, prefix="FrameLevelSynthesizer")
    def forward(
        self, inputs: torch.Tensor, timbre: Optional[torch.Tensor] = None
    ) -> torch.Tensor:
        """Synthesize in frame-level.
        Args:
            inputs [torch.float32; [B, channels, T]]:
                Input features.
            timbre [torch.float32; [B, embed, T]], Optional:
                Time-varying timbre embeddings.
        Returns;
            [torch.float32; [B, channels, T]], outputs.
        """
        # [B, channels, T]
        x: torch.Tensor = self.preconv(inputs)
        # [B, channels, T]
        x = self.convglus(x, timbre)
        # [B, channels, T]
        return self.proj(x)

src/networks/backbone/synthesizers/framelevel.py
MWM-io/nansypp
from typing import Optional, Tuple

import numpy as np
import torch
import torch.nn.functional as F
from torch import nn

from src.networks.backbone.synthesizers.wavenet import WaveNet
from src.utilities.profiling import DO_PROFILING, cuda_synchronized_timer
from src.utilities.types import copy_docstring_and_signature


class SignalGenerator(nn.Module):
    """Additive sinusoidal, subtractive filtered noise signal generator."""

    def __init__(self, scale: int, input_sample_rate: int, output_sample_rate: int):
        """Initializer.
        Args:
            scale: upscaling factor.
            sample_rate: sampling rate.
        """
        super().__init__()
        self.output_sample_rate = output_sample_rate
        self.upsampler = nn.Upsample(
            scale_factor=scale * (output_sample_rate / input_sample_rate), mode="linear"
        )

    @cuda_synchronized_timer(DO_PROFILING, prefix="SignalGenerator")
    def forward(
        self,
        pitch: torch.Tensor,
        p_amp: torch.Tensor,
        ap_amp: torch.Tensor,
        noise: Optional[torch.Tensor] = None,
    ) -> torch.Tensor:
        """Generate the signal.
        Args:
            pitch: [torch.float32; [B, N]], frame-level pitch sequence.
            p_amp: [torch.float32; [B, N]], periodic amplitude.
            ap_amp: [torch.float32; [B, N]], aperiodic amplitude.
            noise: [torch.float32; [B, T]], predefined noise, if provided.
        Returns:
            [torch.float32; [B, T(=N x scale)]], base signal.
        """
        # [B, T]
        pitch = self.upsampler(pitch[:, None]).squeeze(dim=1)
        p_amp = self.upsampler(p_amp[:, None]).squeeze(dim=1)
        # [B, T]
        phase = torch.cumsum(2 * torch.pi * pitch / self.output_sample_rate, dim=-1)
        # [B, T]
        x = p_amp * torch.sin(phase)
        # [B, T]
        ap_amp = self.upsampler(ap_amp[:, None]).squeeze(dim=1)
        if noise is None:
            # [B, T], U[-1, 1] sampled
            noise = torch.rand_like(x) * 2.0 - 1.0
        # [B, T]
        y = ap_amp * noise
        return x + y

    @copy_docstring_and_signature(forward)
    def __call__(self, *args, **kwargs):  # type: ignore[no-untyped-def]
        return super().__call__(*args, **kwargs)


class Synthesizer(nn.Module):
    """Signal-level synthesizer."""

    def __init__(
        self,
        scale: int,
        input_sample_rate: int,
        output_sample_rate: int,
        channels: int,
        aux: int,
        kernels: int,
        dilation_rate: int,
        layers: int,
        cycles: int,
    ):
        """Initializer.
        Args:
            scale: upscaling factor.
            sample_rate: sampling rate.
            channels: size of the hidden channels.
            aux: size of the auxiliary input channels.
            kernels: size of the convolutional kernels.
            dilation_rate: dilation rate.
            layers: the number of the wavenet blocks in single cycle.
            cycles: the number of the cycles.
        """
        super().__init__()
        self.excitation_generator = SignalGenerator(
            scale, input_sample_rate, output_sample_rate
        )

        self.wavenet = WaveNet(channels, aux, kernels, dilation_rate, layers, cycles)

    @cuda_synchronized_timer(DO_PROFILING, prefix="Synthesizer")
    def forward(
        self,
        pitch: torch.Tensor,
        p_amp: torch.Tensor,
        ap_amp: torch.Tensor,
        frame: torch.Tensor,
        noise: Optional[torch.Tensor] = None,
    ) -> Tuple[torch.Tensor, torch.Tensor]:
        """Generate the signal.
        Args:
            pitch: [torch.float32; [B, N]], frame-level pitch sequence.
            p_amp, ap_amp: [torch.float32; [B, N]], periodical and aperiodical amplitudes.
            frame: [torch.float32; [B, aux, N']], frame-level feature map.
            noise: [torch.float32; [B, T]], predefined noise for excitation signal, if provided.
        Returns:
            [torch.float32; [B, T]], excitation signal and generated signal.
        """
        # [B, T]
        excitation = self.excitation_generator(pitch, p_amp, ap_amp, noise=noise)
        # [B, aux, T]
        interp = torch.nn.functional.interpolate(
            frame, size=excitation.shape[-1], mode="linear"
        )
        # [B, T]
        signal = self.wavenet(excitation, interp)
        return excitation, signal

    @copy_docstring_and_signature(forward)
    def __call__(self, *args, **kwargs):  # type: ignore[no-untyped-def]
        return super().__call__(*args, **kwargs)

src/networks/backbone/synthesizers/synthesizer.py
MWM-io/nansypp
from typing import Tuple

import torch
from torch import nn

from src.utilities.profiling import DO_PROFILING, cuda_synchronized_timer
from src.utilities.types import copy_docstring_and_signature


class WaveNetBlock(nn.Module):
    """WaveNet block, dilated convolution and skip connection."""

    def __init__(self, channels: int, aux: int, kernels: int, dilation: int):
        """Initializer.
        Args:
            channels: size of the input channels.
            aux: size of the auxiliary input channels.
            kernels: size of the convolutional kernel.
            dilations: dilation rate.
        """
        super().__init__()
        self.conv = nn.utils.weight_norm(
            nn.Conv1d(
                channels,
                channels * 2,
                kernels,
                padding=(kernels - 1) * dilation // 2,
                dilation=dilation,
            )
        )

        self.proj_aux = nn.utils.weight_norm(
            nn.Conv1d(aux, channels * 2, 1, bias=False)
        )

        self.proj_res = nn.utils.weight_norm(nn.Conv1d(channels, channels, 1))
        self.proj_skip = nn.utils.weight_norm(nn.Conv1d(channels, channels, 1))

    def forward(
        self, inputs: torch.Tensor, aux: torch.Tensor
    ) -> Tuple[torch.Tensor, torch.Tensor]:
        """Pass to the wavenet block.
        Args:
            inputs: [torch.float32; [B, channels, T]], input tensor.
            aux: [torch.float32; [B, aux, T]], auxiliary input tensors.
        Returns:
            residual: [torch.float32; [B, C, T]], residually connected.
            skip: [torch.float32; [B, C, T]], skip connection purposed.
        """
        # [B, C x 2, T]
        x: torch.Tensor = self.conv(inputs) + self.proj_aux(aux)
        # [B, C, T]
        gate, context = x.chunk(2, dim=1)
        # [B, C, T]
        x = torch.sigmoid(gate) * torch.tanh(context)
        # [B, C, T]
        res = (x + self.proj_res(x)) * (2**-0.5)
        return res, self.proj_skip(x)

    @copy_docstring_and_signature(forward)
    def __call__(self, *args, **kwargs):  # type: ignore[no-untyped-def]
        return super().__call__(*args, **kwargs)

    def remove_weight_norm(self):
        nn.utils.remove_weight_norm(self.conv)
        nn.utils.remove_weight_norm(self.proj_aux)
        nn.utils.remove_weight_norm(self.proj_res)
        nn.utils.remove_weight_norm(self.proj_skip)


class WaveNet(nn.Module):
    """WaveNet, Oord et al., 2016."""

    def __init__(
        self,
        channels: int,
        aux: int,
        kernels: int,
        dilation_rate: int,
        layers: int,
        cycles: int,
    ):
        """Initializer.
        Args:
            channels: size of the hidden channels.
            aux: size of the auxiliary input channels.
            kernels: size of the convolutional kernels.
            dilation_rate: base dilation rate.
            layers: the number of the wavenet blocks in single cycle.
            cycles: the number of the cycles.
        """
        super().__init__()
        # channels=64
        self.proj_signal = nn.utils.weight_norm(nn.Conv1d(1, channels, 1))
        # aux=1024, cycles=3, layers=10, dilation_rate=2
        self.blocks = nn.ModuleList(
            [
                WaveNetBlock(channels, aux, kernels, dilation_rate**j)
                for _ in range(cycles)
                for j in range(layers)
            ]
        )

        self.proj_out = nn.Sequential(
            nn.ReLU(),
            nn.utils.weight_norm(nn.Conv1d(channels, channels, 1)),
            nn.ReLU(),
            nn.utils.weight_norm(nn.Conv1d(channels, 1, 1)),
            nn.Tanh(),
        )

    @cuda_synchronized_timer(DO_PROFILING, prefix="WaveNet")
    def forward(self, noise: torch.Tensor, aux: torch.Tensor) -> torch.Tensor:
        """Generate the signal from noise and auxiliary inputs.
        Args:
            noise: [torch.float32; [B, T]], initial noise signal.
            aux: [torch.float32; [B, aux, T]], auxiliary inputs.
        Returns;
            [torch.float32; [B, T]], generated signal.
        """
        # [B, channels, T]
        x: torch.Tensor = self.proj_signal(noise[:, None])
        # (layers x cycles) x [B, channels, T]
        skips = x.new_zeros(1)
        for block in self.blocks:
            # [B, channels, T], [B, channels, T]
            x, skip = block(x, aux)
            skips = skips + skip
        # [B, T]
        return self.proj_out(skips * (len(self.blocks) ** -0.5)).squeeze(dim=1)

    @copy_docstring_and_signature(forward)
    def __call__(self, *args, **kwargs):  # type: ignore[no-untyped-def]
        return super().__call__(*args, **kwargs)

    def remove_weight_norm(self):
        nn.utils.remove_weight_norm(self.proj_signal)
        for block in self.blocks:
            block.remove_weight_norm()
        nn.utils.remove_weight_norm(self.proj_out[1])
        nn.utils.remove_weight_norm(self.proj_out[3])

src/networks/backbone/synthesizers/wavenet.py
MWM-io/nansypp
import torch
from torch import nn


class ConditionalLayerNormalization(nn.Module):
    """Conditional Layer Normalization / Style-Adaptive Layer Normalization
    [1]: Meta-StyleSpeech : Multi-Speaker Adaptive Text-to-Speech Generation,
         Min et al., 2021, arxiv.2106.03153
    """

    def __init__(self, style_dim: int, vector_dim: int):
        """Initializer.
        Args:
            style_dim:
            vector_dim:
        """
        super().__init__()

        self.linear = nn.Linear(in_features=style_dim, out_features=vector_dim)

    def forward(
        self, vector: torch.Tensor, style_embedding: torch.Tensor
    ) -> torch.Tensor:
        """
        Args:
            vector: [B, C, N]
            style_embedding: [B, C_style, 1]
        """
        features = self.linear(style_embedding.squeeze(2))
        mean, std = features.mean(), features.std()
        normalized_vector: torch.Tensor = (vector - mean) / std
        return normalized_vector

src/networks/misc/cln.py
MWM-io/nansypp
from typing import List

import torch
from torch import nn


class ConvReLUNorm(nn.Module):
    """
    Convolutional block with ReLU activations.
    """

    def __init__(
        self,
        in_channels: int,
        out_channels: int,
        kernel_size: int,
        stride: int,
        # normalized_shape: List[int],
        dropout: float,
        **kwargs,
    ):
        """Initializer.
        Args:
            in_channels:
            out_channels:
            kernel_size:
            normalized_shape
            dropout:
        """
        super().__init__()
        # Conv1D(k, d, s)
        self.conv1d = nn.Conv1d(
            in_channels=in_channels,
            out_channels=out_channels,
            kernel_size=kernel_size,
            stride=stride,
            padding=(kernel_size - 1) // 2,
        )
        # LeakyRelu
        self.relu = nn.ReLU()
        # LayerNorm
        self.layer_norm = nn.LayerNorm([out_channels])
        # Dropout
        self.dropout = nn.Dropout(p=dropout)

    def forward(self, embedding: torch.Tensor) -> torch.Tensor:
        """
        Args:
            embedding: [B, C, N]
        Returns:
            output: [B, C, N/S]
        """
        conv_features = self.conv1d(embedding)
        relu = self.relu(conv_features)
        layer_normed = self.layer_norm(relu.transpose(1, 2)).transpose(1, 2)
        dropped_out: torch.Tensor = self.dropout(layer_normed)
        return dropped_out

src/networks/misc/convrelunorm.py
MWM-io/nansypp
from typing import List, Tuple

import torch
import torch.nn.functional as F
from torch import nn

from src.utilities.profiling import DO_PROFILING, cuda_synchronized_timer
from src.utilities.types import copy_docstring_and_signature


class PeriodDiscriminator(nn.Module):
    """Period-aware discriminator."""

    channels: List[int]
    """List of the channel sizes."""
    period: int
    """Size of the unit period."""
    kernels: int
    """Size of the convolutional kernels."""
    stride: int
    """Stride of the convolutions."""
    postkernels: int
    """Size of the postnet convolutional kernels."""
    leak: float
    """Negative slope of leaky ReLUs."""

    def __init__(
        self,
        channels: List[int],
        period: int,
        kernels: int,
        stride: int,
        postkernels: int,
        leak: float,
    ):
        """Initializer.
        Args:
            channels: list of the channel sizes.
            period: size of the unit period.
            kernels: size of the convolutional kernels.
            stride: stride of the convolutions.
            postkernels: size of the postnet convolutional kernels.
            leak: negative slope of leaky ReLU.
        """
        super().__init__()
        self.period = period
        self.convs = nn.ModuleList(
            [
                nn.Sequential(
                    nn.utils.weight_norm(
                        nn.Conv2d(
                            inc,
                            outc,
                            (kernels, 1),
                            (stride, 1),
                            padding=(kernels // 2, 0),
                        )
                    ),
                    nn.LeakyReLU(leak),
                )
                for inc, outc in zip([1] + channels, channels)
            ]
        )

        lastc = channels[-1]
        self.convs.append(
            nn.Sequential(
                nn.utils.weight_norm(
                    nn.Conv2d(lastc, lastc, (kernels, 1), padding=(kernels // 2, 0))
                ),
                nn.LeakyReLU(leak),
            )
        )

        self.postconv = nn.utils.weight_norm(
            nn.Conv2d(lastc, 1, (postkernels, 1), padding=(postkernels // 2, 0))
        )

    def forward(self, inputs: torch.Tensor) -> Tuple[torch.Tensor, List[torch.Tensor]]:
        """Discriminate the inputs in multiple periods.
        Args:
            x: [torch.float32; [B, T]], input audio signal.
        Returns:
            outputs: [torch.float32; [B, S]], logits.
            feature_maps: [torch.float32; [B, C, F, P]], list of feature maps.
        """
        # B, T
        bsize, timestep = inputs.shape
        if timestep % self.period != 0:
            # padding for foldability
            padsize = self.period - timestep % self.period
            # [B, T + R]
            inputs = F.pad(inputs[:, None], (0, padsize), "reflect").squeeze(1)
            # T + R
            timestep = timestep + padsize
        # [B, 1, F(=T // P), P]
        x = inputs.view(bsize, 1, timestep // self.period, self.period)

        # period-aware discriminator
        feature_maps: List[torch.Tensor] = []
        for conv in self.convs:
            x = conv(x)
            feature_maps.append(x)
        # [B, 1, S', P]
        x = self.postconv(x)
        feature_maps.append(x)

        # [B, S]
        return x.view(bsize, -1), feature_maps


class MultiPeriodDiscriminator(nn.Module):
    """MPD: Multi-period discriminator."""

    channels: List[int]
    """List of the channel sizes. (Common to all individual discriminators.)"""
    periods: List[int]
    """Sizes of the respective unit period for each individual discriminator."""
    kernels: int
    """Size of the convolutional kernel. (Common to all individual discriminators.)"""
    stride: int
    """Stride of the convolution. (Common to all individual discriminators.)"""
    postkernels: int
    """Size of the postnet convolutional kernel. (Common to all individual discriminators.)"""
    leak: float
    """Negative slope of leaky ReLUs. (Common to all individual discriminators.)"""

    def __init__(
        self,
        periods: List[int],
        channels: List[int],
        kernels: int,
        stride: int,
        postkernels: int,
        leak: float,
    ):
        """Initializer."""
        super().__init__()
        self.periods = periods
        self.channels = channels
        self.kernels = kernels
        self.stride = stride
        self.postkernels = postkernels
        self.leak = leak

        self.discriminators = nn.ModuleList(
            [
                PeriodDiscriminator(
                    channels=list(channels),
                    period=period,
                    kernels=kernels,
                    stride=stride,
                    postkernels=postkernels,
                    leak=leak,
                )
                for period in list(periods)
            ]
        )

    @cuda_synchronized_timer(DO_PROFILING, prefix="MultiPeriodDiscriminator")
    def forward(
        self, inputs: torch.Tensor
    ) -> Tuple[List[torch.Tensor], List[List[torch.Tensor]]]:
        """Discriminate the samples from real or fake.
        Args:
            x: [B, T], audio sample.
        Returns:
            multiple discriminating results and feature maps.
        """
        results: List[torch.Tensor] = []
        features_maps_per_period: List[List[torch.Tensor]] = []

        for discriminator in self.discriminators:
            result, features_maps = discriminator(inputs)
            results.append(result)
            features_maps_per_period.append(features_maps)

        return results, features_maps_per_period

    @copy_docstring_and_signature(forward)
    def __call__(self, *args, **kwargs):  # type: ignore[no-untyped-def]
        return super().__call__(*args, **kwargs)

src/networks/misc/mpd.py
MWM-io/nansypp
import numpy as np
import torch
from torch import nn


class ParametricEqualizer(nn.Module):
    """Fast-parametric equalizer for approximation of Biquad IIR filter."""

    def __init__(self, sample_rate: int, window_length: int):
        """Initializer.
        Args:
            sr: sample rate.
            window_length: size of the fft window.
        """
        super().__init__()
        self.sample_rate = sample_rate
        self.window_length = window_length

    def forward(self) -> None:
        raise NotImplementedError("Use the various filtering methods")

    def biquad(self, a: torch.Tensor, b: torch.Tensor) -> torch.Tensor:
        """Construct frequency level biquad filter.
        Args:
            a: [torch.float32; [..., 3]], recursive filter, iir.
            b: [torch.float32; [..., 3]], finite impulse filter.
        Returns:
            [torch.float32; [..., window_length // 2 + 1]], biquad filter.
        """
        iir = torch.fft.rfft(a, self.window_length, dim=-1)
        fir = torch.fft.rfft(b, self.window_length, dim=-1)
        return fir / iir

    def low_shelving(
        self, cutoff: float, gain: torch.Tensor, q: torch.Tensor
    ) -> torch.Tensor:
        """Frequency level low-shelving filter.
        Args:
            cutoff: cutoff frequency.
            gain: [torch.float32; [...]], boost of attenutation in decibel.
            q: [torch.float32; [...]], quality factor.
        Returns:
            [torch.float32; [..., window_length // 2 + 1]], frequency filter.
        """
        # ref: torchaudio.functional.lowpass_biquad
        w0 = 2 * np.pi * cutoff / self.sample_rate
        cos_w0 = np.cos(w0)
        # [B]
        alpha = np.sin(w0) / 2 / q
        cos_w0 = torch.full_like(alpha, np.cos(w0))
        A = (gain / 40.0 * np.log(10)).exp()
        # [...], fir
        b0 = A * ((A + 1) - (A - 1) * cos_w0 + 2 * A.sqrt() * alpha)
        b1 = 2 * A * ((A - 1) - (A + 1) * cos_w0)
        b2 = A * ((A + 1) - (A - 1) * cos_w0 - 2 * A.sqrt() * alpha)
        # [...], iir
        a0 = (A + 1) + (A - 1) * cos_w0 + 2 * A.sqrt() * alpha
        a1 = -2 * ((A - 1) + (A + 1) * cos_w0)
        a2 = (A + 1) + (A - 1) * cos_w0 - 2 * A.sqrt() * alpha
        # [..., window_length // 2 + 1]
        return self.biquad(
            a=torch.stack([a0, a1, a2], dim=-1), b=torch.stack([b0, b1, b2], dim=-1)
        )

    def high_shelving(
        self, cutoff: float, gain: torch.Tensor, q: torch.Tensor
    ) -> torch.Tensor:
        """Frequency level high-shelving filter.
        Args:
            cutoff: cutoff frequency.
            gain: [torch.float32; [...]], boost of attenutation in decibel.
            q: [torch.float32; [...]], quality factor.
        Returns:
            [torch.float32; [..., window_length // 2 + 1]], frequency filter.
        """
        # ref: torchaudio.functional.highpass_biquad
        w0 = 2 * np.pi * cutoff / self.sample_rate
        # [...]
        alpha = np.sin(w0) / 2 / q
        cos_w0 = torch.full_like(alpha, np.cos(w0))
        A = (gain / 40.0 * np.log(10)).exp()
        # [...], fir
        b0 = A * ((A + 1) + (A - 1) * cos_w0 + 2 * A.sqrt() * alpha)
        b1 = -2 * A * ((A - 1) + (A + 1) * cos_w0)
        b2 = A * ((A + 1) + (A - 1) * cos_w0 - 2 * A.sqrt() * alpha)
        # [...], iir
        a0 = (A + 1) - (A - 1) * cos_w0 + 2 * A.sqrt() * alpha
        a1 = 2 * ((A - 1) - (A + 1) * cos_w0)
        a2 = (A + 1) - (A - 1) * cos_w0 - 2 * A.sqrt() * alpha
        # [..., window_length // 2 + 1]
        return self.biquad(
            a=torch.stack([a0, a1, a2], dim=-1), b=torch.stack([b0, b1, b2], dim=-1)
        )

    def peaking_equalizer(
        self, center: torch.Tensor, gain: torch.Tensor, q: torch.Tensor
    ) -> torch.Tensor:
        """Frequency level peaking equalizer.
        Args:
            center: [torch.float32; [...]], center frequency.
            gain: [torch.float32; [...]], boost or attenuation in decibel.
            q: [torch.float32; [...]], quality factor.
        Returns:
            [torch.float32; [..., window_length // 2 + 1]], frequency filter.
        """
        # ref: torchaudio.functional.highpass_biquad
        # [...]
        w0 = 2 * np.pi * center / self.sample_rate
        # [...]
        alpha = torch.sin(w0) / 2 / q
        cos_w0 = torch.cos(w0)
        A = (gain / 40.0 * np.log(10)).exp()
        # [..., window_length // 2 + 1]
        return self.biquad(
            a=torch.stack([1 + alpha / A, -2 * cos_w0, 1 - alpha / A], dim=-1),
            b=torch.stack([1 + alpha * A, -2 * cos_w0, 1 - alpha * A], dim=-1),
        )

src/networks/misc/peq.py
MWM-io/nansypp
from typing import Optional, Tuple

import numpy as np
import torch
from torch import nn

from src.data.preprocessing.augmentation import RandomBackgroundNoise
from src.networks.misc.peq import ParametricEqualizer
from src.networks.misc.praat import PraatAugment
from src.utilities.profiling import DO_PROFILING, cuda_synchronized_timer
from src.utilities.types import copy_docstring_and_signature


class InformationPerturbator(nn.Module):
    """Waveform information perturbator.

    To be applied before linguistic features extraction in
    the training stage of the backbone architecture.
    """

    formant_shift: float
    pitch_shift: float
    pitch_range: float
    cutoff_lowpass: float
    cutoff_highpass: float
    q_min: float
    q_max: float
    num_peaks: int
    gain_range: float
    window: torch.Tensor
    peak_centers: torch.Tensor

    def __init__(
        self,
        formant_shift: float,
        pitch_shift: float,
        pitch_range: float,
        cutoff_lowpass: float,
        cutoff_highpass: float,
        q_min: float,
        q_max: float,
        num_peaks: int,
        gain_range: float,
        stft_window_length: int,
        stft_hop_length: int,
        praat_augment: PraatAugment,
        parametric_equalizer: ParametricEqualizer,
        additive_noise: Optional[RandomBackgroundNoise] = None,
    ):
        super().__init__()

        self.num_peaks = num_peaks
        self.formant_shift = formant_shift
        self.pitch_shift = pitch_shift
        self.pitch_range = pitch_range
        self.cutoff_lowpass = cutoff_lowpass
        self.cutoff_highpass = cutoff_highpass
        self.q_min = q_min
        self.q_max = q_max
        self.gain_range = gain_range
        self.stft_window_length = stft_window_length
        self.stft_hop_length = stft_hop_length

        self.praat_augment = praat_augment
        self.parametric_equalizer = parametric_equalizer
        self.additive_noise = additive_noise

        self.register_buffer(
            "window",
            torch.hann_window(self.stft_window_length),
            persistent=False,
        )
        f_min, f_max, peaks = (
            self.cutoff_lowpass,
            self.cutoff_highpass,
            self.num_peaks,
        )
        # peaks except frequency min and max
        self.register_buffer(
            "peak_centers",
            f_min * (f_max / f_min) ** (torch.arange(peaks + 2)[1:-1] / (peaks + 1)),
            persistent=False,
        )

    def augment(
        self,
        wavs: torch.Tensor,
        pitch_shift: Optional[torch.Tensor] = None,
        pitch_range: Optional[torch.Tensor] = None,
        formant_shift: Optional[torch.Tensor] = None,
        quality_power: Optional[torch.Tensor] = None,
        gain: Optional[torch.Tensor] = None,
        noises: Optional[torch.Tensor] = None,
    ) -> torch.Tensor:
        """Augment the audio signal, random pitch, formant shift and PEQ.

        Augmentations are symmetric around the identity, ensuring that the unperturbed
        data remains within the model's training distribution.

        Args:
            wavs: [torch.float32; [B, T]], audio signal.
            pitch_shift: [torch.float32; [B]], pitch shifts.
            pitch_range: [torch.float32; [B]], pitch ranges.
            formant_shift: [torch.float32; [B]], formant shifts.
            quality_power: [torch.float32; [B, num_peaks + 2]],
                exponents of quality factor, for PEQ.
            gain: [torch.float32; [B, num_peaks + 2]], gain in decibel.
        Returns:
            [torch.float32; [B, T]], augmented.
        """
        # B
        bsize, _ = wavs.shape
        # [B, F, T / S], complex64
        fft = torch.stft(
            wavs,
            self.stft_window_length,
            self.stft_hop_length,
            self.stft_window_length,
            self.window,
            return_complex=True,
        )

        # PEQ
        if quality_power is not None:
            # [B, num_peaks + 2]
            q = self.q_min * (self.q_max / self.q_min) ** quality_power
            if gain is None:
                # [B, num_peaks]
                gain = torch.zeros_like(q[:, :-2])
            # [B, num_peaks]
            center = self.peak_centers[None].repeat(bsize, 1)
            # [B, F]
            peaks = torch.prod(
                self.parametric_equalizer.peaking_equalizer(
                    center, gain[:, :-2], q[:, :-2]
                ),
                dim=1,
            )
            # [B, F]
            lowpass = self.parametric_equalizer.low_shelving(
                self.cutoff_lowpass, gain[:, -2], q[:, -2]
            )
            highpass = self.parametric_equalizer.high_shelving(
                self.cutoff_highpass, gain[:, -1], q[:, -1]
            )
            # [B, F]
            filters = peaks * highpass * lowpass
            # [B, F, T / S]
            fft = fft * filters[..., None]

        # [B, T]
        out = torch.istft(
            fft,
            self.stft_window_length,
            self.stft_hop_length,
            self.stft_window_length,
            self.window,
        ).clamp(-1.0, 1.0)
        # max value normalization
        out = out / out.abs().amax(dim=-1, keepdim=True).clamp_min(1e-7)

        if (
            formant_shift is None
            and pitch_shift is None
            and pitch_range is None
            and noises is None
        ):
            return out

        # praat-based augmentation
        if formant_shift is None:
            formant_shift = torch.ones(bsize)
        if pitch_shift is None:
            pitch_shift = torch.ones(bsize)
        if pitch_range is None:
            pitch_range = torch.ones(bsize)
        out = torch.tensor(
            np.stack(
                [
                    self.praat_augment.augment(o, fs.item(), ps.item(), pr.item())
                    for o, fs, ps, pr in zip(
                        out.cpu().numpy(),
                        formant_shift.cpu().numpy(),
                        pitch_shift.cpu().numpy(),
                        pitch_range.cpu().numpy(),
                    )
                ],
                axis=0,
            ),
            device=out.device,
            dtype=torch.float32,
        )
        if self.additive_noise is not None:
            out = self.additive_noise(out, noises=noises)
        return out

    def sample_like(
        self, signal: torch.Tensor
    ) -> Tuple[
        torch.Tensor,
        torch.Tensor,
        torch.Tensor,
        torch.Tensor,
        torch.Tensor,
        torch.Tensor,
    ]:
        """Sample augmentation parameters.
        Args:
            signal: [torch.float32; [B, T]], speech signal.
        Returns:
            augmentation parameters.
        """
        # [B]
        batch_size, audio_length = signal.shape

        def sampler(ratio: float) -> torch.Tensor:
            shifts = torch.rand(batch_size, device=signal.device) * (ratio - 1.0) + 1.0
            # flip
            flip = torch.rand(batch_size) < 0.5
            shifts[flip] = shifts[flip] ** -1
            return shifts

        # sample shifts
        formant_shifts = sampler(self.formant_shift)
        pitch_shifts = sampler(self.pitch_shift)
        pitch_ranges = sampler(self.pitch_range)
        # parametric equalizer
        peaks = self.num_peaks
        # quality factor
        power = torch.rand(batch_size, peaks + 2, device=signal.device)
        # gains
        g_min, g_max = -self.gain_range, self.gain_range
        gain = (
            torch.rand(batch_size, peaks + 2, device=signal.device) * (g_max - g_min)
            + g_min
        )
        # additive noise
        noise = None
        if self.additive_noise is not None:
            noise = self.additive_noise.sample(batch_size, audio_length).to(
                signal.device
            )
        return formant_shifts, pitch_shifts, pitch_ranges, power, gain, noise

    @cuda_synchronized_timer(DO_PROFILING, prefix="Synthesizer")
    @torch.no_grad()
    def forward(self, wavs: torch.Tensor) -> torch.Tensor:
        """Augment the speech.
        Args:
            wavs: [torch.float32; [B, T]], segmented speech.
        Returns:
            [torch.float32; [B, T]], speech signal.
        """
        # B
        bsize, _ = wavs.shape
        saves = None
        while saves is None or len(saves) < bsize:
            # [B] x 4
            (
                formant_shifts,
                pitch_shifts,
                pitch_ranges,
                power,
                gain,
                noises,
            ) = self.sample_like(wavs)
            # [B, T]
            out = self.augment(
                wavs, pitch_shifts, pitch_ranges, formant_shifts, power, gain, noises
            )

            # handle unexpected NaNs
            nan = out.isnan().any(dim=-1)
            if not nan.all():
                # save the outputs for not-nan inputs
                if saves is None:
                    saves = out[~nan]
                else:
                    saves = torch.cat([saves, out[~nan]], dim=0)

        # [B, T]
        return saves[:bsize]

    @copy_docstring_and_signature(forward)
    def __call__(self, *args, **kwargs):  # type: ignore[no-untyped-def]
        return super().__call__(*args, **kwargs)

src/networks/misc/perturbator.py
MWM-io/nansypp
import math

import torch
from torch import nn


class PositionalEncoding(nn.Module):
    def __init__(self, d_model: int, dropout: float = 0.1, max_len: int = 5000):
        super().__init__()
        self.dropout = nn.Dropout(p=dropout)

        position = torch.arange(max_len).unsqueeze(1)
        div_term = torch.exp(
            torch.arange(0, d_model, 2) * (-math.log(10000.0) / d_model)
        )
        pe = torch.zeros(max_len, 1, d_model)
        pe[:, 0, 0::2] = torch.sin(position * div_term)
        pe[:, 0, 1::2] = torch.cos(position * div_term)
        self.register_buffer("pe", pe)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        """
        Arguments:
            x: Tensor, shape ``[batch_size, seq_len, embedding_dim]``
        """
        x = x.transpose(0, 1)
        x = x + self.pe[: x.size(0)]
        x = x.transpose(0, 1)
        return self.dropout(x)

src/networks/misc/positional_encoding.py
MWM-io/nansypp
from typing import Union

import numpy as np
import parselmouth


class PraatAugment:
    """Praat-based pitch, formant and duration augmentations."""

    def __init__(
        self,
        sample_rate: int,
        pitch_steps: float = 0.01,
        pitch_floor: float = 75,
        pitch_ceil: float = 600,
    ):
        """Initializer.
        Args:
            config: configurations.
            pitch_steps: pitch measurement intervals.
            pitch_floor: minimum pitch.
            pitch_ceil: maximum pitch.
        """
        self.sample_rate = sample_rate
        self.pitch_steps = pitch_steps
        self.pitch_floor = pitch_floor
        self.pitch_ceil = pitch_ceil

    def augment(
        self,
        snd: Union[parselmouth.Sound, np.ndarray],
        formant_shift: float = 1.0,
        pitch_shift: float = 1.0,
        pitch_range: float = 1.0,
        duration_factor: float = 1.0,
    ) -> np.ndarray:
        """Augment the sound signal with praat."""
        if not isinstance(snd, parselmouth.Sound):
            snd = parselmouth.Sound(snd, sampling_frequency=self.sample_rate)
        pitch = parselmouth.praat.call(
            snd, "To Pitch", self.pitch_steps, self.pitch_floor, self.pitch_ceil
        )
        ndpit = pitch.selected_array["frequency"]
        # if all unvoiced
        nonzero = ndpit > 1e-5
        if nonzero.sum() == 0:
            return snd.values[0]
        # if voiced
        median, minp = np.median(ndpit[nonzero]).item(), ndpit[nonzero].min().item()
        # scale
        updated = median * pitch_shift
        scaled = updated + (minp * pitch_shift - updated) * pitch_range
        # for preventing infinite loop of `Change gender`
        # ref:https://github.com/praat/praat/issues/1926
        if scaled < 0.0:
            pitch_range = 1.0
        (out,) = parselmouth.praat.call(
            (snd, pitch),
            "Change gender",
            formant_shift,
            median * pitch_shift,
            pitch_range,
            duration_factor,
        ).values
        return out

src/networks/misc/praat.py
MWM-io/nansypp
"""
origin: https://github.com/KinWaiCheuk/nnAudio
commit: dc222b438ecac32492354a2df0b4faa24f473cc2
(last commited in Oct 9, 2022)

---
MIT License

Copyright (c) 2019 KinWaiCheuk

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.
"""
import warnings
from typing import Optional, Tuple

import librosa
import numpy as np
import torch
import torch.nn.functional as F
from scipy import signal
from torch import nn


def nextpow2(A: float) -> int:
    """A helper function to calculate the next nearest number to the power of 2.
    Parameters
    ----------
    A : float
        A float number that is going to be rounded up to the nearest power of 2
    Returns
    -------
    int
        The nearest power of 2 to the input number ``A``
    Examples
    --------
    >>> nextpow2(6)
    3
    """

    return int(np.ceil(np.log2(A)))


def create_lowpass_filter(
    band_center: float = 0.5,
    kernel_length: int = 256,
    transition_bandwidth: float = 0.03,
) -> np.ndarray:
    """
    Calculate the highest frequency we need to preserve and the lowest frequency we allow
    to pass through.
    Note that frequency is on a scale from 0 to 1 where 0 is 0 and 1 is Nyquist frequency of
    the signal BEFORE downsampling.
    """

    # transitionBandwidth = 0.03
    passband_max = band_center / (1 + transition_bandwidth)
    stopband_min = band_center * (1 + transition_bandwidth)

    # Unlike the filter tool we used online yesterday, this tool does
    # not allow us to specify how closely the filter matches our
    # specifications. Instead, we specify the length of the kernel.
    # The longer the kernel is, the more precisely it will match.
    # kernelLength = 256

    # We specify a list of key frequencies for which we will require
    # that the filter match a specific output gain.
    # From [0.0 to passbandMax] is the frequency range we want to keep
    # untouched and [stopbandMin, 1.0] is the range we want to remove
    key_frequencies = [0.0, passband_max, stopband_min, 1.0]

    # We specify a list of output gains to correspond to the key
    # frequencies listed above.
    # The first two gains are 1.0 because they correspond to the first
    # two key frequencies. the second two are 0.0 because they
    # correspond to the stopband frequencies
    gain_at_key_frequencies = [1.0, 1.0, 0.0, 0.0]

    # This command produces the filter kernel coefficients
    filter_kernel = signal.firwin2(
        kernel_length, key_frequencies, gain_at_key_frequencies
    )

    return filter_kernel.astype(np.float32)


def create_cqt_kernels(
    Q: float,
    fs: float,
    fmin: float,
    n_bins: int = 84,
    bins_per_octave: int = 12,
    norm: float = 1,
    window: str = "hann",
    fmax: Optional[float] = None,
    topbin_check: bool = True,
    gamma: float = 0,
    pad_fft: bool = True,
) -> Tuple[np.ndarray, int, torch.Tensor, np.ndarray]:
    """
    Automatically create CQT kernels in time domain
    """

    fft_len = 2 ** nextpow2(np.ceil(Q * fs / fmin))
    # minWin = 2**nextpow2(np.ceil(Q * fs / fmax))

    if (fmax is not None) and (n_bins is None):
        n_bins = np.ceil(
            bins_per_octave * np.log2(fmax / fmin)
        )  # Calculate the number of bins
        freqs = fmin * 2.0 ** (np.r_[0:n_bins] / np.float(bins_per_octave))

    elif (fmax is None) and (n_bins is not None):
        freqs = fmin * 2.0 ** (np.r_[0:n_bins] / np.float(bins_per_octave))

    else:
        warnings.warn("If fmax is given, n_bins will be ignored", SyntaxWarning)
        n_bins = np.ceil(
            bins_per_octave * np.log2(fmax / fmin)
        )  # Calculate the number of bins
        freqs = fmin * 2.0 ** (np.r_[0:n_bins] / np.float(bins_per_octave))

    if np.max(freqs) > fs / 2 and topbin_check:
        raise ValueError(
            "The top bin {}Hz has exceeded the Nyquist frequency, \
                          please reduce the n_bins".format(
                np.max(freqs)
            )
        )

    alpha = 2.0 ** (1.0 / bins_per_octave) - 1.0
    lengths = np.ceil(Q * fs / (freqs + gamma / alpha))

    # get max window length depending on gamma value
    max_len = int(max(lengths))
    fft_len = int(2 ** (np.ceil(np.log2(max_len))))

    temp_kernel = np.zeros((int(n_bins), int(fft_len)), dtype=np.complex64)
    spec_kernel = np.zeros((int(n_bins), int(fft_len)), dtype=np.complex64)

    for k in range(0, int(n_bins)):
        freq = freqs[k]
        l = lengths[k]

        # Centering the kernels
        if l % 2 == 1:  # pad more zeros on RHS
            start = int(np.ceil(fft_len / 2.0 - l / 2.0)) - 1
        else:
            start = int(np.ceil(fft_len / 2.0 - l / 2.0))

        window_dispatch = librosa.filters.get_window(window, int(l), fftbins=True)
        sig = (
            window_dispatch
            * np.exp(np.r_[-l // 2 : l // 2] * 1j * 2 * np.pi * freq / fs)
            / l
        )

        if norm:  # Normalizing the filter # Trying to normalize like librosa
            temp_kernel[k, start : start + int(l)] = sig / np.linalg.norm(sig, norm)
        else:
            temp_kernel[k, start : start + int(l)] = sig
        # specKernel[k, :] = fft(tempKernel[k])

    # return specKernel[:,:fftLen//2+1], fftLen, torch.tensor(lenghts).float()
    return temp_kernel, fft_len, torch.tensor(lengths).float(), freqs


# The following two downsampling count functions are obtained from librosa CQT
# They are used to determine the number of pre resamplings if the starting and ending frequency
# are both in low frequency regions.
def early_downsample_count(
    nyquist: float, filter_cutoff: float, hop_length: int, n_octaves: int
) -> int:
    """Compute the number of early downsampling operations"""

    downsample_count1 = max(
        0, int(np.ceil(np.log2(0.85 * nyquist / filter_cutoff)) - 1) - 1
    )
    # print("downsample_count1 = ", downsample_count1)
    num_twos = nextpow2(hop_length)
    downsample_count2 = max(0, num_twos - n_octaves + 1)
    # print("downsample_count2 = ",downsample_count2)

    return min(downsample_count1, downsample_count2)


def early_downsample(
    sr: float, hop_length: int, n_octaves: int, nyquist: float, filter_cutoff: float
) -> Tuple[float, int, int]:
    """Return new sampling rate and hop length after early dowansampling"""
    downsample_count = early_downsample_count(
        nyquist, filter_cutoff, hop_length, n_octaves
    )
    # print("downsample_count = ", downsample_count)
    downsample_factor = 2 ** (downsample_count)

    hop_length //= downsample_factor  # Getting new hop_length
    new_sr = sr / float(downsample_factor)  # Getting new sampling rate
    sr = new_sr

    return sr, hop_length, downsample_factor


def get_early_downsample_params(
    sr: float, hop_length: int, fmax_t: float, Q: float, n_octaves: int, verbose: bool
) -> Tuple[float, int, int, Optional[torch.Tensor], bool]:
    """Used in CQT2010 and CQT2010v2"""

    window_bandwidth = 1.5  # for hann window
    filter_cutoff = fmax_t * (1 + 0.5 * window_bandwidth / Q)
    sr, hop_length, downsample_factor = early_downsample(
        sr, hop_length, n_octaves, sr // 2, filter_cutoff
    )
    if downsample_factor != 1:
        if verbose:
            print("Can do early downsample, factor = ", downsample_factor)
        earlydownsample = True
        # print("new sr = ", sr)
        # print("new hop_length = ", hop_length)
        early_downsample_filter_np = create_lowpass_filter(
            band_center=1 / downsample_factor,
            kernel_length=256,
            transition_bandwidth=0.03,
        )
        early_downsample_filter = torch.tensor(early_downsample_filter_np)[
            None, None, :
        ]
    else:
        if verbose:
            print(
                "No early downsampling is required, downsample_factor = ",
                downsample_factor,
            )
        early_downsample_filter = None
        earlydownsample = False

    return sr, hop_length, downsample_factor, early_downsample_filter, earlydownsample


def get_cqt_complex(
    x: torch.Tensor,
    cqt_kernels_real: torch.Tensor,
    cqt_kernels_imag: torch.Tensor,
    hop_length: int,
    padding: nn.Module,
) -> torch.Tensor:
    """Multiplying the STFT result with the cqt_kernel, check out the 1992 CQT paper [1]
    for how to multiple the STFT result with the CQT kernel
    [2] Brown, Judith C.C. and Miller Puckette. “An efficient algorithm for the calculation of
    a constant Q transform.” (1992)."""

    # STFT, converting the audio input from time domain to frequency domain
    try:
        x = padding(
            x
        )  # When center is True, we need padding at the beginning and ending
    except:
        warnings.warn(
            f"\ninput size = {x.shape}\tkernel size = {cqt_kernels_real.shape[-1]}\n"
            "padding with reflection mode might not be the best choice, using constant padding instead here",
            UserWarning,
        )
        x = torch.nn.functional.pad(
            x, (cqt_kernels_real.shape[-1] // 2, cqt_kernels_real.shape[-1] // 2)
        )
    CQT_real = F.conv1d(x, cqt_kernels_real, stride=hop_length)
    CQT_imag = -F.conv1d(x, cqt_kernels_imag, stride=hop_length)

    return torch.stack((CQT_real, CQT_imag), -1)


def downsampling_by_n(
    x: torch.Tensor, filterKernel: torch.Tensor, n: int
) -> torch.Tensor:
    """A helper function that downsamples the audio by a arbitary factor n.
    It is used in CQT2010 and CQT2010v2.
    Parameters
    ----------
    x : torch.Tensor
        The input waveform in ``torch.Tensor`` type with shape ``(batch, 1, len_audio)``
    filterKernel : torch.Tensor
        Filter kernel in ``torch.Tensor`` type with shape ``(1, 1, len_kernel)``
    n : int
        The downsampling factor
    Returns
    -------
    torch.Tensor
        The downsampled waveform
    Examples
    --------
    >>> x_down = downsampling_by_n(x, filterKernel)
    """

    x = F.conv1d(x, filterKernel, stride=n, padding=(filterKernel.shape[-1] - 1) // 2)
    return x


def downsampling_by_2(x: torch.Tensor, filterKernel: torch.Tensor) -> torch.Tensor:
    """A helper function that downsamples the audio by half. It is used in CQT2010 and CQT2010v2
    Parameters
    ----------
    x : torch.Tensor
        The input waveform in ``torch.Tensor`` type with shape ``(batch, 1, len_audio)``
    filterKernel : torch.Tensor
        Filter kernel in ``torch.Tensor`` type with shape ``(1, 1, len_kernel)``
    Returns
    -------
    torch.Tensor
        The downsampled waveform
    Examples
    --------
    >>> x_down = downsampling_by_2(x, filterKernel)
    """

    x = F.conv1d(x, filterKernel, stride=2, padding=(filterKernel.shape[-1] - 1) // 2)
    return x


class CQT2010v2(nn.Module):
    """This function is to calculate the CQT of the input signal.
    Input signal should be in either of the following shapes.\n
    1. ``(len_audio)``\n
    2. ``(num_audio, len_audio)``\n
    3. ``(num_audio, 1, len_audio)``
    The correct shape will be inferred autommatically if the input follows these 3 shapes.
    Most of the arguments follow the convention from librosa.
    This class inherits from ``nn.Module``, therefore, the usage is same as ``nn.Module``.
    This alogrithm uses the resampling method proposed in [1].
    Instead of convoluting the STFT results with a gigantic CQT kernel covering the full frequency
    spectrum, we make a small CQT kernel covering only the top octave. Then we keep downsampling the
    input audio by a factor of 2 to convoluting it with the small CQT kernel.
    Everytime the input audio is downsampled, the CQT relative to the downsampled input is equivalent
    to the next lower octave.
    The kernel creation process is still same as the 1992 algorithm. Therefore, we can reuse the
    code from the 1992 alogrithm [2]
    [1] Schörkhuber, Christian. “CONSTANT-Q TRANSFORM TOOLBOX FOR MUSIC PROCESSING.” (2010).
    [2] Brown, Judith C.C. and Miller Puckette. “An efficient algorithm for the calculation of a
    constant Q transform.” (1992).
    Early downsampling factor is to downsample the input audio to reduce the CQT kernel size.
    The result with and without early downsampling are more or less the same except in the very low
    frequency region where freq < 40Hz.
    Parameters
    ----------
    sr : int
        The sampling rate for the input audio. It is used to calculate the correct ``fmin`` and ``fmax``.
        Setting the correct sampling rate is very important for calculating the correct frequency.
    hop_length : int
        The hop (or stride) size. Default value is 512.
    fmin : float
        The frequency for the lowest CQT bin. Default is 32.70Hz, which coresponds to the note C0.
    fmax : float
        The frequency for the highest CQT bin. Default is ``None``, therefore the higest CQT bin is
        inferred from the ``n_bins`` and ``bins_per_octave``.  If ``fmax`` is not ``None``, then the
        argument ``n_bins`` will be ignored and ``n_bins`` will be calculated automatically.
        Default is ``None``
    n_bins : int
        The total numbers of CQT bins. Default is 84. Will be ignored if ``fmax`` is not ``None``.
    bins_per_octave : int
        Number of bins per octave. Default is 12.
    norm : bool
        Normalization for the CQT result.
    basis_norm : int
        Normalization for the CQT kernels. ``1`` means L1 normalization, and ``2`` means L2 normalization.
        Default is ``1``, which is same as the normalization used in librosa.
    window : str
        The windowing function for CQT. It uses ``scipy.signal.get_window``, please refer to
        scipy documentation for possible windowing functions. The default value is 'hann'
    pad_mode : str
        The padding method. Default value is 'reflect'.
    trainable : bool
        Determine if the CQT kernels are trainable or not. If ``True``, the gradients for CQT kernels
        will also be caluclated and the CQT kernels will be updated during model training.
        Default value is ``False``
    output_format : str
        Determine the return type.
        'Magnitude' will return the magnitude of the STFT result, shape = ``(num_samples, freq_bins, time_steps)``;
        'Complex' will return the STFT result in complex number, shape = ``(num_samples, freq_bins, time_steps, 2)``;
        'Phase' will return the phase of the STFT reuslt, shape = ``(num_samples, freq_bins,time_steps, 2)``.
        The complex number is stored as ``(real, imag)`` in the last axis. Default value is 'Magnitude'.
    verbose : bool
        If ``True``, it shows layer information. If ``False``, it suppresses all prints.
    Returns
    -------
    spectrogram : torch.tensor
    It returns a tensor of spectrograms.
    shape = ``(num_samples, freq_bins,time_steps)`` if ``output_format='Magnitude'``;
    shape = ``(num_samples, freq_bins,time_steps, 2)`` if ``output_format='Complex' or 'Phase'``;
    Examples
    --------
    >>> spec_layer = Spectrogram.CQT2010v2()
    >>> specs = spec_layer(x)
    """

    # To DO:
    # need to deal with the filter and other tensors

    lengths: torch.Tensor
    cqt_kernels_real: torch.Tensor
    cqt_kernels_imag: torch.Tensor
    early_downsample_filter: torch.Tensor
    lowpass_filter: torch.Tensor

    def __init__(
        self,
        sr: float = 22050,
        hop_length: int = 512,
        fmin: float = 32.70,
        fmax: Optional[float] = None,
        n_bins: int = 84,
        filter_scale: int = 1,
        bins_per_octave: int = 12,
        norm: bool = True,
        basis_norm: int = 1,
        window: str = "hann",
        pad_mode: str = "reflect",
        earlydownsample: bool = True,
        trainable: bool = False,
        output_format: str = "Magnitude",
    ) -> None:
        super().__init__()

        self.norm = (
            norm  # Now norm is used to normalize the final CQT result by dividing n_fft
        )
        # basis_norm is for normalizing basis
        self.hop_length = hop_length
        self.pad_mode = pad_mode
        self.n_bins = n_bins
        self.earlydownsample = (
            earlydownsample  # We will activate early downsampling later if possible
        )
        self.trainable = trainable
        self.output_format = output_format

        # It will be used to calculate filter_cutoff and creating CQT kernels
        Q = float(filter_scale) / (2 ** (1 / bins_per_octave) - 1)

        # Creating lowpass filter and make it a torch tensor
        lowpass_filter = torch.tensor(
            create_lowpass_filter(
                band_center=0.50, kernel_length=256, transition_bandwidth=0.001
            )
        )

        # Broadcast the tensor to the shape that fits conv1d
        self.register_buffer("lowpass_filter", lowpass_filter[None, None, :])

        # Caluate num of filter requires for the kernel
        # n_octaves determines how many resampling requires for the CQT
        n_filters = min(bins_per_octave, n_bins)
        self.n_octaves = int(np.ceil(float(n_bins) / bins_per_octave))

        # Calculate the lowest frequency bin for the top octave kernel
        self.fmin_t = fmin * 2 ** (self.n_octaves - 1)
        remainder = n_bins % bins_per_octave
        # print("remainder = ", remainder)

        if remainder == 0:
            # Calculate the top bin frequency
            fmax_t = self.fmin_t * 2 ** ((bins_per_octave - 1) / bins_per_octave)
        else:
            # Calculate the top bin frequency
            fmax_t = self.fmin_t * 2 ** ((remainder - 1) / bins_per_octave)

        self.fmin_t = fmax_t / 2 ** (
            1 - 1 / bins_per_octave
        )  # Adjusting the top minium bins
        if fmax_t > sr / 2:
            raise ValueError(
                "The top bin {}Hz has exceeded the Nyquist frequency, \
                            please reduce the n_bins".format(
                    fmax_t
                )
            )

        if self.earlydownsample:
            (
                sr,
                self.hop_length,
                self.downsample_factor,
                early_downsample_filter,
                self.earlydownsample,
            ) = get_early_downsample_params(
                sr, hop_length, fmax_t, Q, self.n_octaves, False
            )
            self.register_buffer("early_downsample_filter", early_downsample_filter)
        else:
            self.downsample_factor = 1

        # Preparing CQT kernels
        basis, self.n_fft, lengths, _ = create_cqt_kernels(
            Q,
            sr,
            self.fmin_t,
            n_filters,
            bins_per_octave,
            norm=basis_norm,
            topbin_check=False,
        )
        # For normalization in the end
        # The freqs returned by create_cqt_kernels cannot be used
        # Since that returns only the top octave bins
        # We need the information for all freq bin
        freqs = fmin * 2.0 ** (np.r_[0:n_bins] / np.float(bins_per_octave))
        self.frequencies = freqs

        lengths = np.ceil(Q * sr / freqs)
        lengths = torch.tensor(lengths).float()
        self.register_buffer("lengths", lengths)

        self.basis = basis
        # These cqt_kernel is already in the frequency domain
        cqt_kernels_real = torch.tensor(basis.real).unsqueeze(1)
        cqt_kernels_imag = torch.tensor(basis.imag).unsqueeze(1)

        if trainable:
            cqt_kernels_real = nn.Parameter(cqt_kernels_real, requires_grad=trainable)
            cqt_kernels_imag = nn.Parameter(cqt_kernels_imag, requires_grad=trainable)
            self.register_parameter("cqt_kernels_real", cqt_kernels_real)
            self.register_parameter("cqt_kernels_imag", cqt_kernels_imag)
        else:
            self.register_buffer("cqt_kernels_real", cqt_kernels_real)
            self.register_buffer("cqt_kernels_imag", cqt_kernels_imag)

        # If center==True, the STFT window will be put in the middle, and paddings at the beginning
        # and ending are required.
        if self.pad_mode == "constant":
            self.padding = nn.ConstantPad1d(self.n_fft // 2, 0)
        elif self.pad_mode == "reflect":
            self.padding = nn.ReflectionPad1d(self.n_fft // 2)

    def forward(
        self,
        x: torch.Tensor,
        output_format: Optional[str] = None,
        normalization_type: str = "librosa",
    ) -> torch.Tensor:
        """
        Convert a batch of waveforms to CQT spectrograms.
        Parameters
        ----------
        x : torch tensor
            Input signal should be in either of the following shapes.\n
            1. ``(len_audio)``\n
            2. ``(num_audio, len_audio)``\n
            3. ``(num_audio, 1, len_audio)``
            It will be automatically broadcast to the right shape
        """
        output_format = output_format or self.output_format

        if self.earlydownsample:
            x = downsampling_by_n(
                x, self.early_downsample_filter, self.downsample_factor
            )
        hop = self.hop_length
        CQT = get_cqt_complex(
            x, self.cqt_kernels_real, self.cqt_kernels_imag, hop, self.padding
        )  # Getting the top octave CQT

        x_down = x  # Preparing a new variable for downsampling

        for _ in range(self.n_octaves - 1):
            hop = hop // 2
            x_down = downsampling_by_2(x_down, self.lowpass_filter)
            CQT1 = get_cqt_complex(
                x_down, self.cqt_kernels_real, self.cqt_kernels_imag, hop, self.padding
            )
            CQT = torch.cat((CQT1, CQT), 1)

        CQT = CQT[:, -self.n_bins :, :]  # Removing unwanted bottom bins
        # print("downsample_factor = ",self.downsample_factor)
        # print(CQT.shape)
        # print(self.lenghts.view(-1,1).shape)

        # Normalizing the output with the downsampling factor, 2**(self.n_octaves-1) is make it
        # same mag as 1992
        CQT = CQT * self.downsample_factor
        # Normalize again to get same result as librosa
        if normalization_type == "librosa":
            CQT = CQT * torch.sqrt(self.lengths.view(-1, 1, 1))
        elif normalization_type == "convolutional":
            pass
        elif normalization_type == "wrap":
            CQT *= 2
        else:
            raise ValueError(
                "The normalization_type %r is not part of our current options."
                % normalization_type
            )

        if output_format == "Magnitude":
            if not self.trainable:
                # Getting CQT Amplitude
                return torch.sqrt(CQT.pow(2).sum(-1))
            else:
                return torch.sqrt(CQT.pow(2).sum(-1) + 1e-8)

        elif output_format == "Complex":
            return CQT

        elif output_format == "Phase":
            phase_real = torch.cos(torch.atan2(CQT[:, :, :, 1], CQT[:, :, :, 0]))
            phase_imag = torch.sin(torch.atan2(CQT[:, :, :, 1], CQT[:, :, :, 0]))
            return torch.stack((phase_real, phase_imag), -1)

        else:
            raise ValueError(
                "The output format %r is not part of our current options."
                % output_format
            )

src/networks/misc/rough_cqt.py
MWM-io/nansypp
from typing import Optional

import torch
import torchaudio
from torch import nn

from src.networks.misc.rough_cqt import CQT2010v2
from src.utilities.profiling import DO_PROFILING, cuda_synchronized_timer
from src.utilities.types import copy_docstring_and_signature


class ConstantQTransform(nn.Module):
    """Constant Q-Transform."""

    hop_length: int
    """The number of samples between adjacent frame."""

    fmin: float
    """The minimum frequency."""

    bins: int
    """The number of output bins."""

    bins_per_octave: int
    """The number of frequency bins per octave."""

    sample_rate: int
    """The sampling rate."""

    def __init__(
        self,
        hop_length: int,
        fmin: float,
        bins: int,
        bins_per_octave: int,
        sample_rate: int,
    ):
        """Initializer.
        Args:
            hop_length: The number of samples between adjacent frame.
            fmin: The minimum frequency.
            bins: The number of output bins.
            bins_per_octave: The number of frequency bins per octave.
            sample_rate: The sampling rate.
        """
        super().__init__()
        # unknown `hop_length`
        # , since linguistic information is 50fps, hop_length could be 441
        self.hop_length = hop_length
        # fmin=32.7(C0)
        self.fmin = fmin
        # bins=191, bins_per_octave=24
        # , fmax = 2 ** (bins / bins_per_octave) * fmin
        #        = 2 ** (191 / 24) * 32.7
        #        = 8132.89

        self.bins = bins
        self.bins_per_octave = bins_per_octave
        self.sample_rate = sample_rate

        self.cqt = CQT2010v2(
            sample_rate,
            hop_length,
            fmin,
            n_bins=bins,
            bins_per_octave=bins_per_octave,
            trainable=False,
            output_format="Magnitude",
        )

    def forward(self, inputs: torch.Tensor) -> torch.Tensor:
        """Apply CQT on inputs.
        Args:
            inputs: [torch.float32; [B, T]], input speech signal.
        Returns:
            [torch.float32; [B, bins, T / hop_length]], CQT magnitudes.
        """
        return self.cqt(inputs[:, None])

    @copy_docstring_and_signature(forward)
    def __call__(self, *args, **kwargs):  # type: ignore[no-untyped-def]
        return super().__call__(*args, **kwargs)


class MelSpectrogram(nn.Module):
    """log-Mel scale spectrogram."""

    hann: torch.Tensor
    melfilter: torch.Tensor

    def __init__(
        self,
        hop_length: int,
        window_length: int,
        mel: int,
        fmin: float,
        sample_rate: int,
        fmax: Optional[float] = None,
    ):
        """Initializer.
        Args:
            hop_length: Hop length, the number of frames between adjacent windows.
            window_length: The length of the window.
            mel: The size of the mel filterbanks.
            fmin: The minimum frequency.
            fmax: The maximum frequency,
                if None, uses half of the sample rate as default.
            sample_rate: The sample rate.
        """
        super().__init__()
        self.hop_length = hop_length
        self.window_length = window_length

        if fmax is None:
            fmax = sample_rate / 2
        # [mel, window_length // 2 + 1]
        # use slaney-scale mel filterbank for `librosa.filters.mel` compatibility.
        melfilter = torchaudio.functional.melscale_fbanks(
            window_length // 2 + 1,
            fmin,
            fmax,
            mel,
            sample_rate,
            norm="slaney",
            mel_scale="slaney",
        ).T

        self.register_buffer(
            "melfilter",
            melfilter,
            persistent=False,
        )
        # [window_length], use hann window
        self.register_buffer("hann", torch.hann_window(window_length), persistent=False)

    @cuda_synchronized_timer(do_profiling=DO_PROFILING, prefix="MelSpectrogram")
    def forward(self, audio: torch.Tensor) -> torch.Tensor:
        """Generate the log-mel scale spectrogram.
        Args:
            audio: [torch.float32; [B, T]], audio signal, [-1, 1]-ranged.
        Returns:
            [torch.float32; [B, mel, T / hop_length]], log-mel spectrogram
        """
        # [B, window_length // 2 + 1, T / hop_length, 2]
        fft = torch.view_as_real(
            torch.stft(
                audio,
                self.window_length,
                self.hop_length,
                window=self.hann,
                center=True,
                pad_mode="reflect",
                return_complex=True,
            )
        )
        # [B, window_length // 2 + 1, T / hop_length]
        mag = torch.sqrt(fft.square().sum(dim=-1) + 1e-7)
        # [B, mel, T / hop_length]
        return torch.log(torch.matmul(self.melfilter, mag) + 1e-7)

    @copy_docstring_and_signature(forward)
    def __call__(self, *args, **kwargs):  # type: ignore[no-untyped-def]
        return super().__call__(*args, **kwargs)

src/networks/misc/transform.py
MWM-io/nansypp
import torch
from torch import nn

from src.networks.misc.cln import ConditionalLayerNormalization
from src.networks.misc.positional_encoding import PositionalEncoding


class TransformerBlock(nn.Module):
    """Transformer block."""

    def __init__(
        self,
        channels: int,
        conditionning_channels: int,
        nhead: int,
        num_encoder_layers: int,
        num_decoder_layers: int,
        dim_feedforward: int,
        dropout: float,
    ) -> None:
        """Initializer."""
        super().__init__()
        self.pos_encoder = PositionalEncoding(channels, dropout=dropout)
        self.transformer = nn.Transformer(
            channels,
            nhead,
            num_encoder_layers,
            num_decoder_layers,
            dim_feedforward,
        )
        self.dropout = nn.Dropout(dropout)
        self.cond1 = ConditionalLayerNormalization(conditionning_channels, channels)
        self.linear = nn.Linear(channels, channels)
        self.cond2 = ConditionalLayerNormalization(conditionning_channels, channels)

    def forward(
        self, features: torch.Tensor, style_embedding: torch.Tensor
    ) -> torch.Tensor:
        """
        Args:
            features: [B, C, N]
            style_embedding: [B, C_style, 1]
        """
        x_transposed = features.transpose(1, 2)
        x_transposed = self.pos_encoder(x_transposed)
        x = self.transformer(x_transposed, x_transposed).transpose(1, 2)
        x = self.dropout(x)
        x = self.cond1(x, style_embedding)
        x = self.linear(x.transpose(1, 2)).transpose(1, 2)
        x = self.cond2(x, style_embedding)
        return x + features

src/networks/misc/transformer.py
MWM-io/nansypp
from typing import Tuple

import torch
from torch import nn

from src.networks.misc.cln import ConditionalLayerNormalization
from src.networks.misc.convrelunorm import ConvReLUNorm


class AmplitudeDecoder(nn.Module):
    """Amplitude decoder."""

    def __init__(
        self,
        conv_in_channels: int,
        conv_out_channels: int,
        conv_kernel_size: int,
        conv_stride: int,
        style_dim: int,
        dropout: float,
        n_convs: int,
        conv_middle_channels: int,
    ):
        """Initializer.
        Args:
            ConvReLUNorm layers params,
            ConditionalLayerNormalization params,
            Linear layer params.
        """
        super().__init__()
        self.convs = nn.ModuleList(
            [
                nn.ModuleDict(
                    {
                        "conv": ConvReLUNorm(
                            conv_in_channels if i == 0 else conv_middle_channels,
                            conv_out_channels
                            if i == n_convs - 1
                            else conv_middle_channels,
                            conv_kernel_size,
                            conv_stride,
                            dropout,
                        ),
                        "cln": ConditionalLayerNormalization(
                            style_dim,
                            conv_out_channels
                            if i == n_convs - 1
                            else conv_middle_channels,
                        ),
                    }
                )
                for i in range(n_convs)
            ]
        )
        self.linear = nn.Linear(conv_out_channels, 2)

    def forward(
        self, phoneme_features: torch.Tensor, style_embedding: torch.Tensor
    ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
        """
        Args:
            phoneme_features: [torch.float32, [B, 128, N]], upsampled phoneme features.
            style_embedding: [torch.float32, [B, 128, N]].
        Returns:
            p_amp: [torch.float32; [B, 1, N]], periodic amplitude.
            ap_amp: [torch.float32; [B, 1, N]], aperiodic amplitude,
            hiddens: [torch.float32; [B, 128, N]], amplitude hidden features.
        """
        hiddens: torch.Tensor = phoneme_features
        for conv in self.convs:
            hiddens = conv["conv"](hiddens)
            hiddens = conv["cln"](hiddens, style_embedding)
        linear_out: torch.Tensor = self.linear(hiddens.transpose(1, 2)).transpose(1, 2)
        p_amp, ap_amp = torch.split(linear_out, [1, 1], dim=1)
        return (
            p_amp.sigmoid().squeeze(dim=-1),
            ap_amp.sigmoid().squeeze(dim=-1),
            hiddens,
        )

src/networks/tts/decoders/amplitude.py
MWM-io/nansypp
import torch
from torch import nn

from src.networks.misc.cln import ConditionalLayerNormalization
from src.networks.misc.convrelunorm import ConvReLUNorm


class DurationPredictor(nn.Module):
    """Duration predictor."""

    def __init__(
        self,
        conv_in_channels: int,
        conv_out_channels: int,
        conv_kernel_size: int,
        conv_stride: int,
        style_dim: int,
        dropout: float,
        n_convs: int,
        conv_middle_channels: int,
    ):
        """Initializer.
        Args:
            ConvReLUNorm layers params,
            ConditionalLayerNormalization layers params,
            Conv1d params.
        """
        super().__init__()
        self.convs = nn.ModuleList(
            [
                nn.ModuleDict(
                    {
                        "conv": ConvReLUNorm(
                            conv_in_channels if i == 0 else conv_middle_channels,
                            conv_out_channels
                            if i == n_convs - 1
                            else conv_middle_channels,
                            conv_kernel_size,
                            conv_stride,
                            # normalized_shape,
                            dropout,
                        ),
                        "cln": ConditionalLayerNormalization(
                            style_dim,
                            conv_out_channels
                            if i == n_convs - 1
                            else conv_middle_channels,
                        ),
                    }
                )
                for i in range(n_convs)
            ]
        )
        # k=1, s=1
        out_dim = 1
        self.conv1d = nn.Conv1d(conv_out_channels, out_dim, 1, 1)

    def forward(
        self, phoneme_features: torch.Tensor, style_embedding: torch.Tensor
    ) -> torch.Tensor:
        """
        Args:
            phoneme_features: [torch.float32, [B, 128, Ntext]], upsampled phoneme features.
            style_embedding: [torch.float32, [B, 128, N]]
        Returns:
            duration: [torch.float32, [B, 1, Ntext]]
        """
        loop_variable = phoneme_features
        for conv in self.convs:
            loop_variable = conv["conv"](loop_variable)
            loop_variable = conv["cln"](loop_variable, style_embedding)
        conv_out: torch.Tensor = self.conv1d(loop_variable)
        return conv_out.relu()

src/networks/tts/decoders/duration.py
MWM-io/nansypp
import torch
from torch import nn

from src.networks.misc.cln import ConditionalLayerNormalization
from src.networks.misc.convrelunorm import ConvReLUNorm


class F0Decoder(nn.Module):
    """
    F0 decoder.
    """

    def __init__(
        self,
        conv_in_channels: int,
        conv_out_channels: int,
        conv_kernel_size: int,
        conv_stride: int,
        style_dim: int,
        dropout: float,
        n_convs: int,
        gru_hidden_size: int,
        gru_num_layers: int,
        conv_middle_channels: int,
    ):
        """Initializer.
        Args:
            ConvReLUNorm layers params,
            ConditionalLayerNormalization layers params,
            GRU params,
            Linear layer params.
        """
        super().__init__()
        self.convs = nn.ModuleList(
            [
                nn.ModuleDict(
                    {
                        "conv": ConvReLUNorm(
                            conv_in_channels if i == 0 else conv_middle_channels,
                            conv_out_channels
                            if i == n_convs - 1
                            else conv_middle_channels,
                            conv_kernel_size,
                            conv_stride,
                            dropout,
                        ),
                        "cln": ConditionalLayerNormalization(
                            style_dim,
                            conv_out_channels
                            if i == n_convs - 1
                            else conv_middle_channels,
                        ),
                    }
                )
                for i in range(n_convs)
            ]
        )
        self.gru = nn.GRU(
            conv_out_channels,
            gru_hidden_size,
            gru_num_layers,
        )
        self.linear = nn.Linear(gru_hidden_size, 1)

    def forward(
        self,
        phoneme_features: torch.Tensor,
        amplitude_hiddens: torch.Tensor,
        style_embedding: torch.Tensor,
    ) -> torch.Tensor:
        """
        Args:
            phoneme_features: [torch.float32; [B, 128, N]], upsampled phoneme features.
            amplitude_hiddens: [torch.float32; [B, 128, N]].
        Returns:
            f0_contour: [torch.float32; [B, 1, N]] predicted fundamental frequency.
        """
        loop_variable = phoneme_features + amplitude_hiddens
        for conv in self.convs:
            loop_variable = conv["conv"](loop_variable)
            loop_variable = conv["cln"](loop_variable, style_embedding)
        gru_out, _ = self.gru(loop_variable.transpose(1, 2))
        f0_contour: torch.Tensor = self.linear(gru_out).transpose(1, 2)
        return f0_contour.sigmoid()

src/networks/tts/decoders/f0.py
MWM-io/nansypp
import torch
from torch import nn

from src.networks.misc.cln import ConditionalLayerNormalization
from src.networks.misc.convrelunorm import ConvReLUNorm
from src.networks.misc.transformer import TransformerBlock


class LinguisticDecoder(nn.Module):
    """Linguistic decoder."""

    def __init__(
        self,
        first_conv_in_channels: int,
        first_conv_out_channels: int,
        first_conv_kernel_size: int,
        first_conv_stride: int,
        dropout: float,
        n_transformer_blocks: int,
        transformer_nhead: int,
        transformer_num_encoder_layers: int,
        transformer_num_decoder_layers: int,
        transformer_dim_feedforward: int,
        cln_convs_middle_channels: int,
        cln_convs_out_channels: int,
        cln_convs_kernel_size: int,
        cln_convs_stride: int,
        style_dim: int,
        n_cln_convs: int,
    ):
        """Initializer.
        Args:
            First ConvReLUNorm layers params,
            Transformer Blocks params,
            Second ConvReLUNorm layers params,
            ConditionalLayerNormalization layers params,
            Conv1d params.
        """
        super().__init__()
        self.first_conv = ConvReLUNorm(
            first_conv_in_channels,
            first_conv_out_channels,
            first_conv_kernel_size,
            first_conv_stride,
            dropout,
        )
        self.transformer_blocks = nn.ModuleList(
            [
                TransformerBlock(
                    first_conv_out_channels,
                    style_dim,
                    transformer_nhead,
                    transformer_num_encoder_layers,
                    transformer_num_decoder_layers,
                    transformer_dim_feedforward,
                    dropout,
                )
                for _ in range(n_transformer_blocks)
            ]
        )
        self.cln_convs = nn.ModuleList(
            [
                nn.ModuleDict(
                    {
                        "conv": ConvReLUNorm(
                            first_conv_out_channels
                            if i == 0
                            else cln_convs_middle_channels
                            if i == n_cln_convs - 1
                            else cln_convs_middle_channels,
                            cln_convs_out_channels,
                            cln_convs_kernel_size,
                            cln_convs_stride,
                            dropout,
                        ),
                        "cln": ConditionalLayerNormalization(
                            style_dim,
                            cln_convs_out_channels
                            if i == n_cln_convs - 1
                            else cln_convs_middle_channels,
                        ),
                    }
                )
                for i in range(n_cln_convs)
            ]
        )
        self.conv1d = nn.Conv1d(cln_convs_out_channels, 128, 1, 1)

    def forward(
        self, phoneme_features: torch.Tensor, style_embedding: torch.Tensor
    ) -> torch.Tensor:
        """
        Args:
            phoneme_features: [torch.float32, [B, 128, N]], upsampled phoneme features.
            style_embedding: [torch.float32, [B, 128, N]].
        Returns:
            linguistic_features: [torch.float32, [B, 128, N]].
        """
        embeddings = self.first_conv(phoneme_features)
        for transformer_block in self.transformer_blocks:
            embeddings = transformer_block(embeddings, style_embedding)
        for conv in self.cln_convs:
            embeddings = conv["conv"](embeddings)
            embeddings = conv["cln"](embeddings, style_embedding)
        linguistic_features: torch.Tensor = self.conv1d(embeddings)
        return linguistic_features

src/networks/tts/decoders/linguistic.py
MWM-io/nansypp
import torch
from torch import nn

from src.networks.misc.convrelunorm import ConvReLUNorm
from src.networks.misc.transformer import TransformerBlock


class PhonemeEncoder(nn.Module):
    """
    Phoneme encoder.
    """

    def __init__(
        self,
        num_labels: int,
        conv_in_channels: int,
        conv_middle_channels: int,
        conv_out_channels: int,
        conv_kernel_size: int,
        conv_stride: int,
        dropout: float,
        n_convs: int,
        conditionning_channels: int,
        nhead: int,
        num_encoder_layers: int,
        num_decoder_layers: int,
        dim_feedforward: int,
        n_transformers: int,
        negative_slope: float,
        out_linear: int,
    ):
        """Initializer.
        Args:
            num_labels: size of phonemtic dictionary,
            ConvReLUNorm layers params,
            Transformer Blocks params,
            Linear layer params,
            LeakyReLU params.
        """
        super().__init__()
        # Lookup embedding table
        self.lookup_embedder = nn.Embedding(num_labels, conv_in_channels)
        # 3 ConvReLUNorm
        self.convs = nn.ModuleList(
            [
                ConvReLUNorm(
                    conv_in_channels if i == 0 else conv_middle_channels,
                    conv_out_channels if i == n_convs - 1 else conv_middle_channels,
                    conv_kernel_size,
                    conv_stride,
                    dropout,
                )
                for i in range(n_convs)
            ]
        )
        # 3 transformer blocks
        self.transformers = nn.ModuleList(
            [
                TransformerBlock(
                    conv_out_channels,
                    conditionning_channels,
                    nhead,
                    num_encoder_layers,
                    num_decoder_layers,
                    dim_feedforward,
                    dropout,
                )
                for _ in range(n_transformers)
            ]
        )
        # Linear with LeakyReLU
        self.linear = nn.Linear(in_features=conv_out_channels, out_features=out_linear)
        self.leakyrelu = nn.LeakyReLU(negative_slope)

    def forward(self, phoneme: torch.Tensor, style_embs: torch.Tensor) -> torch.Tensor:
        """
        Args:
            phoneme: [torch.float32; [B, 1, Ntext]]
            style_embs: [torch.float32; [B, Nstyle, 1]]
        Returns:
            phoneme_features: [torch.float32; [B, 128, Ntext]]
        """
        # [B, C, N]
        embedding = self.lookup_embedder(phoneme).transpose(1, 2)
        for layer in self.convs:
            embedding = layer(embedding)
        for transformer in self.transformers:
            embedding = transformer(embedding, style_embs)
        linear_out: torch.Tensor = self.linear(embedding.transpose(1, 2)).transpose(
            1, 2
        )
        result = self.leakyrelu(linear_out)
        return result

src/networks/tts/encoders/phoneme.py
MWM-io/nansypp
import torch
from torch import nn

from src.networks.backbone.encoders.timbre import TTSAttentiveStatisticsPooling
from src.networks.backbone.encoders.wav2vec2 import Wav2Vec2Wrapper
from src.networks.misc.convrelunorm import ConvReLUNorm


class StyleEncoder(nn.Module):
    """Style encoder."""

    def __init__(
        self,
        conv_in_channels: int,
        conv_middle_channels: int,
        conv_out_channels: int,
        conv_kernel_size: int,
        conv_stride: int,
        dropout: float,
        n_convs: int,
        bottleneck: int,
        wav2vec2: Wav2Vec2Wrapper,
        out_linear: int,
    ):
        """Initializer.
        Args:
            ConvReLUNorm layers params,
            AttentiveStatisticsPooling params,
            Linear layer params.
        """
        super().__init__()
        # Wav2Vec2 first 3 layers
        self.w2v_3rd = wav2vec2
        # ConvReLUNorm
        self.convs = nn.ModuleList(
            [
                ConvReLUNorm(
                    conv_in_channels if i == 0 else conv_middle_channels,
                    conv_out_channels if i == n_convs - 1 else conv_middle_channels,
                    conv_kernel_size,
                    conv_stride,
                    # normalized_shape,
                    dropout,
                )
                for i in range(n_convs)
            ]
        )
        # LeakyRelu
        self.asp = TTSAttentiveStatisticsPooling(conv_out_channels, bottleneck)
        # Dropout
        self.linear = nn.Linear(
            in_features=2 * conv_out_channels, out_features=out_linear
        )

    def forward(self, waveform: torch.Tensor) -> torch.Tensor:
        """
        Args:
            waveform: [torch.float32; [B, N]]
        Returns:
            style_embedding: [torch.float32; [B, 128, 1]]
        """
        w2v_features = self.w2v_3rd(waveform).transpose(1, 2)
        convs_hidden = w2v_features
        for conv in self.convs:
            convs_hidden = conv(convs_hidden)
        asp_hidden = self.asp(convs_hidden)
        linear_out: torch.Tensor = self.linear(asp_hidden).unsqueeze(2)
        return linear_out

src/networks/tts/encoders/style.py
MWM-io/nansypp
import json
from typing import Dict, Optional, Tuple, TypedDict

import torch
from torch import nn
from torch.nn.utils.rnn import pad_sequence

from src.data.datamodule.tts import PitchStats
from src.data.text.text_processor import TextProcessor
from src.networks.tts.decoders.amplitude import AmplitudeDecoder
from src.networks.tts.decoders.duration import DurationPredictor
from src.networks.tts.decoders.f0 import F0Decoder
from src.networks.tts.decoders.linguistic import LinguisticDecoder
from src.networks.tts.encoders.phoneme import PhonemeEncoder
from src.networks.tts.encoders.style import StyleEncoder


class EncodingFeatures(TypedDict):
    style_embedding: torch.Tensor
    duration: torch.Tensor
    upsampled_linguistic_phoneme_features: torch.Tensor
    upsampled_pitch_phoneme_features: torch.Tensor


class DecodingFeatures(TypedDict):
    linguistic: torch.Tensor
    p_amp: torch.Tensor
    ap_amp: torch.Tensor
    pitch: torch.Tensor


class TextToSpeechNetwork(nn.Module):
    """NANSY-TTS network."""

    def __init__(
        self,
        phoneme_encoder: PhonemeEncoder,
        style_encoder: StyleEncoder,
        duration_predictor: DurationPredictor,
        linguistic_decoder: LinguisticDecoder,
        amplitude_decoder: AmplitudeDecoder,
        f0_decoder: F0Decoder,
        pitch_stats: PitchStats,
        text_processor: TextProcessor,
        sample_rate: int,
        backbone_info_file: str,
    ):
        """Initializer."""
        super().__init__()
        self.text_processor = text_processor
        self.phoneme_encoder = phoneme_encoder
        self.style_encoder = style_encoder
        self.duration_predictor = duration_predictor
        self.linguistic_decoder = linguistic_decoder
        self.amplitude_decoder = amplitude_decoder
        self.f0_decoder = f0_decoder
        self.pitch_stats = pitch_stats
        self.sample_rate = sample_rate
        with open(backbone_info_file, "r") as backbone_info:
            self.backbone_info = json.load(backbone_info)

    def encode(
        self,
        phoneme_sequence: torch.Tensor,
        waveform: torch.Tensor,
        alignment: Optional[torch.Tensor],
        upsample_lens: Optional[Dict[str, torch.Tensor]] = None,
    ) -> EncodingFeatures:
        """
        Extract TTS features from input audio and phoneme sequence
        Args:
            phoneme_sequence: [torch.float32; [B, T_text]], encoded text.
            waveform: [torch.float32; [B, T]], speech example audio.
            alignment: Optional[torch.float32; [B, T_text]], duration of each phoneme.
            upsample_lens: Optional[Dict[str, [torch.int64; B]]], lengths of features to generate.
        Returns:
            encoding_features: EncodingFeatures.
        """
        # Encoding
        style_embedding = self.style_encoder(waveform)
        phoneme_features = self.phoneme_encoder(phoneme_sequence, style_embedding)
        # Duration prediction and upsampling
        duration = self.duration_predictor(phoneme_features, style_embedding)
        if alignment is not None and upsample_lens is not None:
            align_to_upsample = alignment.unsqueeze(1)
            cqt_len = upsample_lens["cqt_len"]
            linguistic_len = upsample_lens["linguistic_len"]
        else:
            align_to_upsample = duration
            align_to_upsample[align_to_upsample <= 0] = 1
            source_duration = align_to_upsample.sum(axis=-1).unique()
            linguistic_len = (
                source_duration * self.text_processor.linguistic_features_length_ratio
            ).int()
            cqt_len = (
                source_duration * self.text_processor.pitch_features_length_ratio
            ).int()
        upsampled_linguistic_phoneme_features = (
            self.upsample_padded(  # upsample_features
                phoneme_features,
                align_to_upsample,
                linguistic_len,
            )
        )
        upsampled_pitch_phoneme_features = self.upsample_padded(  # upsample_features
            phoneme_features,
            align_to_upsample,
            cqt_len,
        )
        return {
            "style_embedding": style_embedding,
            "duration": duration,
            "upsampled_linguistic_phoneme_features": upsampled_linguistic_phoneme_features,
            "upsampled_pitch_phoneme_features": upsampled_pitch_phoneme_features,
        }

    def decode(
        self,
        upsampled_linguistic_phoneme_features: torch.Tensor,
        upsampled_pitch_phoneme_features: torch.Tensor,
        style_embedding: torch.Tensor,
    ) -> DecodingFeatures:
        """
        Generate linguistic and pitch featrues for audio generation.
        Args:
            upsampled_linguistic_phoneme_features:
                [torch.float32; [B, tts_linguistic_hidden_channels, T_ling]],
                encoded linguistic features.
            upsampled_pitch_phoneme_features:
                [torch.float32; [B, tts_pitch_hidden_channels, T_pitch]],
                encoded pitch features.
            style_embedding:
                [torch.float32; [B, N_style, 1]],
                embeddings of speech example audio.
        Returns:
            decoding_features: DecodingFeatures.
        """
        linguistic_features = self.linguistic_decoder(
            upsampled_linguistic_phoneme_features, style_embedding
        )
        p_amp_norm, ap_amp_norm, amplitude_hiddens = self.amplitude_decoder(
            upsampled_pitch_phoneme_features, style_embedding
        )
        f0_contour_norm = self.f0_decoder(
            upsampled_pitch_phoneme_features, amplitude_hiddens, style_embedding
        )

        p_amp = self.rescale(p_amp_norm, self.pitch_stats["p_amp"])
        ap_amp = self.rescale(ap_amp_norm, self.pitch_stats["ap_amp"])
        f0_contour = self.rescale(f0_contour_norm, self.pitch_stats["pitch"])

        return {
            "linguistic": linguistic_features,
            "p_amp": p_amp,
            "ap_amp": ap_amp,
            "pitch": f0_contour,
        }

    def forward(
        self,
        phoneme_sequence: torch.Tensor,
        waveform: torch.Tensor,
        alignment: Optional[torch.Tensor] = None,
        upsample_lens: Optional[Dict[str, torch.Tensor]] = None,
    ) -> Tuple[EncodingFeatures, DecodingFeatures]:
        """
        Extract TTS features from input audio and phoneme sequence
        and generate linguistic and pitch featrues for audio generation.
        Args:
            phoneme_sequence:
                [torch.float32; [B, T_text]], encoded text.
            waveform:
                [torch.float32; [B, T]], speech example audio.
            alignment:
                Optional[torch.float32; [B, T_text]], duration of each phoneme.
            upsample_lens:
                Optional[Dict[str, [torch.int64; B]]], lengths of features to generate.
        Returns:
            encoding_features: EncodingFeatures.
            decoding_features: DecodingFeatures.
        """
        encoding_features = self.encode(
            phoneme_sequence, waveform, alignment, upsample_lens
        )
        # Decoding
        decoding_features = self.decode(
            encoding_features["upsampled_linguistic_phoneme_features"],
            encoding_features["upsampled_pitch_phoneme_features"],
            encoding_features["style_embedding"],
        )
        return encoding_features, decoding_features

    @staticmethod
    def rescale(value: torch.Tensor, stats: Dict[str, torch.Tensor]) -> torch.Tensor:
        """
        Min-max normalization for A_p, A_ap and F0.
        """
        return value * (stats["max"] - stats["min"]) + stats["min"]

    @staticmethod
    def resample_duration(duration: torch.Tensor, target_duration: int):
        """
        Upsampling predicted phoneme duration to lengths of linguistic or pitch features.
        """
        source_duration = duration.sum(axis=-1).unique().item()
        length = duration.shape[-1]
        integer_division = (duration / source_duration * target_duration).type(
            torch.int
        )
        modulo = (duration / source_duration * target_duration) % 1.0
        to_add = target_duration - integer_division.sum(axis=1)
        add_from = length - to_add
        _, order = torch.sort(modulo, dim=1)
        for i, value in enumerate(add_from):
            integer_division[i][order[i, value:]] += 1
        assert (len(integer_division.sum(axis=1).unique()) == 1) and (
            integer_division.sum(axis=1).unique().item() == target_duration
        )
        return integer_division

    @staticmethod
    def upsample_features(
        features: torch.Tensor,
        duration: torch.Tensor,
        upsample_len: Optional[int],
    ) -> torch.Tensor:
        """Upsample features given predicted duration.
        Args:
            features: [B, 128, Ntext]
            duration: [B, 1, Ntext]
        Returns:
            upsampled_features: [B, 128, N]
        """
        if len(duration.shape) > 2:
            duration = duration.squeeze(1)
        assert len(duration.sum(axis=-1).unique()) == 1

        if upsample_len is not None:
            duration = TextToSpeechNetwork.resample_duration(duration, upsample_len)

        upsampled_features = torch.stack(
            [
                torch.repeat_interleave(curr_feature, curr_duration, dim=-1)
                for curr_feature, curr_duration in zip(features, duration)
            ]
        )
        return upsampled_features

    @staticmethod
    def upsample_padded(features, duration, upsample_len):
        """
        Upsampling padded features.
        """
        upsampled_features = []
        for i in range(features.shape[0]):
            features_i = features[i, :, duration[i, 0] != 0]
            duration_i = duration[i, :, duration[i, 0] != 0]
            upsample_len_i = upsample_len[i]

            upsampled_features_i = TextToSpeechNetwork.upsample_features(
                features_i.unsqueeze(0), duration_i.unsqueeze(0), upsample_len_i
            )
            upsampled_features.append(upsampled_features_i)

        padded_upsampled_features = pad_sequence(
            [i[0].T for i in upsampled_features],
            batch_first=True,
            padding_value=0,
        ).transpose(1, 2)
        return padded_upsampled_features

src/networks/tts/tts.py
MWM-io/nansypp
import copy
from typing import Optional

import autoroot  # PYTHON_PATH-setup, do not remove # pylint: disable=unused-import
import hydra
import hydra.utils as hu
import pytorch_lightning as pl
import torch
from pytorch_lightning import Trainer
from pytorch_lightning.callbacks import ModelCheckpoint

from src.data.datamodule.backbone import BackboneDataModule
from src.dataclasses.backbone import Config
from src.losses.backbone.discriminator import LeastSquaresDiscriminatorLoss
from src.losses.backbone.generator import GeneratorLoss
from src.models.backbone import Backbone
from src.networks.backbone.generator import Generator
from src.networks.misc.mpd import MultiPeriodDiscriminator


@hydra.main(
    config_path="../../configs/backbone",
    config_name="baseline.yaml",
    version_base="1.3",
)
def main(cfg: Optional[Config] = None) -> None:
    if cfg is None:
        raise RuntimeError(
            "Did not receive a valid Hydra configuration, ",
            "did you mistakingly call the main function directly from another module?",
        )

    if cfg.seed is not None:
        pl.seed_everything(cfg.seed)

    generator: Generator = hu.instantiate(cfg.generator)

    compute_augmentations_in_dataloader = True
    if compute_augmentations_in_dataloader:
        information_perturbator = copy.deepcopy(generator.information_perturbator)
        cqt = copy.deepcopy(generator.cqt)
        mel_spec = copy.deepcopy(generator.mel_spectrogram_transform)
    else:
        mel_spec = cqt = information_perturbator = None

    datamodule: BackboneDataModule = hu.instantiate(
        cfg.datamodule,
        information_perturbator=information_perturbator,
        cqt=cqt,
        mel_spec=mel_spec,
    )
    discriminator: MultiPeriodDiscriminator = hu.instantiate(cfg.discriminator)
    generator_optimizer: torch.optim.Optimizer = hu.instantiate(
        cfg.generator_optimizer,
        filter(lambda p: p.requires_grad, generator.parameters()),
    )
    discriminator_optimizer: torch.optim.Optimizer = hu.instantiate(
        cfg.discriminator_optimizer,
        filter(lambda p: p.requires_grad, discriminator.parameters()),
    )
    if cfg.resume is not None and cfg.resume.checkpoint_path is not None:
        optimizer_step = torch.load(cfg.resume.checkpoint_path, map_location="cpu")[
            "global_step"
        ]
        initial_training_step = optimizer_step // 2
    else:
        initial_training_step = 0
    generator_loss: GeneratorLoss = hu.instantiate(
        cfg.generator_loss,
        generator=generator,
        discriminator=discriminator,
        generator_optimizer=generator_optimizer,
        discriminator_optimizer=discriminator_optimizer,
        initial_training_step=initial_training_step,
        max_training_step=cfg.trainer.max_steps // 2,
    )
    discriminator_loss: LeastSquaresDiscriminatorLoss = hu.instantiate(
        cfg.discriminator_loss
    )

    model: Backbone = hu.instantiate(
        cfg.model,
        generator=generator,
        discriminator=discriminator,
        generator_optimizer=generator_optimizer,
        discriminator_optimizer=discriminator_optimizer,
        generator_loss=generator_loss,
        discriminator_loss=discriminator_loss,
    )

    model_checkpoint: ModelCheckpoint = hu.instantiate(cfg.model_checkpoint)
    callbacks = []
    logger = hu.instantiate(cfg.logger)
    for _, cb_cfg in cfg.callbacks.items():
        callbacks.append(hu.instantiate(cb_cfg))
    callbacks.append(model_checkpoint)

    trainer: Trainer = hu.instantiate(
        cfg.trainer,
        callbacks=callbacks,
        strategy="ddp_find_unused_parameters_true",
        logger=logger,
    )
    trainer.fit(
        model=model,
        datamodule=datamodule,
        ckpt_path=cfg.resume.checkpoint_path if cfg.resume is not None else None,
    )


if __name__ == "__main__":
    main()

src/train/backbone.py
MWM-io/nansypp
from typing import Optional

import autoroot
import hydra
import hydra.utils as hu
import pytorch_lightning as pl
import torch
from pytorch_lightning import Trainer
from pytorch_lightning.callbacks import ModelCheckpoint

from src.data.datamodule.tts import TextToSpeechDataModule
from src.dataclasses.tts import Config
from src.losses.tts import TextToSpeechLoss
from src.models.tts import TextToSpeechModel
from src.networks.tts.tts import TextToSpeechNetwork


@hydra.main(
    config_path="../../configs/tts", config_name="baseline.yaml", version_base="1.3"
)
def main(cfg: Optional[Config] = None) -> None:
    if cfg is None:
        raise RuntimeError(
            "Did not receive a valid Hydra configuration, ",
            "did you mistakingly call the main function directly from another module?",
        )

    if cfg.seed is not None:
        pl.seed_everything(cfg.seed)
    datamodule: TextToSpeechDataModule = hu.instantiate(cfg.datamodule)
    network: TextToSpeechNetwork = hu.instantiate(cfg.network)
    optimizer: torch.optim.Optimizer = hu.instantiate(
        cfg.optimizer, filter(lambda p: p.requires_grad, network.parameters())
    )
    scheduler: torch.optim.lr_scheduler.LRScheduler = hu.instantiate(
        cfg.scheduler, optimizer
    )
    loss: TextToSpeechLoss = hu.instantiate(cfg.loss)
    model: TextToSpeechModel = hu.instantiate(
        cfg.model,
        network=network,
        loss=loss,
        optimizer=optimizer,
        scheduler=scheduler,
    )

    model_checkpoint: ModelCheckpoint = hu.instantiate(cfg.model_checkpoint)
    callbacks = []
    logger = hu.instantiate(cfg.logger)
    for _, cb_cfg in cfg.callbacks.items():
        callbacks.append(hu.instantiate(cb_cfg))
    callbacks.append(model_checkpoint)

    trainer: Trainer = hu.instantiate(
        cfg.trainer,
        callbacks=callbacks,
        strategy="ddp_find_unused_parameters_true",
        logger=logger,
    )
    trainer.fit(
        model=model,
        datamodule=datamodule,
        ckpt_path=cfg.resume.checkpoint_path if cfg.resume is not None else None,
    )


if __name__ == "__main__":
    main()

src/train/tts.py
MWM-io/nansypp
from typing import Optional

import torch
import torchcrepe


@torch.no_grad()
def crepe_pitch_estimation(
    audio: torch.Tensor,
    sample_rate: int,
    hop_len: int,
    device: torch.device,
    model: str = "tiny",
    fmin: int = 50,
    fmax: int = 550,
    batch_size: int = 2048,
) -> Optional[torch.Tensor]:
    try:
        result = torchcrepe.predict(
            audio,
            sample_rate,
            hop_len,
            fmin,
            fmax,
            model,
            batch_size=batch_size,
            device=device,
        )
    except:
        print("Crepe crashed.")
        result = None
    return result

src/utilities/crepe.py
MWM-io/nansypp
import os
import zipfile

import pyrootutils
import requests

root = pyrootutils.setup_root(__file__, dotenv=True, pythonpath=True, cwd=False)


def download_file_from_google_drive(id, destination):
    URL = "https://docs.google.com/uc?export=download&confirm=1"

    session = requests.Session()

    response = session.get(URL, params={"id": id}, stream=True)
    token = get_confirm_token(response)

    if token:
        params = {"id": id, "confirm": token}
        response = session.get(URL, params=params, stream=True)

    save_response_content(response, destination)


def get_confirm_token(response):
    for key, value in response.cookies.items():
        if key.startswith("download_warning"):
            return value

    return None


def save_response_content(response, destination):
    CHUNK_SIZE = 32768

    with open(destination, "wb") as f:
        for chunk in response.iter_content(CHUNK_SIZE):
            if chunk:  # filter out keep-alive new chunks
                f.write(chunk)


def download_backbone(destination_dir):
    print("Downloading backbone..")
    file_id = "1ULSjTnK4wJYFiXzFCPbaZunKu9XGVb0u"
    tmp_dir = os.path.join(root, "static/tmp")
    os.makedirs(tmp_dir, exist_ok=True)
    destination_zip = os.path.join(tmp_dir, "backbone_hifitts.zip")
    download_file_from_google_drive(file_id, destination_zip)
    print("Unzipping...")
    os.makedirs(destination_dir, exist_ok=True)
    with zipfile.ZipFile(destination_zip, "r") as zip_ref:
        zip_ref.extractall(destination_dir)
    os.remove(destination_zip)


def download_tts(destination_dir):
    print("Downloading TTS...")
    file_id = "1oECzHoWyT98Ef9wJeCwnz77koLPhWZf6"
    tmp_dir = os.path.join(root, "static/tmp")
    os.makedirs(tmp_dir, exist_ok=True)
    destination_zip = os.path.join(tmp_dir, "tts_hifitts.zip")
    download_file_from_google_drive(file_id, destination_zip)
    print("Unzipping...")
    os.makedirs(destination_dir, exist_ok=True)
    with zipfile.ZipFile(destination_zip, "r") as zip_ref:
        zip_ref.extractall(destination_dir)
    os.remove(destination_zip)


def main():
    backbone_destination_dir = os.path.join(root, "static/runs/runs_backbone/hifitts")
    if not os.path.isdir(backbone_destination_dir):
        download_backbone(backbone_destination_dir)
    tts_destination_dir = os.path.join(root, "static/runs/runs_tts/hifitts")
    if not os.path.isdir(tts_destination_dir):
        download_tts(tts_destination_dir)


if __name__ == "__main__":
    main()

src/utilities/download_checkpoints.py
MWM-io/nansypp
import functools
import time
from contextlib import contextmanager
from typing import Any, Callable, Generator, Literal, Optional, TypeVar, Union

import torch
from pytorch_lightning import LightningModule, Trainer
from pytorch_lightning.callbacks.timer import Timer

from src.utilities.types import P

DO_PROFILING = False


def synchronize_and_time() -> float:
    torch.cuda.synchronize()
    return time.time()


@contextmanager
def cuda_synchronized_profile_context_manager(
    do_profiling: bool = DO_PROFILING,
    prefix: str = "",
    time_unit: Union[Literal["s"], Literal["ms"]] = "s",
) -> Generator[None, None, None]:
    do_cuda_synchronize = True
    valid_time_units = ["s", "ms"]
    assert (
        valid_time_units.count(time_unit) > 0
    ), f"Unrecognized time_unit option {time_unit}, valid values are: {valid_time_units}"

    if do_profiling and torch.cuda.is_initialized():
        if do_cuda_synchronize:
            torch.cuda.synchronize()
        start_time = time.perf_counter()
    else:
        start_time = None
    try:
        yield
    finally:
        if do_profiling and torch.cuda.is_initialized():
            assert start_time is not None
            if do_cuda_synchronize:
                torch.cuda.synchronize()
            end_time = time.perf_counter()
            duration_seconds = end_time - start_time
            print(
                prefix,
                f"{duration_seconds * (1000 if time_unit == 'ms' else 1):.3f}{time_unit}",
            )


T = TypeVar("T")

TIMER_PREFIX_WIDTH = 40


def cuda_synchronized_timer(
    do_profiling: bool,
    prefix: str = "",
    time_unit: Union[Literal["s"], Literal["ms"]] = "ms",
    before_timing: Callable[[], None] = torch.cuda.synchronize,
    format_specifier: str = "8.2f",
) -> Callable[[Callable[P, T]], Callable[P, T]]:
    def decorator(func: Callable[P, T]) -> Callable[P, T]:
        @functools.wraps(func)
        def wrapper(*args: P.args, **kwargs: P.kwargs) -> T:
            start: Optional[float] = None
            if do_profiling and torch.cuda.is_initialized():
                before_timing()
                start = time.perf_counter()

            func_value = func(*args, **kwargs)

            if start is not None and do_profiling and torch.cuda.is_initialized():
                before_timing()
                run_time_seconds = time.perf_counter() - start
                full_prefix = (
                    f"[{prefix}{'::' if len(prefix)>0 else ''}{func.__name__}] "
                )
                print(
                    full_prefix.ljust(TIMER_PREFIX_WIDTH)
                    + f"{run_time_seconds * (1000 if time_unit == 'ms' else 1):{format_specifier}}{time_unit}",
                )
            return func_value

        return wrapper

    return decorator


class LightningBatchTimer(Timer):
    """A Pytorch-Lightning timer Callback to measure the duration of epochs and batches"""

    batch_time: Optional[float] = None
    time_elapsed_epoch: float = 0
    average_batch_duration: float = 0
    remaining_warmup_batches: int = 0
    num_warmup_batches: int

    def __init__(
        self,
        num_batches_per_epoch: Optional[int],
        num_warmup_batches: int,
        *args: Any,
    ) -> None:
        super().__init__(*args)
        self.num_batches_per_epoch = num_batches_per_epoch
        self.num_warmup_batches = num_warmup_batches
        self.num_effective_batches_per_epoch = (
            (self.num_batches_per_epoch - self.num_warmup_batches)
            if self.num_batches_per_epoch is not None
            else None
        )

    def on_train_epoch_start(
        self, trainer: Trainer, pl_module: LightningModule
    ) -> None:
        torch.cuda.synchronize()
        self.time_elapsed_epoch = self.time_elapsed()
        self.average_batch_duration = 0
        self.remaining_warmup_batches = self.num_warmup_batches

    def on_train_batch_start(
        self,
        trainer: Trainer,
        pl_module: LightningModule,
        batch: Any,
        batch_idx: int,
    ) -> None:
        prev_time = self.batch_time
        torch.cuda.synchronize()
        new_time = self.time_elapsed()
        if prev_time is not None and self.num_effective_batches_per_epoch is not None:
            batch_duration = new_time - prev_time
            # print(f"[MyTimer] Batch: {batch_duration:.3f}s")
            if self.remaining_warmup_batches == 0:
                self.average_batch_duration += batch_duration / (
                    self.num_effective_batches_per_epoch
                )
        self.batch_time = new_time
        self.remaining_warmup_batches = max(0, self.remaining_warmup_batches - 1)

    def on_train_epoch_end(self, trainer: Trainer, *args: Any, **kwargs: Any) -> None:
        torch.cuda.synchronize()
        print(
            f"[MyTimer] Time elapsed on epoch: {self.time_elapsed() - self.time_elapsed_epoch:.3f}s",
            flush=True,
        )
        print(
            f"[MyTimer] Average batch: {self.average_batch_duration:.3f}s",
            flush=True,
        )

src/utilities/profiling.py
MWM-io/nansypp
import json
import os
import string
from datetime import datetime, timedelta
from typing import Dict, List, Tuple

import autoroot
import torch
import torchaudio
import torchaudio.transforms as T
from evaluate import load
from tqdm import tqdm
from transformers import WhisperForConditionalGeneration, WhisperProcessor

from src.inference.tts import TextToSpeechInferencer


class TTSEvaluationModule:
    def __init__(
        self,
        evaluation_data: str,
    ):
        self.tts_output_sample_rate = None

        self.processor = WhisperProcessor.from_pretrained("openai/whisper-medium")
        self.stt_model = WhisperForConditionalGeneration.from_pretrained(
            "openai/whisper-medium"
        )
        self.stt_model_sample_rate = 16000
        self.stt_model.config.forced_decoder_ids = None

        self.resampler = None

        if not os.path.isabs(evaluation_data):
            evaluation_data = os.path.join(autoroot.root, evaluation_data)
        with open(evaluation_data, "r") as file:
            realtive_evaluation_data = json.load(file)
        self.evaluation_data = {}
        for path, text in realtive_evaluation_data.items():
            if not os.path.isabs(path):
                path = os.path.join(autoroot.root, path)
            self.evaluation_data[path] = text
        self.style_audios = list(self.evaluation_data.keys())
        self.phrases = list(self.evaluation_data.values())

        self.cer = load("cer")

    @staticmethod
    def _process_text(text: str) -> str:
        text = " ".join(text.lower().split())
        text = "".join(list(filter(lambda i: i in string.ascii_lowercase + " ", text)))
        return text

    def compute_score(self, transcriptions: Dict[Tuple[str, str], str]) -> float:
        predictions = []
        references = []
        for (_, reference), prediction in transcriptions.items():
            references.append(self._process_text(reference))
            predictions.append(self._process_text(prediction))
        return self.cer.compute(
            predictions=predictions,
            references=references,
        )

    def generate_transcriptions(
        self, generated_audios: Dict[Tuple[str, str], torch.Tensor]
    ) -> Dict[Tuple[str, str], str]:
        self.resampler = T.Resample(
            self.tts_output_sample_rate, self.stt_model_sample_rate
        )
        transcriptions = {}
        for key, audio in tqdm(
            generated_audios.items(), total=len(generated_audios), dynamic_ncols=True
        ):
            audio = self.resampler(audio.detach().cpu())
            sample = audio.detach().cpu().numpy().reshape(-1)
            input_features = self.processor(
                sample, sampling_rate=self.stt_model_sample_rate, return_tensors="pt"
            ).input_features
            predicted_ids = self.stt_model.generate(input_features)
            transcription = self.processor.batch_decode(
                predicted_ids, skip_special_tokens=True
            )[0]
            transcriptions[key] = transcription
        return transcriptions

    def generate_audios(
        self,
        tts_inferencer: TextToSpeechInferencer,
        get_time: bool = False,
    ) -> Dict[Tuple[str, str], torch.Tensor]:
        self.tts_output_sample_rate = tts_inferencer.generator.output_sample_rate
        generated_audios = {}
        start_time = datetime.now()
        for audio_file, text in tqdm(
            self.evaluation_data.items(),
            total=len(self.evaluation_data),
            dynamic_ncols=True,
        ):
            initial_audio, audio_sample_rate = torchaudio.load(audio_file)
            style_audio_tensor = T.Resample(
                audio_sample_rate, tts_inferencer.generator.input_sample_rate
            )(initial_audio).to(tts_inferencer.device)
            generated_audio = tts_inferencer(text, style_audio_tensor)
            generated_audios[(audio_file, text)] = generated_audio.detach().cpu()
            torch.cuda.empty_cache()
        end_time = datetime.now()
        exec_time = end_time - start_time
        if not get_time:
            return generated_audios
        else:
            return generated_audios, exec_time

src/utilities/tts_evaluation.py
MWM-io/nansypp
"""
A decorator that copies a function's docstring and signature to another.

Taken from https://stackoverflow.com/a/75898999/5798793.
Useful to enable typing of nn.Module().__call__() methods by copying
the relevant information from the nn.Module().forward() method.

Usage:

```
from typing import Tuple

from torch import Tensor

class ExampleModule(nn.Module):
    def forward(
        self, x: Tensor, y: Tensor, some_flag: boolean = True
    ) -> Tuple[Tensor, Tensor]:
        ...

    @copy_docstring_and_signature(forward)
    # mypy will complain that this declaration is untyped,
    # hence the type-system 
    def __call__(self, *args, **kwargs):  # type: ignore[no-untyped-def]
        return super().__call__(*args, **kwargs)
```

In the context of nn.Modules' forward / __call__, an alternative
would be to explicitly declare the __call__() method and add type-hints
there copying the ones from the forward() method, but this has the
downside of necessitating manual maintenance to ensure that the
types in __call__() are always up-to-date if making modifications to
the arguments or return type of the forward() method.
This is discussed here:
https://github.com/pytorch/pytorch/issues/44605#issuecomment-692344479
"""

import sys
from functools import partial
from typing import Any, Callable, TypeVar, cast

from hydra.utils import instantiate

if sys.version_info < (3, 10):
    from typing_extensions import ParamSpec
else:
    from typing import ParamSpec


P = ParamSpec("P")
T = TypeVar("T")


def copy_docstring_and_signature(
    from_func: Callable[P, T]
) -> Callable[[Callable], Callable[P, T]]:
    """
    A typed implementation of functools.wraps.

    Taken from https://stackoverflow.com/a/75898999/5798793.
    """

    def decorator(func: Callable) -> Callable[P, T]:
        func.__doc__ = from_func.__doc__
        return cast(Callable[P, T], func)

    return decorator


def typed_instantiate(_: Callable[P, T], *args: Any, **kwargs: Any) -> Callable[P, T]:
    """Return a typed hydra instantiate lambda with the provided function's call signature"""
    return cast(Callable[P, T], partial(instantiate, *args, **kwargs))

src/utilities/types.py
MWM-io/nansypp