chore: import upstream snapshot with attribution

basics/base_augmentation.py

@@ -0,0 +1,28 @@
+from utils.hparams import hparams
+
+
+class BaseAugmentation:
+    """
+    Base class for data augmentation.
+    All methods of this class should be thread-safe.
+    1. *process_item*:
+        Apply augmentation to one piece of data.
+    """
+    def __init__(self, data_dirs: list, augmentation_args: dict):
+        self.raw_data_dirs = data_dirs
+        self.augmentation_args = augmentation_args
+        self.timestep = hparams['hop_size'] / hparams['audio_sample_rate']
+
+    def process_item(self, item: dict, **kwargs) -> dict:
+        raise NotImplementedError()
+
+
+def require_same_keys(func):
+    def run(*args, **kwargs):
+        item: dict = args[1]
+        res: dict = func(*args, **kwargs)
+        assert set(item.keys()) == set(res.keys()), 'Item keys mismatch after augmentation.\n' \
+                                                    f'Before: {sorted(item.keys())}\n' \
+                                                    f'After: {sorted(res.keys())}'
+        return res
+    return run

basics/base_binarizer.py

@@ -0,0 +1,386 @@
+import json
+import pathlib
+import pickle
+import random
+import shutil
+import warnings
+from copy import deepcopy
+
+import numpy as np
+import torch
+from tqdm import tqdm
+
+from utils.hparams import hparams
+from utils.indexed_datasets import IndexedDatasetBuilder
+from utils.multiprocess_utils import chunked_multiprocess_run
+from utils.phoneme_utils import load_phoneme_dictionary
+from utils.plot import distribution_to_figure
+
+
+class BinarizationError(Exception):
+    pass
+
+
+class BaseBinarizer:
+    """
+        Base class for data processing.
+        1. *process* and *process_data_split*:
+            process entire data, generate the train-test split (support parallel processing);
+        2. *process_item*:
+            process singe piece of data;
+        3. *get_pitch*:
+            infer the pitch using some algorithm;
+        4. *get_align*:
+            get the alignment using 'mel2ph' format (see https://arxiv.org/abs/1905.09263).
+        5. phoneme encoder, voice encoder, etc.
+
+        Subclasses should define:
+        1. *load_metadata*:
+            how to read multiple datasets from files;
+        2. *train_item_names*, *valid_item_names*, *test_item_names*:
+            how to split the dataset;
+        3. load_ph_set:
+            the phoneme set.
+    """
+
+    def __init__(self, datasets=None, data_attrs=None):
+        if datasets is None:
+            datasets = hparams['datasets']
+        self.datasets = datasets
+        self.raw_data_dirs = [pathlib.Path(ds['raw_data_dir']) for ds in self.datasets]
+        self.binary_data_dir = pathlib.Path(hparams['binary_data_dir'])
+        self.data_attrs = [] if data_attrs is None else data_attrs
+
+        self.binarization_args = hparams['binarization_args']
+        self.augmentation_args = hparams.get('augmentation_args', {})
+        self.device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
+
+        self.spk_map = {}
+        self.spk_ids = None
+        self.build_spk_map()
+
+        self.lang_map = {}
+        self.dictionaries = hparams['dictionaries']
+        self.build_lang_map()
+
+        self.items = {}
+        self.item_names: list = None
+        self._train_item_names: list = None
+        self._valid_item_names: list = None
+
+        self.phoneme_dictionary = load_phoneme_dictionary()
+        self.timestep = hparams['hop_size'] / hparams['audio_sample_rate']
+
+    def build_spk_map(self):
+        spk_ids = [ds.get('spk_id') for ds in self.datasets]
+        assigned_spk_ids = {spk_id for spk_id in spk_ids if spk_id is not None}
+        idx = 0
+        for i in range(len(spk_ids)):
+            if spk_ids[i] is not None:
+                continue
+            while idx in assigned_spk_ids:
+                idx += 1
+            spk_ids[i] = idx
+            assigned_spk_ids.add(idx)
+        assert max(spk_ids) < hparams['num_spk'], \
+            f'Index in spk_id sequence {spk_ids} is out of range. All values should be smaller than num_spk.'
+
+        for spk_id, dataset in zip(spk_ids, self.datasets):
+            spk_name = dataset['speaker']
+            if spk_name in self.spk_map and self.spk_map[spk_name] != spk_id:
+                raise ValueError(f'Invalid speaker ID assignment. Name \'{spk_name}\' is assigned '
+                                 f'with different speaker IDs: {self.spk_map[spk_name]} and {spk_id}.')
+            self.spk_map[spk_name] = spk_id
+        self.spk_ids = spk_ids
+
+        print("| spk_map: ", self.spk_map)
+
+    def build_lang_map(self):
+        assert len(self.dictionaries.keys()) <= hparams['num_lang'], \
+            'Number of languages must not be greater than num_lang!'
+        for dataset in self.datasets:
+            assert dataset['language'] in self.dictionaries, f'Unrecognized language name: {dataset["language"]}'
+
+        for lang_id, lang_name in enumerate(sorted(self.dictionaries.keys()), start=1):
+            self.lang_map[lang_name] = lang_id
+
+        print("| lang_map: ", self.lang_map)
+
+    def load_meta_data(self, raw_data_dir: pathlib.Path, ds_id, spk, lang) -> dict:
+        raise NotImplementedError()
+
+    def split_train_valid_set(self, prefixes: list):
+        """
+        Split the dataset into training set and validation set.
+        :return: train_item_names, valid_item_names
+        """
+        prefixes = {str(pr): 1 for pr in prefixes}
+        valid_item_names = {}
+        # Add prefixes that specified speaker index and matches exactly item name to test set
+        for prefix in deepcopy(prefixes):
+            if prefix in self.item_names:
+                valid_item_names[prefix] = 1
+                prefixes.pop(prefix)
+        # Add prefixes that exactly matches item name without speaker id to test set
+        for prefix in deepcopy(prefixes):
+            matched = False
+            for name in self.item_names:
+                if name.split(':')[-1] == prefix:
+                    valid_item_names[name] = 1
+                    matched = True
+            if matched:
+                prefixes.pop(prefix)
+        # Add names with one of the remaining prefixes to test set
+        for prefix in deepcopy(prefixes):
+            matched = False
+            for name in self.item_names:
+                if name.startswith(prefix):
+                    valid_item_names[name] = 1
+                    matched = True
+            if matched:
+                prefixes.pop(prefix)
+        for prefix in deepcopy(prefixes):
+            matched = False
+            for name in self.item_names:
+                if name.split(':')[-1].startswith(prefix):
+                    valid_item_names[name] = 1
+                    matched = True
+            if matched:
+                prefixes.pop(prefix)
+
+        if len(prefixes) != 0:
+            warnings.warn(
+                f'The following rules in test_prefixes have no matching names in the dataset: {", ".join(prefixes.keys())}',
+                category=UserWarning
+            )
+            warnings.filterwarnings('default')
+
+        valid_item_names = list(valid_item_names.keys())
+        assert len(valid_item_names) > 0, 'Validation set is empty!'
+        train_item_names = [x for x in self.item_names if x not in set(valid_item_names)]
+        assert len(train_item_names) > 0, 'Training set is empty!'
+
+        return train_item_names, valid_item_names
+
+    @property
+    def train_item_names(self):
+        return self._train_item_names
+
+    @property
+    def valid_item_names(self):
+        return self._valid_item_names
+
+    def meta_data_iterator(self, prefix):
+        if prefix == 'train':
+            item_names = self.train_item_names
+        else:
+            item_names = self.valid_item_names
+        for item_name in item_names:
+            meta_data = self.items[item_name]
+            yield item_name, meta_data
+
+    def process(self):
+        # load each dataset
+        test_prefixes = []
+        for ds_id, dataset in enumerate(self.datasets):
+            items = self.load_meta_data(
+                pathlib.Path(dataset['raw_data_dir']),
+                ds_id=ds_id, spk=dataset['speaker'], lang=dataset['language']
+            )
+            self.items.update(items)
+            test_prefixes.extend(
+                f'{ds_id}:{prefix}'
+                for prefix in dataset.get('test_prefixes', [])
+            )
+        self.item_names = sorted(list(self.items.keys()))
+        self._train_item_names, self._valid_item_names = self.split_train_valid_set(test_prefixes)
+
+        if self.binarization_args['shuffle']:
+            random.shuffle(self.item_names)
+
+        self.binary_data_dir.mkdir(parents=True, exist_ok=True)
+
+        # Copy spk_map, lang_map and dictionary to binary data dir
+        spk_map_fn = self.binary_data_dir / 'spk_map.json'
+        with open(spk_map_fn, 'w', encoding='utf-8') as f:
+            json.dump(self.spk_map, f, ensure_ascii=False)
+        lang_map_fn = self.binary_data_dir / 'lang_map.json'
+        with open(lang_map_fn, 'w', encoding='utf-8') as f:
+            json.dump(self.lang_map, f, ensure_ascii=False)
+        for lang, dict_path in hparams['dictionaries'].items():
+            shutil.copy(dict_path, self.binary_data_dir / f'dictionary-{lang}.txt')
+        self.check_coverage()
+
+        # Process valid set and train set
+        try:
+            self.process_dataset('valid')
+            self.process_dataset(
+                'train',
+                num_workers=int(self.binarization_args['num_workers']),
+                apply_augmentation=any(args['enabled'] for args in self.augmentation_args.values())
+            )
+        except KeyboardInterrupt:
+            exit(-1)
+
+    def check_coverage(self):
+        # Group by phonemes in the dictionary.
+        ph_idx_required = set(range(1, len(self.phoneme_dictionary)))
+        ph_idx_occurred = set()
+        ph_idx_count_map = {
+            idx: 0
+            for idx in ph_idx_required
+        }
+
+        # Load and count those phones that appear in the actual data
+        for item_name in self.items:
+            ph_idx_occurred.update(self.items[item_name]['ph_seq'])
+            for idx in self.items[item_name]['ph_seq']:
+                ph_idx_count_map[idx] += 1
+        ph_count_map = {
+            self.phoneme_dictionary.decode_one(idx, scalar=False): count
+            for idx, count in ph_idx_count_map.items()
+        }
+
+        def display_phoneme(phoneme):
+            if isinstance(phoneme, tuple):
+                return f'({", ".join(phoneme)})'
+            return phoneme
+
+        print('===== Phoneme Distribution Summary =====')
+        keys = sorted(ph_count_map.keys(), key=lambda v: v[0] if isinstance(v, tuple) else v)
+        for i, key in enumerate(keys):
+            if i == len(ph_count_map) - 1:
+                end = '\n'
+            elif i % 10 == 9:
+                end = ',\n'
+            else:
+                end = ', '
+            key_disp = display_phoneme(key)
+            print(f'{key_disp}: {ph_count_map[key]}', end=end)
+
+        # Draw graph.
+        xs = [display_phoneme(k) for k in keys]
+        ys = [ph_count_map[k] for k in keys]
+        plt = distribution_to_figure(
+            title='Phoneme Distribution Summary',
+            x_label='Phoneme', y_label='Number of occurrences',
+            items=xs, values=ys, rotate=len(self.dictionaries) > 1
+        )
+        filename = self.binary_data_dir / 'phoneme_distribution.jpg'
+        plt.savefig(fname=filename,
+                    bbox_inches='tight',
+                    pad_inches=0.25)
+        print(f'| save summary to \'{filename}\'')
+
+        # Check unrecognizable or missing phonemes
+        if ph_idx_occurred != ph_idx_required:
+            missing_phones = sorted({
+                self.phoneme_dictionary.decode_one(idx, scalar=False)
+                for idx in ph_idx_required.difference(ph_idx_occurred)
+            }, key=lambda v: v[0] if isinstance(v, tuple) else v)
+            raise BinarizationError(
+                f'The following phonemes are not covered in transcriptions: {missing_phones}'
+            )
+
+    def process_dataset(self, prefix, num_workers=0, apply_augmentation=False):
+        args = []
+        builder = IndexedDatasetBuilder(self.binary_data_dir, prefix=prefix, allowed_attr=self.data_attrs)
+        total_sec = {k: 0.0 for k in self.spk_map}
+        total_raw_sec = {k: 0.0 for k in self.spk_map}
+        extra_info = {'names': {}, 'ph_texts': {}, 'spk_ids': {}, 'spk_names': {}, 'lengths': {}}
+        max_no = -1
+
+        for item_name, meta_data in self.meta_data_iterator(prefix):
+            args.append([item_name, meta_data, self.binarization_args])
+
+        aug_map = self.arrange_data_augmentation(self.meta_data_iterator(prefix)) if apply_augmentation else {}
+
+        def postprocess(_item):
+            nonlocal total_sec, total_raw_sec, extra_info, max_no
+            if _item is None:
+                return
+            item_no = builder.add_item(_item)
+            max_no = max(max_no, item_no)
+            for k, v in _item.items():
+                if isinstance(v, np.ndarray):
+                    if k not in extra_info:
+                        extra_info[k] = {}
+                    extra_info[k][item_no] = v.shape[0]
+            extra_info['names'][item_no] = _item['name'].split(':', 1)[-1]
+            extra_info['ph_texts'][item_no] = _item['ph_text']
+            extra_info['spk_ids'][item_no] = _item['spk_id']
+            extra_info['spk_names'][item_no] = _item['spk_name']
+            extra_info['lengths'][item_no] = _item['length']
+            total_raw_sec[_item['spk_name']] += _item['seconds']
+            total_sec[_item['spk_name']] += _item['seconds']
+
+            for task in aug_map.get(_item['name'], []):
+                aug_item = task['func'](_item, **task['kwargs'])
+                aug_item_no = builder.add_item(aug_item)
+                max_no = max(max_no, aug_item_no)
+                for k, v in aug_item.items():
+                    if isinstance(v, np.ndarray):
+                        if k not in extra_info:
+                            extra_info[k] = {}
+                        extra_info[k][aug_item_no] = v.shape[0]
+                extra_info['names'][aug_item_no] = aug_item['name'].split(':', 1)[-1]
+                extra_info['ph_texts'][aug_item_no] = aug_item['ph_text']
+                extra_info['spk_ids'][aug_item_no] = aug_item['spk_id']
+                extra_info['spk_names'][aug_item_no] = aug_item['spk_name']
+                extra_info['lengths'][aug_item_no] = aug_item['length']
+                total_sec[aug_item['spk_name']] += aug_item['seconds']
+
+        try:
+            if num_workers > 0:
+                # code for parallel processing
+                for item in tqdm(
+                        chunked_multiprocess_run(self.process_item, args, num_workers=num_workers),
+                        total=len(list(self.meta_data_iterator(prefix)))
+                ):
+                    postprocess(item)
+            else:
+                # code for single cpu processing
+                for a in tqdm(args):
+                    item = self.process_item(*a)
+                    postprocess(item)
+            for k in extra_info:
+                assert set(extra_info[k]) == set(range(max_no + 1)), f'Item numbering is not consecutive.'
+                extra_info[k] = list(map(lambda x: x[1], sorted(extra_info[k].items(), key=lambda x: x[0])))
+        except KeyboardInterrupt:
+            builder.finalize()
+            raise
+
+        builder.finalize()
+        if prefix == "train":
+            extra_info.pop("names")
+            extra_info.pop('ph_texts')
+            extra_info.pop("spk_names")
+        with open(self.binary_data_dir / f"{prefix}.meta", "wb") as f:
+            # noinspection PyTypeChecker
+            pickle.dump(extra_info, f)
+        if apply_augmentation:
+            print(f"| {prefix} total duration (before augmentation): {sum(total_raw_sec.values()):.2f}s")
+            print(
+                f"| {prefix} respective duration (before augmentation): "
+                + ', '.join(f'{k}={v:.2f}s' for k, v in total_raw_sec.items())
+            )
+            print(
+                f"| {prefix} total duration (after augmentation): "
+                f"{sum(total_sec.values()):.2f}s ({sum(total_sec.values()) / sum(total_raw_sec.values()):.2f}x)"
+            )
+            print(
+                f"| {prefix} respective duration (after augmentation): "
+                + ', '.join(f'{k}={v:.2f}s' for k, v in total_sec.items())
+            )
+        else:
+            print(f"| {prefix} total duration: {sum(total_raw_sec.values()):.2f}s")
+            print(f"| {prefix} respective duration: " + ', '.join(f'{k}={v:.2f}s' for k, v in total_raw_sec.items()))
+
+    def arrange_data_augmentation(self, data_iterator):
+        """
+        Code for all types of data augmentation should be added here.
+        """
+        raise NotImplementedError()
+
+    def process_item(self, item_name, meta_data, binarization_args):
+        raise NotImplementedError()

basics/base_dataset.py

@@ -0,0 +1,58 @@
+import os
+import pickle
+
+import torch
+from torch.utils.data import Dataset
+
+from utils.hparams import hparams
+from utils.indexed_datasets import IndexedDataset
+
+
+class BaseDataset(Dataset):
+    """
+        Base class for datasets.
+        1. *sizes*:
+            clipped length if "max_frames" is set;
+        2. *num_frames*:
+            unclipped length.
+
+        Subclasses should define:
+        1. *collate*:
+            take the longest data, pad other data to the same length;
+        2. *__getitem__*:
+            the index function.
+    """
+
+    def __init__(self, prefix, size_key='lengths', preload=False):
+        super().__init__()
+        self.prefix = prefix
+        self.data_dir = hparams['binary_data_dir']
+        with open(os.path.join(self.data_dir, f'{self.prefix}.meta'), 'rb') as f:
+            self.metadata = pickle.load(f)
+        self.sizes = self.metadata[size_key]
+        self._indexed_ds = IndexedDataset(self.data_dir, self.prefix)
+        if preload:
+            self.indexed_ds = [self._indexed_ds[i] for i in range(len(self._indexed_ds))]
+            del self._indexed_ds
+        else:
+            self.indexed_ds = self._indexed_ds
+
+    def __getitem__(self, index):
+        return {'_idx': index, **self.indexed_ds[index]}
+
+    def __len__(self):
+        return len(self.sizes)
+
+    def num_frames(self, index):
+        return self.sizes[index]
+
+    def size(self, index):
+        """Return an example's size as a float or tuple. This value is used when
+        filtering a dataset with ``--max-positions``."""
+        return self.sizes[index]
+
+    def collater(self, samples):
+        return {
+            'size': len(samples),
+            'indices': torch.LongTensor([s['_idx'] for s in samples])
+        }

basics/base_exporter.py

@@ -0,0 +1,86 @@
+import json
+import pathlib
+import shutil
+from pathlib import Path
+from typing import Union
+
+import torch
+import torch.nn as nn
+
+from utils.hparams import hparams
+
+
+class BaseExporter:
+    def __init__(
+            self,
+            device: Union[str, torch.device] = None,
+            cache_dir: Path = None,
+            **kwargs
+    ):
+        self.device = device if device is not None else torch.device('cuda' if torch.cuda.is_available() else 'cpu')
+        self.cache_dir: Path = cache_dir.resolve() if cache_dir is not None \
+            else Path(__file__).parent.parent / 'deployment' / 'cache'
+        self.cache_dir.mkdir(parents=True, exist_ok=True)
+
+    # noinspection PyMethodMayBeStatic
+    def build_spk_map(self) -> dict:
+        if hparams['use_spk_id']:
+            with open(Path(hparams['work_dir']) / 'spk_map.json', 'r', encoding='utf8') as f:
+                spk_map = json.load(f)
+            assert isinstance(spk_map, dict) and len(spk_map) > 0, 'Invalid or empty speaker map!'
+            assert len(spk_map) == len(set(spk_map.values())), 'Duplicate speaker id in speaker map!'
+            return spk_map
+        else:
+            return {}
+
+    # noinspection PyMethodMayBeStatic
+    def build_lang_map(self) -> dict:
+        lang_map_fn = pathlib.Path(hparams['work_dir']) / 'lang_map.json'
+        if lang_map_fn.exists():
+            with open(lang_map_fn, 'r', encoding='utf8') as f:
+                lang_map = json.load(f)
+            assert isinstance(lang_map, dict) and len(lang_map) > 0, 'Invalid or empty language map!'
+            assert len(lang_map) == len(set(lang_map.values())), 'Duplicate language id in language map!'
+            return lang_map
+        else:
+            return {}
+
+    def build_model(self) -> nn.Module:
+        """
+        Creates an instance of nn.Module and load its state dict on the target device.
+        """
+        raise NotImplementedError()
+
+    def export_model(self, path: Path):
+        """
+        Exports the model to ONNX format.
+        :param path: the target model path
+        """
+        raise NotImplementedError()
+
+    # noinspection PyMethodMayBeStatic
+    def export_dictionaries(self, path: Path):
+        dicts = hparams.get('dictionaries')
+        if dicts is not None:
+            for lang in dicts.keys():
+                fn = f'dictionary-{lang}.txt'
+                shutil.copy(pathlib.Path(hparams['work_dir']) / fn, path)
+                print(f'| export dictionary => {path / fn}')
+        else:
+            fn = 'dictionary.txt'
+            shutil.copy(pathlib.Path(hparams['work_dir']) / fn, path)
+            print(f'| export dictionary => {path / fn}')
+
+    def export_attachments(self, path: Path):
+        """
+        Exports related files and configs (e.g. the dictionary) to the target directory.
+        :param path: the target directory
+        """
+        raise NotImplementedError()
+
+    def export(self, path: Path):
+        """
+        Exports all the artifacts to the target directory.
+        :param path: the target directory
+        """
+        raise NotImplementedError()

basics/base_module.py

@@ -0,0 +1,18 @@
+from torch import nn
+
+
+class CategorizedModule(nn.Module):
+    @property
+    def category(self):
+        raise NotImplementedError()
+
+    def check_category(self, category):
+        if category is None:
+            raise RuntimeError('Category is not specified in this checkpoint.\n'
+                               'If this is a checkpoint in the old format, please consider '
+                               'migrating it to the new format via the following command:\n'
+                               'python scripts/migrate.py ckpt <INPUT_CKPT> <OUTPUT_CKPT>')
+        elif category != self.category:
+            raise RuntimeError('Category mismatches!\n'
+                               f'This checkpoint is of the category \'{category}\', '
+                               f'but a checkpoint of category \'{self.category}\' is required.')

basics/base_pe.py

@@ -0,0 +1,7 @@
+class BasePE:
+    def get_pitch(
+            self, waveform, samplerate, length,
+            *, hop_size, f0_min=65, f0_max=1100,
+            speed=1, interp_uv=False
+    ):
+        raise NotImplementedError()

basics/base_svs_infer.py

@@ -0,0 +1,136 @@
+# coding=utf8
+import numpy as np
+import torch
+from torch import Tensor
+from typing import Tuple, Dict
+
+from utils.hparams import hparams
+from utils.infer_utils import resample_align_curve
+
+
+class BaseSVSInfer:
+    """
+        Base class for SVS inference models.
+        Subclasses should define:
+        1. *build_model*:
+            how to build the model;
+        2. *run_model*:
+            how to run the model (typically, generate a mel-spectrogram and
+            pass it to the pre-built vocoder);
+        3. *preprocess_input*:
+            how to preprocess user input.
+        4. *infer_once*
+            infer from raw inputs to the final outputs
+    """
+
+    def __init__(self, device=None):
+        if device is None:
+            device = 'cuda' if torch.cuda.is_available() else 'cpu'
+        self.device = device
+        self.timestep = hparams['hop_size'] / hparams['audio_sample_rate']
+        self.spk_map = {}
+        self.lang_map = {}
+        self.model: torch.nn.Module = None
+
+    def build_model(self, ckpt_steps=None) -> torch.nn.Module:
+        raise NotImplementedError()
+
+    def load_speaker_mix(self, param_src: dict, summary_dst: dict,
+                         mix_mode: str = 'frame', mix_length: int = None) -> Tuple[Tensor, Tensor]:
+        """
+
+        :param param_src: param dict
+        :param summary_dst: summary dict
+        :param mix_mode: 'token' or 'frame'
+        :param mix_length: total tokens or frames to mix
+        :return: spk_mix_id [B=1, 1, N], spk_mix_value [B=1, T, N]
+        """
+        assert mix_mode == 'token' or mix_mode == 'frame'
+        param_key = 'spk_mix' if mix_mode == 'frame' else 'ph_spk_mix'
+        summary_solo_key = 'spk' if mix_mode == 'frame' else 'ph_spk'
+        spk_mix_map = param_src.get(param_key)  # { spk_name: value } or { spk_name: "value value value ..." }
+        dynamic = False
+        if spk_mix_map is None:
+            assert len(self.spk_map) == 1, (
+                "This is a multi-speaker model. "
+                "Please specify a speaker or speaker mix by --spk option."
+            )
+            # Get the only speaker
+            for name in self.spk_map.keys():
+                spk_mix_map = {name: 1.0}
+                break
+        else:
+            for name in spk_mix_map:
+                assert name in self.spk_map, f'Speaker \'{name}\' not found.'
+        if len(spk_mix_map) == 1:
+            summary_dst[summary_solo_key] = list(spk_mix_map.keys())[0]
+        elif any([isinstance(val, str) for val in spk_mix_map.values()]):
+            print_mix = '|'.join(spk_mix_map.keys())
+            summary_dst[param_key] = f'dynamic({print_mix})'
+            dynamic = True
+        else:
+            print_mix = '|'.join([f'{n}:{"%.3f" % spk_mix_map[n]}' for n in spk_mix_map])
+            summary_dst[param_key] = f'static({print_mix})'
+        spk_mix_id_list = []
+        spk_mix_value_list = []
+        if dynamic:
+            for name, values in spk_mix_map.items():
+                spk_mix_id_list.append(self.spk_map[name])
+                if isinstance(values, str):
+                    # this speaker has a variable proportion
+                    if mix_mode == 'token':
+                        cur_spk_mix_value = values.split()
+                        assert len(cur_spk_mix_value) == mix_length, \
+                            'Speaker mix checks failed. In dynamic token-level mix, ' \
+                            'number of proportion values must equal number of tokens.'
+                        cur_spk_mix_value = torch.from_numpy(
+                            np.array(cur_spk_mix_value, 'float32')
+                        ).to(self.device)[None]  # => [B=1, T]
+                    else:
+                        cur_spk_mix_value = torch.from_numpy(resample_align_curve(
+                            np.array(values.split(), 'float32'),
+                            original_timestep=float(param_src['spk_mix_timestep']),
+                            target_timestep=self.timestep,
+                            align_length=mix_length
+                        )).to(self.device)[None]  # => [B=1, T]
+                    assert torch.all(cur_spk_mix_value >= 0.), \
+                        f'Speaker mix checks failed.\n' \
+                        f'Proportions of speaker \'{name}\' on some {mix_mode}s are negative.'
+                else:
+                    # this speaker has a constant proportion
+                    assert values >= 0., f'Speaker mix checks failed.\n' \
+                                         f'Proportion of speaker \'{name}\' is negative.'
+                    cur_spk_mix_value = torch.full(
+                        (1, mix_length), fill_value=values,
+                        dtype=torch.float32, device=self.device
+                    )
+                spk_mix_value_list.append(cur_spk_mix_value)
+            spk_mix_id = torch.LongTensor(spk_mix_id_list).to(self.device)[None, None]  # => [B=1, 1, N]
+            spk_mix_value = torch.stack(spk_mix_value_list, dim=2)  # [B=1, T] => [B=1, T, N]
+            spk_mix_value_sum = torch.sum(spk_mix_value, dim=2, keepdim=True)  # => [B=1, T, 1]
+            assert torch.all(spk_mix_value_sum > 0.), \
+                f'Speaker mix checks failed.\n' \
+                f'Proportions of speaker mix on some frames sum to zero.'
+            spk_mix_value /= spk_mix_value_sum  # normalize
+        else:
+            for name, value in spk_mix_map.items():
+                spk_mix_id_list.append(self.spk_map[name])
+                assert value >= 0., f'Speaker mix checks failed.\n' \
+                                    f'Proportion of speaker \'{name}\' is negative.'
+                spk_mix_value_list.append(value)
+            spk_mix_id = torch.LongTensor(spk_mix_id_list).to(self.device)[None, None]  # => [B=1, 1, N]
+            spk_mix_value = torch.FloatTensor(spk_mix_value_list).to(self.device)[None, None]  # => [B=1, 1, N]
+            spk_mix_value_sum = spk_mix_value.sum()
+            assert spk_mix_value_sum > 0., f'Speaker mix checks failed.\n' \
+                                           f'Proportions of speaker mix sum to zero.'
+            spk_mix_value /= spk_mix_value_sum  # normalize
+        return spk_mix_id, spk_mix_value
+
+    def preprocess_input(self, param: dict, idx=0) -> Dict[str, torch.Tensor]:
+        raise NotImplementedError()
+
+    def forward_model(self, sample: Dict[str, torch.Tensor]):
+        raise NotImplementedError()
+
+    def run_inference(self, params, **kwargs):
+        raise NotImplementedError()

basics/base_task.py

@@ -0,0 +1,514 @@
+import logging
+import os
+import pathlib
+import shutil
+import sys
+from typing import Dict
+
+import matplotlib
+
+import utils
+
+matplotlib.use('Agg')
+
+import torch.utils.data
+from torchmetrics import Metric, MeanMetric
+import lightning.pytorch as pl
+from lightning.pytorch.utilities.rank_zero import rank_zero_debug, rank_zero_info, rank_zero_only
+
+from basics.base_module import CategorizedModule
+from utils.hparams import hparams
+from utils.training_utils import (
+    DsModelCheckpoint, DsTQDMProgressBar,
+    DsBatchSampler, DsTensorBoardLogger,
+    get_latest_checkpoint_path, get_strategy
+)
+from utils.phoneme_utils import load_phoneme_dictionary
+
+torch.multiprocessing.set_sharing_strategy(os.getenv('TORCH_SHARE_STRATEGY', 'file_system'))
+
+log_format = '%(asctime)s %(message)s'
+logging.basicConfig(stream=sys.stdout, level=logging.INFO,
+                    format=log_format, datefmt='%m/%d %I:%M:%S %p')
+
+
+class BaseTask(pl.LightningModule):
+    """
+        Base class for training tasks.
+        1. *load_ckpt*:
+            load checkpoint;
+        2. *training_step*:
+            record and log the loss;
+        3. *optimizer_step*:
+            run backwards step;
+        4. *start*:
+            load training configs, backup code, log to tensorboard, start training;
+        5. *configure_ddp* and *init_ddp_connection*:
+            start parallel training.
+
+        Subclasses should define:
+        1. *build_model*, *build_optimizer*, *build_scheduler*:
+            how to build the model, the optimizer and the training scheduler;
+        2. *_training_step*:
+            one training step of the model;
+        3. *on_validation_end* and *_on_validation_end*:
+            postprocess the validation output.
+    """
+
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+        self.max_batch_frames = hparams['max_batch_frames']
+        self.max_batch_size = hparams['max_batch_size']
+        self.max_val_batch_frames = hparams['max_val_batch_frames']
+        if self.max_val_batch_frames == -1:
+            hparams['max_val_batch_frames'] = self.max_val_batch_frames = self.max_batch_frames
+        self.max_val_batch_size = hparams['max_val_batch_size']
+        if self.max_val_batch_size == -1:
+            hparams['max_val_batch_size'] = self.max_val_batch_size = self.max_batch_size
+
+        self.training_sampler = None
+        self.skip_immediate_validation = False
+        self.skip_immediate_ckpt_save = False
+
+        self.phoneme_dictionary = load_phoneme_dictionary()
+        self.build_model()
+
+        self.valid_losses: Dict[str, Metric] = {}
+        self.valid_metrics: Dict[str, Metric] = {}
+
+    def _finish_init(self):
+        self.register_validation_loss('total_loss')
+        self.build_losses_and_metrics()
+        assert len(self.valid_losses) > 0, "No validation loss registered. Please check your configuration file."
+
+    ###########
+    # Training, validation and testing
+    ###########
+    def setup(self, stage):
+        self.train_dataset = self.dataset_cls('train')
+        self.valid_dataset = self.dataset_cls('valid')
+        self.num_replicas = (self.trainer.distributed_sampler_kwargs or {}).get('num_replicas', 1)
+
+    def get_need_freeze_state_dict_key(self, model_state_dict) -> list:
+        key_list = []
+        for i in hparams['frozen_params']:
+            for j in model_state_dict:
+                if j.startswith(i):
+                    key_list.append(j)
+        return list(set(key_list))
+
+    def freeze_params(self) -> None:
+        model_state_dict = self.state_dict().keys()
+        freeze_key = self.get_need_freeze_state_dict_key(model_state_dict=model_state_dict)
+
+        for i in freeze_key:
+            params=self.get_parameter(i)
+
+            params.requires_grad = False
+
+    def unfreeze_all_params(self) -> None:
+        for i in self.model.parameters():
+            i.requires_grad = True
+
+    def load_finetune_ckpt(
+            self, state_dict
+    ) -> None:
+        adapt_shapes = hparams['finetune_strict_shapes']
+        if not adapt_shapes:
+            cur_model_state_dict = self.state_dict()
+            unmatched_keys = []
+            for key, param in state_dict.items():
+                if key in cur_model_state_dict:
+                    new_param = cur_model_state_dict[key]
+                    if new_param.shape != param.shape:
+                        unmatched_keys.append(key)
+                        print('| Unmatched keys: ', key, new_param.shape, param.shape)
+            for key in unmatched_keys:
+                del state_dict[key]
+        self.load_state_dict(state_dict, strict=False)
+
+    def load_pre_train_model(self):
+        pre_train_ckpt_path = hparams['finetune_ckpt_path']
+        blacklist = hparams['finetune_ignored_params']
+        # whitelist=hparams['pre_train_whitelist']
+        if blacklist is None:
+            blacklist = []
+        # if whitelist is  None:
+        #     raise RuntimeError("")
+
+        if pre_train_ckpt_path is not None:
+            ckpt = torch.load(pre_train_ckpt_path)
+            # if ckpt.get('category') is None:
+            #     raise RuntimeError("")
+
+            if isinstance(self.model, CategorizedModule):
+                self.model.check_category(ckpt.get('category'))
+
+            state_dict = {}
+            for i in ckpt['state_dict']:
+                # if 'diffusion' in i:
+                # if i in rrrr:
+                #     continue
+                skip = False
+                for b in blacklist:
+                    if i.startswith(b):
+                        skip = True
+                        break
+
+                if skip:
+                    continue
+
+                state_dict[i] = ckpt['state_dict'][i]
+                print(i)
+            return state_dict
+        else:
+            raise RuntimeError("")
+
+    def _build_model(self):
+        raise NotImplementedError()
+
+    def build_model(self):
+        self.model = self._build_model()
+        # utils.load_warp(self)
+        self.unfreeze_all_params()
+        if hparams['freezing_enabled']:
+            self.freeze_params()
+        if hparams['finetune_enabled'] and get_latest_checkpoint_path(pathlib.Path(hparams['work_dir'])) is None:
+            self.load_finetune_ckpt(self.load_pre_train_model())
+        self.print_arch()
+
+    @rank_zero_only
+    def print_arch(self):
+        utils.print_arch(self.model)
+
+    def build_losses_and_metrics(self):
+        raise NotImplementedError()
+
+    def register_validation_metric(self, name: str, metric: Metric):
+        assert isinstance(metric, Metric)
+        self.valid_metrics[name] = metric
+
+    def register_validation_loss(self, name: str, Aggregator: Metric = MeanMetric):
+        assert issubclass(Aggregator, Metric)
+        self.valid_losses[name] = Aggregator()
+
+    def run_model(self, sample, infer=False):
+        """
+        steps:
+            1. run the full model
+            2. calculate losses if not infer
+        """
+        raise NotImplementedError()
+
+    def on_train_epoch_start(self):
+        if self.training_sampler is not None:
+            self.training_sampler.set_epoch(self.current_epoch)
+
+    def _training_step(self, sample):
+        """
+        :return: total loss: torch.Tensor, loss_log: dict, other_log: dict
+        """
+        losses = self.run_model(sample)
+        total_loss = sum(losses.values())
+        return total_loss, {**losses, 'batch_size': float(sample['size'])}
+
+    def training_step(self, sample, batch_idx):
+        total_loss, log_outputs = self._training_step(sample)
+
+        # logs to progress bar
+        self.log_dict(log_outputs, prog_bar=True, logger=False, on_step=True, on_epoch=False)
+        self.log('lr', self.lr_schedulers().get_last_lr()[0], prog_bar=True, logger=False, on_step=True, on_epoch=False)
+        # logs to tensorboard
+        if self.global_step % hparams['log_interval'] == 0:
+            tb_log = {f'training/{k}': v for k, v in log_outputs.items()}
+            tb_log['training/lr'] = self.lr_schedulers().get_last_lr()[0]
+            self.logger.log_metrics(tb_log, step=self.global_step)
+
+        return total_loss
+
+    # def on_before_optimizer_step(self, *args, **kwargs):
+    #     self.log_dict(grad_norm(self, norm_type=2))
+
+    def _on_validation_start(self):
+        pass
+
+    def on_validation_start(self):
+        if self.skip_immediate_validation:
+            rank_zero_debug("Skip validation")
+            return
+        self._on_validation_start()
+        for metric in self.valid_losses.values():
+            metric.to(self.device)
+            metric.reset()
+        for metric in self.valid_metrics.values():
+            metric.to(self.device)
+            metric.reset()
+
+    def _validation_step(self, sample, batch_idx):
+        """
+
+        :param sample:
+        :param batch_idx:
+        :return: loss_log: dict, weight: int
+        """
+        raise NotImplementedError()
+
+    def validation_step(self, sample, batch_idx):
+        """
+
+        :param sample:
+        :param batch_idx:
+        """
+        if self.skip_immediate_validation:
+            rank_zero_debug("Skip validation")
+            return
+        if sample['size'] > 0:
+            with torch.autocast(self.device.type, enabled=False):
+                losses, weight = self._validation_step(sample, batch_idx)
+            losses = {
+                'total_loss': sum(losses.values()),
+                **losses
+            }
+            for k, v in losses.items():
+                self.valid_losses[k].update(v, weight=weight)
+
+    def _on_validation_epoch_end(self):
+        pass
+
+    def on_validation_epoch_end(self):
+        if self.skip_immediate_validation:
+            self.skip_immediate_validation = False
+            self.skip_immediate_ckpt_save = True
+            return
+        self._on_validation_epoch_end()
+        loss_vals = {k: v.compute() for k, v in self.valid_losses.items()}
+        metric_vals = {k: v.compute() for k, v in self.valid_metrics.items()}
+        self.log('val_loss', loss_vals['total_loss'], on_epoch=True, prog_bar=True, logger=False, sync_dist=True)
+        self.logger.log_metrics({f'validation/{k}': v for k, v in loss_vals.items()}, step=self.global_step)
+        self.logger.log_metrics({f'metrics/{k}': v for k, v in metric_vals.items()}, step=self.global_step)
+
+    # noinspection PyMethodMayBeStatic
+    def build_scheduler(self, optimizer):
+        from utils import build_lr_scheduler_from_config
+
+        scheduler_args = hparams['lr_scheduler_args']
+        assert scheduler_args['scheduler_cls'] != ''
+        scheduler = build_lr_scheduler_from_config(optimizer, scheduler_args)
+        return scheduler
+
+    # noinspection PyMethodMayBeStatic
+    def build_optimizer(self, model):
+        from utils import build_object_from_class_name
+
+        optimizer_args = hparams['optimizer_args']
+        assert optimizer_args['optimizer_cls'] != ''
+        if 'beta1' in optimizer_args and 'beta2' in optimizer_args and 'betas' not in optimizer_args:
+            optimizer_args['betas'] = (optimizer_args['beta1'], optimizer_args['beta2'])
+        optimizer = build_object_from_class_name(
+            optimizer_args['optimizer_cls'],
+            torch.optim.Optimizer,
+            model if optimizer_args['optimizer_cls'] == 'modules.optimizer.muon.Muon_AdamW' else model.parameters(),
+            **optimizer_args
+        )
+        return optimizer
+
+    def configure_optimizers(self):
+        optm = self.build_optimizer(self.model)
+        scheduler = self.build_scheduler(optm)
+        if scheduler is None:
+            return optm
+        return {
+            "optimizer": optm,
+            "lr_scheduler": {
+                "scheduler": scheduler,
+                "interval": "step",
+                "frequency": 1
+            }
+        }
+
+    def train_dataloader(self):
+        self.training_sampler = DsBatchSampler(
+            self.train_dataset,
+            max_batch_frames=self.max_batch_frames,
+            max_batch_size=self.max_batch_size,
+            num_replicas=self.num_replicas,
+            rank=self.global_rank,
+            sort_by_similar_size=hparams['sort_by_len'],
+            size_reversed=True,
+            required_batch_count_multiple=hparams['accumulate_grad_batches'],
+            shuffle_sample=True,
+            shuffle_batch=True
+        )
+        return torch.utils.data.DataLoader(
+            self.train_dataset,
+            collate_fn=self.train_dataset.collater,
+            batch_sampler=self.training_sampler,
+            num_workers=hparams['ds_workers'],
+            prefetch_factor=hparams['dataloader_prefetch_factor'],
+            pin_memory=True,
+            persistent_workers=True
+        )
+
+    def val_dataloader(self):
+        sampler = DsBatchSampler(
+            self.valid_dataset,
+            max_batch_frames=self.max_val_batch_frames,
+            max_batch_size=self.max_val_batch_size,
+            num_replicas=self.num_replicas,
+            rank=self.global_rank,
+            shuffle_sample=False,
+            shuffle_batch=False,
+            disallow_empty_batch=False,
+            pad_batch_assignment=False
+        )
+        return torch.utils.data.DataLoader(
+            self.valid_dataset,
+            collate_fn=self.valid_dataset.collater,
+            batch_sampler=sampler,
+            num_workers=hparams['ds_workers'],
+            prefetch_factor=hparams['dataloader_prefetch_factor'],
+            persistent_workers=True
+        )
+
+    def test_dataloader(self):
+        return self.val_dataloader()
+
+    def on_test_start(self):
+        self.on_validation_start()
+
+    def test_step(self, sample, batch_idx):
+        return self.validation_step(sample, batch_idx)
+
+    def on_test_end(self):
+        return self.on_validation_end()
+
+    ###########
+    # Running configuration
+    ###########
+
+    @classmethod
+    def start(cls):
+        task = cls()
+
+        # if pre_train is not None:
+        #     task.load_state_dict(pre_train,strict=False)
+        #     print("load success-------------------------------------------------------------------")
+
+        work_dir = pathlib.Path(hparams['work_dir'])
+        trainer = pl.Trainer(
+            accelerator=hparams['pl_trainer_accelerator'],
+            devices=hparams['pl_trainer_devices'],
+            num_nodes=hparams['pl_trainer_num_nodes'],
+            strategy=get_strategy(
+                hparams['pl_trainer_devices'],
+                hparams['pl_trainer_num_nodes'],
+                hparams['pl_trainer_accelerator'],
+                hparams['pl_trainer_strategy'],
+                hparams['pl_trainer_precision'],
+            ),
+            precision=hparams['pl_trainer_precision'],
+            callbacks=[
+                DsModelCheckpoint(
+                    dirpath=work_dir,
+                    filename='model_ckpt_steps_{step}',
+                    auto_insert_metric_name=False,
+                    monitor='step',
+                    mode='max',
+                    save_last=False,
+                    # every_n_train_steps=hparams['val_check_interval'],
+                    save_top_k=hparams['num_ckpt_keep'],
+                    permanent_ckpt_start=hparams['permanent_ckpt_start'],
+                    permanent_ckpt_interval=hparams['permanent_ckpt_interval'],
+                    verbose=True
+                ),
+                # LearningRateMonitor(logging_interval='step'),
+                DsTQDMProgressBar(),
+            ],
+            logger=DsTensorBoardLogger(
+                save_dir=str(work_dir),
+                name='lightning_logs',
+                version='latest'
+            ),
+            gradient_clip_val=hparams['clip_grad_norm'],
+            val_check_interval=hparams['val_check_interval'] * hparams['accumulate_grad_batches'],
+            # so this is global_steps
+            check_val_every_n_epoch=None,
+            log_every_n_steps=1,
+            max_steps=hparams['max_updates'],
+            use_distributed_sampler=False,
+            num_sanity_val_steps=hparams['num_sanity_val_steps'],
+            accumulate_grad_batches=hparams['accumulate_grad_batches']
+        )
+        if not hparams['infer']:  # train
+            @rank_zero_only
+            def train_payload_copy():
+                # Copy files to work_dir
+                binary_dir = pathlib.Path(hparams['binary_data_dir'])
+                spk_map_dst = work_dir / 'spk_map.json'
+                spk_map_src = binary_dir / 'spk_map.json'
+                shutil.copy(spk_map_src, spk_map_dst)
+                print(f'| Copied spk map to {spk_map_dst}.')
+                lang_map_dst = work_dir / 'lang_map.json'
+                lang_map_src = binary_dir / 'lang_map.json'
+                shutil.copy(lang_map_src, lang_map_dst)
+                print(f'| Copied lang map to {lang_map_dst}.')
+                for lang in hparams['dictionaries'].keys():
+                    dict_dst = work_dir / f'dictionary-{lang}.txt'
+                    dict_src = binary_dir / f'dictionary-{lang}.txt'
+                    shutil.copy(dict_src, dict_dst)
+                    print(f'| Copied dictionary for language \'{lang}\' to {dict_dst}.')
+
+            train_payload_copy()
+            trainer.fit(task, ckpt_path=get_latest_checkpoint_path(work_dir))
+        else:
+            trainer.test(task)
+
+    def on_save_checkpoint(self, checkpoint):
+        if isinstance(self.model, CategorizedModule):
+            checkpoint['category'] = self.model.category
+        checkpoint['trainer_stage'] = self.trainer.state.stage.value
+
+    def on_load_checkpoint(self, checkpoint):
+        from lightning.pytorch.trainer.states import RunningStage
+        from utils import simulate_lr_scheduler
+        if checkpoint.get('trainer_stage', '') == RunningStage.VALIDATING.value:
+            self.skip_immediate_validation = True
+
+        optimizer_args = hparams['optimizer_args']
+        scheduler_args = hparams['lr_scheduler_args']
+
+        if 'beta1' in optimizer_args and 'beta2' in optimizer_args and 'betas' not in optimizer_args:
+            optimizer_args['betas'] = (optimizer_args['beta1'], optimizer_args['beta2'])
+
+        if checkpoint.get('optimizer_states', None):
+            opt_states = checkpoint['optimizer_states']
+            assert len(opt_states) == 1  # only support one optimizer
+            opt_state = opt_states[0]
+            for param_group in opt_state['param_groups']:
+                for k, v in optimizer_args.items():
+                    if k in param_group and param_group[k] != v:
+                        if 'lr_schedulers' in checkpoint and checkpoint['lr_schedulers'] and k == 'lr':
+                            continue
+                        rank_zero_info(f'| Overriding optimizer parameter {k} from checkpoint: {param_group[k]} -> {v}')
+                        param_group[k] = v
+                if 'initial_lr' in param_group and param_group['initial_lr'] != optimizer_args['lr']:
+                    rank_zero_info(
+                        f'| Overriding optimizer parameter initial_lr from checkpoint: {param_group["initial_lr"]} -> {optimizer_args["lr"]}'
+                    )
+                    param_group['initial_lr'] = optimizer_args['lr']
+
+        if checkpoint.get('lr_schedulers', None):
+            assert checkpoint.get('optimizer_states', False)
+            assert len(checkpoint['lr_schedulers']) == 1  # only support one scheduler
+            checkpoint['lr_schedulers'][0] = simulate_lr_scheduler(
+                optimizer_args, scheduler_args,
+                step_count=checkpoint['global_step'],
+                num_param_groups=len(checkpoint['optimizer_states'][0]['param_groups'])
+            )
+            for param_group, new_lr in zip(
+                checkpoint['optimizer_states'][0]['param_groups'],
+                checkpoint['lr_schedulers'][0]['_last_lr'],
+            ):
+                if param_group['lr'] != new_lr:
+                    rank_zero_info(f'| Overriding optimizer parameter lr from checkpoint: {param_group["lr"]} -> {new_lr}')
+                    param_group['lr'] = new_lr

basics/base_vocoder.py

@@ -0,0 +1,23 @@
+class BaseVocoder:
+    def to_device(self, device):
+        """
+        
+        :param device: torch.device or str
+        """
+        raise NotImplementedError()
+
+    def get_device(self):
+        """
+        
+        :return: device: torch.device or str
+        """
+        raise NotImplementedError()
+    
+    def spec2wav(self, mel, **kwargs):
+        """
+
+        :param mel: [T, 80]
+        :return: wav: [T']
+        """
+
+        raise NotImplementedError()