Add audio files with Git LFS support

- Add .wav, .ogg, .mp3, .flac to .gitattributes for LFS tracking
- Migrate existing audio files to Git LFS
- This resolves the binary file rejection from Hugging Face Spaces

.gitattributes

@@ -33,3 +33,8 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 *.zip filter=lfs diff=lfs merge=lfs -text
 *.zst filter=lfs diff=lfs merge=lfs -text
 *tfevents* filter=lfs diff=lfs merge=lfs -text
+# Audio files
+*.wav filter=lfs diff=lfs merge=lfs -text
+*.ogg filter=lfs diff=lfs merge=lfs -text
+*.mp3 filter=lfs diff=lfs merge=lfs -text
+*.flac filter=lfs diff=lfs merge=lfs -text

.gitignore

@@ -0,0 +1,219 @@
+# pretrained models
+models/**/*.pt
+
+# Byte-compiled / optimized / DLL files
+__pycache__/
+*.py[codz]
+*$py.class
+
+# C extensions
+*.so
+
+# Distribution / packaging
+.Python
+build/
+develop-eggs/
+dist/
+downloads/
+eggs/
+.eggs/
+lib/
+lib64/
+parts/
+sdist/
+var/
+wheels/
+share/python-wheels/
+*.egg-info/
+.installed.cfg
+*.egg
+MANIFEST
+
+# PyInstaller
+#   Usually these files are written by a python script from a template
+#   before PyInstaller builds the exe, so as to inject date/other infos into it.
+*.manifest
+*.spec
+
+# Installer logs
+pip-log.txt
+pip-delete-this-directory.txt
+
+# Unit test / coverage reports
+htmlcov/
+.tox/
+.nox/
+.coverage
+.coverage.*
+.cache
+nosetests.xml
+coverage.xml
+*.cover
+*.py.cover
+.hypothesis/
+.pytest_cache/
+cover/
+
+# Translations
+*.mo
+*.pot
+
+# Django stuff:
+*.log
+local_settings.py
+db.sqlite3
+db.sqlite3-journal
+
+# Flask stuff:
+instance/
+.webassets-cache
+
+# Scrapy stuff:
+.scrapy
+
+# Sphinx documentation
+docs/_build/
+
+# PyBuilder
+.pybuilder/
+target/
+
+# Jupyter Notebook
+.ipynb_checkpoints
+
+# IPython
+profile_default/
+ipython_config.py
+
+# pyenv
+#   For a library or package, you might want to ignore these files since the code is
+#   intended to run in multiple environments; otherwise, check them in:
+# .python-version
+
+# pipenv
+#   According to pypa/pipenv#598, it is recommended to include Pipfile.lock in version control.
+#   However, in case of collaboration, if having platform-specific dependencies or dependencies
+#   having no cross-platform support, pipenv may install dependencies that don't work, or not
+#   install all needed dependencies.
+# Pipfile.lock
+
+# UV
+#   Similar to Pipfile.lock, it is generally recommended to include uv.lock in version control.
+#   This is especially recommended for binary packages to ensure reproducibility, and is more
+#   commonly ignored for libraries.
+# uv.lock
+
+# poetry
+#   Similar to Pipfile.lock, it is generally recommended to include poetry.lock in version control.
+#   This is especially recommended for binary packages to ensure reproducibility, and is more
+#   commonly ignored for libraries.
+#   https://python-poetry.org/docs/basic-usage/#commit-your-poetrylock-file-to-version-control
+# poetry.lock
+# poetry.toml
+
+# pdm
+#   Similar to Pipfile.lock, it is generally recommended to include pdm.lock in version control.
+#   pdm recommends including project-wide configuration in pdm.toml, but excluding .pdm-python.
+#   https://pdm-project.org/en/latest/usage/project/#working-with-version-control
+# pdm.lock
+# pdm.toml
+.pdm-python
+.pdm-build/
+
+# pixi
+#   Similar to Pipfile.lock, it is generally recommended to include pixi.lock in version control.
+# pixi.lock
+#   Pixi creates a virtual environment in the .pixi directory, just like venv module creates one
+#   in the .venv directory. It is recommended not to include this directory in version control.
+.pixi
+
+# PEP 582; used by e.g. github.com/David-OConnor/pyflow and github.com/pdm-project/pdm
+__pypackages__/
+
+# Celery stuff
+celerybeat-schedule
+celerybeat.pid
+
+# Redis
+*.rdb
+*.aof
+*.pid
+
+# RabbitMQ
+mnesia/
+rabbitmq/
+rabbitmq-data/
+
+# ActiveMQ
+activemq-data/
+
+# SageMath parsed files
+*.sage.py
+
+# Environments
+.env
+.envrc
+.venv
+env/
+venv/
+ENV/
+env.bak/
+venv.bak/
+
+# Spyder project settings
+.spyderproject
+.spyproject
+
+# Rope project settings
+.ropeproject
+
+# mkdocs documentation
+/site
+
+# mypy
+.mypy_cache/
+.dmypy.json
+dmypy.json
+
+# Pyre type checker
+.pyre/
+
+# pytype static type analyzer
+.pytype/
+
+# Cython debug symbols
+cython_debug/
+
+# PyCharm
+#   JetBrains specific template is maintained in a separate JetBrains.gitignore that can
+#   be found at https://github.com/github/gitignore/blob/main/Global/JetBrains.gitignore
+#   and can be added to the global gitignore or merged into this file.  For a more nuclear
+#   option (not recommended) you can uncomment the following to ignore the entire idea folder.
+# .idea/
+
+# Abstra
+#   Abstra is an AI-powered process automation framework.
+#   Ignore directories containing user credentials, local state, and settings.
+#   Learn more at https://abstra.io/docs
+.abstra/
+
+# Visual Studio Code
+#   Visual Studio Code specific template is maintained in a separate VisualStudioCode.gitignore
+#   that can be found at https://github.com/github/gitignore/blob/main/Global/VisualStudioCode.gitignore
+#   and can be added to the global gitignore or merged into this file. However, if you prefer,
+#   you could uncomment the following to ignore the entire vscode folder
+# .vscode/
+
+# Ruff stuff:
+.ruff_cache/
+
+# PyPI configuration file
+.pypirc
+
+# Marimo
+marimo/_static/
+marimo/_lsp/
+__marimo__/
+
+# Streamlit
+.streamlit/secrets.toml

Makefile

@@ -0,0 +1,5 @@
+run:
+	PYTHONPATH=. uv run gradio app.py
+
+export-requirements:
+	uv export --format requirements-txt > requirements.txt

README.md

@@ -4,7 +4,9 @@ emoji: 🐨
 colorFrom: purple
 colorTo: pink
 sdk: gradio
+python_version: 3.13
 sdk_version: 5.49.1
+suggested_hardware: g4
 app_file: app.py
 pinned: false
 ---

app.py

@@ -0,0 +1,18 @@
+import gradio as gr
+
+from apps.audio_cloning.main import main as audio_cloning
+from apps.dev.main import main as dev
+
+with gr.Blocks(theme=gr.themes.Soft()) as demo:
+    audio_cloning()
+
+
+with demo.route(name="Dev", path="/dev"):
+    dev()
+
+
+if __name__ == "__main__":
+    # demo.queue(max_size=2, concurrency_limit=2, concurrency_id="gpu_queue")
+    # auth = ("charaxim", "chrmx-demo-wordpass")
+    # demo.launch(share=False, auth=auth)
+    demo.launch(share=False)

apps/audio_cloning/main.py

@@ -0,0 +1,15 @@
+import gradio as gr
+
+from .vallex.main import main as vallex
+
+
+def main():
+    gr.Markdown("# Charamix Audio Cloning Prototype")
+
+    # zero-shot audio cloning
+    with gr.Tab("Zero-shot Audio Cloning with VALL-E-X"):
+        vallex()
+
+    # fine-tuning audio cloning
+    # with gr.Tab("Fine-tuning Audio Cloning"):
+    #     gr.Markdown("TODO")

apps/audio_cloning/vallex/data/__init__.py

@@ -0,0 +1,3 @@
+# from .datamodule import *
+# from .tokenizer import *
+from .collation import *

apps/audio_cloning/vallex/data/collation.py

@@ -0,0 +1,107 @@
+from typing import List, Tuple
+
+import numpy as np
+import torch
+
+
+class TextTokenCollater:
+    """Collate list of text tokens
+
+    Map sentences to integers. Sentences are padded to equal length.
+    Beginning and end-of-sequence symbols can be added.
+
+    Example:
+        >>> token_collater = TextTokenCollater(text_tokens)
+        >>> tokens_batch, tokens_lens = token_collater(text)
+
+    Returns:
+        tokens_batch: IntTensor of shape (B, L)
+            B: batch dimension, number of input sentences
+            L: length of the longest sentence
+        tokens_lens: IntTensor of shape (B,)
+            Length of each sentence after adding <eos> and <bos>
+            but before padding.
+    """
+
+    def __init__(
+        self,
+        text_tokens: List[str],
+        add_eos: bool = True,
+        add_bos: bool = True,
+        pad_symbol: str = "<pad>",
+        bos_symbol: str = "<bos>",
+        eos_symbol: str = "<eos>",
+    ):
+        self.pad_symbol = pad_symbol
+
+        self.add_eos = add_eos
+        self.add_bos = add_bos
+
+        self.bos_symbol = bos_symbol
+        self.eos_symbol = eos_symbol
+
+        unique_tokens = (
+            [pad_symbol]
+            + ([bos_symbol] if add_bos else [])
+            + ([eos_symbol] if add_eos else [])
+            + sorted(text_tokens)
+        )
+
+        self.token2idx = {token: idx for idx, token in enumerate(unique_tokens)}
+        self.idx2token = [token for token in unique_tokens]
+
+    def index(self, tokens_list: List[str]) -> Tuple[torch.Tensor, torch.Tensor]:
+        seqs, seq_lens = [], []
+        for tokens in tokens_list:
+            assert all([True if s in self.token2idx else False for s in tokens]) is True
+            seq = (
+                ([self.bos_symbol] if self.add_bos else [])
+                + list(tokens)
+                + ([self.eos_symbol] if self.add_eos else [])
+            )
+            seqs.append(seq)
+            seq_lens.append(len(seq))
+
+        max_len = max(seq_lens)
+        for k, (seq, seq_len) in enumerate(zip(seqs, seq_lens)):
+            seq.extend([self.pad_symbol] * (max_len - seq_len))
+
+        tokens = torch.from_numpy(
+            np.array(
+                [[self.token2idx[token] for token in seq] for seq in seqs],
+                dtype=np.int64,
+            )
+        )
+        tokens_lens = torch.IntTensor(seq_lens)
+
+        return tokens, tokens_lens
+
+    def __call__(self, texts: List[str]) -> Tuple[torch.Tensor, torch.Tensor]:
+        tokens_seqs = [[p for p in text] for text in texts]
+        max_len = len(max(tokens_seqs, key=len))
+
+        seqs = [
+            ([self.bos_symbol] if self.add_bos else [])
+            + list(seq)
+            + ([self.eos_symbol] if self.add_eos else [])
+            + [self.pad_symbol] * (max_len - len(seq))
+            for seq in tokens_seqs
+        ]
+
+        tokens_batch = torch.from_numpy(
+            np.array(
+                [seq for seq in seqs],
+                dtype=np.int64,
+            )
+        )
+
+        tokens_lens = torch.IntTensor(
+            [len(seq) + int(self.add_eos) + int(self.add_bos) for seq in tokens_seqs]
+        )
+
+        return tokens_batch, tokens_lens
+
+
+def get_text_token_collater() -> TextTokenCollater:
+    collater = TextTokenCollater(["0"], add_bos=False, add_eos=False)
+    return collater

apps/audio_cloning/vallex/data/datamodule.py

@@ -0,0 +1,419 @@
+# Copyright      2023                          (authors: Feiteng Li)
+#
+# See ../../../../LICENSE for clarification regarding multiple authors
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+
+import argparse
+import inspect
+import logging
+from functools import lru_cache
+from pathlib import Path
+from typing import Any, Dict, Optional
+
+import torch
+# from icefall.utils import str2bool
+# from lhotse import CutSet, load_manifest_lazy
+# from lhotse.dataset import (
+#     CutConcatenate,
+#     DynamicBucketingSampler,
+#     PrecomputedFeatures,
+#     SingleCutSampler,
+#     SpecAugment,
+# )
+# from lhotse.dataset.input_strategies import OnTheFlyFeatures
+# from lhotse.utils import fix_random_seed
+from torch.utils.data import DataLoader
+
+from data.collation import get_text_token_collater
+# from data.dataset import SpeechSynthesisDataset
+from data.fbank import get_fbank_extractor
+from data.input_strategies import PromptedPrecomputedFeatures
+
+# PrecomputedFeatures = PrecomputedFeatures
+
+
+class _SeedWorkers:
+    def __init__(self, seed: int):
+        self.seed = seed
+
+    def __call__(self, worker_id: int):
+        fix_random_seed(self.seed + worker_id)
+
+
+def _get_input_strategy(input_strategy, dataset, cuts):
+    if input_strategy == "PromptedPrecomputedFeatures":
+        return PromptedPrecomputedFeatures(dataset, cuts)
+
+    return eval(input_strategy)()
+
+
+class TtsDataModule:
+    """
+    DataModule for VALL-E TTS experiments.
+    It assumes there is always one train and valid dataloader.
+
+    It contains all the common data pipeline modules used in TTS
+    experiments, e.g.:
+    - dynamic batch size,
+    - bucketing samplers,
+    - cut concatenation[not used & tested yet],
+    - augmentation[not used & tested yet],
+    - on-the-fly feature extraction[not used & tested yet]
+
+    This class should be derived for specific corpora used in TTS tasks.
+    """
+
+    def __init__(self, args: argparse.Namespace):
+        self.args = args
+
+    @classmethod
+    def add_arguments(cls, parser: argparse.ArgumentParser):
+        group = parser.add_argument_group(
+            title="TTS data related options",
+            description="These options are used for the preparation of "
+            "PyTorch DataLoaders from Lhotse CutSet's -- they control the "
+            "effective batch sizes, sampling strategies, applied data "
+            "augmentations, etc.",
+        )
+        group.add_argument(
+            "--manifest-dir",
+            type=Path,
+            default=Path("data/tokenized"),
+            help="Path to directory with train/valid/test cuts.",
+        )
+        group.add_argument(
+            "--max-duration",
+            type=int,
+            default=40.0,
+            help="Maximum pooled recordings duration (seconds) in a "
+            "single batch. You can reduce it if it causes CUDA OOM.",
+        )
+        group.add_argument(
+            "--bucketing-sampler",
+            type=str2bool,
+            default=True,
+            help="When enabled, the batches will come from buckets of "
+            "similar duration (saves padding frames).",
+        )
+        group.add_argument(
+            "--num-buckets",
+            type=int,
+            default=10,
+            help="The number of buckets for the DynamicBucketingSampler"
+            "(you might want to increase it for larger datasets).",
+        )
+        group.add_argument(
+            "--concatenate-cuts",
+            type=str2bool,
+            default=False,
+            help="When enabled, utterances (cuts) will be concatenated "
+            "to minimize the amount of padding.",
+        )
+        group.add_argument(
+            "--duration-factor",
+            type=float,
+            default=1.0,
+            help="Determines the maximum duration of a concatenated cut "
+            "relative to the duration of the longest cut in a batch.",
+        )
+        group.add_argument(
+            "--gap",
+            type=float,
+            default=0.1,
+            help="The amount of padding (in seconds) inserted between "
+            "concatenated cuts. This padding is filled with noise when "
+            "noise augmentation is used.",
+        )
+        group.add_argument(
+            "--on-the-fly-feats",
+            type=str2bool,
+            default=False,
+            help="When enabled, use on-the-fly cut mixing and feature "
+            "extraction. Will drop existing precomputed feature manifests "
+            "if available.",
+        )
+        group.add_argument(
+            "--shuffle",
+            type=str2bool,
+            default=True,
+            help="When enabled (=default), the examples will be "
+            "shuffled for each epoch.",
+        )
+        group.add_argument(
+            "--drop-last",
+            type=str2bool,
+            default=False,
+            help="Whether to drop last batch. Used by sampler.",
+        )
+        group.add_argument(
+            "--return-cuts",
+            type=str2bool,
+            default=True,
+            help="When enabled, each batch will have the "
+            "field: batch['supervisions']['cut'] with the cuts that "
+            "were used to construct it.",
+        )
+
+        group.add_argument(
+            "--num-workers",
+            type=int,
+            default=8,
+            help="The number of training dataloader workers that "
+            "collect the batches.",
+        )
+
+        group.add_argument(
+            "--enable-spec-aug",
+            type=str2bool,
+            default=False,
+            help="When enabled, use SpecAugment for training dataset.",
+        )
+
+        group.add_argument(
+            "--spec-aug-time-warp-factor",
+            type=int,
+            default=80,
+            help="Used only when --enable-spec-aug is True. "
+            "It specifies the factor for time warping in SpecAugment. "
+            "Larger values mean more warping. "
+            "A value less than 1 means to disable time warp.",
+        )
+
+        group.add_argument(
+            "--input-strategy",
+            type=str,
+            default="PrecomputedFeatures",
+            help="AudioSamples or PrecomputedFeatures or PromptedPrecomputedFeatures",
+        )
+
+        group.add_argument(
+            "--dataset",
+            type=str,
+            default="ljspeech",
+            help="--input-strategy PromptedPrecomputedFeatures needs dataset name to prepare prompts.",
+        )
+
+        parser.add_argument(
+            "--text-tokens",
+            type=str,
+            default="data/tokenized/unique_text_tokens.k2symbols",
+            help="Path to the unique text tokens file",
+        )
+
+        parser.add_argument(
+            "--sampling-rate",
+            type=int,
+            default=24000,
+            help="""Audio sampling rate.""",
+        )
+
+    def train_dataloaders(
+        self,
+        cuts_train: CutSet,
+        sampler_state_dict: Optional[Dict[str, Any]] = None,
+    ) -> DataLoader:
+        """
+        Args:
+          cuts_train:
+            CutSet for training.
+          sampler_state_dict:
+            The state dict for the training sampler.
+        """
+        transforms = []
+
+        if self.args.concatenate_cuts:
+            logging.info(
+                f"Using cut concatenation with duration factor "
+                f"{self.args.duration_factor} and gap {self.args.gap}."
+            )
+            # Cut concatenation should be the first transform in the list,
+            # so that if we e.g. mix noise in, it will fill the gaps between
+            # different utterances.
+            transforms = [
+                CutConcatenate(
+                    duration_factor=self.args.duration_factor, gap=self.args.gap
+                )
+            ] + transforms
+
+        input_transforms = []
+        if self.args.enable_spec_aug:
+            logging.info("Enable SpecAugment")
+            logging.info(
+                f"Time warp factor: {self.args.spec_aug_time_warp_factor}"
+            )
+            # Set the value of num_frame_masks according to Lhotse's version.
+            # In different Lhotse's versions, the default of num_frame_masks is
+            # different.
+            num_frame_masks = 10
+            num_frame_masks_parameter = inspect.signature(
+                SpecAugment.__init__
+            ).parameters["num_frame_masks"]
+            if num_frame_masks_parameter.default == 1:
+                num_frame_masks = 2
+            logging.info(f"Num frame mask: {num_frame_masks}")
+            input_transforms.append(
+                SpecAugment(
+                    time_warp_factor=self.args.spec_aug_time_warp_factor,
+                    num_frame_masks=num_frame_masks,
+                    features_mask_size=27,
+                    num_feature_masks=2,
+                    frames_mask_size=100,
+                )
+            )
+        else:
+            logging.info("Disable SpecAugment")
+
+        logging.info("About to create train dataset")
+        if self.args.on_the_fly_feats:
+            # NOTE: the PerturbSpeed transform should be added only if we
+            # remove it from data prep stage.
+            # Add on-the-fly speed perturbation; since originally it would
+            # have increased epoch size by 3, we will apply prob 2/3 and use
+            # 3x more epochs.
+            # Speed perturbation probably should come first before
+            # concatenation, but in principle the transforms order doesn't have
+            # to be strict (e.g. could be randomized)
+            # transforms = [PerturbSpeed(factors=[0.9, 1.1], p=2/3)] + transforms   # noqa
+            # Drop feats to be on the safe side.
+            train = SpeechSynthesisDataset(
+                get_text_token_collater(self.args.text_tokens),
+                cut_transforms=transforms,
+                feature_input_strategy=OnTheFlyFeatures(get_fbank_extractor()),
+                feature_transforms=input_transforms,
+            )
+        else:
+            train = SpeechSynthesisDataset(
+                get_text_token_collater(self.args.text_tokens),
+                feature_input_strategy=_get_input_strategy(
+                    self.args.input_strategy, self.args.dataset, cuts_train
+                ),
+                cut_transforms=transforms,
+                feature_transforms=input_transforms,
+            )
+
+        if self.args.bucketing_sampler:
+            logging.info("Using DynamicBucketingSampler")
+            train_sampler = DynamicBucketingSampler(
+                cuts_train,
+                max_duration=self.args.max_duration,
+                shuffle=self.args.shuffle,
+                num_buckets=self.args.num_buckets,
+                drop_last=self.args.drop_last,
+            )
+        else:
+            logging.info(
+                "Using SingleCutSampler and sort by duraton(ascending=True)."
+            )
+            cuts_train = cuts_train.to_eager().sort_by_duration(ascending=True)
+            train_sampler = SingleCutSampler(
+                cuts_train,
+                max_duration=self.args.max_duration,
+                shuffle=self.args.shuffle,
+            )
+        logging.info("About to create train dataloader")
+
+        if sampler_state_dict is not None:
+            logging.info("Loading sampler state dict")
+            train_sampler.load_state_dict(sampler_state_dict)
+
+        # 'seed' is derived from the current random state, which will have
+        # previously been set in the main process.
+        seed = torch.randint(0, 100000, ()).item()
+        worker_init_fn = _SeedWorkers(seed)
+
+        train_dl = DataLoader(
+            train,
+            sampler=train_sampler,
+            batch_size=None,
+            num_workers=self.args.num_workers,
+            persistent_workers=False,
+            worker_init_fn=worker_init_fn,
+        )
+
+        return train_dl
+
+    def valid_dataloaders(self, cuts_valid: CutSet) -> DataLoader:
+        logging.info("About to create dev dataset")
+        if self.args.on_the_fly_feats:
+            validate = SpeechSynthesisDataset(
+                get_text_token_collater(self.args.text_tokens),
+                feature_input_strategy=OnTheFlyFeatures(get_fbank_extractor()),
+                cut_transforms=[],
+            )
+        else:
+            validate = SpeechSynthesisDataset(
+                get_text_token_collater(self.args.text_tokens),
+                feature_input_strategy=_get_input_strategy(
+                    self.args.input_strategy, self.args.dataset, cuts_valid
+                ),
+                cut_transforms=[],
+            )
+        valid_sampler = DynamicBucketingSampler(
+            cuts_valid,
+            max_duration=self.args.max_duration,
+            shuffle=False,
+        )
+        logging.info("About to create dev dataloader")
+        valid_dl = DataLoader(
+            validate,
+            sampler=valid_sampler,
+            batch_size=None,
+            num_workers=4,
+            persistent_workers=False,
+        )
+
+        return valid_dl
+
+    def test_dataloaders(self, cuts: CutSet) -> DataLoader:
+        logging.debug("About to create test dataset")
+        test = SpeechSynthesisDataset(
+            get_text_token_collater(self.args.text_tokens),
+            feature_input_strategy=OnTheFlyFeatures(get_fbank_extractor())
+            if self.args.on_the_fly_feats
+            else _get_input_strategy(
+                self.args.input_strategy, self.args.dataset, cuts
+            ),
+            cut_transforms=[],
+        )
+        sampler = DynamicBucketingSampler(
+            cuts,
+            max_duration=self.args.max_duration,
+            shuffle=False,
+        )
+        logging.debug("About to create test dataloader")
+        test_dl = DataLoader(
+            test,
+            batch_size=None,
+            sampler=sampler,
+            num_workers=self.args.num_workers,
+        )
+        return test_dl
+
+    @lru_cache()
+    def train_cuts(self) -> CutSet:
+        logging.info("About to get train cuts")
+        return load_manifest_lazy(
+            self.args.manifest_dir / "cuts_train.jsonl.gz"
+        )
+
+    @lru_cache()
+    def dev_cuts(self) -> CutSet:
+        logging.info("About to get dev cuts")
+        return load_manifest_lazy(self.args.manifest_dir / "cuts_dev.jsonl.gz")
+
+    @lru_cache()
+    def test_cuts(self) -> CutSet:
+        logging.info("About to get test cuts")
+        return load_manifest_lazy(self.args.manifest_dir / "cuts_test.jsonl.gz")

apps/audio_cloning/vallex/data/dataset.py

@@ -0,0 +1,242 @@
+# Copyright      2023                           (authors: Feiteng Li)
+#
+# See ../../../../LICENSE for clarification regarding multiple authors
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""
+modified from lhoste.dataset.speech_synthesis.py
+"""
+
+import torch
+import math
+import h5py
+from tokenizers import Tokenizer
+from typing import Union, List
+import numpy as np
+from tqdm import tqdm
+
+_pad        = '_'
+_punctuation = ',.!?-~…'
+_letters = 'NQabdefghijklmnopstuvwxyzɑæʃʑçɯɪɔɛɹðəɫɥɸʊɾʒθβŋɦ⁼ʰ`^#*=ˈˌ→↓↑ '
+symbols = [_pad] + list(_punctuation) + list(_letters)
+
+language_dict = {
+    'en': 0,
+    'zh': 1,
+    'ja': 2,
+}
+def seq2phone(tokens: Union[List, np.ndarray]):
+    """
+    Convert tokenized phoneme ID sequence back to phoneme string
+    :param tokens: phoneme tokens
+    :return: recovered phoneme sequence
+    """
+    phones = "".join([symbols[i] for i in tokens])
+    return phones
+
+class DynamicBatchSampler(torch.utils.data.Sampler):
+    def __init__(self, sampler, num_tokens_fn, num_buckets=100, min_size=0, max_size=1000,
+                 max_tokens=None, max_sentences=None, drop_last=False):
+        """
+
+        :param sampler:
+        :param num_tokens_fn: 根据idx返回样本的长度的函数
+        :param num_buckets: 利用桶原理将相似长度的样本放在一个batchsize中，桶的数量
+        :param min_size: 最小长度的样本， 小于这个值的样本会被过滤掉。 依据这个值来创建样桶
+        :param max_size: 最大长度的样本
+        :param max_sentences: batch_size, 但是这里可以通过max_sentences 和 max_tokens 共同控制最终的大小
+        """
+        super(DynamicBatchSampler, self).__init__(sampler)
+        self.sampler = sampler
+        self.num_tokens_fn = num_tokens_fn
+        self.num_buckets = num_buckets
+
+        self.min_size = min_size
+        self.max_size = max_size
+
+        assert max_size <= max_tokens, "max_size should be smaller than max tokens"
+        assert max_tokens is not None or max_sentences is not None, \
+            "max_tokens and max_sentences should not be null at the same time, please specify one parameter at least"
+        self.max_tokens = max_tokens if max_tokens is not None else float('Inf')
+        self.max_sentences = max_sentences if max_sentences is not None else float('Inf')
+        self.drop_last = drop_last
+
+    def set_epoch(self, epoch):
+        self.sampler.set_epoch(epoch)
+    def is_batch_full(self, num_tokens, batch):
+        if len(batch) == 0:
+            return False
+        if len(batch) == self.max_sentences:
+            return True
+        if num_tokens > self.max_tokens:
+            return True
+        return False
+
+    def __iter__(self):
+        buckets = [[] for _ in range(self.num_buckets)]
+        sample_len = [0] * self.num_buckets
+
+        for idx in self.sampler:
+            idx_length = self.num_tokens_fn(idx)
+            if not (self.min_size <= idx_length <= self.max_size):
+                print("sentence at index {} of size {} exceeds max_tokens, the sentence is ignored".format(idx, idx_length))
+                continue
+
+            index_buckets = math.floor((idx_length - self.min_size) / (self.max_size - self.min_size + 1)
+                                       * self.num_buckets)
+            sample_len[index_buckets] = max(sample_len[index_buckets], idx_length)
+
+            num_tokens = (len(buckets[index_buckets]) + 1) * sample_len[index_buckets]
+            if self.is_batch_full(num_tokens, buckets[index_buckets]):
+                # yield this batch
+                yield buckets[index_buckets]
+                buckets[index_buckets] = []
+                sample_len[index_buckets] = 0
+
+            buckets[index_buckets].append(idx)
+
+        # process left-over
+        leftover_batch = []
+        leftover_sample_len = 0
+        leftover = [idx for bucket in buckets for idx in bucket]
+        for idx in leftover:
+            idx_length = self.num_tokens_fn(idx)
+            leftover_sample_len = max(leftover_sample_len, idx_length)
+            num_tokens = (len(leftover_batch) + 1) * leftover_sample_len
+            if self.is_batch_full(num_tokens, leftover_batch):
+                yield leftover_batch
+                leftover_batch = []
+                leftover_sample_len = 0
+            leftover_batch.append(idx)
+
+        if len(leftover_batch) > 0 and not self.drop_last:
+            yield leftover_batch
+
+    def __len__(self):
+        # we do not know the exactly batch size, so do not call len(dataloader)
+        pass
+
+
+class AudioDataset(torch.utils.data.Dataset):
+    def __init__(self, h5_path, ann_path, tokenizer_path):
+        self.h5_path = h5_path
+        with open(ann_path, 'r', encoding='utf-8') as f:
+            lines = f.readlines()
+        ls = [l.split("|") for l in lines]
+        ls_T = list(zip(*ls))
+        del ls_T[-1]
+        self.h5_paths, self.durations, self.langs, self.texts = \
+            list(ls_T[0]), list(ls_T[1]), list(ls_T[2]), list(ls_T[3])
+        self.durations = [float(dur) for dur in self.durations]
+        self.tokenizer = Tokenizer.from_file(tokenizer_path)
+
+        self._archive = None
+
+    def __len__(self):
+        return len(self.h5_paths)
+
+    def get_dur(self, idx):
+        return self.durations[idx]
+
+    @property
+    def archive(self):
+        if self._archive is None:  # lazy loading here!
+            self._archive = h5py.File(self.h5_path, "r")
+        return self._archive
+    def __getitem__(self, idx):
+        archive = self.archive
+        h5_path = self.h5_paths[idx]
+        sub = archive[h5_path]
+        audio_tokens = sub['audio'][()]
+        phone_tokens = sub['text'][()]
+        dur = self.durations[idx]
+        lang = self.langs[idx]
+        text = self.texts[idx]
+        # tokenization should be done within dataloader
+        phones = seq2phone(phone_tokens)
+        phones = phones.replace(" ", "_")
+        if not len(phones):
+            cptpho_tokens = self.tokenizer.encode(text).ids
+        else:
+            cptpho_tokens = self.tokenizer.encode(phones).ids
+        assert len(cptpho_tokens)
+        return {
+            'utt_id': h5_path,
+            'text': text,
+            'audio': None,
+            'audio_lens': None,
+            'audio_features': audio_tokens,
+            'audio_features_lens': len(audio_tokens.T),
+            'text_tokens': np.array(cptpho_tokens),
+            'text_tokens_lens': len(cptpho_tokens),
+            'language': language_dict[lang],
+        }
+
+def collate(batch):
+    utt_id_s = [b['utt_id'] for b in batch]
+    text_s = [b['text'] for b in batch]
+
+    audio_s = [b['audio'] for b in batch]
+    audio_lens_s = [b['audio_lens'] for b in batch]
+
+    audio_features_lens_s = [b['audio_features_lens'] for b in batch]
+    # create an empty tensor with maximum audio feature length
+    audio_features_s = torch.zeros([len(batch), max(audio_features_lens_s), 8], dtype=torch.int64) - 1 # audio pad with -1
+
+    text_tokens_lens_s = [b['text_tokens_lens'] for b in batch]
+    # create an empty tensor with maximum text tokens length
+    text_tokens_s = torch.zeros([len(batch), max(text_tokens_lens_s)], dtype=torch.int64) + 3 # [PAD] token id 3
+
+    language_s = [b['language'] for b in batch]
+
+    for i, b in enumerate(batch):
+        audio_features = b['audio_features']
+        audio_features_lens = b['audio_features_lens']
+        audio_features_s[i, :audio_features_lens, :] = torch.LongTensor(audio_features.T)
+
+        text_tokens = b['text_tokens']
+        text_tokens_lens = b['text_tokens_lens']
+        text_tokens_s[i, :text_tokens_lens] = torch.LongTensor(text_tokens)
+
+    batch = {
+        'utt_id': utt_id_s,
+        'text': text_s,
+        'audio': audio_s,
+        'audio_lens': audio_lens_s,
+        'audio_features': audio_features_s,
+        'audio_features_lens': torch.LongTensor(np.array(audio_features_lens_s)),
+        'text_tokens': text_tokens_s,
+        'text_tokens_lens': torch.LongTensor(np.array(text_tokens_lens_s)),
+        'languages': torch.LongTensor(np.array(language_s)),
+    }
+    return batch
+
+def create_dataloader(data_dir="/root/valle/egs/mix", n_gpus=1, rank=0, num_workers=0, num_buckets=10, max_duration=120):
+    train_dataset = AudioDataset(h5_path=f"{data_dir}/audio_sum.hdf5",
+                                 ann_path=f"{data_dir}/audio_ann_sum.txt",
+                                 tokenizer_path=f"{data_dir}/bpe_69.json")
+    ran_sampler = torch.utils.data.distributed.DistributedSampler(
+            train_dataset,
+            num_replicas=n_gpus,
+            rank=rank,
+            shuffle=True,
+        )
+    dynamic_sampler = DynamicBatchSampler(ran_sampler, train_dataset.get_dur, num_buckets=num_buckets, max_size=20,
+                                          max_tokens=max_duration)
+
+
+    train_loader = torch.utils.data.DataLoader(train_dataset, num_workers=num_workers, collate_fn=collate,
+                                               batch_sampler=dynamic_sampler)
+
+    return train_loader

apps/audio_cloning/vallex/data/fbank.py

@@ -0,0 +1,212 @@
+# Copyright      2023                          (authors: Feiteng Li)
+#
+# See ../../../../LICENSE for clarification regarding multiple authors
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+
+from dataclasses import asdict, dataclass
+from typing import Any, Dict, Optional, Union
+
+import numpy as np
+import torch
+# from lhotse.features.base import FeatureExtractor
+# from lhotse.utils import EPSILON, Seconds, compute_num_frames
+from librosa.filters import mel as librosa_mel_fn
+
+
+@dataclass
+class BigVGANFbankConfig:
+    # Spectogram-related part
+    # Note that frame_length and frame_shift will be converted to milliseconds before torchaudio/Kaldi sees them
+    frame_length: Seconds = 1024 / 24000.0
+    frame_shift: Seconds = 256 / 24000.0
+    remove_dc_offset: bool = True
+    round_to_power_of_two: bool = True
+
+    # Fbank-related part
+    low_freq: float = 0.0
+    high_freq: float = 12000.0
+    num_mel_bins: int = 100
+    use_energy: bool = False
+
+    def to_dict(self) -> Dict[str, Any]:
+        return asdict(self)
+
+    @staticmethod
+    def from_dict(data: Dict[str, Any]) -> "BigVGANFbankConfig":
+        return BigVGANFbankConfig(**data)
+
+
+def dynamic_range_compression_torch(x, C=1, clip_val=1e-5):
+    return torch.log(torch.clamp(x, min=clip_val) * C)
+
+
+def spectral_normalize_torch(magnitudes):
+    output = dynamic_range_compression_torch(magnitudes)
+    return output
+
+
+# https://github.com/NVIDIA/BigVGAN
+# bigvgan_24khz_100band https://drive.google.com/drive/folders/1EpxX6AsxjCbbk0mmAhE0td6eYiABr8Oz
+class BigVGANFbank(FeatureExtractor):
+    name = "fbank"
+    config_type = BigVGANFbankConfig
+
+    def __init__(self, config: Optional[Any] = None):
+        super(BigVGANFbank, self).__init__(config)
+        sampling_rate = 24000
+        self.mel_basis = torch.from_numpy(
+            librosa_mel_fn(
+                sampling_rate,
+                1024,
+                self.config.num_mel_bins,
+                self.config.low_freq,
+                self.config.high_freq,
+            ).astype(np.float32)
+        )
+        self.hann_window = torch.hann_window(1024)
+
+    def _feature_fn(self, samples, **kwargs):
+        win_length, n_fft = 1024, 1024
+        hop_size = 256
+        if True:
+            sampling_rate = 24000
+            duration = round(samples.shape[-1] / sampling_rate, ndigits=12)
+            expected_num_frames = compute_num_frames(
+                duration=duration,
+                frame_shift=self.frame_shift,
+                sampling_rate=sampling_rate,
+            )
+            pad_size = (
+                (expected_num_frames - 1) * hop_size
+                + win_length
+                - samples.shape[-1]
+            )
+            assert pad_size >= 0
+
+            y = torch.nn.functional.pad(
+                samples,
+                (0, pad_size),
+                mode="constant",
+            )
+        else:
+            y = torch.nn.functional.pad(
+                samples,
+                (int((n_fft - hop_size) / 2), int((n_fft - hop_size) / 2)),
+                mode="reflect",
+            )
+
+        y = y.squeeze(1)
+
+        # complex tensor as default, then use view_as_real for future pytorch compatibility
+        spec = torch.stft(
+            y,
+            n_fft,
+            hop_length=hop_size,
+            win_length=win_length,
+            window=self.hann_window,
+            center=False,
+            pad_mode="reflect",
+            normalized=False,
+            onesided=True,
+            return_complex=True,
+        )
+        spec = torch.view_as_real(spec)
+        spec = torch.sqrt(spec.pow(2).sum(-1) + (1e-9))
+
+        spec = torch.matmul(self.mel_basis, spec)
+        spec = spectral_normalize_torch(spec)
+
+        return spec.transpose(2, 1).squeeze(0)
+
+    def extract(
+        self, samples: Union[np.ndarray, torch.Tensor], sampling_rate: int
+    ) -> np.ndarray:
+        assert sampling_rate == 24000
+        params = asdict(self.config)
+        params.update({"sample_frequency": sampling_rate, "snip_edges": False})
+        params["frame_shift"] *= 1000.0
+        params["frame_length"] *= 1000.0
+        if not isinstance(samples, torch.Tensor):
+            samples = torch.from_numpy(samples)
+        # Torchaudio Kaldi feature extractors expect the channel dimension to be first.
+        if len(samples.shape) == 1:
+            samples = samples.unsqueeze(0)
+        features = self._feature_fn(samples, **params).to(torch.float32)
+        return features.numpy()
+
+    @property
+    def frame_shift(self) -> Seconds:
+        return self.config.frame_shift
+
+    def feature_dim(self, sampling_rate: int) -> int:
+        return self.config.num_mel_bins
+
+    @staticmethod
+    def mix(
+        features_a: np.ndarray,
+        features_b: np.ndarray,
+        energy_scaling_factor_b: float,
+    ) -> np.ndarray:
+        return np.log(
+            np.maximum(
+                # protection against log(0); max with EPSILON is adequate since these are energies (always >= 0)
+                EPSILON,
+                np.exp(features_a)
+                + energy_scaling_factor_b * np.exp(features_b),
+            )
+        )
+
+    @staticmethod
+    def compute_energy(features: np.ndarray) -> float:
+        return float(np.sum(np.exp(features)))
+
+
+def get_fbank_extractor() -> BigVGANFbank:
+    return BigVGANFbank(BigVGANFbankConfig())
+
+
+if __name__ == "__main__":
+    extractor = BigVGANFbank(BigVGANFbankConfig())
+
+    samples = torch.from_numpy(np.random.random([1000]).astype(np.float32))
+    samples = torch.clip(samples, -1.0, 1.0)
+    fbank = extractor.extract(samples, 24000.0)
+    print(f"fbank {fbank.shape}")
+
+    from scipy.io.wavfile import read
+
+    MAX_WAV_VALUE = 32768.0
+
+    sampling_rate, samples = read(
+        "egs/libritts/prompts/5639_40744_000000_000002.wav"
+    )
+    print(f"samples: [{samples.min()}, {samples.max()}]")
+    fbank = extractor.extract(samples.astype(np.float32) / MAX_WAV_VALUE, 24000)
+    print(f"fbank {fbank.shape}")
+
+    import matplotlib.pyplot as plt
+
+    _ = plt.figure(figsize=(18, 10))
+    plt.imshow(
+        X=fbank.transpose(1, 0),
+        cmap=plt.get_cmap("jet"),
+        aspect="auto",
+        interpolation="nearest",
+    )
+    plt.gca().invert_yaxis()
+    plt.savefig("egs/libritts/prompts/5639_40744_000000_000002.png")
+    plt.close()
+
+    print("fbank test PASS!")

apps/audio_cloning/vallex/data/input_strategies.py

@@ -0,0 +1,159 @@
+import random
+from collections import defaultdict
+from concurrent.futures import ThreadPoolExecutor
+from typing import Tuple, Type
+
+# from lhotse import CutSet
+# from lhotse.dataset.collation import collate_features
+# from lhotse.dataset.input_strategies import (
+#     ExecutorType,
+#     PrecomputedFeatures,
+#     _get_executor,
+# )
+# from lhotse.utils import fastcopy
+
+
+class PromptedFeatures:
+    def __init__(self, prompts, features):
+        self.prompts = prompts
+        self.features = features
+
+    def to(self, device):
+        return PromptedFeatures(
+            self.prompts.to(device), self.features.to(device)
+        )
+
+    def sum(self):
+        return self.features.sum()
+
+    @property
+    def ndim(self):
+        return self.features.ndim
+
+    @property
+    def data(self):
+        return (self.prompts, self.features)
+
+
+# class PromptedPrecomputedFeatures(PrecomputedFeatures):
+#     """
+#     :class:`InputStrategy` that reads pre-computed features, whose manifests
+#     are attached to cuts, from disk.
+#
+#     It automatically pads the feature matrices with pre or post feature.
+#
+#     .. automethod:: __call__
+#     """
+#
+#     def __init__(
+#         self,
+#         dataset: str,
+#         cuts: CutSet,
+#         num_workers: int = 0,
+#         executor_type: Type[ExecutorType] = ThreadPoolExecutor,
+#     ) -> None:
+#         super(PromptedPrecomputedFeatures, self).__init__(
+#             num_workers, executor_type
+#         )
+#
+#         self.utt2neighbors = defaultdict(lambda: [])
+#
+#         if dataset.lower() == "libritts":
+#             # 909_131041_000013_000002
+#             # 909_131041_000013_000003
+#             speaker2utts = defaultdict(lambda: [])
+#
+#             utt2cut = {}
+#             for cut in cuts:
+#                 speaker = cut.supervisions[0].speaker
+#                 speaker2utts[speaker].append(cut.id)
+#                 utt2cut[cut.id] = cut
+#
+#             for spk in speaker2utts:
+#                 uttids = sorted(speaker2utts[spk])
+#                 # Using the property of sorted keys to find previous utterance
+#                 # The keys has structure speaker_book_x_y e.g. 1089_134691_000004_000001
+#                 if len(uttids) == 1:
+#                     self.utt2neighbors[uttids[0]].append(utt2cut[uttids[0]])
+#                     continue
+#
+#                 utt2prevutt = dict(zip(uttids, [uttids[1]] + uttids[:-1]))
+#                 utt2postutt = dict(zip(uttids[:-1], uttids[1:]))
+#
+#                 for utt in utt2prevutt:
+#                     self.utt2neighbors[utt].append(utt2cut[utt2prevutt[utt]])
+#
+#                 for utt in utt2postutt:
+#                     self.utt2neighbors[utt].append(utt2cut[utt2postutt[utt]])
+#         elif dataset.lower() == "ljspeech":
+#             utt2cut = {}
+#             uttids = []
+#             for cut in cuts:
+#                 uttids.append(cut.id)
+#                 utt2cut[cut.id] = cut
+#
+#             if len(uttids) == 1:
+#                 self.utt2neighbors[uttids[0]].append(utt2cut[uttids[0]])
+#             else:
+#                 # Using the property of sorted keys to find previous utterance
+#                 # The keys has structure: LJ001-0010
+#                 utt2prevutt = dict(zip(uttids, [uttids[1]] + uttids[:-1]))
+#                 utt2postutt = dict(zip(uttids[:-1], uttids[1:]))
+#
+#                 for utt in utt2postutt:
+#                     postutt = utt2postutt[utt]
+#                     if utt[:5] == postutt[:5]:
+#                         self.utt2neighbors[utt].append(utt2cut[postutt])
+#
+#                 for utt in utt2prevutt:
+#                     prevutt = utt2prevutt[utt]
+#                     if utt[:5] == prevutt[:5] or not self.utt2neighbors[utt]:
+#                         self.utt2neighbors[utt].append(utt2cut[prevutt])
+#         else:
+#             raise ValueError
+#
+#     def __call__(
+#         self, cuts: CutSet
+#     ) -> Tuple[PromptedFeatures, PromptedFeatures]:
+#         """
+#         Reads the pre-computed features from disk/other storage.
+#         The returned shape is``(B, T, F) => (batch_size, num_frames, num_features)``.
+#
+#         :return: a tensor with collated features, and a tensor of ``num_frames`` of each cut before padding.
+#         """
+#         features, features_lens = collate_features(
+#             cuts,
+#             executor=_get_executor(
+#                 self.num_workers, executor_type=self._executor_type
+#             ),
+#         )
+#
+#         prompts_cuts = []
+#         for k, cut in enumerate(cuts):
+#             prompts_cut = random.choice(self.utt2neighbors[cut.id])
+#             prompts_cuts.append(fastcopy(prompts_cut, id=f"{cut.id}-{str(k)}"))
+#
+#         mini_duration = min([cut.duration for cut in prompts_cuts] + [3.0])
+#         # prompts_cuts = CutSet.from_cuts(prompts_cuts).truncate(
+#         #     max_duration=mini_duration,
+#         #     offset_type="random",
+#         #     preserve_id=True,
+#         # )
+#         prompts_cuts = CutSet(
+#             cuts={k: cut for k, cut in enumerate(prompts_cuts)}
+#         ).truncate(
+#             max_duration=mini_duration,
+#             offset_type="random",
+#             preserve_id=False,
+#         )
+#
+#         prompts, prompts_lens = collate_features(
+#             prompts_cuts,
+#             executor=_get_executor(
+#                 self.num_workers, executor_type=self._executor_type
+#             ),
+#         )
+#
+#         return PromptedFeatures(prompts, features), PromptedFeatures(
+#             prompts_lens, features_lens
+#         )

apps/audio_cloning/vallex/data/symbol_table.py

@@ -0,0 +1,289 @@
+# Copyright      2020  Mobvoi Inc.        (authors: Fangjun Kuang)
+#
+# See ../../../LICENSE for clarification regarding multiple authors
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from dataclasses import dataclass, field
+from typing import Dict, Generic, List, Optional, TypeVar, Union
+
+Symbol = TypeVar("Symbol")
+
+
+# Disable __repr__ otherwise it could freeze e.g. Jupyter.
+@dataclass(repr=False)
+class SymbolTable(Generic[Symbol]):
+    """SymbolTable that maps symbol IDs, found on the FSA arcs to
+    actual objects. These objects can be arbitrary Python objects
+    that can serve as keys in a dictionary (i.e. they need to be
+    hashable and immutable).
+
+    The SymbolTable can only be read to/written from disk if the
+    symbols are strings.
+    """
+
+    _id2sym: Dict[int, Symbol] = field(default_factory=dict)
+    """Map an integer to a symbol.
+    """
+
+    _sym2id: Dict[Symbol, int] = field(default_factory=dict)
+    """Map a symbol to an integer.
+    """
+
+    _next_available_id: int = 1
+    """A helper internal field that helps adding new symbols
+    to the table efficiently.
+    """
+
+    eps: Symbol = "<eps>"
+    """Null symbol, always mapped to index 0.
+    """
+
+    def __post_init__(self):
+        for idx, sym in self._id2sym.items():
+            assert self._sym2id[sym] == idx
+            assert idx >= 0
+
+        for sym, idx in self._sym2id.items():
+            assert idx >= 0
+            assert self._id2sym[idx] == sym
+
+        if 0 not in self._id2sym:
+            self._id2sym[0] = self.eps
+            self._sym2id[self.eps] = 0
+        else:
+            assert self._id2sym[0] == self.eps
+            assert self._sym2id[self.eps] == 0
+
+        self._next_available_id = max(self._id2sym) + 1
+
+    @staticmethod
+    def from_str(s: str) -> "SymbolTable":
+        """Build a symbol table from a string.
+
+        The string consists of lines. Every line has two fields separated
+        by space(s), tab(s) or both. The first field is the symbol and the
+        second the integer id of the symbol.
+
+        Args:
+          s:
+            The input string with the format described above.
+        Returns:
+          An instance of :class:`SymbolTable`.
+        """
+        id2sym: Dict[int, str] = dict()
+        sym2id: Dict[str, int] = dict()
+
+        for line in s.split("\n"):
+            fields = line.split()
+            if len(fields) == 0:
+                continue  # skip empty lines
+            assert len(fields) == 2, (
+                f"Expect a line with 2 fields. Given: {len(fields)}"
+            )
+            sym, idx = fields[0], int(fields[1])
+            assert sym not in sym2id, f"Duplicated symbol {sym}"
+            assert idx not in id2sym, f"Duplicated id {idx}"
+            id2sym[idx] = sym
+            sym2id[sym] = idx
+
+        eps = id2sym.get(0, "<eps>")
+
+        return SymbolTable(_id2sym=id2sym, _sym2id=sym2id, eps=eps)
+
+    @staticmethod
+    def from_file(filename: str) -> "SymbolTable":
+        """Build a symbol table from file.
+
+        Every line in the symbol table file has two fields separated by
+        space(s), tab(s) or both. The following is an example file:
+
+        .. code-block::
+
+            <eps> 0
+            a 1
+            b 2
+            c 3
+
+        Args:
+          filename:
+            Name of the symbol table file. Its format is documented above.
+
+        Returns:
+          An instance of :class:`SymbolTable`.
+
+        """
+        with open(filename, "r", encoding="utf-8") as f:
+            return SymbolTable.from_str(f.read().strip())
+
+    def to_str(self) -> str:
+        """
+        Returns:
+          Return a string representation of this object. You can pass
+          it to the method ``from_str`` to recreate an identical object.
+        """
+        s = ""
+        for idx, symbol in sorted(self._id2sym.items()):
+            s += f"{symbol} {idx}\n"
+        return s
+
+    def to_file(self, filename: str):
+        """Serialize the SymbolTable to a file.
+
+        Every line in the symbol table file has two fields separated by
+        space(s), tab(s) or both. The following is an example file:
+
+        .. code-block::
+
+            <eps> 0
+            a 1
+            b 2
+            c 3
+
+        Args:
+          filename:
+            Name of the symbol table file. Its format is documented above.
+        """
+        with open(filename, "w", encoding="utf-8") as f:
+            for idx, symbol in sorted(self._id2sym.items()):
+                print(symbol, idx, file=f)
+
+    def add(self, symbol: Symbol, index: Optional[int] = None) -> int:
+        """Add a new symbol to the SymbolTable.
+
+        Args:
+            symbol:
+                The symbol to be added.
+            index:
+                Optional int id to which the symbol should be assigned.
+                If it is not available, a ValueError will be raised.
+
+        Returns:
+            The int id to which the symbol has been assigned.
+        """
+        # Already in the table? Return its ID.
+        if symbol in self._sym2id:
+            return self._sym2id[symbol]
+        # Specific ID not provided - use next available.
+        if index is None:
+            index = self._next_available_id
+        # Specific ID provided but not available.
+        if index in self._id2sym:
+            raise ValueError(
+                f"Cannot assign id '{index}' to '{symbol}' - "
+                f"already occupied by {self._id2sym[index]}"
+            )
+        self._sym2id[symbol] = index
+        self._id2sym[index] = symbol
+
+        # Update next available ID if needed
+        if self._next_available_id <= index:
+            self._next_available_id = index + 1
+
+        return index
+
+    def get(self, k: Union[int, Symbol]) -> Union[Symbol, int]:
+        """Get a symbol for an id or get an id for a symbol
+
+        Args:
+          k:
+            If it is an id, it tries to find the symbol corresponding
+            to the id; if it is a symbol, it tries to find the id
+            corresponding to the symbol.
+
+        Returns:
+          An id or a symbol depending on the given `k`.
+        """
+        if isinstance(k, int):
+            return self._id2sym[k]
+        else:
+            return self._sym2id[k]
+
+    def merge(self, other: "SymbolTable") -> "SymbolTable":
+        """Create a union of two SymbolTables.
+        Raises an AssertionError if the same IDs are occupied by
+        different symbols.
+
+        Args:
+            other:
+                A symbol table to merge with ``self``.
+
+        Returns:
+            A new symbol table.
+        """
+        self._check_compatible(other)
+
+        id2sym = {**self._id2sym, **other._id2sym}
+        sym2id = {**self._sym2id, **other._sym2id}
+
+        return SymbolTable(_id2sym=id2sym, _sym2id=sym2id, eps=self.eps)
+
+    def _check_compatible(self, other: "SymbolTable") -> None:
+        # Epsilon compatibility
+        assert self.eps == other.eps, (
+            f"Mismatched epsilon symbol: {self.eps} != {other.eps}"
+        )
+        # IDs compatibility
+        common_ids = set(self._id2sym).intersection(other._id2sym)
+        for idx in common_ids:
+            assert self[idx] == other[idx], (
+                f"ID conflict for id: {idx}, "
+                f'self[idx] = "{self[idx]}", '
+                f'other[idx] = "{other[idx]}"'
+            )
+        # Symbols compatibility
+        common_symbols = set(self._sym2id).intersection(other._sym2id)
+        for sym in common_symbols:
+            assert self[sym] == other[sym], (
+                f"ID conflict for id: {sym}, "
+                f'self[sym] = "{self[sym]}", '
+                f'other[sym] = "{other[sym]}"'
+            )
+
+    def __getitem__(self, item: Union[int, Symbol]) -> Union[Symbol, int]:
+        return self.get(item)
+
+    def __contains__(self, item: Union[int, Symbol]) -> bool:
+        if isinstance(item, int):
+            return item in self._id2sym
+        else:
+            return item in self._sym2id
+
+    def __len__(self) -> int:
+        return len(self._id2sym)
+
+    def __eq__(self, other: "SymbolTable") -> bool:
+        if len(self) != len(other):
+            return False
+
+        for s in self.symbols:
+            if self[s] != other[s]:
+                return False
+
+        return True
+
+    @property
+    def ids(self) -> List[int]:
+        """Returns a list of integer IDs corresponding to the symbols."""
+        ans = list(self._id2sym.keys())
+        ans.sort()
+        return ans
+
+    @property
+    def symbols(self) -> List[Symbol]:
+        """Returns a list of symbols (e.g., strings) corresponding to
+        the integer IDs.
+        """
+        ans = list(self._sym2id.keys())
+        ans.sort()
+        return ans

apps/audio_cloning/vallex/data/tokenizer.py

@@ -0,0 +1,121 @@
+#!/usr/bin/env python3
+# Copyright    2023                            (authors: Feiteng Li)
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from typing import Any
+
+import numpy as np
+import torch
+import torchaudio
+from encodec import EncodecModel
+from encodec.utils import convert_audio
+
+try:
+    pass
+except Exception:
+    pass
+
+
+def remove_encodec_weight_norm(model):
+    from encodec.modules import SConv1d
+    from encodec.modules.seanet import SConvTranspose1d, SEANetResnetBlock
+    from torch.nn.utils import remove_weight_norm
+
+    encoder = model.encoder.model
+    for key in encoder._modules:
+        if isinstance(encoder._modules[key], SEANetResnetBlock):
+            remove_weight_norm(encoder._modules[key].shortcut.conv.conv)
+            block_modules = encoder._modules[key].block._modules
+            for skey in block_modules:
+                if isinstance(block_modules[skey], SConv1d):
+                    remove_weight_norm(block_modules[skey].conv.conv)
+        elif isinstance(encoder._modules[key], SConv1d):
+            remove_weight_norm(encoder._modules[key].conv.conv)
+
+    decoder = model.decoder.model
+    for key in decoder._modules:
+        if isinstance(decoder._modules[key], SEANetResnetBlock):
+            remove_weight_norm(decoder._modules[key].shortcut.conv.conv)
+            block_modules = decoder._modules[key].block._modules
+            for skey in block_modules:
+                if isinstance(block_modules[skey], SConv1d):
+                    remove_weight_norm(block_modules[skey].conv.conv)
+        elif isinstance(decoder._modules[key], SConvTranspose1d):
+            remove_weight_norm(decoder._modules[key].convtr.convtr)
+        elif isinstance(decoder._modules[key], SConv1d):
+            remove_weight_norm(decoder._modules[key].conv.conv)
+
+
+class AudioTokenizer:
+    """EnCodec audio."""
+
+    def __init__(
+        self,
+        device: Any = None,
+    ) -> None:
+        # Instantiate a pretrained EnCodec model
+        model = EncodecModel.encodec_model_24khz()
+        model.set_target_bandwidth(6.0)
+        remove_encodec_weight_norm(model)
+
+        if not device:
+            device = torch.device("cpu")
+            if torch.cuda.is_available():
+                device = torch.device("cuda:0")
+            if torch.backends.mps.is_available():
+                device = torch.device("mps")
+
+        self._device = device
+
+        self.codec = model.to(device)
+        self.sample_rate = model.sample_rate
+        self.channels = model.channels
+
+    @property
+    def device(self):
+        return self._device
+
+    def encode(self, wav: torch.Tensor) -> torch.Tensor:
+        return self.codec.encode(wav.to(self.device))
+
+    def decode(self, frames: torch.Tensor) -> torch.Tensor:
+        return self.codec.decode(frames)
+
+
+def tokenize_audio(tokenizer: AudioTokenizer, audio):
+    # Load and pre-process the audio waveform
+    if isinstance(audio, str):
+        wav, sr = torchaudio.load(audio)
+    else:
+        wav, sr = audio
+    wav = convert_audio(wav, sr, tokenizer.sample_rate, tokenizer.channels)
+    wav = wav.unsqueeze(0)
+
+    # Extract discrete codes from EnCodec
+    with torch.no_grad():
+        encoded_frames = tokenizer.encode(wav)
+    return encoded_frames
+
+
+if __name__ == "__main__":
+    model = EncodecModel.encodec_model_24khz()
+    model.set_target_bandwidth(6.0)
+
+    samples = torch.from_numpy(np.random.random([4, 1, 1600])).type(torch.float32)
+    codes_raw = model.encode(samples)
+
+    remove_encodec_weight_norm(model)
+    codes_norm = model.encode(samples)
+
+    assert torch.allclose(codes_raw[0][0], codes_norm[0][0])

apps/audio_cloning/vallex/descriptions.py

@@ -0,0 +1,34 @@
+top_md_org = """
+# VALL-E X
+VALL-E X can synthesize high-quality personalized speech with only a 3-second enrolled recording of
+an unseen speaker as an acoustic prompt, even in another language for a monolingual speaker.<br>
+This implementation supports zero-shot, mono-lingual/cross-lingual text-to-speech functionality of three languages (English, Chinese, Japanese)<br>
+See this [demo](https://plachtaa.github.io/) page for more details.
+"""
+
+top_ja_md = """
+# VALL-E X
+
+VALL-E X は、未学習の話者でも 3 秒間の音声プロンプトだけで高品質なパーソナライズ音声を合成できます。<br>
+単一言語話者であっても別の言語による音声合成が可能です。<br>
+本実装は英語・中国語・日本語のゼロショット単言語/クロス言語テキスト読み上げをサポートしています。
+
+## Reference
+
+- [github.com/Plachtaa/VALL-E-X](https://github.com/Plachtaa/VALL-E-X/tree/master#readme)
+- [github.com/gemelo-ai/vocos](https://github.com/gemelo-ai/vocos)
+"""
+
+infer_from_audio_md_org = """
+Upload a speech of 3~10 seconds as the audio prompt and type in the text you'd like to synthesize.<br>
+The model will synthesize speech of given text with the same voice of your audio prompt.<br>
+The model also tends to preserve the emotion & acoustic environment of your given speech.<br>
+For faster inference, please use **"Make prompt"** to get a `.npz` file as the encoded audio prompt, and use it by **"Infer from prompt"**
+"""
+
+infer_from_audio_ja_md = """
+3〜10 秒程度の音声をプロンプトとしてアップロードし、合成したいテキストを入力してください。<br>
+モデルは、プロンプトと同じ声質でテキストを読み上げる音声を生成します。<br>
+元の音声に含まれる感情や音響環境も比較的保持されます。<br>
+推論を高速化したい場合は **"Make prompt"** で `.npz` ファイルを作成し、 **"Infer from prompt"** で利用してください。
+"""

apps/audio_cloning/vallex/examples.py

@@ -0,0 +1,108 @@
+_prompts_dir = "apps/audio_cloning/vallex/prompts"
+
+infer_from_audio_examples = [
+    [
+        "私のクローンに騙されないでください。",
+        "日本語",
+        "no-accent",
+        f"{_prompts_dir}/ja-okuwaki.wav",
+        None,
+        "こんにちは、私の名前はオクワキヨウスケです。",
+    ],
+    [
+        "ぼくのクローンに騙されないでくれなのだ。",
+        "日本語",
+        "no-accent",
+        f"{_prompts_dir}/ja-zundamon.wav",
+        None,
+        "はじめまして、ずんだもんなのだ",
+    ],
+    [
+        "私のクローンに騙されないでください。",
+        "日本語",
+        "no-accent",
+        f"{_prompts_dir}/ja-okuwaki-long.wav",
+        None,
+        "こんにちは、私の名前はオクワキヨウスケです。これは音声クローニング用のサンプルです。",
+    ],
+    [
+        "私の声を真似するのはそんなに面白いですか？",
+        "日本語",
+        "no-accent",
+        f"{_prompts_dir}/ja-2.ogg",
+        None,
+        "初めまして、朝武よしのです。",
+    ],
+    [
+        "This is how this machine has taken my voice.",
+        "English",
+        "no-accent",
+        f"{_prompts_dir}/en-2.wav",
+        None,
+        "Wow, look at that! That's no ordinary Teddy bear!",
+    ],
+    [
+        "我喜欢抽电子烟，尤其是锐刻五代。",
+        "中文",
+        "no-accent",
+        f"{_prompts_dir}/zh-1.wav",
+        None,
+        "今天我很荣幸，",
+    ],
+    [
+        "你可以听得出来我有多困。",
+        "中文",
+        "no-accent",
+        f"{_prompts_dir}/en-1.wav",
+        None,
+        "",
+    ],
+    [
+        "この文は、クロスリンガル合成の例です。",
+        "日本語",
+        "no-accent",
+        f"{_prompts_dir}/zh-2.wav",
+        None,
+        "",
+    ],
+    [
+        "Actually, I can't speak English, but this machine helped me do it.",
+        "English",
+        "no-accent",
+        f"{_prompts_dir}/ja-1.wav",
+        None,
+        "",
+    ],
+]
+
+make_npz_prompt_examples = [
+    [
+        "Gem-trader",
+        f"{_prompts_dir}/en-2.wav",
+        None,
+        "Wow, look at that! That's no ordinary Teddy bear!",
+    ],
+    ["Ding Zhen", f"{_prompts_dir}/zh-1.wav", None, "今天我很荣幸，"],
+    ["Yoshino", f"{_prompts_dir}/ja-2.ogg", None, "初めまして、朝武よしのです。"],
+    ["Sleepy-woman", f"{_prompts_dir}/en-1.wav", None, ""],
+    ["Yae", f"{_prompts_dir}/zh-2.wav", None, ""],
+    ["Cafe", f"{_prompts_dir}/ja-1.wav", None, ""],
+]
+
+infer_from_prompt_examples = [
+    [
+        "A prompt contains voice, prosody and emotion information of a certain speaker.",
+        "English",
+        "no-accent",
+        f"{_prompts_dir}/vctk_1",
+        None,
+    ],
+    [
+        "This prompt is made with an audio of three seconds.",
+        "English",
+        "no-accent",
+        f"{_prompts_dir}/librispeech_1",
+        None,
+    ],
+    ["This prompt is made with Chinese speech", "English", "no-accent", "seel", None],
+]

apps/audio_cloning/vallex/g2p/__init__.py

@@ -0,0 +1,84 @@
+"""from https://github.com/keithito/tacotron"""
+
+# import utils.g2p.cleaners
+from tokenizers import Tokenizer
+
+import apps.audio_cloning.vallex.g2p.cleaners as cleaners
+
+from .symbols import 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)}
+
+TOKENIZER_PATH = "./apps/audio_cloning/vallex/g2p/bpe_1024.json"
+
+
+class PhonemeBpeTokenizer:
+    def __init__(self, tokenizer_path=TOKENIZER_PATH):
+        print(f"Initializing PhonemeBpeTokenizer with tokenizer path: {tokenizer_path}")
+        self.tokenizer = Tokenizer.from_file(tokenizer_path)
+
+    def tokenize(self, text):
+        # 1. convert text to phoneme
+        phonemes, langs = _clean_text(text, ["cje_cleaners"])
+        # 2. replace blank space " " with "_"
+        phonemes = phonemes.replace(" ", "_")
+        # 3. tokenize phonemes
+        phoneme_tokens = self.tokenizer.encode(phonemes).ids
+        assert len(phoneme_tokens) == len(langs)
+        if not len(phoneme_tokens):
+            raise ValueError("Empty text is given")
+        return phoneme_tokens, langs
+
+
+def text_to_sequence(text, cleaner_names):
+    """Converts a string of text to a sequence of IDs corresponding to the symbols in the text.
+    Args:
+      text: string to convert to a sequence
+      cleaner_names: names of the cleaner functions to run the text through
+    Returns:
+      List of integers corresponding to the symbols in the text
+    """
+    sequence = []
+    symbol_to_id = {s: i for i, s in enumerate(symbols)}
+    clean_text = _clean_text(text, cleaner_names)
+    for symbol in clean_text:
+        if symbol not in symbol_to_id.keys():
+            continue
+        symbol_id = symbol_to_id[symbol]
+        sequence += [symbol_id]
+    return sequence
+
+
+def cleaned_text_to_sequence(cleaned_text):
+    """Converts a string of text to a sequence of IDs corresponding to the symbols in the text.
+    Args:
+      text: string to convert to a sequence
+    Returns:
+      List of integers corresponding to the symbols in the text
+    """
+    sequence = [
+        _symbol_to_id[symbol]
+        for symbol in cleaned_text
+        if symbol in _symbol_to_id.keys()
+    ]
+    return sequence
+
+
+def sequence_to_text(sequence):
+    """Converts a sequence of IDs back to a string"""
+    result = ""
+    for symbol_id in sequence:
+        s = _id_to_symbol[symbol_id]
+        result += s
+    return result
+
+
+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, langs = cleaner(text)
+    return text, langs

apps/audio_cloning/vallex/g2p/bpe_1024.json

@@ -0,0 +1,2049 @@
+{
+  "version": "1.0",
+  "truncation": null,
+  "padding": null,
+  "added_tokens": [
+    {
+      "id": 0,
+      "content": "[UNK]",
+      "single_word": false,
+      "lstrip": false,
+      "rstrip": false,
+      "normalized": false,
+      "special": true
+    },
+    {
+      "id": 1,
+      "content": "[CLS]",
+      "single_word": false,
+      "lstrip": false,
+      "rstrip": false,
+      "normalized": false,
+      "special": true
+    },
+    {
+      "id": 2,
+      "content": "[SEP]",
+      "single_word": false,
+      "lstrip": false,
+      "rstrip": false,
+      "normalized": false,
+      "special": true
+    },
+    {
+      "id": 3,
+      "content": "[PAD]",
+      "single_word": false,
+      "lstrip": false,
+      "rstrip": false,
+      "normalized": false,
+      "special": true
+    },
+    {
+      "id": 4,
+      "content": "[MASK]",
+      "single_word": false,
+      "lstrip": false,
+      "rstrip": false,
+      "normalized": false,
+      "special": true
+    }
+  ],
+  "normalizer": null,
+  "pre_tokenizer": {
+    "type": "Whitespace"
+  },
+  "post_processor": null,
+  "decoder": null,
+  "model": {
+    "type": "BPE",
+    "dropout": null,
+    "unk_token": "[UNK]",
+    "continuing_subword_prefix": null,
+    "end_of_word_suffix": null,
+    "fuse_unk": false,
+    "byte_fallback": false,
+    "vocab": {
+      "[UNK]": 0,
+      "[CLS]": 1,
+      "[SEP]": 2,
+      "[PAD]": 3,
+      "[MASK]": 4,
+      "!": 5,
+      "#": 6,
+      "*": 7,
+      ",": 8,
+      "-": 9,
+      ".": 10,
+      "=": 11,
+      "?": 12,
+      "N": 13,
+      "Q": 14,
+      "^": 15,
+      "_": 16,
+      "`": 17,
+      "a": 18,
+      "b": 19,
+      "d": 20,
+      "e": 21,
+      "f": 22,
+      "g": 23,
+      "h": 24,
+      "i": 25,
+      "j": 26,
+      "k": 27,
+      "l": 28,
+      "m": 29,
+      "n": 30,
+      "o": 31,
+      "p": 32,
+      "s": 33,
+      "t": 34,
+      "u": 35,
+      "v": 36,
+      "w": 37,
+      "x": 38,
+      "y": 39,
+      "z": 40,
+      "~": 41,
+      "æ": 42,
+      "ç": 43,
+      "ð": 44,
+      "ŋ": 45,
+      "ɑ": 46,
+      "ɔ": 47,
+      "ə": 48,
+      "ɛ": 49,
+      "ɥ": 50,
+      "ɪ": 51,
+      "ɫ": 52,
+      "ɯ": 53,
+      "ɸ": 54,
+      "ɹ": 55,
+      "ɾ": 56,
+      "ʃ": 57,
+      "ʊ": 58,
+      "ʑ": 59,
+      "ʒ": 60,
+      "ʰ": 61,
+      "ˈ": 62,
+      "ˌ": 63,
+      "θ": 64,
+      "…": 65,
+      "⁼": 66,
+      "↑": 67,
+      "→": 68,
+      "↓": 69,
+      "_t": 70,
+      "↓↑": 71,
+      "_ˈ": 72,
+      "ən": 73,
+      "_s": 74,
+      "aɪ": 75,
+      "əɹ": 76,
+      "eɪ": 77,
+      "oʊ": 78,
+      "_k": 79,
+      "ʃi": 80,
+      "_w": 81,
+      "_ð": 82,
+      "ts": 83,
+      "tʃ": 84,
+      "_ts": 85,
+      "_h": 86,
+      "_ə": 87,
+      "_m": 88,
+      "an": 89,
+      "_n": 90,
+      "_ðə": 91,
+      "ɛn": 92,
+      "ɑʊ": 93,
+      "ɑŋ": 94,
+      "`⁼": 95,
+      "_p": 96,
+      "_i": 97,
+      "_ɪ": 98,
+      "_tʃ": 99,
+      "_l": 100,
+      "jɛn": 101,
+      "_d": 102,
+      "_f": 103,
+      "_j": 104,
+      "wo": 105,
+      "_b": 106,
+      "ta": 107,
+      "`↓": 108,
+      "te": 109,
+      "ənd": 110,
+      "_ʃi": 111,
+      "wa": 112,
+      "ka": 113,
+      "ɪŋ": 114,
+      "in": 115,
+      "st": 116,
+      "li": 117,
+      "ʊŋ": 118,
+      "_tɪ": 119,
+      "to": 120,
+      "weɪ": 121,
+      "_ənd": 122,
+      "ʰi": 123,
+      "_əv": 124,
+      "əŋ": 125,
+      "no": 126,
+      "_x": 127,
+      "ɾɯ": 128,
+      "na": 129,
+      "_a": 130,
+      "_ɹ": 131,
+      "ɪn": 132,
+      "ga": 133,
+      "de": 134,
+      "joʊ": 135,
+      "æn": 136,
+      "kɯ": 137,
+      "ɾe": 138,
+      "ma": 139,
+      "_ðə_ˈ": 140,
+      "ɾa": 141,
+      "ɛɹ": 142,
+      "mo": 143,
+      "ɔɹ": 144,
+      "əɫ": 145,
+      "_g": 146,
+      "da": 147,
+      "*↑": 148,
+      "ɪˈ": 149,
+      "_o": 150,
+      "_ʃ": 151,
+      "iŋ": 152,
+      "ja": 153,
+      "əm": 154,
+      "_ˌ": 155,
+      "aʊ": 156,
+      "_əˈ": 157,
+      "`↑": 158,
+      "ət": 159,
+      "_aɪ": 160,
+      "oo": 161,
+      "sɯ": 162,
+      "↓.": 163,
+      "_ɪn": 164,
+      "_hi": 165,
+      "_wɪ": 166,
+      "ɪz": 167,
+      "_na": 168,
+      "wan": 169,
+      "_ko": 170,
+      "_wo": 171,
+      "ɪd": 172,
+      "ɾi": 173,
+      "_ju": 174,
+      "mə": 175,
+      "_lə": 176,
+      "_hæ": 177,
+      "_ðət": 178,
+      "ɑɹ": 179,
+      "tʰ": 180,
+      "ki": 181,
+      "……": 182,
+      "ɑz": 183,
+      "_ɔ": 184,
+      "_mi": 185,
+      "_wɑz": 186,
+      "_ˈs": 187,
+      "↓,": 188,
+      "_tʰ": 189,
+      "əˈ": 190,
+      "dʑ": 191,
+      "ɪt": 192,
+      "_kʰ": 193,
+      "iɛ": 194,
+      "_ma": 195,
+      "ɪs": 196,
+      "tsɯ": 197,
+      "_ni": 198,
+      "_ɪt": 199,
+      "ke": 200,
+      "iɑʊ": 201,
+      "_ka": 202,
+      "_əɹ": 203,
+      "nd": 204,
+      "_ˈp": 205,
+      "ko": 206,
+      "jo": 207,
+      "ɹi": 208,
+      "mən": 209,
+      "ʊd": 210,
+      "_ˈm": 211,
+      "_fəɹ": 212,
+      "tʃʰi": 213,
+      "sa": 214,
+      "ʰɥ": 215,
+      "kʰ": 216,
+      "ˈs": 217,
+      "ɑt": 218,
+      "ɛd": 219,
+      "se": 220,
+      "tʃi": 221,
+      "ɛɫ": 222,
+      "_ˈk": 223,
+      "_joʊ": 224,
+      "təɹ": 225,
+      "ɛz": 226,
+      "--": 227,
+      "vəɹ": 228,
+      "`→": 229,
+      "ʃən": 230,
+      "_ɪz": 231,
+      "_meɪ": 232,
+      "_æ": 233,
+      "dʒ": 234,
+      "_ki": 235,
+      "_hɪz": 236,
+      "_bi": 237,
+      "uɑŋ": 238,
+      "_ˈf": 239,
+      "↓↑.": 240,
+      "_wɪθ": 241,
+      "ju": 242,
+      "iɑŋ": 243,
+      "→.": 244,
+      "_so": 245,
+      "_həɹ": 246,
+      "↑.": 247,
+      "ni": 248,
+      "_mo": 249,
+      "_maɪ": 250,
+      "laɪ": 251,
+      "ɥɛ": 252,
+      "_ta": 253,
+      "ənt": 254,
+      "_tʃʰi": 255,
+      "_sɯ": 256,
+      "_θ": 257,
+      "_ɛz": 258,
+      "wən": 259,
+      "me": 260,
+      "mi": 261,
+      "_hæd": 262,
+      "_ha": 263,
+      "əs": 264,
+      "_ˈl": 265,
+      "_st": 266,
+      "ðəɹ": 267,
+      "oʊn": 268,
+      "_wa": 269,
+      "ʰəŋ": 270,
+      "_nɑt": 271,
+      "*.": 272,
+      "kt": 273,
+      "_ˈh": 274,
+      "do": 275,
+      "ɥæn": 276,
+      "ne": 277,
+      "_to": 278,
+      "_wən": 279,
+      "_no": 280,
+      "_laɪ": 281,
+      "_wəɹ": 282,
+      "↑,": 283,
+      "→,": 284,
+      "ɛs": 285,
+      "↓↑,": 286,
+      "_ɔn": 287,
+      "ʰu": 288,
+      "so": 289,
+      "_ˈb": 290,
+      "ɫd": 291,
+      "ɪk": 292,
+      "ɪst": 293,
+      "_fɹ": 294,
+      "_ðɛɹ": 295,
+      "_weɪ": 296,
+      "kaɾa": 297,
+      "_ˈd": 298,
+      "_hæv": 299,
+      "tsʰ": 300,
+      "waɪ": 301,
+      "ɾo": 302,
+      "ɛm": 303,
+      "_æt": 304,
+      "ʊɹ": 305,
+      "_ˈw": 306,
+      "ba": 307,
+      "_noʊ": 308,
+      "ʰjɛn": 309,
+      "ɹeɪ": 310,
+      "_jo": 311,
+      "ɸɯ": 312,
+      "_sa": 313,
+      "_ɹɪˈ": 314,
+      "_ˈn": 315,
+      "ai": 316,
+      "_bət": 317,
+      "ɪɹ": 318,
+      "tʃʰɥ": 319,
+      "_dʑ": 320,
+      "əˌ": 321,
+      "_ðɪs": 322,
+      "..": 323,
+      "xwa": 324,
+      "_ɪm": 325,
+      "_dɪˈ": 326,
+      "_kən": 327,
+      "dʑi": 328,
+      "*,": 329,
+      "ɑn": 330,
+      "_ʃiɑŋ": 331,
+      "_kɯ": 332,
+      "ʃin": 333,
+      "_soʊ": 334,
+      "bi": 335,
+      "tʰjɛn": 336,
+      "te_i": 337,
+      "_tsʰ": 338,
+      "_ɯ": 339,
+      "aɪt": 340,
+      "ʰiŋ": 341,
+      "ðə": 342,
+      "_ɔɫ": 343,
+      "_ˈɹ": 344,
+      "nai": 345,
+      "əɹd": 346,
+      "_ˈt": 347,
+      "_ən": 348,
+      "_tʃʰɥ": 349,
+      "_iɛ": 350,
+      "leɪ": 351,
+      "ɛɹi": 352,
+      "ˈt": 353,
+      "ha": 354,
+      "ʃiŋ": 355,
+      "ɛvəɹ": 356,
+      "zɯ": 357,
+      "_wi": 358,
+      "_ja": 359,
+      "ɛk": 360,
+      "ʰɑŋ": 361,
+      "_tsɯ": 362,
+      "_əv_ðə": 363,
+      "taʃi": 364,
+      "_sɛd": 365,
+      "_xə": 366,
+      "_li": 367,
+      "_si": 368,
+      "desɯ": 369,
+      "_ˌɪn": 370,
+      "ʃjɛn": 371,
+      "_baɪ": 372,
+      "on": 373,
+      "_xɑʊ": 374,
+      "_ðeɪ": 375,
+      "_xaɪ": 376,
+      "`↓↑": 377,
+      "xweɪ": 378,
+      "hi": 379,
+      "_se": 380,
+      "ə_s": 381,
+      "_fɹəm": 382,
+      "ʊt": 383,
+      "di": 384,
+      "aʊt": 385,
+      "əb": 386,
+      "sɹ": 387,
+      "əz": 388,
+      "_xweɪ": 389,
+      "_kʰə": 390,
+      "ɹu": 391,
+      "_u": 392,
+      "_de": 393,
+      "aɪd": 394,
+      "ɪv": 395,
+      "bɯ": 396,
+      "_ho": 397,
+      "əɹz": 398,
+      "joo": 399,
+      "_bɪˈ": 400,
+      "_tʰa": 401,
+      "ɛt": 402,
+      "en": 403,
+      "ɛni": 404,
+      "əst": 405,
+      "æk": 406,
+      "ə_ts": 407,
+      "_ˈɪn": 408,
+      "ti": 409,
+      "ɥn": 410,
+      "_dʒ": 411,
+      "xɑʊ": 412,
+      "_ˈv": 413,
+      "ʃiɑŋ": 414,
+      "pʰ": 415,
+      "_wɪtʃ": 416,
+      "eɪm": 417,
+      "oʊz": 418,
+      "əðəɹ": 419,
+      "fɑŋ": 420,
+      "_ˈg": 421,
+      "_do": 422,
+      "_ʃiɑʊ": 423,
+      "_ˈæ": 424,
+      "_jʊɹ": 425,
+      "_ðɛm": 426,
+      "ɪm": 427,
+      "ɛst": 428,
+      "ænd": 429,
+      "_du": 430,
+      "ɯɯ": 431,
+      "kan": 432,
+      "_da": 433,
+      "ino": 434,
+      "_e": 435,
+      "_wʊd": 436,
+      "ɛnd": 437,
+      "meɪ": 438,
+      "θɪŋ": 439,
+      "_ʃjɛn": 440,
+      "iz": 441,
+      "aɪm": 442,
+      "_hu": 443,
+      "_əˈb": 444,
+      "əns": 445,
+      "_wɪɫ": 446,
+      "tʰi": 447,
+      "go": 448,
+      "ɛnt": 449,
+      "fu": 450,
+      "æp": 451,
+      "xoʊ": 452,
+      "eɪk": 453,
+      "ʊk": 454,
+      "əɹˈ": 455,
+      "_θɪŋ": 456,
+      "əl": 457,
+      "pɹ": 458,
+      "ətʃ": 459,
+      "nt": 460,
+      "_ɸɯ": 461,
+      "lu": 462,
+      "_ˈɔ": 463,
+      "_iɑʊ": 464,
+      "lə": 465,
+      "tu": 466,
+      "_dʑi": 467,
+      "eɪt": 468,
+      "_ʃin": 469,
+      "nna": 470,
+      "_ˈpɹ": 471,
+      "fən": 472,
+      "_əp": 473,
+      "njɛn": 474,
+      "_aʊt": 475,
+      "fɔɹ": 476,
+      "_tu": 477,
+      "eɪʃən": 478,
+      "ɪɫ": 479,
+      "_wət": 480,
+      "_ɪf": 481,
+      "_ɥ": 482,
+      "_fa": 483,
+      "ˈw": 484,
+      "tʃʰjɛn": 485,
+      "_wɪn": 486,
+      "oʊɫd": 487,
+      "_əˈp": 488,
+      "aʊnd": 489,
+      "san": 490,
+      "he": 491,
+      "_bɪn": 492,
+      "fa": 493,
+      "ɪf": 494,
+      "ɔŋ": 495,
+      "ge": 496,
+      "_ɪn_ðə": 497,
+      "miŋ": 498,
+      "_pɹ": 499,
+      "ina": 500,
+      "ano": 501,
+      "əbəɫ": 502,
+      "kˈs": 503,
+      "_ˈɛni": 504,
+      "nəŋ": 505,
+      "əd": 506,
+      "_əv_ðə_ˈ": 507,
+      "_waɪ": 508,
+      "_taɪm": 509,
+      "ˈsɛɫ": 510,
+      "ʃiɛ": 511,
+      "_kəm": 512,
+      "æst": 513,
+      "_goʊ": 514,
+      "mɯ": 515,
+      "ˈp": 516,
+      "_ˈst": 517,
+      "ə_t": 518,
+      "pt": 519,
+      "_pʰ": 520,
+      "ʰɹ": 521,
+      "ʃja": 522,
+      "iwa": 523,
+      "ɪl": 524,
+      "bət": 525,
+      "_fɑŋ": 526,
+      "ho": 527,
+      "iv": 528,
+      "loʊ": 529,
+      "be": 530,
+      "_laɪk": 531,
+      "ɪʃ": 532,
+      "_fu": 533,
+      "ze": 534,
+      "ə_tʃ": 535,
+      "ɑɹt": 536,
+      "ɔɹd": 537,
+      "tʃʰiŋ": 538,
+      "mp": 539,
+      "_ðə_s": 540,
+      "_əˈbaʊt": 541,
+      "_ˈoʊ": 542,
+      "kʰə": 543,
+      "d_tɪ": 544,
+      "ŋga": 545,
+      "əli": 546,
+      "_kʰan": 547,
+      "çi": 548,
+      "_ˈju": 549,
+      "_kʊd": 550,
+      "ɔɫ": 551,
+      "ɔt": 552,
+      "_ɪts": 553,
+      "_san": 554,
+      "tʃa": 555,
+      "i_na": 556,
+      "xə": 557,
+      "ɛkt": 558,
+      "_mɔɹ": 559,
+      "te_kɯ": 560,
+      "ɪdʒ": 561,
+      "jʊŋ": 562,
+      "_wan": 563,
+      "æt": 564,
+      "kat": 565,
+      "ˈsɛɫf": 566,
+      "_ke": 567,
+      "aɪnd": 568,
+      "it": 569,
+      "_ɑɹ": 570,
+      "sp": 571,
+      "oʊnt": 572,
+      "_tʃi": 573,
+      "tsʰɹ": 574,
+      "_xən": 575,
+      "_əˈg": 576,
+      "ə_k": 577,
+      "to_i": 578,
+      "_tʰi": 579,
+      "_iŋ": 580,
+      "aʊn": 581,
+      "gɯ": 582,
+      "_ɪkˈs": 583,
+      "ɛv": 584,
+      "gi": 585,
+      "ks": 586,
+      "_səm": 587,
+      "ana": 588,
+      "ɪtəɫ": 589,
+      "nan": 590,
+      "_ˈɪntu": 591,
+      "_hiɹ": 592,
+      "_te": 593,
+      "_naʊ": 594,
+      "ʃiɑʊ": 595,
+      "ʃo": 596,
+      "ɹe": 597,
+      "xaɪ": 598,
+      "_tʃʰiŋ": 599,
+      "_sɹ": 600,
+      "_haʊ": 601,
+      "?.": 602,
+      "_feɪ": 603,
+      "liŋ": 604,
+      "_ʃja": 605,
+      "_ˈdʒ": 606,
+      "_seɪ": 607,
+      "ˈn": 608,
+      "soʊ": 609,
+      "tʰʊŋ": 610,
+      "_ljoʊ": 611,
+      "maɪ": 612,
+      "_bɹ": 613,
+      "ɹeɪt": 614,
+      "_nəŋ": 615,
+      "ʰə": 616,
+      "æns": 617,
+      "_ˈɔl": 618,
+      "tatʃi": 619,
+      "nto": 620,
+      "_ˌɪnˈ": 621,
+      "le": 622,
+      "nde": 623,
+      "_ˈvɛɹi": 624,
+      "mənt": 625,
+      "ɾima": 626,
+      "_ðɛn": 627,
+      "_həz": 628,
+      "_ɹi": 629,
+      "ftəɹ": 630,
+      "_sp": 631,
+      "ɾewa": 632,
+      "ga_a": 633,
+      "z_əv": 634,
+      "_miŋ": 635,
+      "_tɪ_ðə": 636,
+      "ɹaɪ": 637,
+      "ɛl": 638,
+      "ɹæ": 639,
+      "_hoʊ": 640,
+      "xu": 641,
+      "oʊnli": 642,
+      "ŋk": 643,
+      "i_i": 644,
+      "_dɪd": 645,
+      "_dʒɪst": 646,
+      "ing": 647,
+      "kai": 648,
+      "_mæn": 649,
+      "_in": 650,
+      "zo": 651,
+      "əf": 652,
+      "dake": 653,
+      "_ˈsəm": 654,
+      "ɾɯ_no": 655,
+      "_go": 656,
+      "tʃəɹ": 657,
+      "ite": 658,
+      "`↓.": 659,
+      "_kʰaɪ": 660,
+      "sk": 661,
+      "ɔɹs": 662,
+      "_tʰiŋ": 663,
+      "_nə": 664,
+      "pəɫ": 665,
+      "_tɪ_bi": 666,
+      "ˈfɔɹ": 667,
+      "mu": 668,
+      "su": 669,
+      "aa": 670,
+      "ɪstəɹ": 671,
+      "ʰan": 672,
+      "pəɹ": 673,
+      "ə_p": 674,
+      "liɑŋ": 675,
+      "_v": 676,
+      "oʊst": 677,
+      "_əˈgɛn": 678,
+      "ənz": 679,
+      "No": 680,
+      "ɔɹt": 681,
+      "_səˈ": 682,
+      "_mɯ": 683,
+      "tʃʰ": 684,
+      "_ˈlɪtəɫ": 685,
+      "_xwo": 686,
+      "_ˌbi": 687,
+      "_ˈoʊvəɹ": 688,
+      "_çi": 689,
+      "_deɪ": 690,
+      "aɪn": 691,
+      "_ʃiŋ": 692,
+      "i_ʃi": 693,
+      "_tsʰaɪ": 694,
+      "ʃoo": 695,
+      "ɾoo": 696,
+      "bəɹ": 697,
+      "ʰa": 698,
+      "ˈɛs": 699,
+      "_ɪn_ðə_ˈ": 700,
+      "Nwa": 701,
+      "_ðən": 702,
+      "saɪ": 703,
+      "_ˈjuˈɛs": 704,
+      "nda": 705,
+      "_pleɪ": 706,
+      "ɪŋ_tɪ": 707,
+      "ɪti": 708,
+      "_me": 709,
+      "_ʃʊd": 710,
+      "_nu": 711,
+      "_ðə_k": 712,
+      "za": 713,
+      "_ˈɛvəɹ": 714,
+      "əɹn": 715,
+      "æd": 716,
+      "ˈm": 717,
+      "_doʊnt": 718,
+      "_məst": 719,
+      "jɯɯ": 720,
+      "ɑɹd": 721,
+      "_jɛn": 722,
+      "ʃɥ": 723,
+      "_ˈoʊnli": 724,
+      "_ʃo": 725,
+      "_liŋ": 726,
+      "ss": 727,
+      "ɑl": 728,
+      "dea": 729,
+      "ɾeta": 730,
+      "mjɛn": 731,
+      "_gʊd": 732,
+      "_wɔ": 733,
+      "imo": 734,
+      "no_ko": 735,
+      "_ɥæn": 736,
+      "ndʒ": 737,
+      "ɪʃən": 738,
+      "o_ʃi": 739,
+      "_θɪŋk": 740,
+      "_nan": 741,
+      "to_o": 742,
+      "_tʰʊŋ": 743,
+      "ljoʊ": 744,
+      "tai": 745,
+      "mə_s": 746,
+      "_jɯ": 747,
+      "_uɑŋ": 748,
+      "_ˌbiˈfɔɹ": 749,
+      "æs": 750,
+      "_tʃʰjɛn": 751,
+      "ik": 752,
+      "_bæk": 753,
+      "_ˈiv": 754,
+      "eɪn": 755,
+      "un": 756,
+      "la": 757,
+      "ˈk": 758,
+      "_daʊn": 759,
+      "anai": 760,
+      "_lɛ": 761,
+      "əɹt": 762,
+      "ðɛɹ": 763,
+      "_ˈæftəɹ": 764,
+      "dat": 765,
+      "fan": 766,
+      "bəɫ": 767,
+      "temo": 768,
+      "tʰa": 769,
+      "ɾɯ_ko": 770,
+      "ˈv": 771,
+      "feɪ": 772,
+      "_mətʃ": 773,
+      "xwo": 774,
+      "ɹoʊ": 775,
+      "_ba": 776,
+      "_ˈnɛvəɹ": 777,
+      "_meɪd": 778,
+      "_jʊŋ": 779,
+      "_əˈpɑn": 780,
+      "!?": 781,
+      "_ˈʃ": 782,
+      "_ðə_ˈk": 783,
+      "ft": 784,
+      "_bo": 785,
+      "_ɪn_ə": 786,
+      "tʃʰɥæn": 787,
+      "ˈz": 788,
+      "`↓,": 789,
+      "_bɪˈk": 790,
+      "ɪg": 791,
+      "kin": 792,
+      "_kl": 793,
+      "ɾɯ_n": 794,
+      "_lɑʊ": 795,
+      "----": 796,
+      "ika": 797,
+      "_ɹaɪt": 798,
+      "zd": 799,
+      "z_ənd": 800,
+      "_kjo": 801,
+      "xwan": 802,
+      "too": 803,
+      "_gɪt": 804,
+      "_liɑŋ": 805,
+      "ta_n": 806,
+      "_keɪm": 807,
+      "_ˈəðəɹ": 808,
+      "_wɛɫ": 809,
+      "teki": 810,
+      "see": 811,
+      "jɯ": 812,
+      "i_o": 813,
+      "to_ʃi": 814,
+      "fəɫ": 815,
+      "bo": 816,
+      "ˌt": 817,
+      "ɪp": 818,
+      "ane": 819,
+      "_tʰjɛn": 820,
+      "_tʃo": 821,
+      "ɾjo": 822,
+      "ɪns": 823,
+      "_he": 824,
+      "ŋka": 825,
+      "ʃɥɛ": 826,
+      "dʑa": 827,
+      "vd": 828,
+      "ʰwan": 829,
+      "_gɹeɪt": 830,
+      "_əv_ə": 831,
+      "əndəɹ": 832,
+      "kedo": 833,
+      "_ðə_b": 834,
+      "ək": 835,
+      "_teɪk": 836,
+      "kʰan": 837,
+      "_ˈɔlˌ": 838,
+      "swo": 839,
+      "_ɪt_wɑz": 840,
+      "_ʃɥ": 841,
+      "_sim": 842,
+      "_ˈfɑ": 843,
+      "min": 844,
+      "i_a": 845,
+      "soo": 846,
+      "ɛns": 847,
+      "_sətʃ": 848,
+      "tʰaɪ": 849,
+      "_ga": 850,
+      "i_ka": 851,
+      "koo": 852,
+      "_fəɹst": 853,
+      "_ˈtʃ": 854,
+      "nno": 855,
+      "ə_ɹ": 856,
+      "taɾa": 857,
+      "tʃʰjoʊ": 858,
+      "_æm": 859,
+      "_mu": 860,
+      "_meɪk": 861,
+      "↓…": 862,
+      "ɪˈθ": 863,
+      "ɑb": 864,
+      "ɹa": 865,
+      "_wɛɹ": 866,
+      "_ðə_ˈs": 867,
+      "_əˈl": 868,
+      "_oʊɫd": 869,
+      "æl": 870,
+      "_ˈpi": 871,
+      "_lɔŋ": 872,
+      "dʑo": 873,
+      "_tʰaɪ": 874,
+      "ɔɹn": 875,
+      "əɫz": 876,
+      "_təˈ": 877,
+      "_əˈweɪ": 878,
+      "pa": 879,
+      "_ðiz": 880,
+      "_ˈsp": 881,
+      "nn": 882,
+      "mae": 883,
+      "towa": 884,
+      "ta_no": 885,
+      "_an": 886,
+      "kʰaɪ": 887,
+      "ɾaɾe": 888,
+      "eɪs": 889,
+      "ɑd": 890,
+      "_wɪˈθ": 891,
+      "_ˈivɪn": 892,
+      "_lu": 893,
+      "ɔɪ": 894,
+      "lɪŋ": 895,
+      "əti": 896,
+      "_ðə_f": 897,
+      "oʃi": 898,
+      "_la": 899,
+      "si": 900,
+      "tɪd": 901,
+      "haʊ": 902,
+      "pʰin": 903,
+      "ˈst": 904,
+      "_ˈpəɹ": 905,
+      "eɹ": 906,
+      "*!": 907,
+      "_ˈmɪstəɹ": 908,
+      "ʃa": 909,
+      "_ˌɪm": 910,
+      "ˌθɪŋ": 911,
+      "_neɪ": 912,
+      "_nɥ": 913,
+      "ɑk": 914,
+      "_ɹu": 915,
+      "_ʃɯ": 916,
+      "_ðə_ˈm": 917,
+      "demo": 918,
+      "_dɹ": 919,
+      "dʑoo": 920,
+      "_stɪɫ": 921,
+      "_pʰiŋ": 922,
+      "ə_i": 923,
+      "_ɪkˈsp": 924,
+      "_wɛnt": 925,
+      "ɪɹi": 926,
+      "əˈm": 927,
+      "o_ka": 928,
+      "_əˈk": 929,
+      "ɔk": 930,
+      "_ɥɛ": 931,
+      "_lʊk": 932,
+      "ˈd": 933,
+      "kaʃi": 934,
+      "_wɪθ_ə": 935,
+      "ljɛn": 936,
+      "ɔn": 937,
+      "_ljɛn": 938,
+      "_hɛɫ": 939,
+      "uɹ": 940,
+      "_tʰoʊ": 941,
+      "_tʃʰɥæn": 942,
+      "_sk": 943,
+      "tsʰaɪ": 944,
+      "ɛtəɹ": 945,
+      "_min": 946,
+      "noʊ": 947,
+      "ʃɯ": 948,
+      "_θɹu": 949,
+      "_θɔt": 950,
+      "dajo": 951,
+      "wi": 952,
+      "i_ko": 953,
+      "_tɹ": 954,
+      "_fan": 955,
+      "ɹɛ": 956,
+      "saN": 957,
+      "_hi_wɑz": 958,
+      "_ɾe": 959,
+      "_əm": 960,
+      "te_ki": 961,
+      "_xoʊ": 962,
+      "ˈl": 963,
+      "ˈg": 964,
+      "ga_i": 965,
+      "_ɔn_ðə": 966,
+      "_xwa": 967,
+      "vɪŋ": 968,
+      "man": 969,
+      "fəɹ": 970,
+      "_oʊn": 971,
+      "ˈɹ": 972,
+      "_kɹ": 973,
+      "te_o": 974,
+      "ɪli": 975,
+      "_ʃɥɛ": 976,
+      "_fəŋ": 977,
+      "æɫ": 978,
+      "ɑp": 979,
+      "_ˈɛv": 980,
+      "eɪndʒ": 981,
+      "iɫ": 982,
+      "wət": 983,
+      "ɛðəɹ": 984,
+      "_fən": 985,
+      "ɾee": 986,
+      "_hi_hæd": 987,
+      "_maɪt": 988,
+      "_ge": 989,
+      "ækt": 990,
+      "ɪts": 991,
+      "_hɪm": 992,
+      "_ze": 993,
+      "ii": 994,
+      "_N": 995,
+      "_əv_hɪz": 996,
+      "_gɹ": 997,
+      "ænt": 998,
+      "ɪˌ": 999,
+      "_hɪmˈsɛɫf": 1000,
+      "wa_na": 1001,
+      "aɪəɹ": 1002,
+      "dʑanai": 1003,
+      "kana": 1004,
+      "aɪz": 1005,
+      "_ɪt_ɪz": 1006,
+      "mase": 1007,
+      "wɪn": 1008,
+      "əθɪŋ": 1009,
+      "_pɹəˈ": 1010,
+      "kɯn": 1011,
+      "ˈju": 1012,
+      "_fɔɹ": 1013,
+      "pʰi": 1014,
+      "pʰiŋ": 1015,
+      "o_i": 1016,
+      "vz": 1017,
+      "ɔɪn": 1018,
+      "tʰiŋ": 1019,
+      "_ne": 1020,
+      "gəɹ": 1021,
+      "æts": 1022,
+      "_ˈɹi": 1023
+    },
+    "merges": [
+      "_ t",
+      "↓ ↑",
+      "_ ˈ",
+      "ə n",
+      "_ s",
+      "a ɪ",
+      "ə ɹ",
+      "e ɪ",
+      "o ʊ",
+      "_ k",
+      "ʃ i",
+      "_ w",
+      "_ ð",
+      "t s",
+      "t ʃ",
+      "_t s",
+      "_ h",
+      "_ ə",
+      "_ m",
+      "a n",
+      "_ n",
+      "_ð ə",
+      "ɛ n",
+      "ɑ ʊ",
+      "ɑ ŋ",
+      "` ⁼",
+      "_ p",
+      "_ i",
+      "_ ɪ",
+      "_t ʃ",
+      "_ l",
+      "j ɛn",
+      "_ d",
+      "_ f",
+      "_ j",
+      "w o",
+      "_ b",
+      "t a",
+      "` ↓",
+      "t e",
+      "ən d",
+      "_ ʃi",
+      "w a",
+      "k a",
+      "ɪ ŋ",
+      "i n",
+      "s t",
+      "l i",
+      "ʊ ŋ",
+      "_t ɪ",
+      "t o",
+      "w eɪ",
+      "_ ənd",
+      "ʰ i",
+      "_ə v",
+      "ə ŋ",
+      "n o",
+      "_ x",
+      "ɾ ɯ",
+      "n a",
+      "_ a",
+      "_ ɹ",
+      "ɪ n",
+      "g a",
+      "d e",
+      "j oʊ",
+      "æ n",
+      "k ɯ",
+      "ɾ e",
+      "m a",
+      "_ðə _ˈ",
+      "ɾ a",
+      "ɛ ɹ",
+      "m o",
+      "ɔ ɹ",
+      "ə ɫ",
+      "_ g",
+      "d a",
+      "* ↑",
+      "ɪ ˈ",
+      "_ o",
+      "_ ʃ",
+      "i ŋ",
+      "j a",
+      "ə m",
+      "_ ˌ",
+      "a ʊ",
+      "_ə ˈ",
+      "` ↑",
+      "ə t",
+      "_ aɪ",
+      "o o",
+      "s ɯ",
+      "↓ .",
+      "_ɪ n",
+      "_h i",
+      "_w ɪ",
+      "ɪ z",
+      "_n a",
+      "w an",
+      "_k o",
+      "_w o",
+      "ɪ d",
+      "ɾ i",
+      "_j u",
+      "m ə",
+      "_l ə",
+      "_h æ",
+      "_ðə t",
+      "ɑ ɹ",
+      "t ʰ",
+      "k i",
+      "… …",
+      "ɑ z",
+      "_ ɔ",
+      "_m i",
+      "_w ɑz",
+      "_ˈ s",
+      "↓ ,",
+      "_t ʰ",
+      "ə ˈ",
+      "d ʑ",
+      "ɪ t",
+      "_k ʰ",
+      "i ɛ",
+      "_m a",
+      "ɪ s",
+      "ts ɯ",
+      "_n i",
+      "_ɪ t",
+      "k e",
+      "i ɑʊ",
+      "_k a",
+      "_ əɹ",
+      "n d",
+      "_ˈ p",
+      "k o",
+      "j o",
+      "ɹ i",
+      "m ən",
+      "ʊ d",
+      "_ˈ m",
+      "_f əɹ",
+      "tʃ ʰi",
+      "s a",
+      "ʰ ɥ",
+      "k ʰ",
+      "ˈ s",
+      "ɑ t",
+      "ɛ d",
+      "s e",
+      "t ʃi",
+      "ɛ ɫ",
+      "_ˈ k",
+      "_j oʊ",
+      "t əɹ",
+      "ɛ z",
+      "- -",
+      "v əɹ",
+      "` →",
+      "ʃ ən",
+      "_ɪ z",
+      "_m eɪ",
+      "_ æ",
+      "d ʒ",
+      "_k i",
+      "_h ɪz",
+      "_b i",
+      "u ɑŋ",
+      "_ˈ f",
+      "↓↑ .",
+      "_wɪ θ",
+      "j u",
+      "i ɑŋ",
+      "→ .",
+      "_s o",
+      "_h əɹ",
+      "↑ .",
+      "n i",
+      "_m o",
+      "_m aɪ",
+      "l aɪ",
+      "ɥ ɛ",
+      "_t a",
+      "ən t",
+      "_tʃ ʰi",
+      "_s ɯ",
+      "_ θ",
+      "_ ɛz",
+      "w ən",
+      "m e",
+      "m i",
+      "_hæ d",
+      "_h a",
+      "ə s",
+      "_ˈ l",
+      "_s t",
+      "ð əɹ",
+      "oʊ n",
+      "_w a",
+      "ʰ əŋ",
+      "_n ɑt",
+      "* .",
+      "k t",
+      "_ˈ h",
+      "d o",
+      "ɥ æn",
+      "n e",
+      "_t o",
+      "_w ən",
+      "_n o",
+      "_l aɪ",
+      "_w əɹ",
+      "↑ ,",
+      "→ ,",
+      "ɛ s",
+      "↓↑ ,",
+      "_ɔ n",
+      "ʰ u",
+      "s o",
+      "_ˈ b",
+      "ɫ d",
+      "ɪ k",
+      "ɪ st",
+      "_f ɹ",
+      "_ð ɛɹ",
+      "_w eɪ",
+      "ka ɾa",
+      "_ˈ d",
+      "_hæ v",
+      "ts ʰ",
+      "w aɪ",
+      "ɾ o",
+      "ɛ m",
+      "_æ t",
+      "ʊ ɹ",
+      "_ˈ w",
+      "b a",
+      "_n oʊ",
+      "ʰ jɛn",
+      "ɹ eɪ",
+      "_j o",
+      "ɸ ɯ",
+      "_s a",
+      "_ɹ ɪˈ",
+      "_ˈ n",
+      "a i",
+      "_b ət",
+      "ɪ ɹ",
+      "tʃ ʰɥ",
+      "_d ʑ",
+      "ə ˌ",
+      "_ð ɪs",
+      ". .",
+      "x wa",
+      "_ɪ m",
+      "_d ɪˈ",
+      "_k ən",
+      "dʑ i",
+      "* ,",
+      "ɑ n",
+      "_ʃi ɑŋ",
+      "_k ɯ",
+      "ʃi n",
+      "_s oʊ",
+      "b i",
+      "tʰ jɛn",
+      "te _i",
+      "_ts ʰ",
+      "_ ɯ",
+      "aɪ t",
+      "ʰi ŋ",
+      "ð ə",
+      "_ɔ ɫ",
+      "_ˈ ɹ",
+      "na i",
+      "əɹ d",
+      "_ˈ t",
+      "_ ən",
+      "_tʃ ʰɥ",
+      "_i ɛ",
+      "l eɪ",
+      "ɛɹ i",
+      "ˈ t",
+      "h a",
+      "ʃi ŋ",
+      "ɛ vəɹ",
+      "z ɯ",
+      "_w i",
+      "_j a",
+      "ɛ k",
+      "ʰ ɑŋ",
+      "_ts ɯ",
+      "_əv _ðə",
+      "ta ʃi",
+      "_s ɛd",
+      "_x ə",
+      "_l i",
+      "_s i",
+      "de sɯ",
+      "_ˌ ɪn",
+      "ʃ jɛn",
+      "_b aɪ",
+      "o n",
+      "_x ɑʊ",
+      "_ð eɪ",
+      "_x aɪ",
+      "` ↓↑",
+      "x weɪ",
+      "h i",
+      "_s e",
+      "ə _s",
+      "_fɹ əm",
+      "ʊ t",
+      "d i",
+      "aʊ t",
+      "ə b",
+      "s ɹ",
+      "ə z",
+      "_x weɪ",
+      "_kʰ ə",
+      "ɹ u",
+      "_ u",
+      "_d e",
+      "aɪ d",
+      "ɪ v",
+      "b ɯ",
+      "_h o",
+      "əɹ z",
+      "j oo",
+      "_b ɪˈ",
+      "_tʰ a",
+      "ɛ t",
+      "e n",
+      "ɛn i",
+      "ə st",
+      "æ k",
+      "ə _ts",
+      "_ˈ ɪn",
+      "t i",
+      "ɥ n",
+      "_d ʒ",
+      "x ɑʊ",
+      "_ˈ v",
+      "ʃi ɑŋ",
+      "p ʰ",
+      "_wɪ tʃ",
+      "eɪ m",
+      "oʊ z",
+      "ə ðəɹ",
+      "f ɑŋ",
+      "_ˈ g",
+      "_d o",
+      "_ʃi ɑʊ",
+      "_ˈ æ",
+      "_j ʊɹ",
+      "_ð ɛm",
+      "ɪ m",
+      "ɛ st",
+      "æn d",
+      "_d u",
+      "ɯ ɯ",
+      "k an",
+      "_d a",
+      "in o",
+      "_ e",
+      "_w ʊd",
+      "ɛn d",
+      "m eɪ",
+      "θ ɪŋ",
+      "_ʃ jɛn",
+      "i z",
+      "aɪ m",
+      "_h u",
+      "_əˈ b",
+      "ən s",
+      "_wɪ ɫ",
+      "t ʰi",
+      "g o",
+      "ɛn t",
+      "f u",
+      "æ p",
+      "x oʊ",
+      "eɪ k",
+      "ʊ k",
+      "əɹ ˈ",
+      "_θ ɪŋ",
+      "ə l",
+      "p ɹ",
+      "ə tʃ",
+      "n t",
+      "_ ɸɯ",
+      "l u",
+      "_ˈ ɔ",
+      "_i ɑʊ",
+      "l ə",
+      "t u",
+      "_dʑ i",
+      "eɪ t",
+      "_ʃi n",
+      "n na",
+      "_ˈp ɹ",
+      "f ən",
+      "_ə p",
+      "n jɛn",
+      "_a ʊt",
+      "f ɔɹ",
+      "_t u",
+      "eɪ ʃən",
+      "ɪ ɫ",
+      "_w ət",
+      "_ɪ f",
+      "_ ɥ",
+      "_f a",
+      "ˈ w",
+      "tʃ ʰjɛn",
+      "_w ɪn",
+      "oʊ ɫd",
+      "_əˈ p",
+      "aʊ nd",
+      "s an",
+      "h e",
+      "_b ɪn",
+      "f a",
+      "ɪ f",
+      "ɔ ŋ",
+      "g e",
+      "_ɪn _ðə",
+      "m iŋ",
+      "_p ɹ",
+      "in a",
+      "an o",
+      "əb əɫ",
+      "k ˈs",
+      "_ˈ ɛni",
+      "n əŋ",
+      "ə d",
+      "_əv _ðə_ˈ",
+      "_w aɪ",
+      "_t aɪm",
+      "ˈs ɛɫ",
+      "ʃi ɛ",
+      "_k əm",
+      "æ st",
+      "_g oʊ",
+      "m ɯ",
+      "ˈ p",
+      "_ˈ st",
+      "ə _t",
+      "p t",
+      "_p ʰ",
+      "ʰ ɹ",
+      "ʃ ja",
+      "i wa",
+      "ɪ l",
+      "b ət",
+      "_f ɑŋ",
+      "h o",
+      "i v",
+      "l oʊ",
+      "b e",
+      "_laɪ k",
+      "ɪ ʃ",
+      "_f u",
+      "z e",
+      "ə _tʃ",
+      "ɑɹ t",
+      "ɔɹ d",
+      "tʃʰi ŋ",
+      "m p",
+      "_ðə _s",
+      "_əˈb aʊt",
+      "_ˈ oʊ",
+      "kʰ ə",
+      "d _tɪ",
+      "ŋ ga",
+      "ə li",
+      "_kʰ an",
+      "ç i",
+      "_ˈ ju",
+      "_k ʊd",
+      "ɔ ɫ",
+      "ɔ t",
+      "_ɪ ts",
+      "_s an",
+      "tʃ a",
+      "i _na",
+      "x ə",
+      "ɛ kt",
+      "_m ɔɹ",
+      "te _kɯ",
+      "ɪd ʒ",
+      "j ʊŋ",
+      "_w an",
+      "æ t",
+      "ka t",
+      "ˈsɛɫ f",
+      "_k e",
+      "aɪ nd",
+      "i t",
+      "_ ɑɹ",
+      "s p",
+      "oʊn t",
+      "_t ʃi",
+      "tsʰ ɹ",
+      "_x ən",
+      "_əˈ g",
+      "ə _k",
+      "to _i",
+      "_t ʰi",
+      "_i ŋ",
+      "aʊ n",
+      "g ɯ",
+      "_ɪ kˈs",
+      "ɛ v",
+      "g i",
+      "k s",
+      "_s əm",
+      "an a",
+      "ɪt əɫ",
+      "n an",
+      "_ˈɪn tu",
+      "_hi ɹ",
+      "_t e",
+      "_n aʊ",
+      "ʃi ɑʊ",
+      "ʃ o",
+      "ɹ e",
+      "x aɪ",
+      "_tʃʰi ŋ",
+      "_s ɹ",
+      "_h aʊ",
+      "? .",
+      "_f eɪ",
+      "li ŋ",
+      "_ʃ ja",
+      "_ˈ dʒ",
+      "_s eɪ",
+      "ˈ n",
+      "s oʊ",
+      "tʰ ʊŋ",
+      "_l joʊ",
+      "m aɪ",
+      "_b ɹ",
+      "ɹeɪ t",
+      "_n əŋ",
+      "ʰ ə",
+      "æn s",
+      "_ˈɔ l",
+      "ta tʃi",
+      "n to",
+      "_ˌɪn ˈ",
+      "l e",
+      "n de",
+      "_ˈv ɛɹi",
+      "mən t",
+      "ɾi ma",
+      "_ð ɛn",
+      "_h əz",
+      "_ɹ i",
+      "f təɹ",
+      "_s p",
+      "ɾe wa",
+      "ga _a",
+      "z _əv",
+      "_m iŋ",
+      "_tɪ _ðə",
+      "ɹ aɪ",
+      "ɛ l",
+      "ɹ æ",
+      "_h oʊ",
+      "x u",
+      "oʊn li",
+      "ŋ k",
+      "i _i",
+      "_d ɪd",
+      "_dʒ ɪst",
+      "in g",
+      "ka i",
+      "_m æn",
+      "_i n",
+      "z o",
+      "ə f",
+      "da ke",
+      "_ˈs əm",
+      "ɾɯ _no",
+      "_g o",
+      "tʃ əɹ",
+      "i te",
+      "`↓ .",
+      "_kʰ aɪ",
+      "s k",
+      "ɔɹ s",
+      "_t ʰiŋ",
+      "_n ə",
+      "p əɫ",
+      "_tɪ _bi",
+      "ˈ fɔɹ",
+      "m u",
+      "s u",
+      "a a",
+      "ɪst əɹ",
+      "ʰ an",
+      "p əɹ",
+      "ə _p",
+      "li ɑŋ",
+      "_ v",
+      "oʊ st",
+      "_əˈg ɛn",
+      "ən z",
+      "N o",
+      "ɔɹ t",
+      "_s əˈ",
+      "_m ɯ",
+      "tʃ ʰ",
+      "_ˈl ɪtəɫ",
+      "_x wo",
+      "_ˌ bi",
+      "_ˈoʊ vəɹ",
+      "_ çi",
+      "_d eɪ",
+      "aɪ n",
+      "_ʃi ŋ",
+      "i _ʃi",
+      "_tsʰ aɪ",
+      "ʃ oo",
+      "ɾ oo",
+      "b əɹ",
+      "ʰ a",
+      "ˈ ɛs",
+      "_ɪn _ðə_ˈ",
+      "N wa",
+      "_ð ən",
+      "s aɪ",
+      "_ˈju ˈɛs",
+      "n da",
+      "_p leɪ",
+      "ɪŋ _tɪ",
+      "ɪt i",
+      "_m e",
+      "_ʃ ʊd",
+      "_n u",
+      "_ðə _k",
+      "z a",
+      "_ˈ ɛvəɹ",
+      "əɹ n",
+      "æ d",
+      "ˈ m",
+      "_d oʊnt",
+      "_m əst",
+      "j ɯɯ",
+      "ɑɹ d",
+      "_ jɛn",
+      "ʃ ɥ",
+      "_ˈ oʊnli",
+      "_ʃ o",
+      "_l iŋ",
+      "s s",
+      "ɑ l",
+      "de a",
+      "ɾe ta",
+      "m jɛn",
+      "_g ʊd",
+      "_w ɔ",
+      "i mo",
+      "no _ko",
+      "_ ɥæn",
+      "nd ʒ",
+      "ɪ ʃən",
+      "o _ʃi",
+      "_θɪŋ k",
+      "_n an",
+      "to _o",
+      "_tʰ ʊŋ",
+      "l joʊ",
+      "ta i",
+      "mə _s",
+      "_j ɯ",
+      "_ uɑŋ",
+      "_ˌbi ˈfɔɹ",
+      "æ s",
+      "_tʃ ʰjɛn",
+      "i k",
+      "_b æk",
+      "_ˈ iv",
+      "eɪ n",
+      "u n",
+      "l a",
+      "ˈ k",
+      "_d aʊn",
+      "an ai",
+      "_l ɛ",
+      "əɹ t",
+      "ð ɛɹ",
+      "_ˈæ ftəɹ",
+      "da t",
+      "f an",
+      "b əɫ",
+      "te mo",
+      "tʰ a",
+      "ɾɯ _ko",
+      "ˈ v",
+      "f eɪ",
+      "_m ətʃ",
+      "x wo",
+      "ɹ oʊ",
+      "_b a",
+      "_ˈn ɛvəɹ",
+      "_meɪ d",
+      "_j ʊŋ",
+      "_əˈp ɑn",
+      "! ?",
+      "_ˈ ʃ",
+      "_ðə_ˈ k",
+      "f t",
+      "_b o",
+      "_ɪn _ə",
+      "tʃʰɥ æn",
+      "ˈ z",
+      "`↓ ,",
+      "_bɪˈ k",
+      "ɪ g",
+      "k in",
+      "_k l",
+      "ɾɯ _n",
+      "_l ɑʊ",
+      "-- --",
+      "i ka",
+      "_ɹ aɪt",
+      "z d",
+      "z _ənd",
+      "_k jo",
+      "x wan",
+      "to o",
+      "_g ɪt",
+      "_l iɑŋ",
+      "ta _n",
+      "_k eɪm",
+      "_ˈ əðəɹ",
+      "_w ɛɫ",
+      "te ki",
+      "se e",
+      "j ɯ",
+      "i _o",
+      "to _ʃi",
+      "f əɫ",
+      "b o",
+      "ˌ t",
+      "ɪ p",
+      "an e",
+      "_tʰ jɛn",
+      "_tʃ o",
+      "ɾ jo",
+      "ɪn s",
+      "_h e",
+      "ŋ ka",
+      "ʃ ɥɛ",
+      "dʑ a",
+      "v d",
+      "ʰ wan",
+      "_g ɹeɪt",
+      "_əv _ə",
+      "ənd əɹ",
+      "ke do",
+      "_ðə _b",
+      "ə k",
+      "_t eɪk",
+      "kʰ an",
+      "_ˈɔl ˌ",
+      "s wo",
+      "_ɪt _wɑz",
+      "_ʃ ɥ",
+      "_si m",
+      "_ˈf ɑ",
+      "m in",
+      "i _a",
+      "s oo",
+      "ɛn s",
+      "_s ətʃ",
+      "tʰ aɪ",
+      "_ ga",
+      "i _ka",
+      "k oo",
+      "_fəɹ st",
+      "_ˈ tʃ",
+      "n no",
+      "ə _ɹ",
+      "ta ɾa",
+      "tʃʰ joʊ",
+      "_æ m",
+      "_m u",
+      "_meɪ k",
+      "↓ …",
+      "ɪˈ θ",
+      "ɑ b",
+      "ɹ a",
+      "_w ɛɹ",
+      "_ðə_ˈ s",
+      "_əˈ l",
+      "_ oʊɫd",
+      "æ l",
+      "_ˈp i",
+      "_l ɔŋ",
+      "dʑ o",
+      "_tʰ aɪ",
+      "ɔɹ n",
+      "əɫ z",
+      "_t əˈ",
+      "_əˈ weɪ",
+      "p a",
+      "_ð iz",
+      "_ˈs p",
+      "n n",
+      "ma e",
+      "to wa",
+      "ta _no",
+      "_ an",
+      "kʰ aɪ",
+      "ɾa ɾe",
+      "eɪ s",
+      "ɑ d",
+      "_w ɪˈθ",
+      "_ˈiv ɪn",
+      "_l u",
+      "ɔ ɪ",
+      "l ɪŋ",
+      "ət i",
+      "_ðə _f",
+      "o ʃi",
+      "_l a",
+      "s i",
+      "t ɪd",
+      "h aʊ",
+      "pʰ in",
+      "ˈ st",
+      "_ˈp əɹ",
+      "e ɹ",
+      "* !",
+      "_ˈm ɪstəɹ",
+      "ʃ a",
+      "_ˌ ɪm",
+      "ˌ θɪŋ",
+      "_n eɪ",
+      "_n ɥ",
+      "ɑ k",
+      "_ɹ u",
+      "_ʃ ɯ",
+      "_ðə_ˈ m",
+      "de mo",
+      "_d ɹ",
+      "dʑ oo",
+      "_st ɪɫ",
+      "_p ʰiŋ",
+      "ə _i",
+      "_ɪkˈs p",
+      "_w ɛnt",
+      "ɪ ɹi",
+      "əˈ m",
+      "o _ka",
+      "_əˈ k",
+      "ɔ k",
+      "_ ɥɛ",
+      "_l ʊk",
+      "ˈ d",
+      "ka ʃi",
+      "_wɪθ _ə",
+      "l jɛn",
+      "ɔ n",
+      "_l jɛn",
+      "_h ɛɫ",
+      "u ɹ",
+      "_tʰ oʊ",
+      "_tʃʰɥ æn",
+      "_s k",
+      "tsʰ aɪ",
+      "ɛ təɹ",
+      "_m in",
+      "n oʊ",
+      "ʃ ɯ",
+      "_θ ɹu",
+      "_θ ɔt",
+      "da jo",
+      "w i",
+      "i _ko",
+      "_t ɹ",
+      "_f an",
+      "ɹ ɛ",
+      "sa N",
+      "_hi _wɑz",
+      "_ ɾe",
+      "_ə m",
+      "te _ki",
+      "_x oʊ",
+      "ˈ l",
+      "ˈ g",
+      "ga _i",
+      "_ɔn _ðə",
+      "_x wa",
+      "v ɪŋ",
+      "m an",
+      "f əɹ",
+      "_ oʊn",
+      "ˈ ɹ",
+      "_k ɹ",
+      "te _o",
+      "ɪ li",
+      "_ʃ ɥɛ",
+      "_f əŋ",
+      "æ ɫ",
+      "ɑ p",
+      "_ˈ ɛv",
+      "eɪ ndʒ",
+      "i ɫ",
+      "w ət",
+      "ɛ ðəɹ",
+      "_f ən",
+      "ɾe e",
+      "_hi _hæd",
+      "_maɪ t",
+      "_g e",
+      "æ kt",
+      "ɪ ts",
+      "_h ɪm",
+      "_ ze",
+      "i i",
+      "_ N",
+      "_əv _hɪz",
+      "_g ɹ",
+      "æn t",
+      "ɪ ˌ",
+      "_hɪm ˈsɛɫf",
+      "wa _na",
+      "aɪ əɹ",
+      "dʑ anai",
+      "kan a",
+      "aɪ z",
+      "_ɪt _ɪz",
+      "ma se",
+      "w ɪn",
+      "ə θɪŋ",
+      "_pɹ əˈ",
+      "kɯ n",
+      "ˈ ju",
+      "_f ɔɹ",
+      "p ʰi",
+      "p ʰiŋ",
+      "o _i",
+      "v z",
+      "ɔ ɪn",
+      "t ʰiŋ",
+      "_n e",
+      "g əɹ",
+      "æ ts",
+      "_ˈ ɹi"
+    ]
+  }
+}

apps/audio_cloning/vallex/g2p/bpe_69.json

@@ -0,0 +1,141 @@
+{
+  "version": "1.0",
+  "truncation": null,
+  "padding": null,
+  "added_tokens": [
+    {
+      "id": 0,
+      "content": "[UNK]",
+      "single_word": false,
+      "lstrip": false,
+      "rstrip": false,
+      "normalized": false,
+      "special": true
+    },
+    {
+      "id": 1,
+      "content": "[CLS]",
+      "single_word": false,
+      "lstrip": false,
+      "rstrip": false,
+      "normalized": false,
+      "special": true
+    },
+    {
+      "id": 2,
+      "content": "[SEP]",
+      "single_word": false,
+      "lstrip": false,
+      "rstrip": false,
+      "normalized": false,
+      "special": true
+    },
+    {
+      "id": 3,
+      "content": "[PAD]",
+      "single_word": false,
+      "lstrip": false,
+      "rstrip": false,
+      "normalized": false,
+      "special": true
+    },
+    {
+      "id": 4,
+      "content": "[MASK]",
+      "single_word": false,
+      "lstrip": false,
+      "rstrip": false,
+      "normalized": false,
+      "special": true
+    }
+  ],
+  "normalizer": null,
+  "pre_tokenizer": {
+    "type": "Whitespace"
+  },
+  "post_processor": null,
+  "decoder": null,
+  "model": {
+    "type": "BPE",
+    "dropout": null,
+    "unk_token": "[UNK]",
+    "continuing_subword_prefix": null,
+    "end_of_word_suffix": null,
+    "fuse_unk": false,
+    "byte_fallback": false,
+    "vocab": {
+      "[UNK]": 0,
+      "[CLS]": 1,
+      "[SEP]": 2,
+      "[PAD]": 3,
+      "[MASK]": 4,
+      "!": 5,
+      "#": 6,
+      "*": 7,
+      ",": 8,
+      "-": 9,
+      ".": 10,
+      "=": 11,
+      "?": 12,
+      "N": 13,
+      "Q": 14,
+      "^": 15,
+      "_": 16,
+      "`": 17,
+      "a": 18,
+      "b": 19,
+      "d": 20,
+      "e": 21,
+      "f": 22,
+      "g": 23,
+      "h": 24,
+      "i": 25,
+      "j": 26,
+      "k": 27,
+      "l": 28,
+      "m": 29,
+      "n": 30,
+      "o": 31,
+      "p": 32,
+      "s": 33,
+      "t": 34,
+      "u": 35,
+      "v": 36,
+      "w": 37,
+      "x": 38,
+      "y": 39,
+      "z": 40,
+      "~": 41,
+      "æ": 42,
+      "ç": 43,
+      "ð": 44,
+      "ŋ": 45,
+      "ɑ": 46,
+      "ɔ": 47,
+      "ə": 48,
+      "ɛ": 49,
+      "ɥ": 50,
+      "ɪ": 51,
+      "ɫ": 52,
+      "ɯ": 53,
+      "ɸ": 54,
+      "ɹ": 55,
+      "ɾ": 56,
+      "ʃ": 57,
+      "ʊ": 58,
+      "ʑ": 59,
+      "ʒ": 60,
+      "ʰ": 61,
+      "ˈ": 62,
+      "ˌ": 63,
+      "θ": 64,
+      "…": 65,
+      "⁼": 66,
+      "↑": 67,
+      "→": 68,
+      "↓": 69
+    },
+    "merges": [
+    ]
+  }
+}

apps/audio_cloning/vallex/g2p/cleaners.py

@@ -0,0 +1,76 @@
+import re
+
+from .english import english_to_ipa2
+from .japanese import japanese_to_ipa2, japanese_to_romaji_with_accent
+from .mandarin import (
+    chinese_to_bopomofo,
+    chinese_to_ipa,
+    latin_to_bopomofo,
+    number_to_chinese,
+)
+
+patterns = [r"\[EN\](.*?)\[EN\]", r"\[ZH\](.*?)\[ZH\]", r"\[JA\](.*?)\[JA\]"]
+
+
+def japanese_cleaners(text):
+    text = japanese_to_romaji_with_accent(text)
+    text = re.sub(r"([A-Za-z])$", r"\1.", text)
+    return text
+
+
+def japanese_cleaners2(text):
+    return japanese_cleaners(text).replace("ts", "ʦ").replace("...", "…")
+
+
+def chinese_cleaners(text):
+    """Pipeline for Chinese text"""
+    text = number_to_chinese(text)
+    text = chinese_to_bopomofo(text)
+    text = latin_to_bopomofo(text)
+    text = re.sub(r"([ˉˊˇˋ˙])$", r"\1。", text)
+    return text
+
+
+def cje_cleaners(text):
+    matches = []
+    for pattern in patterns:
+        matches.extend(re.finditer(pattern, text))
+
+    matches.sort(key=lambda x: x.start())  # Sort matches by their start positions
+
+    outputs = ""
+    output_langs = []
+
+    for match in matches:
+        text_segment = text[match.start() : match.end()]
+        phon = clean_one(text_segment)
+        if "[EN]" in text_segment:
+            lang = "en"
+        elif "[ZH]" in text_segment:
+            lang = "zh"
+        elif "[JA]" in text_segment:
+            lang = "ja"
+        else:
+            raise ValueError("If you see this error, please report this bug to issues.")
+        outputs += phon
+        output_langs += [lang] * len(phon)
+    assert len(outputs) == len(output_langs)
+    return outputs, output_langs
+
+
+def clean_one(text):
+    if text.find("[ZH]") != -1:
+        text = re.sub(
+            r"\[ZH\](.*?)\[ZH\]", lambda x: chinese_to_ipa(x.group(1)) + " ", text
+        )
+    if text.find("[JA]") != -1:
+        text = re.sub(
+            r"\[JA\](.*?)\[JA\]", lambda x: japanese_to_ipa2(x.group(1)) + " ", text
+        )
+    if text.find("[EN]") != -1:
+        text = re.sub(
+            r"\[EN\](.*?)\[EN\]", lambda x: english_to_ipa2(x.group(1)) + " ", text
+        )
+    text = re.sub(r"\s+$", "", text)
+    text = re.sub(r"([^\.,!\?\-…~])$", r"\1.", text)
+    return text

apps/audio_cloning/vallex/g2p/english.py

@@ -0,0 +1,197 @@
+import re
+
+import inflect
+from unidecode import unidecode
+
+"""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).
+"""
+
+
+_inflect = inflect.engine()
+_comma_number_re = re.compile(r"([0-9][0-9\,]+[0-9])")
+_decimal_number_re = re.compile(r"([0-9]+\.[0-9]+)")
+_pounds_re = re.compile(r"£([0-9\,]*[0-9]+)")
+_dollars_re = re.compile(r"\$([0-9\.\,]*[0-9]+)")
+_ordinal_re = re.compile(r"[0-9]+(st|nd|rd|th)")
+_number_re = re.compile(r"[0-9]+")
+
+# 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"),
+    ]
+]
+
+
+# List of (ipa, lazy ipa) pairs:
+_lazy_ipa = [
+    (re.compile("%s" % x[0]), x[1])
+    for x in [
+        ("r", "ɹ"),
+        ("æ", "e"),
+        ("ɑ", "a"),
+        ("ɔ", "o"),
+        ("ð", "z"),
+        ("θ", "s"),
+        ("ɛ", "e"),
+        ("ɪ", "i"),
+        ("ʊ", "u"),
+        ("ʒ", "ʥ"),
+        ("ʤ", "ʥ"),
+        ("ˈ", "↓"),
+    ]
+]
+
+# List of (ipa, lazy ipa2) pairs:
+_lazy_ipa2 = [
+    (re.compile("%s" % x[0]), x[1])
+    for x in [
+        ("r", "ɹ"),
+        ("ð", "z"),
+        ("θ", "s"),
+        ("ʒ", "ʑ"),
+        ("ʤ", "dʑ"),
+        ("ˈ", "↓"),
+    ]
+]
+
+# List of (ipa, ipa2) pairs
+_ipa_to_ipa2 = [
+    (re.compile("%s" % x[0]), x[1]) for x in [("r", "ɹ"), ("ʤ", "dʒ"), ("ʧ", "tʃ")]
+]
+
+
+def expand_abbreviations(text):
+    for regex, replacement in _abbreviations:
+        text = re.sub(regex, replacement, text)
+    return text
+
+
+def collapse_whitespace(text):
+    return re.sub(r"\s+", " ", text)
+
+
+def _remove_commas(m):
+    return m.group(1).replace(",", "")
+
+
+def _expand_decimal_point(m):
+    return m.group(1).replace(".", " point ")
+
+
+def _expand_dollars(m):
+    match = m.group(1)
+    parts = match.split(".")
+    if len(parts) > 2:
+        return match + " dollars"  # Unexpected format
+    dollars = int(parts[0]) if parts[0] else 0
+    cents = int(parts[1]) if len(parts) > 1 and parts[1] else 0
+    if dollars and cents:
+        dollar_unit = "dollar" if dollars == 1 else "dollars"
+        cent_unit = "cent" if cents == 1 else "cents"
+        return "%s %s, %s %s" % (dollars, dollar_unit, cents, cent_unit)
+    elif dollars:
+        dollar_unit = "dollar" if dollars == 1 else "dollars"
+        return "%s %s" % (dollars, dollar_unit)
+    elif cents:
+        cent_unit = "cent" if cents == 1 else "cents"
+        return "%s %s" % (cents, cent_unit)
+    else:
+        return "zero dollars"
+
+
+def _expand_ordinal(m):
+    return _inflect.number_to_words(m.group(0))
+
+
+def _expand_number(m):
+    num = int(m.group(0))
+    if num > 1000 and num < 3000:
+        if num == 2000:
+            return "two thousand"
+        elif num > 2000 and num < 2010:
+            return "two thousand " + _inflect.number_to_words(num % 100)
+        elif num % 100 == 0:
+            return _inflect.number_to_words(num // 100) + " hundred"
+        else:
+            return _inflect.number_to_words(
+                num, andword="", zero="oh", group=2
+            ).replace(", ", " ")
+    else:
+        return _inflect.number_to_words(num, andword="")
+
+
+def normalize_numbers(text):
+    text = re.sub(_comma_number_re, _remove_commas, text)
+    text = re.sub(_pounds_re, r"\1 pounds", text)
+    text = re.sub(_dollars_re, _expand_dollars, text)
+    text = re.sub(_decimal_number_re, _expand_decimal_point, text)
+    text = re.sub(_ordinal_re, _expand_ordinal, text)
+    text = re.sub(_number_re, _expand_number, text)
+    return text
+
+
+def mark_dark_l(text):
+    return re.sub(r"l([^aeiouæɑɔəɛɪʊ ]*(?: |$))", lambda x: "ɫ" + x.group(1), text)
+
+
+def english_to_ipa(text):
+    import eng_to_ipa as ipa
+
+    text = unidecode(text).lower()
+    text = expand_abbreviations(text)
+    text = normalize_numbers(text)
+    phonemes = ipa.convert(text)
+    phonemes = collapse_whitespace(phonemes)
+    return phonemes
+
+
+def english_to_lazy_ipa(text):
+    text = english_to_ipa(text)
+    for regex, replacement in _lazy_ipa:
+        text = re.sub(regex, replacement, text)
+    return text
+
+
+def english_to_ipa2(text):
+    text = english_to_ipa(text)
+    text = mark_dark_l(text)
+    for regex, replacement in _ipa_to_ipa2:
+        text = re.sub(regex, replacement, text)
+    return text.replace("...", "…")
+
+
+def english_to_lazy_ipa2(text):
+    text = english_to_ipa(text)
+    for regex, replacement in _lazy_ipa2:
+        text = re.sub(regex, replacement, text)
+    return text

apps/audio_cloning/vallex/g2p/japanese.py

@@ -0,0 +1,173 @@
+import re
+
+from unidecode import unidecode
+
+# Regular expression matching Japanese without punctuation marks:
+_japanese_characters = re.compile(
+    r"[A-Za-z\d\u3005\u3040-\u30ff\u4e00-\u9fff\uff11-\uff19\uff21-\uff3a\uff41-\uff5a\uff66-\uff9d]"
+)
+
+# Regular expression matching non-Japanese characters or punctuation marks:
+_japanese_marks = re.compile(
+    r"[^A-Za-z\d\u3005\u3040-\u30ff\u4e00-\u9fff\uff11-\uff19\uff21-\uff3a\uff41-\uff5a\uff66-\uff9d]"
+)
+
+# List of (symbol, Japanese) pairs for marks:
+_symbols_to_japanese = [(re.compile("%s" % x[0]), x[1]) for x in [("％", "パーセント")]]
+
+# List of (romaji, ipa) pairs for marks:
+_romaji_to_ipa = [
+    (re.compile("%s" % x[0]), x[1])
+    for x in [
+        ("ts", "ʦ"),
+        ("u", "ɯ"),
+        ("j", "ʥ"),
+        ("y", "j"),
+        ("ni", "n^i"),
+        ("nj", "n^"),
+        ("hi", "çi"),
+        ("hj", "ç"),
+        ("f", "ɸ"),
+        ("I", "i*"),
+        ("U", "ɯ*"),
+        ("r", "ɾ"),
+    ]
+]
+
+# List of (romaji, ipa2) pairs for marks:
+_romaji_to_ipa2 = [
+    (re.compile("%s" % x[0]), x[1])
+    for x in [
+        ("u", "ɯ"),
+        ("ʧ", "tʃ"),
+        ("j", "dʑ"),
+        ("y", "j"),
+        ("ni", "n^i"),
+        ("nj", "n^"),
+        ("hi", "çi"),
+        ("hj", "ç"),
+        ("f", "ɸ"),
+        ("I", "i*"),
+        ("U", "ɯ*"),
+        ("r", "ɾ"),
+    ]
+]
+
+# List of (consonant, sokuon) pairs:
+_real_sokuon = [
+    (re.compile("%s" % x[0]), x[1])
+    for x in [
+        (r"Q([↑↓]*[kg])", r"k#\1"),
+        (r"Q([↑↓]*[tdjʧ])", r"t#\1"),
+        (r"Q([↑↓]*[sʃ])", r"s\1"),
+        (r"Q([↑↓]*[pb])", r"p#\1"),
+    ]
+]
+
+# List of (consonant, hatsuon) pairs:
+_real_hatsuon = [
+    (re.compile("%s" % x[0]), x[1])
+    for x in [
+        (r"N([↑↓]*[pbm])", r"m\1"),
+        (r"N([↑↓]*[ʧʥj])", r"n^\1"),
+        (r"N([↑↓]*[tdn])", r"n\1"),
+        (r"N([↑↓]*[kg])", r"ŋ\1"),
+    ]
+]
+
+
+def symbols_to_japanese(text):
+    for regex, replacement in _symbols_to_japanese:
+        text = re.sub(regex, replacement, text)
+    return text
+
+
+def japanese_to_romaji_with_accent(text):
+    """Reference https://r9y9.github.io/ttslearn/latest/notebooks/ch10_Recipe-Tacotron.html"""
+    import pyopenjtalk
+
+    text = symbols_to_japanese(text)
+    sentences = re.split(_japanese_marks, text)
+    marks = re.findall(_japanese_marks, text)
+    text = ""
+    for i, sentence in enumerate(sentences):
+        if re.match(_japanese_characters, sentence):
+            if text != "":
+                text += " "
+            labels = pyopenjtalk.extract_fullcontext(sentence)
+            for n, label in enumerate(labels):
+                phoneme = re.search(r"\-([^\+]*)\+", label).group(1)
+                if phoneme not in ["sil", "pau"]:
+                    text += (
+                        phoneme.replace("ch", "ʧ").replace("sh", "ʃ").replace("cl", "Q")
+                    )
+                else:
+                    continue
+                # n_moras = int(re.search(r'/F:(\d+)_', label).group(1))
+                a1 = int(re.search(r"/A:(\-?[0-9]+)\+", label).group(1))
+                a2 = int(re.search(r"\+(\d+)\+", label).group(1))
+                a3 = int(re.search(r"\+(\d+)/", label).group(1))
+                if re.search(r"\-([^\+]*)\+", labels[n + 1]).group(1) in ["sil", "pau"]:
+                    a2_next = -1
+                else:
+                    a2_next = int(re.search(r"\+(\d+)\+", labels[n + 1]).group(1))
+                # Accent phrase boundary
+                if a3 == 1 and a2_next == 1:
+                    text += " "
+                # Falling
+                elif a1 == 0 and a2_next == a2 + 1:
+                    text += "↓"
+                # Rising
+                elif a2 == 1 and a2_next == 2:
+                    text += "↑"
+        if i < len(marks):
+            text += unidecode(marks[i]).replace(" ", "")
+    return text
+
+
+def get_real_sokuon(text):
+    for regex, replacement in _real_sokuon:
+        text = re.sub(regex, replacement, text)
+    return text
+
+
+def get_real_hatsuon(text):
+    for regex, replacement in _real_hatsuon:
+        text = re.sub(regex, replacement, text)
+    return text
+
+
+def japanese_to_ipa(text):
+    text = japanese_to_romaji_with_accent(text).replace("...", "…")
+    text = re.sub(
+        r"([aiueo])\1+", lambda x: x.group(0)[0] + "ː" * (len(x.group(0)) - 1), text
+    )
+    text = get_real_sokuon(text)
+    text = get_real_hatsuon(text)
+    for regex, replacement in _romaji_to_ipa:
+        text = re.sub(regex, replacement, text)
+    return text
+
+
+def japanese_to_ipa2(text):
+    text = japanese_to_romaji_with_accent(text).replace("...", "…")
+    text = get_real_sokuon(text)
+    text = get_real_hatsuon(text)
+    for regex, replacement in _romaji_to_ipa2:
+        text = re.sub(regex, replacement, text)
+    return text
+
+
+def japanese_to_ipa3(text):
+    text = (
+        japanese_to_ipa2(text)
+        .replace("n^", "ȵ")
+        .replace("ʃ", "ɕ")
+        .replace("*", "\u0325")
+        .replace("#", "\u031a")
+    )
+    text = re.sub(
+        r"([aiɯeo])\1+", lambda x: x.group(0)[0] + "ː" * (len(x.group(0)) - 1), text
+    )
+    text = re.sub(r"((?:^|\s)(?:ts|tɕ|[kpt]))", r"\1ʰ", text)
+    return text

apps/audio_cloning/vallex/g2p/mandarin.py

@@ -0,0 +1,337 @@
+import re
+
+import cn2an
+import jieba
+
+# List of (Latin alphabet, bopomofo) pairs:
+_latin_to_bopomofo = [
+    (re.compile("%s" % x[0], re.IGNORECASE), x[1])
+    for x in [
+        ("a", "ㄟˉ"),
+        ("b", "ㄅㄧˋ"),
+        ("c", "ㄙㄧˉ"),
+        ("d", "ㄉㄧˋ"),
+        ("e", "ㄧˋ"),
+        ("f", "ㄝˊㄈㄨˋ"),
+        ("g", "ㄐㄧˋ"),
+        ("h", "ㄝˇㄑㄩˋ"),
+        ("i", "ㄞˋ"),
+        ("j", "ㄐㄟˋ"),
+        ("k", "ㄎㄟˋ"),
+        ("l", "ㄝˊㄛˋ"),
+        ("m", "ㄝˊㄇㄨˋ"),
+        ("n", "ㄣˉ"),
+        ("o", "ㄡˉ"),
+        ("p", "ㄆㄧˉ"),
+        ("q", "ㄎㄧㄡˉ"),
+        ("r", "ㄚˋ"),
+        ("s", "ㄝˊㄙˋ"),
+        ("t", "ㄊㄧˋ"),
+        ("u", "ㄧㄡˉ"),
+        ("v", "ㄨㄧˉ"),
+        ("w", "ㄉㄚˋㄅㄨˋㄌㄧㄡˋ"),
+        ("x", "ㄝˉㄎㄨˋㄙˋ"),
+        ("y", "ㄨㄞˋ"),
+        ("z", "ㄗㄟˋ"),
+    ]
+]
+
+# List of (bopomofo, romaji) pairs:
+_bopomofo_to_romaji = [
+    (re.compile("%s" % x[0]), x[1])
+    for x in [
+        ("ㄅㄛ", "p⁼wo"),
+        ("ㄆㄛ", "pʰwo"),
+        ("ㄇㄛ", "mwo"),
+        ("ㄈㄛ", "fwo"),
+        ("ㄅ", "p⁼"),
+        ("ㄆ", "pʰ"),
+        ("ㄇ", "m"),
+        ("ㄈ", "f"),
+        ("ㄉ", "t⁼"),
+        ("ㄊ", "tʰ"),
+        ("ㄋ", "n"),
+        ("ㄌ", "l"),
+        ("ㄍ", "k⁼"),
+        ("ㄎ", "kʰ"),
+        ("ㄏ", "h"),
+        ("ㄐ", "ʧ⁼"),
+        ("ㄑ", "ʧʰ"),
+        ("ㄒ", "ʃ"),
+        ("ㄓ", "ʦ`⁼"),
+        ("ㄔ", "ʦ`ʰ"),
+        ("ㄕ", "s`"),
+        ("ㄖ", "ɹ`"),
+        ("ㄗ", "ʦ⁼"),
+        ("ㄘ", "ʦʰ"),
+        ("ㄙ", "s"),
+        ("ㄚ", "a"),
+        ("ㄛ", "o"),
+        ("ㄜ", "ə"),
+        ("ㄝ", "e"),
+        ("ㄞ", "ai"),
+        ("ㄟ", "ei"),
+        ("ㄠ", "au"),
+        ("ㄡ", "ou"),
+        ("ㄧㄢ", "yeNN"),
+        ("ㄢ", "aNN"),
+        ("ㄧㄣ", "iNN"),
+        ("ㄣ", "əNN"),
+        ("ㄤ", "aNg"),
+        ("ㄧㄥ", "iNg"),
+        ("ㄨㄥ", "uNg"),
+        ("ㄩㄥ", "yuNg"),
+        ("ㄥ", "əNg"),
+        ("ㄦ", "əɻ"),
+        ("ㄧ", "i"),
+        ("ㄨ", "u"),
+        ("ㄩ", "ɥ"),
+        ("ˉ", "→"),
+        ("ˊ", "↑"),
+        ("ˇ", "↓↑"),
+        ("ˋ", "↓"),
+        ("˙", ""),
+        ("，", ","),
+        ("。", "."),
+        ("！", "!"),
+        ("？", "?"),
+        ("—", "-"),
+    ]
+]
+
+# List of (romaji, ipa) pairs:
+_romaji_to_ipa = [
+    (re.compile("%s" % x[0], re.IGNORECASE), x[1])
+    for x in [
+        ("ʃy", "ʃ"),
+        ("ʧʰy", "ʧʰ"),
+        ("ʧ⁼y", "ʧ⁼"),
+        ("NN", "n"),
+        ("Ng", "ŋ"),
+        ("y", "j"),
+        ("h", "x"),
+    ]
+]
+
+# List of (bopomofo, ipa) pairs:
+_bopomofo_to_ipa = [
+    (re.compile("%s" % x[0]), x[1])
+    for x in [
+        ("ㄅㄛ", "p⁼wo"),
+        ("ㄆㄛ", "pʰwo"),
+        ("ㄇㄛ", "mwo"),
+        ("ㄈㄛ", "fwo"),
+        ("ㄅ", "p⁼"),
+        ("ㄆ", "pʰ"),
+        ("ㄇ", "m"),
+        ("ㄈ", "f"),
+        ("ㄉ", "t⁼"),
+        ("ㄊ", "tʰ"),
+        ("ㄋ", "n"),
+        ("ㄌ", "l"),
+        ("ㄍ", "k⁼"),
+        ("ㄎ", "kʰ"),
+        ("ㄏ", "x"),
+        ("ㄐ", "tʃ⁼"),
+        ("ㄑ", "tʃʰ"),
+        ("ㄒ", "ʃ"),
+        ("ㄓ", "ts`⁼"),
+        ("ㄔ", "ts`ʰ"),
+        ("ㄕ", "s`"),
+        ("ㄖ", "ɹ`"),
+        ("ㄗ", "ts⁼"),
+        ("ㄘ", "tsʰ"),
+        ("ㄙ", "s"),
+        ("ㄚ", "a"),
+        ("ㄛ", "o"),
+        ("ㄜ", "ə"),
+        ("ㄝ", "ɛ"),
+        ("ㄞ", "aɪ"),
+        ("ㄟ", "eɪ"),
+        ("ㄠ", "ɑʊ"),
+        ("ㄡ", "oʊ"),
+        ("ㄧㄢ", "jɛn"),
+        ("ㄩㄢ", "ɥæn"),
+        ("ㄢ", "an"),
+        ("ㄧㄣ", "in"),
+        ("ㄩㄣ", "ɥn"),
+        ("ㄣ", "ən"),
+        ("ㄤ", "ɑŋ"),
+        ("ㄧㄥ", "iŋ"),
+        ("ㄨㄥ", "ʊŋ"),
+        ("ㄩㄥ", "jʊŋ"),
+        ("ㄥ", "əŋ"),
+        ("ㄦ", "əɻ"),
+        ("ㄧ", "i"),
+        ("ㄨ", "u"),
+        ("ㄩ", "ɥ"),
+        ("ˉ", "→"),
+        ("ˊ", "↑"),
+        ("ˇ", "↓↑"),
+        ("ˋ", "↓"),
+        ("˙", ""),
+        ("，", ","),
+        ("。", "."),
+        ("！", "!"),
+        ("？", "?"),
+        ("—", "-"),
+    ]
+]
+
+# List of (bopomofo, ipa2) pairs:
+_bopomofo_to_ipa2 = [
+    (re.compile("%s" % x[0]), x[1])
+    for x in [
+        ("ㄅㄛ", "pwo"),
+        ("ㄆㄛ", "pʰwo"),
+        ("ㄇㄛ", "mwo"),
+        ("ㄈㄛ", "fwo"),
+        ("ㄅ", "p"),
+        ("ㄆ", "pʰ"),
+        ("ㄇ", "m"),
+        ("ㄈ", "f"),
+        ("ㄉ", "t"),
+        ("ㄊ", "tʰ"),
+        ("ㄋ", "n"),
+        ("ㄌ", "l"),
+        ("ㄍ", "k"),
+        ("ㄎ", "kʰ"),
+        ("ㄏ", "h"),
+        ("ㄐ", "tɕ"),
+        ("ㄑ", "tɕʰ"),
+        ("ㄒ", "ɕ"),
+        ("ㄓ", "tʂ"),
+        ("ㄔ", "tʂʰ"),
+        ("ㄕ", "ʂ"),
+        ("ㄖ", "ɻ"),
+        ("ㄗ", "ts"),
+        ("ㄘ", "tsʰ"),
+        ("���", "s"),
+        ("ㄚ", "a"),
+        ("ㄛ", "o"),
+        ("ㄜ", "ɤ"),
+        ("ㄝ", "ɛ"),
+        ("ㄞ", "aɪ"),
+        ("ㄟ", "eɪ"),
+        ("ㄠ", "ɑʊ"),
+        ("ㄡ", "oʊ"),
+        ("ㄧㄢ", "jɛn"),
+        ("ㄩㄢ", "yæn"),
+        ("ㄢ", "an"),
+        ("ㄧㄣ", "in"),
+        ("ㄩㄣ", "yn"),
+        ("ㄣ", "ən"),
+        ("ㄤ", "ɑŋ"),
+        ("ㄧㄥ", "iŋ"),
+        ("ㄨㄥ", "ʊŋ"),
+        ("ㄩㄥ", "jʊŋ"),
+        ("ㄥ", "ɤŋ"),
+        ("ㄦ", "əɻ"),
+        ("ㄧ", "i"),
+        ("ㄨ", "u"),
+        ("ㄩ", "y"),
+        ("ˉ", "˥"),
+        ("ˊ", "˧˥"),
+        ("ˇ", "˨˩˦"),
+        ("ˋ", "˥˩"),
+        ("˙", ""),
+        ("，", ","),
+        ("。", "."),
+        ("！", "!"),
+        ("？", "?"),
+        ("—", "-"),
+    ]
+]
+
+
+def number_to_chinese(text):
+    numbers = re.findall(r"\d+(?:\.?\d+)?", text)
+    for number in numbers:
+        text = text.replace(number, cn2an.an2cn(number), 1)
+    return text
+
+
+def chinese_to_bopomofo(text):
+    from pypinyin import BOPOMOFO, lazy_pinyin
+
+    text = text.replace("、", "，").replace("；", "，").replace("：", "，")
+    words = jieba.lcut(text, cut_all=False)
+    text = ""
+    for word in words:
+        bopomofos = lazy_pinyin(word, BOPOMOFO)
+        if not re.search("[\u4e00-\u9fff]", word):
+            text += word
+            continue
+        for i in range(len(bopomofos)):
+            bopomofos[i] = re.sub(r"([\u3105-\u3129])$", r"\1ˉ", bopomofos[i])
+        if text != "":
+            text += " "
+        text += "".join(bopomofos)
+    return text
+
+
+def latin_to_bopomofo(text):
+    for regex, replacement in _latin_to_bopomofo:
+        text = re.sub(regex, replacement, text)
+    return text
+
+
+def bopomofo_to_romaji(text):
+    for regex, replacement in _bopomofo_to_romaji:
+        text = re.sub(regex, replacement, text)
+    return text
+
+
+def bopomofo_to_ipa(text):
+    for regex, replacement in _bopomofo_to_ipa:
+        text = re.sub(regex, replacement, text)
+    return text
+
+
+def bopomofo_to_ipa2(text):
+    for regex, replacement in _bopomofo_to_ipa2:
+        text = re.sub(regex, replacement, text)
+    return text
+
+
+def chinese_to_romaji(text):
+    text = number_to_chinese(text)
+    text = chinese_to_bopomofo(text)
+    text = latin_to_bopomofo(text)
+    text = bopomofo_to_romaji(text)
+    text = re.sub("i([aoe])", r"y\1", text)
+    text = re.sub("u([aoəe])", r"w\1", text)
+    text = re.sub("([ʦsɹ]`[⁼ʰ]?)([→↓↑ ]+|$)", r"\1ɹ`\2", text).replace("ɻ", "ɹ`")
+    text = re.sub("([ʦs][⁼ʰ]?)([→↓↑ ]+|$)", r"\1ɹ\2", text)
+    return text
+
+
+def chinese_to_lazy_ipa(text):
+    text = chinese_to_romaji(text)
+    for regex, replacement in _romaji_to_ipa:
+        text = re.sub(regex, replacement, text)
+    return text
+
+
+def chinese_to_ipa(text):
+    text = number_to_chinese(text)
+    text = chinese_to_bopomofo(text)
+    text = latin_to_bopomofo(text)
+    text = bopomofo_to_ipa(text)
+    text = re.sub("i([aoe])", r"j\1", text)
+    text = re.sub("u([aoəe])", r"w\1", text)
+    text = re.sub("([sɹ]`[⁼ʰ]?)([→↓↑ ]+|$)", r"\1ɹ`\2", text).replace("ɻ", "ɹ`")
+    text = re.sub("([s][⁼ʰ]?)([→↓↑ ]+|$)", r"\1ɹ\2", text)
+    return text
+
+
+def chinese_to_ipa2(text):
+    text = number_to_chinese(text)
+    text = chinese_to_bopomofo(text)
+    text = latin_to_bopomofo(text)
+    text = bopomofo_to_ipa2(text)
+    text = re.sub(r"i([aoe])", r"j\1", text)
+    text = re.sub(r"u([aoəe])", r"w\1", text)
+    text = re.sub(r"([ʂɹ]ʰ?)([˩˨˧˦˥ ]+|$)", r"\1ʅ\2", text)
+    text = re.sub(r"(sʰ?)([˩˨˧˦˥ ]+|$)", r"\1ɿ\2", text)
+    return text

apps/audio_cloning/vallex/g2p/symbols.py

@@ -0,0 +1,76 @@
+'''
+Defines the set of symbols used in text input to the model.
+'''
+
+# japanese_cleaners
+# _pad        = '_'
+# _punctuation = ',.!?-'
+# _letters = 'AEINOQUabdefghijkmnoprstuvwyzʃʧ↓↑ '
+
+
+'''# japanese_cleaners2
+_pad        = '_'
+_punctuation = ',.!?-~…'
+_letters = 'AEINOQUabdefghijkmnoprstuvwyzʃʧʦ↓↑ '
+'''
+
+
+'''# korean_cleaners
+_pad        = '_'
+_punctuation = ',.!?…~'
+_letters = 'ㄱㄴㄷㄹㅁㅂㅅㅇㅈㅊㅋㅌㅍㅎㄲㄸㅃㅆㅉㅏㅓㅗㅜㅡㅣㅐㅔ '
+'''
+
+'''# chinese_cleaners
+_pad        = '_'
+_punctuation = '，。！？—…'
+_letters = 'ㄅㄆㄇㄈㄉㄊㄋㄌㄍㄎㄏㄐㄑㄒㄓㄔㄕㄖㄗㄘㄙㄚㄛㄜㄝㄞㄟㄠㄡㄢㄣㄤㄥㄦㄧㄨㄩˉˊˇˋ˙ '
+'''
+
+# # zh_ja_mixture_cleaners
+# _pad        = '_'
+# _punctuation = ',.!?-~…'
+# _letters = 'AEINOQUabdefghijklmnoprstuvwyzʃʧʦɯɹəɥ⁼ʰ`→↓↑ '
+
+
+'''# sanskrit_cleaners
+_pad        = '_'
+_punctuation = '।'
+_letters = 'ँंःअआइईउऊऋएऐओऔकखगघङचछजझञटठडढणतथदधनपफबभमयरलळवशषसहऽािीुूृॄेैोौ्ॠॢ '
+'''
+
+'''# cjks_cleaners
+_pad        = '_'
+_punctuation = ',.!?-~…'
+_letters = 'NQabdefghijklmnopstuvwxyzʃʧʥʦɯɹəɥçɸɾβŋɦː⁼ʰ`^#*=→↓↑ '
+'''
+
+'''# thai_cleaners
+_pad        = '_'
+_punctuation = '.!? '
+_letters = 'กขฃคฆงจฉชซฌญฎฏฐฑฒณดตถทธนบปผฝพฟภมยรฤลวศษสหฬอฮฯะัาำิีึืุูเแโใไๅๆ็่้๊๋์'
+'''
+
+# # cjke_cleaners2
+_pad        = '_'
+_punctuation = ',.!?-~…'
+_letters = 'NQabdefghijklmnopstuvwxyzɑæʃʑçɯɪɔɛɹðəɫɥɸʊɾʒθβŋɦ⁼ʰ`^#*=ˈˌ→↓↑ '
+
+
+'''# shanghainese_cleaners
+_pad        = '_'
+_punctuation = ',.!?…'
+_letters = 'abdfghiklmnopstuvyzøŋȵɑɔɕəɤɦɪɿʑʔʰ̩̃ᴀᴇ15678 '
+'''
+
+'''# chinese_dialect_cleaners
+_pad        = '_'
+_punctuation = ',.!?~…─'
+_letters = '#Nabdefghijklmnoprstuvwxyzæçøŋœȵɐɑɒɓɔɕɗɘəɚɛɜɣɤɦɪɭɯɵɷɸɻɾɿʂʅʊʋʌʏʑʔʦʮʰʷˀː˥˦˧˨˩̥̩̃̚ᴀᴇ↑↓∅ⱼ '
+'''
+
+# Export all symbols:
+symbols = [_pad] + list(_punctuation) + list(_letters)
+
+# Special symbol ids
+SPACE_ID = symbols.index(" ")

apps/audio_cloning/vallex/macros.py

@@ -0,0 +1,34 @@
+NUM_LAYERS = 12
+NUM_HEAD = 16
+N_DIM = 1024
+PREFIX_MODE = 1
+NUM_QUANTIZERS = 8
+SAMPLE_RATE = 24000
+
+lang2token = {
+    "zh": "[ZH]",
+    "ja": "[JA]",
+    "en": "[EN]",
+    "mix": "",
+}
+
+lang2code = {
+    "zh": 0,
+    "ja": 1,
+    "en": 2,
+}
+
+token2lang = {"[ZH]": "zh", "[JA]": "ja", "[EN]": "en", "": "mix"}
+
+code2lang = {
+    0: "zh",
+    1: "ja",
+    2: "en",
+}
+
+langdropdown2token = {
+    "English": "[EN]",
+    "中文": "[ZH]",
+    "日本語": "[JA]",
+    "Mix": "",
+}

apps/audio_cloning/vallex/main.py

@@ -0,0 +1,461 @@
+import logging
+import multiprocessing
+import os
+import pathlib
+import platform
+import sys
+import tempfile
+import time
+
+import gradio as gr
+import langid
+import nltk
+import numpy as np
+import torch
+import torchaudio
+import whisper
+from vocos import Vocos
+
+from .data.collation import get_text_token_collater
+from .data.tokenizer import (
+    AudioTokenizer,
+    tokenize_audio,
+)
+from .descriptions import infer_from_audio_ja_md, top_ja_md
+from .examples import infer_from_audio_examples
+from .g2p import PhonemeBpeTokenizer
+from .macros import (
+    N_DIM,
+    NUM_HEAD,
+    NUM_LAYERS,
+    NUM_QUANTIZERS,
+    PREFIX_MODE,
+    lang2code,
+    lang2token,
+    langdropdown2token,
+    token2lang,
+)
+from .models.vallex import VALLE
+
+logger = logging.getLogger(__name__)
+
+# set languages
+langid.set_languages(["en", "zh", "ja"])
+
+# set nltk data path
+nltk.data.path = nltk.data.path + [os.path.join(os.getcwd(), "nltk_data")]
+logger.info("nltk_data path: %s", nltk.data.path)
+
+# get encoding
+logger.info(
+    "default encoding is %s,file system encoding is %s",
+    sys.getdefaultencoding(),
+    sys.getfilesystemencoding(),
+)
+
+# check python version
+logger.info("You are using Python version %s", platform.python_version())
+if sys.version_info[0] < 3 or sys.version_info[1] < 7:
+    logger.warning("The Python version is too low and may cause problems")
+if platform.system().lower() == "windows":
+    temp = pathlib.PosixPath
+    pathlib.PosixPath = pathlib.WindowsPath
+else:
+    temp = pathlib.WindowsPath
+    pathlib.WindowsPath = pathlib.PosixPath
+os.environ["PROTOCOL_BUFFERS_PYTHON_IMPLEMENTATION"] = "python"
+
+# set torch threads (guarded for hot-reload)
+thread_count = multiprocessing.cpu_count()
+logger.info("Use %d cpu cores for computing", thread_count)
+if not getattr(torch, "_vallex_threads_configured", False):
+    torch.set_num_threads(thread_count)
+    try:
+        torch.set_num_interop_threads(thread_count)
+    except RuntimeError as err:
+        logger.warning("Skipping set_num_interop_threads: %s", err)
+    torch._C._jit_set_profiling_executor(False)
+    torch._C._jit_set_profiling_mode(False)
+    torch._C._set_graph_executor_optimize(False)
+
+    # gradio のリロード時に torch.set_num_iterop_threads を実行するとエラーになるので、設定済みのフラグをセット
+    setattr(torch, "_vallex_threads_configured", True)
+else:
+    logger.info("Torch threads already configured; skipping reconfiguration")
+
+# set text tokenizer and collater
+logger.info("Setting text tokenizer and collater...")
+tokenizer_path = "./apps/audio_cloning/vallex/g2p/bpe_69.json"
+text_tokenizer = PhonemeBpeTokenizer(tokenizer_path=tokenizer_path)
+text_collater = get_text_token_collater()
+
+# set device
+logger.info("Setting device...")
+device = torch.device("cpu")
+if torch.cuda.is_available():
+    device = torch.device("cuda", 0)
+# if torch.backends.mps.is_available():
+#     device = torch.device("mps")
+logger.info("Device set to %s", device)
+
+# Download VALL-E-X model weights if not exists
+OUTPUT_DIR_CHECKPOINTS = "./models/checkpoints"
+if platform.system().lower() == "linux":
+    # docker（linux）環境では /app/models/checkpoints にする
+    OUTPUT_DIR_CHECKPOINTS = "/app/models/checkpoints"
+
+OUTPUT_FILENAME_CHECKPOINTS = "vallex-checkpoint.pt"
+OUTPUT_PATH_CHECKPOINTS = os.path.join(
+    OUTPUT_DIR_CHECKPOINTS, OUTPUT_FILENAME_CHECKPOINTS
+)
+if not os.path.exists(OUTPUT_DIR_CHECKPOINTS):
+    os.makedirs(OUTPUT_DIR_CHECKPOINTS, exist_ok=True)
+if not os.path.exists(OUTPUT_PATH_CHECKPOINTS):
+    import wget
+
+    logging.info(
+        "Downloading model from https://huggingface.co/Plachta/VALL-E-X/resolve/main/vallex-checkpoint.pt ..."
+    )
+    wget.download(
+        "https://huggingface.co/Plachta/VALL-E-X/resolve/main/vallex-checkpoint.pt",
+        out=OUTPUT_PATH_CHECKPOINTS,
+        bar=wget.bar_adaptive,
+    )
+    raise Exception(
+        "\n Model weights download failed, please go to 'https://huggingface.co/Plachta/VALL-E-X/resolve/main/vallex-checkpoint.pt'"
+        "\n manually download model weights and put it to {} .".format(
+            os.getcwd() + f"{OUTPUT_DIR_CHECKPOINTS}"
+        )
+    )
+
+# initialize VALL-E-X model
+model = VALLE(
+    N_DIM,
+    NUM_HEAD,
+    NUM_LAYERS,
+    norm_first=True,
+    add_prenet=False,
+    prefix_mode=PREFIX_MODE,
+    share_embedding=True,
+    nar_scale_factor=1.0,
+    prepend_bos=True,
+    num_quantizers=NUM_QUANTIZERS,
+)
+checkpoint = torch.load(OUTPUT_PATH_CHECKPOINTS, map_location="cpu", weights_only=False)
+missing_keys, unexpected_keys = model.load_state_dict(checkpoint["model"], strict=True)
+assert not missing_keys
+model.eval()
+
+# Encodec-based tokenizer: converts reference audio into discrete conditioning tokens for VALLE
+logger.info("Initializing Encodec-based tokenizer...")
+audio_tokenizer = AudioTokenizer(device)
+
+# Vocos vocoder: decodes VALLE's discrete acoustic codes back into a 24 kHz waveform
+vocos = Vocos.from_pretrained("charactr/vocos-encodec-24khz").to(device)
+
+# initialize ASR model
+OUTPUT_DIR_WHISPER = "./models/whisper"
+if platform.system().lower() == "linux":
+    OUTPUT_DIR_WHISPER = "/app/models/whisper"
+
+if not os.path.exists(OUTPUT_DIR_WHISPER):
+    os.makedirs(OUTPUT_DIR_WHISPER, exist_ok=True)
+try:
+    logger.info("Loading Whisper model...")
+    model_name = "tiny"
+    whisper_model = whisper.load_model(
+        model_name, download_root=OUTPUT_DIR_WHISPER
+    ).cpu()
+    logger.info("Whisper model loaded successfully")
+except Exception as e:
+    logging.info(e)
+    raise Exception(
+        "\n Whisper download failed or damaged, please go to "
+        "'https://openaipublic.azureedge.net/main/whisper/models/345ae4da62f9b3d59415adc60127b97c714f32e89e936602e85993674d08dcb1/medium.pt'"
+        "\n manually download model and put it to {} .".format(os.getcwd() + "/whisper")
+    )
+
+# Initialize Voice Presets
+logger.info("Initializing Voice Presets...")
+PRESETS_DIR = "apps/audio_cloning/vallex/presets"
+preset_list = os.walk(PRESETS_DIR).__next__()[2]
+preset_list = [preset[:-4] for preset in preset_list if preset.endswith(".npz")]
+
+
+def clear_prompts():
+    try:
+        path = tempfile.gettempdir()
+        for eachfile in os.listdir(path):
+            filename = os.path.join(path, eachfile)
+            if os.path.isfile(filename) and filename.endswith(".npz"):
+                lastmodifytime = os.stat(filename).st_mtime
+                endfiletime = time.time() - 60
+                if endfiletime > lastmodifytime:
+                    os.remove(filename)
+    except Exception as e:
+        logger.error("Error clearing prompts: %s", e)
+        return
+
+
+def transcribe_one(model, audio_path):
+    # load audio and pad/trim it to fit 30 seconds
+    audio = whisper.load_audio(audio_path)
+    audio = whisper.pad_or_trim(audio)
+
+    # make log-Mel spectrogram and move to the same device as the model
+    mel = whisper.log_mel_spectrogram(audio).to(model.device)
+
+    # detect the spoken language
+    _, probs = model.detect_language(mel)
+    print(f"Detected language: {max(probs, key=probs.get)}")
+    lang = max(probs, key=probs.get)
+    # decode the audio
+    options = whisper.DecodingOptions(
+        temperature=1.0,
+        best_of=5,
+        fp16=False if device == torch.device("cpu") else True,
+        sample_len=150,
+    )
+    result = whisper.decode(model, mel, options)
+
+    # print the recognized text
+    print(result.text)
+
+    text_pr = result.text
+    if text_pr.strip(" ")[-1] not in "?!.,。，？！。、":
+        text_pr += "."
+    return lang, text_pr
+
+
+def make_npz_prompt(name, uploaded_audio, recorded_audio, transcript_content):
+    global model, text_collater, text_tokenizer, audio_tokenizer
+    clear_prompts()
+    audio_prompt = uploaded_audio if uploaded_audio is not None else recorded_audio
+    sr, wav_pr = audio_prompt
+    if not isinstance(wav_pr, torch.FloatTensor):
+        wav_pr = torch.FloatTensor(wav_pr)
+    if wav_pr.abs().max() > 1:
+        wav_pr /= wav_pr.abs().max()
+    if wav_pr.size(-1) == 2:
+        wav_pr = wav_pr[:, 0]
+    if wav_pr.ndim == 1:
+        wav_pr = wav_pr.unsqueeze(0)
+    assert wav_pr.ndim and wav_pr.size(0) == 1
+
+    if transcript_content == "":
+        text_pr, lang_pr = make_prompt(name, wav_pr, sr, save=False)
+    else:
+        lang_pr = langid.classify(str(transcript_content))[0]
+        lang_token = lang2token[lang_pr]
+        text_pr = f"{lang_token}{str(transcript_content)}{lang_token}"
+    # tokenize audio
+    encoded_frames = tokenize_audio(audio_tokenizer, (wav_pr, sr))
+    audio_tokens = encoded_frames[0][0].transpose(2, 1).cpu().numpy()
+
+    # tokenize text
+    phonemes, _ = text_tokenizer.tokenize(text=f"{text_pr}".strip())
+    text_tokens, enroll_x_lens = text_collater([phonemes])
+
+    message = f"Detected language: {lang_pr}\n Detected text {text_pr}\n"
+
+    # save as npz file
+    np.savez(
+        os.path.join(tempfile.gettempdir(), f"{name}.npz"),
+        audio_tokens=audio_tokens,
+        text_tokens=text_tokens,
+        lang_code=lang2code[lang_pr],
+    )
+    return message, os.path.join(tempfile.gettempdir(), f"{name}.npz")
+
+
+def make_prompt(name, wav, sr, save=True):
+    global whisper_model
+    whisper_model.to(device)
+    if not isinstance(wav, torch.FloatTensor):
+        wav = torch.tensor(wav)
+    if wav.abs().max() > 1:
+        wav /= wav.abs().max()
+    if wav.size(-1) == 2:
+        wav = wav.mean(-1, keepdim=False)
+    if wav.ndim == 1:
+        wav = wav.unsqueeze(0)
+    assert wav.ndim and wav.size(0) == 1
+    torchaudio.save(f"./prompts/{name}.wav", wav, sr)
+    lang, text = transcribe_one(whisper_model, f"./prompts/{name}.wav")
+    lang_token = lang2token[lang]
+    text = lang_token + text + lang_token
+    with open(f"./prompts/{name}.txt", "w", encoding="utf-8") as f:
+        f.write(text)
+    if not save:
+        os.remove(f"./prompts/{name}.wav")
+        os.remove(f"./prompts/{name}.txt")
+
+    whisper_model.cpu()
+    torch.cuda.empty_cache()
+    return text, lang
+
+
+@torch.no_grad()
+def infer_from_audio(
+    text, language, accent, audio_prompt, record_audio_prompt, transcript_content
+):
+    global model, text_collater, text_tokenizer, audio_tokenizer
+    audio_prompt = audio_prompt if audio_prompt is not None else record_audio_prompt
+    sr, wav_pr = audio_prompt
+    if not isinstance(wav_pr, torch.FloatTensor):
+        wav_pr = torch.FloatTensor(wav_pr)
+    if wav_pr.abs().max() > 1:
+        wav_pr /= wav_pr.abs().max()
+    if wav_pr.size(-1) == 2:
+        wav_pr = wav_pr[:, 0]
+    if wav_pr.ndim == 1:
+        wav_pr = wav_pr.unsqueeze(0)
+    assert wav_pr.ndim and wav_pr.size(0) == 1
+
+    if transcript_content == "":
+        text_pr, lang_pr = make_prompt("dummy", wav_pr, sr, save=False)
+    else:
+        lang_pr = langid.classify(str(transcript_content))[0]
+        lang_token = lang2token[lang_pr]
+        text_pr = f"{lang_token}{str(transcript_content)}{lang_token}"
+
+    if language == "auto-detect":
+        lang_token = lang2token[langid.classify(text)[0]]
+    else:
+        lang_token = langdropdown2token[language]
+    lang = token2lang[lang_token]
+    text = lang_token + text + lang_token
+
+    # onload model
+    model.to(device)
+
+    # tokenize audio
+    encoded_frames = tokenize_audio(audio_tokenizer, (wav_pr, sr))
+    audio_prompts = encoded_frames[0][0].transpose(2, 1).to(device)
+
+    # tokenize text
+    logging.info(f"synthesize text: {text}")
+    phone_tokens, langs = text_tokenizer.tokenize(text=f"_{text}".strip())
+    text_tokens, text_tokens_lens = text_collater([phone_tokens])
+
+    enroll_x_lens = None
+    if text_pr:
+        text_prompts, _ = text_tokenizer.tokenize(text=f"{text_pr}".strip())
+        text_prompts, enroll_x_lens = text_collater([text_prompts])
+    text_tokens = torch.cat([text_prompts, text_tokens], dim=-1)
+    text_tokens_lens += enroll_x_lens
+    lang = lang if accent == "no-accent" else token2lang[langdropdown2token[accent]]
+    encoded_frames = model.inference(
+        text_tokens.to(device),
+        text_tokens_lens.to(device),
+        audio_prompts,
+        enroll_x_lens=enroll_x_lens,
+        top_k=-100,
+        temperature=1,
+        prompt_language=lang_pr,
+        text_language=langs if accent == "no-accent" else lang,
+        best_of=5,
+    )
+    # Decode with Vocos
+    frames = encoded_frames.permute(2, 0, 1)
+    features = vocos.codes_to_features(frames)
+    samples = vocos.decode(features, bandwidth_id=torch.tensor([2], device=device))
+
+    # offload model
+    model.to("cpu")
+    torch.cuda.empty_cache()
+
+    message = f"text prompt: {text_pr}\nsythesized text: {text}"
+    return message, (24000, samples.squeeze(0).cpu().numpy())
+
+
+def main():
+    app = gr.Blocks(title="VALL-E X")
+    with app:
+        gr.Markdown(top_ja_md)
+        with gr.Tab("Infer from audio"):
+            gr.Markdown(infer_from_audio_ja_md)
+            with gr.Row():
+                with gr.Column():
+                    textbox = gr.TextArea(
+                        label="音声合成で喋らせたいテキスト",
+                        # placeholder="Type your sentence here",
+                        placeholder="ここに音声合成で喋らせたいテキストを入力してください。",
+                        value="Welcome back, Master. What can I do for you today?",
+                        elem_id="tts-input",
+                    )
+                    language_dropdown = gr.Dropdown(
+                        choices=["auto-detect", "English", "中文", "日本語"],
+                        value="auto-detect",
+                        label="language",
+                    )
+                    accent_dropdown = gr.Dropdown(
+                        choices=["no-accent", "English", "中文", "日本語"],
+                        value="no-accent",
+                        label="accent",
+                    )
+                    textbox_transcript = gr.TextArea(
+                        label="Transcript",
+                        # placeholder="Write transcript here. (leave empty to use whisper)",
+                        placeholder="アップロードした音声、または録音した音声のテキストを入力してください。(whisper を使用する場合は空のままにしてください。)",
+                        value="",
+                        elem_id="prompt-name",
+                    )
+                    upload_audio_prompt = gr.Audio(
+                        label="音声アップロード",
+                        sources=["upload"],
+                        interactive=True,
+                    )
+                    record_audio_prompt = gr.Audio(
+                        label="音声を録音する",
+                        sources=["microphone"],
+                        interactive=True,
+                    )
+                with gr.Column():
+                    text_output = gr.Textbox(label="Message")
+                    audio_output = gr.Audio(label="Output Audio", elem_id="tts-audio")
+                    btn = gr.Button("音声合成を開始する")
+                    btn.click(
+                        infer_from_audio,
+                        inputs=[
+                            textbox,
+                            language_dropdown,
+                            accent_dropdown,
+                            upload_audio_prompt,
+                            record_audio_prompt,
+                            textbox_transcript,
+                        ],
+                        outputs=[text_output, audio_output],
+                    )
+                    textbox_mp = gr.TextArea(
+                        label="Prompt name",
+                        placeholder="Name your prompt here",
+                        value="prompt_1",
+                        elem_id="prompt-name",
+                    )
+                    btn_mp = gr.Button("Make prompt!")
+                    prompt_output = gr.File(interactive=False)
+                    btn_mp.click(
+                        make_npz_prompt,
+                        inputs=[
+                            textbox_mp,
+                            upload_audio_prompt,
+                            record_audio_prompt,
+                            textbox_transcript,
+                        ],
+                        outputs=[text_output, prompt_output],
+                    )
+            gr.Examples(
+                examples=infer_from_audio_examples,
+                inputs=[
+                    textbox,
+                    language_dropdown,
+                    accent_dropdown,
+                    upload_audio_prompt,
+                    record_audio_prompt,
+                    textbox_transcript,
+                ],
+                outputs=[text_output, audio_output],
+                fn=infer_from_audio,
+                cache_examples=False,
+            )

apps/audio_cloning/vallex/models/__init__.py

@@ -0,0 +1,127 @@
+import argparse
+
+import torch.nn as nn
+
+from .transformer import Transformer
+from .vallex import VALLE, VALLF
+
+
+def add_model_arguments(parser: argparse.ArgumentParser):
+    parser.add_argument(
+        "--model-name",
+        type=str,
+        default="VALL-E",
+        help="VALL-E, VALL-F, Transformer.",
+    )
+    parser.add_argument(
+        "--decoder-dim",
+        type=int,
+        default=1024,
+        help="Embedding dimension in the decoder model.",
+    )
+    parser.add_argument(
+        "--nhead",
+        type=int,
+        default=16,
+        help="Number of attention heads in the Decoder layers.",
+    )
+    parser.add_argument(
+        "--num-decoder-layers",
+        type=int,
+        default=12,
+        help="Number of Decoder layers.",
+    )
+    parser.add_argument(
+        "--scale-factor",
+        type=float,
+        default=1.0,
+        help="Model scale factor which will be assigned different meanings in different models.",
+    )
+    parser.add_argument(
+        "--norm-first",
+        type=bool,
+        default=True,
+        help="Pre or Post Normalization.",
+    )
+    parser.add_argument(
+        "--add-prenet",
+        type=bool,
+        default=False,
+        help="Whether add PreNet after Inputs.",
+    )
+
+    # VALL-E & F
+    parser.add_argument(
+        "--prefix-mode",
+        type=int,
+        default=1,
+        help="The mode for how to prefix VALL-E NAR Decoder, "
+        "0: no prefix, 1: 0 to random, 2: random to random, 4: chunk of pre or post utterance.",
+    )
+    parser.add_argument(
+        "--share-embedding",
+        type=bool,
+        default=True,
+        help="Share the parameters of the output projection layer with the parameters of the acoustic embedding.",
+    )
+    parser.add_argument(
+        "--prepend-bos",
+        type=bool,
+        default=False,
+        help="Whether prepend <BOS> to the acoustic tokens -> AR Decoder inputs.",
+    )
+    parser.add_argument(
+        "--num-quantizers",
+        type=int,
+        default=8,
+        help="Number of Audio/Semantic quantization layers.",
+    )
+
+    # Transformer
+    parser.add_argument(
+        "--scaling-xformers",
+        type=bool,
+        default=False,
+        help="Apply Reworked Conformer scaling on Transformers.",
+    )
+
+
+def get_model(params) -> nn.Module:
+    if params.model_name.lower() in ["vall-f", "vallf"]:
+        model = VALLF(
+            params.decoder_dim,
+            params.nhead,
+            params.num_decoder_layers,
+            norm_first=params.norm_first,
+            add_prenet=params.add_prenet,
+            prefix_mode=params.prefix_mode,
+            share_embedding=params.share_embedding,
+            nar_scale_factor=params.scale_factor,
+            prepend_bos=params.prepend_bos,
+            num_quantizers=params.num_quantizers,
+        )
+    elif params.model_name.lower() in ["vall-e", "valle"]:
+        model = VALLE(
+            params.decoder_dim,
+            params.nhead,
+            params.num_decoder_layers,
+            norm_first=params.norm_first,
+            add_prenet=params.add_prenet,
+            prefix_mode=params.prefix_mode,
+            share_embedding=params.share_embedding,
+            nar_scale_factor=params.scale_factor,
+            prepend_bos=params.prepend_bos,
+            num_quantizers=params.num_quantizers,
+        )
+    else:
+        assert params.model_name in ["Transformer"]
+        model = Transformer(
+            params.decoder_dim,
+            params.nhead,
+            params.num_decoder_layers,
+            norm_first=params.norm_first,
+            add_prenet=params.add_prenet,
+            scaling_xformers=params.scaling_xformers,
+        )
+
+    return model

apps/audio_cloning/vallex/models/macros.py

@@ -0,0 +1,11 @@
+# Text
+NUM_TEXT_TOKENS = 2048
+
+# Audio
+NUM_AUDIO_TOKENS = 1024  # EnCodec RVQ bins
+NUM_MEL_BINS = 100  # BigVGAN bigvgan_24khz_100band
+
+
+# Speaker
+NUM_SPEAKER_CLASSES = 4096
+SPEAKER_EMBEDDING_DIM = 64

apps/audio_cloning/vallex/models/transformer.py

@@ -0,0 +1,386 @@
+# Copyright    2023                             (authors: Feiteng Li)
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from functools import partial
+from typing import Any, Dict, List, Tuple, Union
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+from ..modules.embedding import SinePositionalEmbedding, TokenEmbedding
+from ..modules.scaling import BalancedDoubleSwish, ScaledLinear
+from ..modules.transformer import (
+    BalancedBasicNorm,
+    IdentityNorm,
+    TransformerDecoderLayer,
+    TransformerEncoder,
+    TransformerEncoderLayer,
+)
+from .macros import NUM_MEL_BINS, NUM_TEXT_TOKENS
+
+# from icefall.utils import make_pad_mask
+# from torchmetrics.classification import BinaryAccuracy
+from .vallex import Transpose
+from .visualizer import visualize
+
+IdentityNorm = IdentityNorm
+
+
+class Transformer(nn.Module):
+    """It implements seq2seq Transformer TTS for debug(No StopPredictor and SpeakerEmbeding)
+    Neural Speech Synthesis with Transformer Network
+    https://arxiv.org/abs/1809.08895
+    """
+
+    def __init__(
+        self,
+        d_model: int,
+        nhead: int,
+        num_layers: int,
+        norm_first: bool = True,
+        add_prenet: bool = False,
+        scaling_xformers: bool = False,
+    ):
+        """
+        Args:
+          d_model:
+            The number of expected features in the input (required).
+          nhead:
+            The number of heads in the multiheadattention models (required).
+          num_layers:
+            The number of sub-decoder-layers in the decoder (required).
+        """
+        super().__init__()
+        self.text_embedding = TokenEmbedding(d_model, NUM_TEXT_TOKENS)  # W_x
+
+        if add_prenet:
+            self.encoder_prenet = nn.Sequential(
+                Transpose(),
+                nn.Conv1d(d_model, d_model, kernel_size=5, padding="same"),
+                nn.BatchNorm1d(d_model),
+                nn.ReLU(),
+                nn.Dropout(0.5),
+                nn.Conv1d(d_model, d_model, kernel_size=5, padding="same"),
+                nn.BatchNorm1d(d_model),
+                nn.ReLU(),
+                nn.Dropout(0.5),
+                nn.Conv1d(d_model, d_model, kernel_size=5, padding="same"),
+                nn.BatchNorm1d(d_model),
+                nn.ReLU(),
+                nn.Dropout(0.5),
+                Transpose(),
+                nn.Linear(d_model, d_model),
+            )
+
+            self.decoder_prenet = nn.Sequential(
+                nn.Linear(NUM_MEL_BINS, 256),
+                nn.ReLU(),
+                nn.Dropout(0.5),
+                nn.Linear(256, 256),
+                nn.ReLU(),
+                nn.Dropout(0.5),
+                nn.Linear(256, d_model),
+            )
+
+            assert scaling_xformers is False  # TODO: update this block
+        else:
+            self.encoder_prenet = nn.Identity()
+            if scaling_xformers:
+                self.decoder_prenet = ScaledLinear(NUM_MEL_BINS, d_model)
+            else:
+                self.decoder_prenet = nn.Linear(NUM_MEL_BINS, d_model)
+
+        self.encoder_position = SinePositionalEmbedding(
+            d_model,
+            dropout=0.1,
+            scale=False,
+        )
+        self.decoder_position = SinePositionalEmbedding(
+            d_model, dropout=0.1, scale=False
+        )
+
+        if scaling_xformers:
+            self.encoder = TransformerEncoder(
+                TransformerEncoderLayer(
+                    d_model,
+                    nhead,
+                    dim_feedforward=d_model * 4,
+                    dropout=0.1,
+                    batch_first=True,
+                    norm_first=norm_first,
+                    linear1_self_attention_cls=ScaledLinear,
+                    linear2_self_attention_cls=partial(
+                        ScaledLinear, initial_scale=0.01
+                    ),
+                    linear1_feedforward_cls=ScaledLinear,
+                    linear2_feedforward_cls=partial(ScaledLinear, initial_scale=0.01),
+                    activation=partial(
+                        BalancedDoubleSwish,
+                        channel_dim=-1,
+                        max_abs=10.0,
+                        min_prob=0.25,
+                    ),
+                    layer_norm_cls=IdentityNorm,
+                ),
+                num_layers=num_layers,
+                norm=BalancedBasicNorm(d_model) if norm_first else None,
+            )
+
+            self.decoder = nn.TransformerDecoder(
+                TransformerDecoderLayer(
+                    d_model,
+                    nhead,
+                    dim_feedforward=d_model * 4,
+                    dropout=0.1,
+                    batch_first=True,
+                    norm_first=norm_first,
+                    linear1_self_attention_cls=ScaledLinear,
+                    linear2_self_attention_cls=partial(
+                        ScaledLinear, initial_scale=0.01
+                    ),
+                    linear1_feedforward_cls=ScaledLinear,
+                    linear2_feedforward_cls=partial(ScaledLinear, initial_scale=0.01),
+                    activation=partial(
+                        BalancedDoubleSwish,
+                        channel_dim=-1,
+                        max_abs=10.0,
+                        min_prob=0.25,
+                    ),
+                    layer_norm_cls=IdentityNorm,
+                ),
+                num_layers=num_layers,
+                norm=BalancedBasicNorm(d_model) if norm_first else None,
+            )
+
+            self.predict_layer = ScaledLinear(d_model, NUM_MEL_BINS)
+            self.stop_layer = nn.Linear(d_model, 1)
+        else:
+            self.encoder = nn.TransformerEncoder(
+                nn.TransformerEncoderLayer(
+                    d_model,
+                    nhead,
+                    dim_feedforward=d_model * 4,
+                    activation=F.relu,
+                    dropout=0.1,
+                    batch_first=True,
+                    norm_first=norm_first,
+                ),
+                num_layers=num_layers,
+                norm=nn.LayerNorm(d_model) if norm_first else None,
+            )
+
+            self.decoder = nn.TransformerDecoder(
+                nn.TransformerDecoderLayer(
+                    d_model,
+                    nhead,
+                    dim_feedforward=d_model * 4,
+                    activation=F.relu,
+                    dropout=0.1,
+                    batch_first=True,
+                    norm_first=norm_first,
+                ),
+                num_layers=num_layers,
+                norm=nn.LayerNorm(d_model) if norm_first else None,
+            )
+
+            self.predict_layer = nn.Linear(d_model, NUM_MEL_BINS)
+            self.stop_layer = nn.Linear(d_model, 1)
+
+        self.stop_accuracy_metric = BinaryAccuracy(
+            threshold=0.5, multidim_average="global"
+        )
+
+    #     self.apply(self._init_weights)
+
+    # def _init_weights(self, module):
+    #     if isinstance(module, (nn.Linear)):
+    #         module.weight.data.normal_(mean=0.0, std=0.02)
+    #         if isinstance(module, nn.Linear) and module.bias is not None:
+    #             module.bias.data.zero_()
+    #     elif isinstance(module, nn.LayerNorm):
+    #         module.bias.data.zero_()
+    #         module.weight.data.fill_(1.0)
+    #     elif isinstance(module, nn.Embedding):
+    #         module.weight.data.normal_(mean=0.0, std=0.02)
+
+    def forward(
+        self,
+        x: torch.Tensor,
+        x_lens: torch.Tensor,
+        y: torch.Tensor,
+        y_lens: torch.Tensor,
+        reduction: str = "sum",
+        train_stage: int = 0,
+        **kwargs,
+    ) -> Tuple[torch.Tensor, Union[torch.Tensor, None]]:
+        """
+        Args:
+          x:
+            A 2-D tensor of shape (N, S).
+          x_lens:
+            A 1-D tensor of shape (N,). It contains the number of tokens in `x`
+            before padding.
+          y:
+            A 3-D tensor of shape (N, T, 8).
+          y_lens:
+            A 1-D tensor of shape (N,). It contains the number of tokens in `x`
+            before padding.
+          train_stage:
+            Not used in this model.
+        Returns:
+          Return the predicted audio code matrix, cross-entropy loss and Top-10 accuracy.
+        """
+        del train_stage
+
+        assert x.ndim == 2, x.shape
+        assert x_lens.ndim == 1, x_lens.shape
+        assert y.ndim == 3, y.shape
+        assert y_lens.ndim == 1, y_lens.shape
+
+        assert torch.all(x_lens > 0)
+
+        # NOTE: x has been padded in TextTokenCollater
+        x_mask = make_pad_mask(x_lens).to(x.device)
+
+        x = self.text_embedding(x)
+        x = self.encoder_prenet(x)
+        x = self.encoder_position(x)
+        x = self.encoder(x, src_key_padding_mask=x_mask)
+
+        total_loss, metrics = 0.0, {}
+
+        y_mask = make_pad_mask(y_lens).to(y.device)
+        y_mask_float = y_mask.type(torch.float32)
+        data_mask = 1.0 - y_mask_float.unsqueeze(-1)
+
+        # Training
+        # AR Decoder
+        def pad_y(y):
+            y = F.pad(y, (0, 0, 1, 0, 0, 0), value=0).detach()
+            # inputs, targets
+            return y[:, :-1], y[:, 1:]
+
+        y, targets = pad_y(y * data_mask)  # mask padding as zeros
+
+        y_emb = self.decoder_prenet(y)
+        y_pos = self.decoder_position(y_emb)
+
+        y_len = y_lens.max()
+        tgt_mask = torch.triu(
+            torch.ones(y_len, y_len, device=y.device, dtype=torch.bool),
+            diagonal=1,
+        )
+        y_dec = self.decoder(
+            y_pos,
+            x,
+            tgt_mask=tgt_mask,
+            memory_key_padding_mask=x_mask,
+        )
+
+        predict = self.predict_layer(y_dec)
+        # loss
+        total_loss = F.mse_loss(predict, targets, reduction=reduction)
+
+        logits = self.stop_layer(y_dec).squeeze(-1)
+        stop_loss = F.binary_cross_entropy_with_logits(
+            logits,
+            y_mask_float.detach(),
+            weight=1.0 + y_mask_float.detach() * 4.0,
+            reduction=reduction,
+        )
+        metrics["stop_loss"] = stop_loss.detach()
+
+        stop_accuracy = self.stop_accuracy_metric(
+            (torch.sigmoid(logits) >= 0.5).type(torch.int64),
+            y_mask.type(torch.int64),
+        )
+        # icefall MetricsTracker.norm_items()
+        metrics["stop_accuracy"] = stop_accuracy.item() * y_lens.sum().type(
+            torch.float32
+        )
+
+        return ((x, predict), total_loss + 100.0 * stop_loss, metrics)
+
+    def inference(
+        self,
+        x: torch.Tensor,
+        x_lens: torch.Tensor,
+        y: Any = None,
+        **kwargs,
+    ) -> torch.Tensor:
+        """
+        Args:
+          x:
+            A 2-D tensor of shape (1, S).
+          x_lens:
+            A 1-D tensor of shape (1,). It contains the number of tokens in `x`
+            before padding.
+        Returns:
+          Return the predicted audio code matrix and cross-entropy loss.
+        """
+        assert x.ndim == 2, x.shape
+        assert x_lens.ndim == 1, x_lens.shape
+
+        assert torch.all(x_lens > 0)
+
+        x_mask = make_pad_mask(x_lens).to(x.device)
+
+        x = self.text_embedding(x)
+        x = self.encoder_prenet(x)
+        x = self.encoder_position(x)
+        x = self.encoder(x, src_key_padding_mask=x_mask)
+
+        x_mask = make_pad_mask(x_lens).to(x.device)
+
+        # AR Decoder
+        # TODO: Managing decoder steps avoid repetitive computation
+        y = torch.zeros(
+            [x.shape[0], 1, NUM_MEL_BINS], dtype=torch.float32, device=x.device
+        )
+        while True:
+            y_emb = self.decoder_prenet(y)
+            y_pos = self.decoder_position(y_emb)
+
+            tgt_mask = torch.triu(
+                torch.ones(y.shape[1], y.shape[1], device=y.device, dtype=torch.bool),
+                diagonal=1,
+            )
+
+            y_dec = self.decoder(
+                y_pos,
+                x,
+                tgt_mask=tgt_mask,
+                memory_mask=None,
+                memory_key_padding_mask=x_mask,
+            )
+            predict = self.predict_layer(y_dec[:, -1:])
+
+            logits = self.stop_layer(y_dec[:, -1:]) > 0  # sigmoid(0.0) = 0.5
+            if y.shape[1] > x_lens.max() * 10 or all(logits.cpu().numpy()):
+                print(f"TransformerTTS EOS [Text {x_lens[0]} -> Audio {y.shape[1]}]")
+                break
+
+            y = torch.concat([y, predict], dim=1)
+
+        return y[:, 1:]
+
+    def visualize(
+        self,
+        predicts: Tuple[torch.Tensor],
+        batch: Dict[str, Union[List, torch.Tensor]],
+        output_dir: str,
+        limit: int = 4,
+    ) -> None:
+        visualize(predicts, batch, output_dir, limit=limit)

apps/audio_cloning/vallex/models/vallex.py

@@ -0,0 +1,823 @@
+# Copyright    2023                             (authors: Feiteng Li)
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import random
+from typing import Dict, Iterator, List, Tuple, Union
+
+import numpy as np
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+# from icefall.utils import make_pad_mask
+# from torchmetrics.classification import MulticlassAccuracy
+from ..data.input_strategies import PromptedFeatures
+from ..modules.embedding import SinePositionalEmbedding, TokenEmbedding
+from ..modules.transformer import (
+    AdaptiveLayerNorm,
+    LayerNorm,
+    TransformerDecoderLayer,
+    TransformerEncoder,
+    TransformerEncoderLayer,
+)
+from .macros import NUM_AUDIO_TOKENS, NUM_TEXT_TOKENS
+
+
+class Transpose(nn.Identity):
+    """(N, T, D) -> (N, D, T)"""
+
+    def forward(self, input: torch.Tensor) -> torch.Tensor:
+        return input.transpose(1, 2)
+
+
+# NOTE: There are two ways to implement the model
+#       1) [VALL-F] standard TransformerDecoder, use x as memory
+#       2) [VALL-E] modified TransformerDecoder like GPT-x(e.g. causal TransformerEncoder),
+#          use x as the prefix of decoder inputs
+class VALLF(nn.Module):
+    """It implements https://arxiv.org/abs/2301.02111
+    "Neural Codec Language Models are Zero-Shot Text to Speech Synthesizers"
+    """
+
+    def __init__(
+        self,
+        d_model: int,
+        nhead: int,
+        num_layers: int,
+        norm_first: bool = True,
+        add_prenet: bool = False,
+        decoder_cls: Union[
+            nn.TransformerDecoder, nn.TransformerEncoder
+        ] = nn.TransformerDecoder,
+        decoder_layer_cls: Union[
+            TransformerDecoderLayer, TransformerEncoderLayer
+        ] = TransformerDecoderLayer,
+        prefix_mode: int = 0,
+        share_embedding: bool = True,
+        nar_scale_factor: float = 1.0,
+        prepend_bos: bool = True,
+        num_quantizers: int = 8,
+    ):
+        """
+        Args:
+          d_model:
+            The number of expected features in the input (required).
+          nhead:
+            The number of heads in the multiheadattention models (required).
+          num_layers:
+            The number of sub-decoder-layers in the decoder (required).
+        """
+        super().__init__()
+        nar_d_model = int(d_model * nar_scale_factor)
+
+        self.ar_text_embedding = TokenEmbedding(d_model, NUM_TEXT_TOKENS)  # W_x
+        self.nar_text_embedding = TokenEmbedding(nar_d_model, NUM_TEXT_TOKENS)
+
+        # ID NUM_AUDIO_TOKENS     -> PAD
+        # ID NUM_AUDIO_TOKENS + 1 -> BOS
+        self.ar_audio_prepend_bos = prepend_bos
+        self.ar_audio_embedding = TokenEmbedding(
+            d_model, NUM_AUDIO_TOKENS + 1 + int(prepend_bos)
+        )
+
+        # PreNet
+        if add_prenet:
+            self.ar_text_prenet = nn.Sequential(
+                Transpose(),
+                nn.Conv1d(d_model, d_model, kernel_size=5, padding="same"),
+                nn.BatchNorm1d(d_model),
+                nn.ReLU(),
+                nn.Dropout(0.5),
+                nn.Conv1d(d_model, d_model, kernel_size=5, padding="same"),
+                nn.BatchNorm1d(d_model),
+                nn.ReLU(),
+                nn.Dropout(0.5),
+                nn.Conv1d(d_model, d_model, kernel_size=5, padding="same"),
+                nn.BatchNorm1d(d_model),
+                nn.ReLU(),
+                nn.Dropout(0.5),
+                Transpose(),
+                nn.Linear(d_model, d_model),
+            )
+
+            self.ar_audio_prenet = nn.Sequential(
+                nn.Linear(d_model, 256),
+                nn.ReLU(),
+                nn.Dropout(0.25),
+                nn.Linear(256, 256),
+                nn.ReLU(),
+                nn.Dropout(0.25),
+                nn.Linear(256, d_model),
+            )
+        else:
+            self.ar_text_prenet = nn.Identity()
+            self.ar_audio_prenet = nn.Identity()
+
+        self.ar_text_position = SinePositionalEmbedding(
+            d_model,
+            dropout=0.1,
+            scale=False,
+            alpha=True,
+        )
+        self.ar_audio_position = SinePositionalEmbedding(
+            d_model,
+            dropout=0.1,
+            scale=False,
+            alpha=True,
+        )
+
+        self.ar_decoder = decoder_cls(
+            decoder_layer_cls(
+                d_model,
+                nhead,
+                dim_feedforward=d_model * 4,
+                dropout=0.1,
+                batch_first=True,
+                norm_first=norm_first,
+            ),
+            num_layers=num_layers,
+            norm=LayerNorm(d_model) if norm_first else None,
+        )
+        self.ar_predict_layer = nn.Linear(d_model, NUM_AUDIO_TOKENS + 1, bias=False)
+
+        self.rng = random.Random(0)
+        self.num_heads = nhead
+        self.prefix_mode = prefix_mode
+        self.num_quantizers = num_quantizers
+
+        assert num_quantizers >= 1
+        if num_quantizers > 1:
+            self.nar_audio_embeddings = nn.ModuleList(
+                [TokenEmbedding(nar_d_model, NUM_AUDIO_TOKENS + 1)]
+                + [
+                    TokenEmbedding(nar_d_model, NUM_AUDIO_TOKENS)
+                    for i in range(num_quantizers - 1)
+                ]
+            )  # W_a
+
+            # PreNet
+            if add_prenet:
+                self.nar_text_prenet = nn.Sequential(
+                    Transpose(),
+                    nn.Conv1d(nar_d_model, nar_d_model, kernel_size=5, padding="same"),
+                    nn.BatchNorm1d(nar_d_model),
+                    nn.ReLU(),
+                    nn.Dropout(0.5),
+                    nn.Conv1d(nar_d_model, nar_d_model, kernel_size=5, padding="same"),
+                    nn.BatchNorm1d(nar_d_model),
+                    nn.ReLU(),
+                    nn.Dropout(0.5),
+                    nn.Conv1d(nar_d_model, nar_d_model, kernel_size=5, padding="same"),
+                    nn.BatchNorm1d(nar_d_model),
+                    nn.ReLU(),
+                    nn.Dropout(0.5),
+                    Transpose(),
+                    nn.Linear(nar_d_model, nar_d_model),
+                )
+                self.nar_audio_prenet = nn.Sequential(
+                    nn.Linear(nar_d_model, 256),
+                    nn.ReLU(),
+                    nn.Dropout(0.25),
+                    nn.Linear(256, 256),
+                    nn.ReLU(),
+                    nn.Dropout(0.25),
+                    nn.Linear(256, nar_d_model),
+                )
+            else:
+                self.nar_text_prenet = nn.Identity()
+                self.nar_audio_prenet = nn.Identity()
+
+            self.nar_text_position = SinePositionalEmbedding(
+                nar_d_model,
+                dropout=0.0,
+                scale=False,
+                alpha=False,
+            )
+            self.nar_audio_position = SinePositionalEmbedding(
+                nar_d_model,
+                dropout=0.1,
+                scale=False,
+                alpha=False,
+            )
+
+            self.nar_decoder = decoder_cls(
+                decoder_layer_cls(
+                    nar_d_model,
+                    int(nhead * nar_scale_factor),
+                    dim_feedforward=nar_d_model * 4,
+                    dropout=0.1,
+                    batch_first=True,
+                    norm_first=norm_first,
+                    adaptive_layer_norm=True,
+                ),
+                num_layers=int(num_layers * nar_scale_factor),
+                norm=AdaptiveLayerNorm(nar_d_model, norm=nn.LayerNorm(nar_d_model))
+                if norm_first
+                else None,
+            )
+            self.nar_predict_layers = nn.ModuleList(
+                [
+                    nn.Linear(nar_d_model, NUM_AUDIO_TOKENS, bias=False)
+                    for i in range(num_quantizers - 1)
+                ]
+            )
+            self.nar_stage_embeddings = nn.ModuleList(
+                [TokenEmbedding(nar_d_model, 1) for i in range(num_quantizers - 1)]
+            )
+
+            if share_embedding:
+                # We share the parameters of the output projection layer with the parameters of the acoustic embedding Wa
+                # NOTE(Feiteng): In the experiment, this undermines accuracy
+                # self.ar_predict_layer.weight = self.ar_audio_embedding.weight
+
+                # We also share the parameters of the acoustic embedding layer and the output prediction layer,
+                # which means the weights of the j-th prediction layer are the same as the (j + 1)-th acoustic embedding layer.
+                for j in range(0, num_quantizers - 2):
+                    self.nar_predict_layers[j].weight = self.nar_audio_embeddings[
+                        j + 2
+                    ].weight
+
+    def stage_parameters(self, stage: int = 1) -> Iterator[nn.Parameter]:
+        assert stage > 0
+        if stage == 1:
+            for name, param in self.named_parameters():
+                if name.startswith("ar_"):
+                    print(f" AR parameter: {name}")
+                    yield param
+
+        if stage == 2:
+            for name, param in self.named_parameters():
+                if name.startswith("nar_"):
+                    print(f"NAR parameter: {name}")
+                    yield param
+
+    def stage_named_parameters(
+        self, stage: int = 1
+    ) -> Iterator[Tuple[str, nn.Parameter]]:
+        assert stage > 0
+        if stage == 1:
+            for pair in self.named_parameters():
+                if pair[0].startswith("ar_"):
+                    yield pair
+
+        if stage == 2:
+            for pair in self.named_parameters():
+                if pair[0].startswith("nar_"):
+                    yield pair
+
+    def pad_y_eos(self, y, y_mask_int, eos_id):
+        targets = F.pad(y, (0, 1), value=0) + eos_id * F.pad(
+            y_mask_int, (0, 1), value=1
+        )
+        # inputs, targets
+        if self.ar_audio_prepend_bos:
+            return (
+                F.pad(targets[:, :-1], (1, 0), value=NUM_AUDIO_TOKENS + 1),
+                targets,
+            )
+
+        return targets[:, :-1], targets[:, 1:]
+
+    def _prepare_prompts(
+        self, y, y_lens, codes, nar_stage, y_prompts_codes, prefix_mode
+    ):
+        # 5.1 For the NAR acoustic prompt tokens, we select a random segment waveform of 3 seconds
+        # from the same utterance.
+        # We implement this differently.
+        if prefix_mode == 0:
+            # no prefix
+            prefix_len = 0
+            y_emb = self.nar_audio_embeddings[0](y)
+            for j in range(1, nar_stage):
+                # Formula (4) (5)
+                y_emb = y_emb + self.nar_audio_embeddings[j](codes[..., j])
+        elif prefix_mode == 1:
+            # prefix at begining
+            int_low = (0.25 * y_lens.min()).type(torch.int64).item()
+            prefix_len = torch.randint(0, int_low * 2, size=()).item()
+            prefix_len = min(prefix_len, 225)  # 24000/320 * 3s = 225 frames
+
+            y_prompts = self.nar_audio_embeddings[0](y[:, :prefix_len])
+            y_emb = self.nar_audio_embeddings[0](y[:, prefix_len:])
+            for j in range(1, self.num_quantizers):
+                y_prompts += self.nar_audio_embeddings[j](codes[:, :prefix_len, j])
+                if j < nar_stage:
+                    y_emb += self.nar_audio_embeddings[j](codes[:, prefix_len:, j])
+            y_emb = torch.concat([y_prompts, y_emb], axis=1)
+        elif prefix_mode in [2, 4]:
+            if prefix_mode == 2:
+                # random prefix
+                prefix_len = min(225, int(0.25 * y_lens.min().item()))
+
+                y_prompts_codes = []
+                for b in range(codes.shape[0]):
+                    start = self.rng.randint(0, y_lens[b].item() - prefix_len)
+                    y_prompts_codes.append(
+                        torch.clone(codes[b, start : start + prefix_len])
+                    )
+                    codes[b, start : start + prefix_len, nar_stage] = NUM_AUDIO_TOKENS
+                y_prompts_codes = torch.stack(y_prompts_codes, dim=0)
+            else:
+                prefix_len = y_prompts_codes.shape[1]
+
+            y_prompts = self.nar_audio_embeddings[0](y_prompts_codes[..., 0])
+            y_emb = self.nar_audio_embeddings[0](y)
+            for j in range(1, self.num_quantizers):
+                y_prompts += self.nar_audio_embeddings[j](y_prompts_codes[..., j])
+                if j < nar_stage:
+                    y_emb += self.nar_audio_embeddings[j](codes[..., j])
+            y_emb = torch.concat([y_prompts, y_emb], axis=1)
+        else:
+            raise ValueError
+
+        return y_emb, prefix_len
+
+    def forward(
+        self,
+        x: torch.Tensor,
+        x_lens: torch.Tensor,
+        y: Union[torch.Tensor, PromptedFeatures],
+        y_lens: Union[torch.Tensor, PromptedFeatures],
+        reduction: str = "sum",
+        train_stage: int = 0,
+        **kwargs,
+    ) -> Tuple[torch.Tensor, Union[torch.Tensor, None]]:
+        raise NotImplementedError
+
+    def inference(
+        self,
+        x: torch.Tensor,
+        x_lens: torch.Tensor,
+        y: torch.Tensor,
+        enroll_x_lens: Union[torch.Tensor, None] = None,
+        top_k: int = -100,
+        temperature: float = 1.0,
+    ) -> torch.Tensor:
+        raise NotImplementedError
+
+    def visualize(
+        self,
+        predicts: Tuple[torch.Tensor],
+        batch: Dict[str, Union[List, torch.Tensor]],
+        output_dir: str,
+        limit: int = 4,
+    ) -> None:
+        raise NotImplementedError
+
+
+class VALLE(VALLF):
+    """It implements https://arxiv.org/abs/2301.02111
+    "Neural Codec Language Models are Zero-Shot Text to Speech Synthesizers"
+    """
+
+    def __init__(
+        self,
+        d_model: int,
+        nhead: int,
+        num_layers: int,
+        norm_first: bool = True,
+        add_prenet: bool = False,
+        prefix_mode: int = 0,
+        share_embedding: bool = True,
+        nar_scale_factor: float = 1.0,
+        **kwargs,
+    ):
+        """
+        Args:
+          d_model:
+            The number of expected features in the input (required).
+          nhead:
+            The number of heads in the multiheadattention models (required).
+          num_layers:
+            The number of sub-decoder-layers in the decoder (required).
+        """
+        super(VALLE, self).__init__(
+            d_model,
+            nhead,
+            num_layers,
+            norm_first=norm_first,
+            add_prenet=add_prenet,
+            decoder_cls=TransformerEncoder,
+            decoder_layer_cls=TransformerEncoderLayer,
+            prefix_mode=prefix_mode,
+            share_embedding=share_embedding,
+            nar_scale_factor=nar_scale_factor,
+            **kwargs,
+        )
+        self.language_ID = {
+            "en": 0,
+            "zh": 1,
+            "ja": 2,
+        }
+        self.ar_language_embedding = TokenEmbedding(d_model, len(self.language_ID))
+        self.nar_language_embedding = TokenEmbedding(d_model, len(self.language_ID))
+
+    def forward(
+        self,
+        x: torch.Tensor,
+        x_lens: torch.Tensor,
+        y: Union[torch.Tensor, PromptedFeatures],
+        y_lens: Union[torch.Tensor, PromptedFeatures],
+        reduction: str = "sum",
+        train_stage: int = 0,
+        **kwargs,
+    ):
+        raise NotImplementedError
+
+    def inference(
+        self,
+        x: torch.Tensor,
+        x_lens: torch.Tensor,
+        y: torch.Tensor,
+        enroll_x_lens: torch.Tensor,
+        top_k: int = -100,
+        temperature: float = 1.0,
+        prompt_language: str = None,
+        text_language: str = None,
+        best_of: int = 1,
+        length_penalty: float = 1.0,
+        return_worst: bool = False,
+    ) -> torch.Tensor:
+        """
+        Args:
+          x:
+            A 2-D tensor of shape (1, S).
+          x_lens:
+            A 1-D tensor of shape (1,). It contains the number of tokens in `x`
+            before padding.
+          y:
+            A 3-D tensor of shape (1, T, 8).
+          top_k: (`optional`) int
+            The number of highest probability tokens to keep for top-k-filtering. Default to -100.
+          temperature: (`optional`) float
+            The value used to module the next token probabilities. Must be strictly positive. Default to 1.0.
+        Returns:
+          Return the predicted audio code matrix.
+        """
+        assert x.ndim == 2, x.shape
+        assert x_lens.ndim == 1, x_lens.shape
+        assert y.ndim == 3, y.shape
+        assert y.shape[0] == 1, y.shape
+
+        assert torch.all(x_lens > 0)
+
+        # NOTE: x has been padded in TextTokenCollater
+        text = x
+        x = self.ar_text_embedding(text)
+        # Add language embedding
+        prompt_language_id = torch.LongTensor(
+            np.array([self.language_ID[prompt_language]])
+        ).to(x.device)
+        if isinstance(text_language, str):
+            text_language_id = torch.LongTensor(
+                np.array([self.language_ID[text_language]])
+            ).to(x.device)
+        elif isinstance(text_language, List):
+            text_language_id = torch.LongTensor(
+                np.array([self.language_ID[tl] for tl in text_language])
+            ).to(x.device)
+        x[:, :enroll_x_lens, :] += self.ar_language_embedding(prompt_language_id)
+        x[:, enroll_x_lens:, :] += self.ar_language_embedding(text_language_id)
+        x = self.ar_text_prenet(x)
+        x = self.ar_text_position(x)
+
+        text_len = x_lens.max()
+        prompts = y
+        prefix_len = y.shape[1]
+
+        # AR Decoder
+        # TODO: Managing decoder steps avoid repetitive computation
+        y = prompts[..., 0]
+        if self.ar_audio_prepend_bos:
+            y = F.pad(y, (1, 0), value=NUM_AUDIO_TOKENS + 1)
+
+        x_len = x_lens.max()
+        x_attn_mask = torch.zeros((x_len, x_len), dtype=torch.bool)
+
+        kv_cache = None
+        use_kv_caching = True
+
+        sum_logprobs = torch.zeros(
+            best_of, device=y.device
+        )  # implement batch decoding here
+        x = x.repeat(best_of, 1, 1)
+        y = y.repeat(best_of, 1)
+        while True:
+            y_emb = self.ar_audio_embedding(y)
+            y_emb = self.ar_audio_prenet(y_emb)
+            y_pos = self.ar_audio_position(y_emb)
+            xy_pos = torch.concat([x, y_pos], dim=1)
+
+            y_len = y.shape[1]
+            x_attn_mask_pad = F.pad(
+                x_attn_mask,
+                (0, y_len),
+                value=True,
+            )
+            y_attn_mask = F.pad(
+                torch.triu(torch.ones(y_len, y_len, dtype=torch.bool), diagonal=1),
+                (x_len, 0),
+                value=False,
+            )
+            xy_attn_mask = torch.concat([x_attn_mask_pad, y_attn_mask], dim=0).to(
+                y.device
+            )
+
+            if use_kv_caching and kv_cache is not None:
+                xy_pos = xy_pos[:, [-1]]
+            else:
+                pass
+
+            xy_dec, kv_cache = self.ar_decoder.infer(
+                xy_pos,
+                mask=xy_attn_mask,
+                past_kv=kv_cache,
+                use_cache=use_kv_caching,
+            )
+            # xy_dec, _ = self.ar_decoder(
+            #     (xy_pos, None),
+            #     mask=xy_attn_mask,
+            # )
+
+            logits = self.ar_predict_layer(xy_dec[:, -1])
+            samples, current_logprobs = topk_sampling(
+                logits, top_k=top_k, top_p=1, temperature=temperature
+            )
+            sum_logprobs += current_logprobs * (y[:, -1] != NUM_AUDIO_TOKENS)
+            samples[y[:, -1] == NUM_AUDIO_TOKENS] = NUM_AUDIO_TOKENS
+            completed = (samples[:, -1] == NUM_AUDIO_TOKENS).all()
+            if completed or (y.shape[1] - prompts.shape[1]) > x_lens.max() * 16:
+                if prompts.shape[1] == y.shape[1]:
+                    raise SyntaxError("well trained model shouldn't reach here.")
+                lengths = torch.sum(y != NUM_AUDIO_TOKENS, dim=1)
+                avg_logprobs = sum_logprobs / lengths**length_penalty
+                # choose the best beam according to sum_logprobs
+                best_beam = y[torch.argmax(avg_logprobs), :]
+                worst_beam = y[torch.argmin(avg_logprobs), :]
+                # strip all eos tokens
+                best_beam = best_beam[best_beam != NUM_AUDIO_TOKENS]
+                worst_beam = worst_beam[worst_beam != NUM_AUDIO_TOKENS]
+                if return_worst:
+                    y = worst_beam.unsqueeze(0)
+                else:
+                    y = best_beam.unsqueeze(0)
+                print(f"VALL-E EOS [{prompts.shape[1]} -> {y.shape[1]}]")
+                break
+
+            y = torch.concat([y, samples], dim=1)
+
+        codes = [y[:, prefix_len + int(self.ar_audio_prepend_bos) :]]
+        if self.num_quantizers == 1:
+            return torch.stack(codes, dim=-1)
+
+        # Non-AR Decoders
+        y_emb = self.nar_audio_embeddings[0](y[:, int(self.ar_audio_prepend_bos) :])
+
+        if self.prefix_mode in [2, 4]:  # Exclude enrolled_phonemes
+            enrolled_len = enroll_x_lens.max().item()
+            # SOS + Synthesis Text + EOS
+            text = torch.concat(
+                [
+                    text[:, :1],
+                    text[:, enrolled_len - 1 :],
+                ],
+                dim=1,
+            )
+            text_len = text_len - (enrolled_len - 2)
+            assert text.shape[0] == 1
+
+        x = self.nar_text_embedding(text)
+        # Add language embedding
+        prompt_language_id = torch.LongTensor(
+            np.array([self.language_ID[prompt_language]])
+        ).to(x.device)
+        if isinstance(text_language, str):
+            text_language_id = torch.LongTensor(
+                np.array([self.language_ID[text_language]])
+            ).to(x.device)
+        elif isinstance(text_language, List):
+            text_language_id = torch.LongTensor(
+                np.array([self.language_ID[tl] for tl in text_language])
+            ).to(x.device)
+        x[:, :enroll_x_lens, :] += self.nar_language_embedding(prompt_language_id)
+        x[:, enroll_x_lens:, :] += self.nar_language_embedding(text_language_id)
+        x = self.nar_text_prenet(x)
+        x = self.nar_text_position(x)
+
+        if self.prefix_mode == 0:
+            for i, (predict_layer, embedding_layer) in enumerate(
+                zip(
+                    self.nar_predict_layers,
+                    self.nar_audio_embeddings[1:],
+                )
+            ):
+                y_pos = self.nar_audio_prenet(y_emb)
+                y_pos = self.nar_audio_position(y_pos)
+                xy_pos = torch.concat([x, y_pos], dim=1)
+
+                xy_dec, _ = self.nar_decoder(
+                    (xy_pos, self.nar_stage_embeddings[i].weight)
+                )
+                logits = predict_layer(xy_dec[:, text_len + prefix_len :])
+
+                samples = torch.argmax(logits, dim=-1)
+                codes.append(samples)
+
+                if i < self.num_quantizers - 2:
+                    y_emb[:, :prefix_len] += embedding_layer(prompts[..., i + 1])
+                    y_emb[:, prefix_len:] += embedding_layer(samples)
+        else:
+            for j in range(1, self.num_quantizers):
+                y_emb[:, :prefix_len] += self.nar_audio_embeddings[j](prompts[..., j])
+
+            for i, (predict_layer, embedding_layer) in enumerate(
+                zip(
+                    self.nar_predict_layers,
+                    self.nar_audio_embeddings[1:],
+                )
+            ):
+                y_pos = self.nar_audio_prenet(y_emb)
+                y_pos = self.nar_audio_position(y_pos)
+                xy_pos = torch.concat([x, y_pos], dim=1)
+
+                xy_dec, _ = self.nar_decoder(
+                    (xy_pos, self.nar_stage_embeddings[i].weight)
+                )
+                logits = predict_layer(xy_dec[:, text_len + prefix_len :])
+
+                samples = torch.argmax(logits, dim=-1)
+                codes.append(samples)
+
+                if i < self.num_quantizers - 2:
+                    y_emb[:, prefix_len:] += embedding_layer(samples)
+
+        assert len(codes) == self.num_quantizers
+        return torch.stack(codes, dim=-1)
+
+    def continual(
+        self,
+        x: torch.Tensor,
+        x_lens: torch.Tensor,
+        y: torch.Tensor,
+    ) -> torch.Tensor:
+        """
+        Args:
+          x:
+            A 2-D tensor of shape (1, S).
+          x_lens:
+            A 1-D tensor of shape (1,). It contains the number of tokens in `x`
+            before padding.
+          y:
+            A 3-D tensor of shape (1, T, 8).
+        Returns:
+          Return the predicted audio code matrix.
+        """
+        assert x.ndim == 2, x.shape
+        assert x_lens.ndim == 1, x_lens.shape
+        assert y.ndim == 3, y.shape
+        assert y.shape[0] == 1, y.shape
+
+        assert torch.all(x_lens > 0)
+        assert self.num_quantizers == 8
+
+        # NOTE: x has been padded in TextTokenCollater
+        text = x
+        x = self.ar_text_embedding(text)
+        x = self.ar_text_prenet(x)
+        x = self.ar_text_position(x)
+
+        text_len = x_lens.max()
+
+        prefix_len = min(int(y.shape[1] * 0.5), 3 * 75)
+
+        # AR Decoder
+        prompts = y[:, :prefix_len]
+
+        codes = [y[:, prefix_len:, 0]]
+        # Non-AR Decoders
+        x = self.nar_text_embedding(text)
+        x = self.nar_text_prenet(x)
+        x = self.nar_text_position(x)
+
+        y_emb = self.nar_audio_embeddings[0](y[..., 0])
+
+        if self.prefix_mode == 0:
+            for i, (predict_layer, embedding_layer) in enumerate(
+                zip(
+                    self.nar_predict_layers,
+                    self.nar_audio_embeddings[1:],
+                )
+            ):
+                y_pos = self.nar_audio_position(y_emb)
+                y_pos = self.nar_audio_prenet(y_pos)
+                xy_pos = torch.concat([x, y_pos], dim=1)
+
+                xy_dec, _ = self.nar_decoder(
+                    (xy_pos, self.nar_stage_embeddings[i].weight)
+                )
+                logits = predict_layer(xy_dec[:, text_len + prefix_len :])
+
+                samples = torch.argmax(logits, dim=-1)
+                codes.append(samples)
+
+                if i < 6:
+                    y_emb[:, :prefix_len] += embedding_layer(prompts[..., i + 1])
+                    y_emb[:, prefix_len:] += embedding_layer(samples)
+        else:
+            for j in range(1, 8):
+                y_emb[:, :prefix_len] += self.nar_audio_embeddings[j](prompts[..., j])
+
+            for i, (predict_layer, embedding_layer) in enumerate(
+                zip(
+                    self.nar_predict_layers,
+                    self.nar_audio_embeddings[1:],
+                )
+            ):
+                y_pos = self.nar_audio_prenet(y_emb)
+                y_pos = self.nar_audio_position(y_pos)
+                xy_pos = torch.concat([x, y_pos], dim=1)
+
+                xy_dec, _ = self.nar_decoder(
+                    (xy_pos, self.nar_stage_embeddings[i].weight)
+                )
+                logits = predict_layer(xy_dec[:, text_len + prefix_len :])
+
+                samples = torch.argmax(logits, dim=-1)
+                codes.append(samples)
+
+                if i < 6:
+                    y_emb[:, prefix_len:] += embedding_layer(samples)
+
+        assert len(codes) == 8
+        return torch.stack(codes, dim=-1)
+
+
+# https://github.com/microsoft/unilm/blob/master/xtune/src/transformers/modeling_utils.py
+def top_k_top_p_filtering(
+    logits, top_k=0, top_p=1.0, filter_value=-float("Inf"), min_tokens_to_keep=1
+):
+    """Filter a distribution of logits using top-k and/or nucleus (top-p) filtering
+    Args:
+        logits: logits distribution shape (batch size, vocabulary size)
+        if top_k > 0: keep only top k tokens with highest probability (top-k filtering).
+        if top_p < 1.0: keep the top tokens with cumulative probability >= top_p (nucleus filtering).
+            Nucleus filtering is described in Holtzman et al. (http://arxiv.org/abs/1904.09751)
+        Make sure we keep at least min_tokens_to_keep per batch example in the output
+    From: https://gist.github.com/thomwolf/1a5a29f6962089e871b94cbd09daf317
+    """
+    if top_k > 0:
+        top_k = min(max(top_k, min_tokens_to_keep), logits.size(-1))  # Safety check
+        # Remove all tokens with a probability less than the last token of the top-k
+        indices_to_remove = logits < torch.topk(logits, top_k)[0][..., -1, None]
+        logits[indices_to_remove] = filter_value
+
+    if top_p < 1.0:
+        sorted_logits, sorted_indices = torch.sort(logits, descending=True)
+        cumulative_probs = torch.cumsum(F.softmax(sorted_logits, dim=-1), dim=-1)
+
+        # Remove tokens with cumulative probability above the threshold (token with 0 are kept)
+        sorted_indices_to_remove = cumulative_probs > top_p
+        if min_tokens_to_keep > 1:
+            # Keep at least min_tokens_to_keep (set to min_tokens_to_keep-1 because we add the first one below)
+            sorted_indices_to_remove[..., :min_tokens_to_keep] = 0
+        # Shift the indices to the right to keep also the first token above the threshold
+        sorted_indices_to_remove[..., 1:] = sorted_indices_to_remove[..., :-1].clone()
+        sorted_indices_to_remove[..., 0] = 0
+
+        # scatter sorted tensors to original indexing
+        indices_to_remove = sorted_indices_to_remove.scatter(
+            1, sorted_indices, sorted_indices_to_remove
+        )
+        logits[indices_to_remove] = filter_value
+    return logits
+
+
+def topk_sampling(logits, top_k=10, top_p=1.0, temperature=1.0):
+    # temperature: (`optional`) float
+    #     The value used to module the next token probabilities. Must be strictly positive. Default to 1.0.
+    # top_k: (`optional`) int
+    #     The number of highest probability vocabulary tokens to keep for top-k-filtering. Between 1 and infinity. Default to 50.
+    # top_p: (`optional`) float
+    #     The cumulative probability of parameter highest probability vocabulary tokens to keep for nucleus sampling. Must be between 0 and 1. Default to 1.
+
+    # Temperature (higher temperature => more likely to sample low probability tokens)
+    if temperature != 1.0:
+        logits = logits / temperature
+    # Top-p/top-k filtering
+    logits = top_k_top_p_filtering(logits, top_k=top_k, top_p=top_p)
+    # Sample
+    token = torch.multinomial(F.softmax(logits, dim=-1), num_samples=1)
+    logprobs = F.log_softmax(logits.float(), dim=-1)
+    current_logprobs = logprobs[torch.arange(logprobs.shape[0]), token.squeeze(1)]
+    return token, current_logprobs

apps/audio_cloning/vallex/models/visualizer.py

@@ -0,0 +1,102 @@
+#!/usr/bin/env python3
+# Copyright    2023                           (authors: Feiteng Li)
+#
+# See ../../../../LICENSE for clarification regarding multiple authors
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+
+from typing import Dict, List, Tuple, Union
+
+import matplotlib.pyplot as plt
+import numpy as np
+import torch
+
+
+def visualize(
+    predicts: Tuple[torch.Tensor],
+    batch: Dict[str, Union[List, torch.Tensor]],
+    output_dir: str,
+    limit: int = 4,
+) -> None:
+    text_tokens = batch["text_tokens"].to("cpu").detach().numpy()
+    text_tokens_lens = batch["text_tokens_lens"].to("cpu").detach().numpy()
+    audio_features = batch["audio_features"].to("cpu").detach().numpy()
+    audio_features_lens = batch["audio_features_lens"].to("cpu").detach().numpy()
+    assert text_tokens.ndim == 2
+
+    utt_ids, texts = batch["utt_id"], batch["text"]
+
+    encoder_outputs = predicts[0].to("cpu").type(torch.float32).detach().numpy()
+    decoder_outputs = predicts[1]
+    if isinstance(decoder_outputs, list):
+        decoder_outputs = decoder_outputs[-1]
+    decoder_outputs = decoder_outputs.to("cpu").type(torch.float32).detach().numpy()
+
+    vmin, vmax = 0, 1024  # Encodec
+    if decoder_outputs.dtype == np.float32:
+        vmin, vmax = -6, 0  # Fbank
+
+    num_figures = 3
+    for b, (utt_id, text) in enumerate(zip(utt_ids[:limit], texts[:limit])):
+        _ = plt.figure(figsize=(14, 8 * num_figures))
+
+        S = text_tokens_lens[b]
+        T = audio_features_lens[b]
+
+        # encoder
+        plt.subplot(num_figures, 1, 1)
+        plt.title(f"Text: {text}")
+        plt.imshow(
+            X=np.transpose(encoder_outputs[b]),
+            cmap=plt.get_cmap("jet"),
+            aspect="auto",
+            interpolation="nearest",
+        )
+        plt.gca().invert_yaxis()
+        plt.axvline(x=S - 0.4, linewidth=2, color="r")
+        plt.xlabel("Encoder Output")
+        plt.colorbar()
+
+        # decoder
+        plt.subplot(num_figures, 1, 2)
+        plt.imshow(
+            X=np.transpose(decoder_outputs[b]),
+            cmap=plt.get_cmap("jet"),
+            aspect="auto",
+            interpolation="nearest",
+            vmin=vmin,
+            vmax=vmax,
+        )
+        plt.gca().invert_yaxis()
+        plt.axvline(x=T - 0.4, linewidth=2, color="r")
+        plt.xlabel("Decoder Output")
+        plt.colorbar()
+
+        # target
+        plt.subplot(num_figures, 1, 3)
+        plt.imshow(
+            X=np.transpose(audio_features[b]),
+            cmap=plt.get_cmap("jet"),
+            aspect="auto",
+            interpolation="nearest",
+            vmin=vmin,
+            vmax=vmax,
+        )
+        plt.gca().invert_yaxis()
+        plt.axvline(x=T - 0.4, linewidth=2, color="r")
+        plt.xlabel("Decoder Target")
+        plt.colorbar()
+
+        plt.savefig(f"{output_dir}/{utt_id}.png")
+        plt.close()

apps/audio_cloning/vallex/modules/activation.py

@@ -0,0 +1,612 @@
+from typing import Optional, Tuple, List
+import math
+
+import torch
+from torch import Tensor
+from torch.nn import Linear, Module
+from torch.nn import functional as F
+from torch.nn.init import constant_, xavier_normal_, xavier_uniform_
+from torch.nn.modules.linear import NonDynamicallyQuantizableLinear
+from torch.nn.parameter import Parameter
+
+def _in_projection_packed(
+    q: Tensor,
+    k: Tensor,
+    v: Tensor,
+    w: Tensor,
+    b: Optional[Tensor] = None,
+) -> List[Tensor]:
+    r"""
+    Performs the in-projection step of the attention operation, using packed weights.
+    Output is a triple containing projection tensors for query, key and value.
+
+    Args:
+        q, k, v: query, key and value tensors to be projected. For self-attention,
+            these are typically the same tensor; for encoder-decoder attention,
+            k and v are typically the same tensor. (We take advantage of these
+            identities for performance if they are present.) Regardless, q, k and v
+            must share a common embedding dimension; otherwise their shapes may vary.
+        w: projection weights for q, k and v, packed into a single tensor. Weights
+            are packed along dimension 0, in q, k, v order.
+        b: optional projection biases for q, k and v, packed into a single tensor
+            in q, k, v order.
+
+    Shape:
+        Inputs:
+        - q: :math:`(..., E)` where E is the embedding dimension
+        - k: :math:`(..., E)` where E is the embedding dimension
+        - v: :math:`(..., E)` where E is the embedding dimension
+        - w: :math:`(E * 3, E)` where E is the embedding dimension
+        - b: :math:`E * 3` where E is the embedding dimension
+
+        Output:
+        - in output list :math:`[q', k', v']`, each output tensor will have the
+            same shape as the corresponding input tensor.
+    """
+    E = q.size(-1)
+    if k is v:
+        if q is k:
+            # self-attention
+            return F.linear(q, w, b).chunk(3, dim=-1)
+        else:
+            # encoder-decoder attention
+            w_q, w_kv = w.split([E, E * 2])
+            if b is None:
+                b_q = b_kv = None
+            else:
+                b_q, b_kv = b.split([E, E * 2])
+            return (F.linear(q, w_q, b_q),) + F.linear(k, w_kv, b_kv).chunk(2, dim=-1)
+    else:
+        w_q, w_k, w_v = w.chunk(3)
+        if b is None:
+            b_q = b_k = b_v = None
+        else:
+            b_q, b_k, b_v = b.chunk(3)
+        return F.linear(q, w_q, b_q), F.linear(k, w_k, b_k), F.linear(v, w_v, b_v)
+
+def _scaled_dot_product_attention(
+    q: Tensor,
+    k: Tensor,
+    v: Tensor,
+    attn_mask: Optional[Tensor] = None,
+    dropout_p: float = 0.0,
+) -> Tuple[Tensor, Tensor]:
+    r"""
+    Computes scaled dot product attention on query, key and value tensors, using
+    an optional attention mask if passed, and applying dropout if a probability
+    greater than 0.0 is specified.
+    Returns a tensor pair containing attended values and attention weights.
+
+    Args:
+        q, k, v: query, key and value tensors. See Shape section for shape details.
+        attn_mask: optional tensor containing mask values to be added to calculated
+            attention. May be 2D or 3D; see Shape section for details.
+        dropout_p: dropout probability. If greater than 0.0, dropout is applied.
+
+    Shape:
+        - q: :math:`(B, Nt, E)` where B is batch size, Nt is the target sequence length,
+            and E is embedding dimension.
+        - key: :math:`(B, Ns, E)` where B is batch size, Ns is the source sequence length,
+            and E is embedding dimension.
+        - value: :math:`(B, Ns, E)` where B is batch size, Ns is the source sequence length,
+            and E is embedding dimension.
+        - attn_mask: either a 3D tensor of shape :math:`(B, Nt, Ns)` or a 2D tensor of
+            shape :math:`(Nt, Ns)`.
+
+        - Output: attention values have shape :math:`(B, Nt, E)`; attention weights
+            have shape :math:`(B, Nt, Ns)`
+    """
+    B, Nt, E = q.shape
+    q = q / math.sqrt(E)
+    # (B, Nt, E) x (B, E, Ns) -> (B, Nt, Ns)
+    if attn_mask is not None:
+        attn = torch.baddbmm(attn_mask, q, k.transpose(-2, -1))
+    else:
+        attn = torch.bmm(q, k.transpose(-2, -1))
+
+    attn = F.softmax(attn, dim=-1)
+    if dropout_p > 0.0:
+        attn = F.dropout(attn, p=dropout_p)
+    # (B, Nt, Ns) x (B, Ns, E) -> (B, Nt, E)
+    output = torch.bmm(attn, v)
+    return output, attn
+
+def multi_head_attention_forward(
+        x,
+        ipw,
+        ipb,
+        opw,
+        opb,
+        n_head,
+        attn_mask,
+        past_kv=None,
+        use_cache=False,
+):
+    # x = x.transpose(1, 0)
+    # tgt_len, bsz, embed_dim = x.shape
+    # head_dim = embed_dim // n_head
+    # q, k, v = _in_projection_packed(x, x, x, ipw, ipb)
+    # q = q.contiguous().view(tgt_len, bsz * n_head, head_dim).transpose(0, 1)
+    # k = k.contiguous().view(k.shape[0], bsz * n_head, head_dim).transpose(0, 1)
+    # v = v.contiguous().view(v.shape[0], bsz * n_head, head_dim).transpose(0, 1)
+
+    # new_attn_mask = torch.zeros_like(attn_mask, dtype=q.dtype)
+    # new_attn_mask.masked_fill_(attn_mask, float("-inf"))
+    # attn_mask = new_attn_mask
+    #
+    # attn_output, attn_output_weights = _scaled_dot_product_attention(q, k, v, attn_mask, 0.0)
+    # attn_output = attn_output.transpose(0, 1).contiguous().view(tgt_len * bsz, embed_dim)
+    # attn_output = torch._C._nn.linear(attn_output, opw, opb)
+    # attn_output = attn_output.view(tgt_len, bsz, attn_output.size(1))
+
+    B, T, C = x.size()
+
+    q, k, v = torch._C._nn.linear(x, ipw, ipb).chunk(3, dim=-1)
+    k = k.view(B, T, n_head, C // n_head).transpose(1, 2)  # (B, nh, T, hs)
+    q = q.view(B, T, n_head, C // n_head).transpose(1, 2)  # (B, nh, T, hs)
+    v = v.view(B, T, n_head, C // n_head).transpose(1, 2)  # (B, nh, T, hs)
+    if past_kv is not None:
+        past_key = past_kv[0]
+        past_value = past_kv[1]
+        k = torch.cat((past_key, k), dim=-2)
+        v = torch.cat((past_value, v), dim=-2)
+
+    FULL_T = k.shape[-2]
+
+    if use_cache is True:
+        present = (k, v)
+    else:
+        present = None
+
+    att = (q @ k.transpose(-2, -1)) * (1.0 / math.sqrt(k.size(-1)))
+    att = att.masked_fill(attn_mask[FULL_T - T:FULL_T, :FULL_T], float('-inf'))
+    att = F.softmax(att, dim=-1)
+    y = att @ v  # (B, nh, T, T) x (B, nh, T, hs) -> (B, nh, T, hs)
+    y = y.transpose(1, 2).contiguous().view(B, T, C)  # re-assemble all head outputs side by side
+    y = torch._C._nn.linear(y, opw, opb)
+    return (y, present)
+
+
+class MultiheadAttention(Module):
+    r"""Allows the model to jointly attend to information
+    from different representation subspaces as described in the paper:
+    `Attention Is All You Need <https://arxiv.org/abs/1706.03762>`_.
+
+    Multi-Head Attention is defined as:
+
+    .. math::
+        \text{MultiHead}(Q, K, V) = \text{Concat}(head_1,\dots,head_h)W^O
+
+    where :math:`head_i = \text{Attention}(QW_i^Q, KW_i^K, VW_i^V)`.
+
+    ``forward()`` will use a special optimized implementation if all of the following
+    conditions are met:
+
+    - self attention is being computed (i.e., ``query``, ``key``, and ``value`` are the same tensor. This
+      restriction will be loosened in the future.)
+    - Either autograd is disabled (using ``torch.inference_mode`` or ``torch.no_grad``) or no tensor argument ``requires_grad``
+    - training is disabled (using ``.eval()``)
+    - dropout is 0
+    - ``add_bias_kv`` is ``False``
+    - ``add_zero_attn`` is ``False``
+    - ``batch_first`` is ``True`` and the input is batched
+    - ``kdim`` and ``vdim`` are equal to ``embed_dim``
+    - at most one of ``key_padding_mask`` or ``attn_mask`` is passed
+    - if a `NestedTensor <https://pytorch.org/docs/stable/nested.html>`_ is passed, neither ``key_padding_mask``
+      nor ``attn_mask`` is passed
+
+    If the optimized implementation is in use, a
+    `NestedTensor <https://pytorch.org/docs/stable/nested.html>`_ can be passed for
+    ``query``/``key``/``value`` to represent padding more efficiently than using a
+    padding mask. In this case, a `NestedTensor <https://pytorch.org/docs/stable/nested.html>`_
+    will be returned, and an additional speedup proportional to the fraction of the input
+    that is padding can be expected.
+
+    Args:
+        embed_dim: Total dimension of the model.
+        num_heads: Number of parallel attention heads. Note that ``embed_dim`` will be split
+            across ``num_heads`` (i.e. each head will have dimension ``embed_dim // num_heads``).
+        dropout: Dropout probability on ``attn_output_weights``. Default: ``0.0`` (no dropout).
+        bias: If specified, adds bias to input / output projection layers. Default: ``True``.
+        add_bias_kv: If specified, adds bias to the key and value sequences at dim=0. Default: ``False``.
+        add_zero_attn: If specified, adds a new batch of zeros to the key and value sequences at dim=1.
+            Default: ``False``.
+        kdim: Total number of features for keys. Default: ``None`` (uses ``kdim=embed_dim``).
+        vdim: Total number of features for values. Default: ``None`` (uses ``vdim=embed_dim``).
+        batch_first: If ``True``, then the input and output tensors are provided
+            as (batch, seq, feature). Default: ``False`` (seq, batch, feature).
+
+    Examples::
+
+        >>> # xdoctest: +SKIP
+        >>> multihead_attn = nn.MultiheadAttention(embed_dim, num_heads)
+        >>> attn_output, attn_output_weights = multihead_attn(query, key, value)
+
+    """
+    __constants__ = ["batch_first"]
+    bias_k: Optional[torch.Tensor]
+    bias_v: Optional[torch.Tensor]
+
+    def __init__(
+            self,
+            embed_dim,
+            num_heads,
+            dropout=0.0,
+            bias=True,
+            add_bias_kv=False,
+            add_zero_attn=False,
+            kdim=None,
+            vdim=None,
+            batch_first=False,
+            linear1_cls=Linear,
+            linear2_cls=Linear,
+            device=None,
+            dtype=None,
+    ) -> None:
+        factory_kwargs = {"device": device, "dtype": dtype}
+        super(MultiheadAttention, self).__init__()
+        self.embed_dim = embed_dim
+        self.kdim = kdim if kdim is not None else embed_dim
+        self.vdim = vdim if vdim is not None else embed_dim
+        self._qkv_same_embed_dim = (
+                self.kdim == embed_dim and self.vdim == embed_dim
+        )
+
+        self.num_heads = num_heads
+        self.dropout = dropout
+        self.batch_first = batch_first
+        self.head_dim = embed_dim // num_heads
+        assert (
+                self.head_dim * num_heads == self.embed_dim
+        ), "embed_dim must be divisible by num_heads"
+
+        if add_bias_kv:
+            self.bias_k = Parameter(
+                torch.empty((1, 1, embed_dim), **factory_kwargs)
+            )
+            self.bias_v = Parameter(
+                torch.empty((1, 1, embed_dim), **factory_kwargs)
+            )
+        else:
+            self.bias_k = self.bias_v = None
+
+        if linear1_cls == Linear:
+            if not self._qkv_same_embed_dim:
+                self.q_proj_weight = Parameter(
+                    torch.empty((embed_dim, embed_dim), **factory_kwargs)
+                )
+                self.k_proj_weight = Parameter(
+                    torch.empty((embed_dim, self.kdim), **factory_kwargs)
+                )
+                self.v_proj_weight = Parameter(
+                    torch.empty((embed_dim, self.vdim), **factory_kwargs)
+                )
+                self.register_parameter("in_proj_weight", None)
+            else:
+                self.in_proj_weight = Parameter(
+                    torch.empty((3 * embed_dim, embed_dim), **factory_kwargs)
+                )
+                self.register_parameter("q_proj_weight", None)
+                self.register_parameter("k_proj_weight", None)
+                self.register_parameter("v_proj_weight", None)
+
+            if bias:
+                self.in_proj_bias = Parameter(
+                    torch.empty(3 * embed_dim, **factory_kwargs)
+                )
+            else:
+                self.register_parameter("in_proj_bias", None)
+            self.out_proj = NonDynamicallyQuantizableLinear(
+                embed_dim, embed_dim, bias=bias, **factory_kwargs
+            )
+
+            self._reset_parameters()
+        else:
+            if not self._qkv_same_embed_dim:
+                raise NotImplementedError
+            else:
+                self.in_proj_linear = linear1_cls(
+                    embed_dim, 3 * embed_dim, bias=bias, **factory_kwargs
+                )
+                self.in_proj_weight = self.in_proj_linear.weight
+
+                self.register_parameter("q_proj_weight", None)
+                self.register_parameter("k_proj_weight", None)
+                self.register_parameter("v_proj_weight", None)
+
+                if bias:
+                    self.in_proj_bias = self.in_proj_linear.bias
+                else:
+                    self.register_parameter("in_proj_bias", None)
+
+            self.out_proj = linear2_cls(
+                embed_dim, embed_dim, bias=bias, **factory_kwargs
+            )
+
+            if self.bias_k is not None:
+                xavier_normal_(self.bias_k)
+            if self.bias_v is not None:
+                xavier_normal_(self.bias_v)
+
+        self.add_zero_attn = add_zero_attn
+
+    def _reset_parameters(self):
+        if self._qkv_same_embed_dim:
+            xavier_uniform_(self.in_proj_weight)
+        else:
+            xavier_uniform_(self.q_proj_weight)
+            xavier_uniform_(self.k_proj_weight)
+            xavier_uniform_(self.v_proj_weight)
+
+        if self.in_proj_bias is not None:
+            constant_(self.in_proj_bias, 0.0)
+            constant_(self.out_proj.bias, 0.0)
+
+        if self.bias_k is not None:
+            xavier_normal_(self.bias_k)
+        if self.bias_v is not None:
+            xavier_normal_(self.bias_v)
+
+    def __setstate__(self, state):
+        # Support loading old MultiheadAttention checkpoints generated by v1.1.0
+        if "_qkv_same_embed_dim" not in state:
+            state["_qkv_same_embed_dim"] = True
+
+        super(MultiheadAttention, self).__setstate__(state)
+
+    def forward(
+            self,
+            query: Tensor,
+            key: Tensor,
+            value: Tensor,
+            key_padding_mask: Optional[Tensor] = None,
+            need_weights: bool = True,
+            attn_mask: Optional[Tensor] = None,
+            average_attn_weights: bool = True,
+    ) -> Tuple[Tensor, Optional[Tensor]]:
+        r"""
+        Args:
+            query: Query embeddings of shape :math:`(L, E_q)` for unbatched input, :math:`(L, N, E_q)` when ``batch_first=False``
+                or :math:`(N, L, E_q)` when ``batch_first=True``, where :math:`L` is the target sequence length,
+                :math:`N` is the batch size, and :math:`E_q` is the query embedding dimension ``embed_dim``.
+                Queries are compared against key-value pairs to produce the output.
+                See "Attention Is All You Need" for more details.
+            key: Key embeddings of shape :math:`(S, E_k)` for unbatched input, :math:`(S, N, E_k)` when ``batch_first=False``
+                or :math:`(N, S, E_k)` when ``batch_first=True``, where :math:`S` is the source sequence length,
+                :math:`N` is the batch size, and :math:`E_k` is the key embedding dimension ``kdim``.
+                See "Attention Is All You Need" for more details.
+            value: Value embeddings of shape :math:`(S, E_v)` for unbatched input, :math:`(S, N, E_v)` when
+                ``batch_first=False`` or :math:`(N, S, E_v)` when ``batch_first=True``, where :math:`S` is the source
+                sequence length, :math:`N` is the batch size, and :math:`E_v` is the value embedding dimension ``vdim``.
+                See "Attention Is All You Need" for more details.
+            key_padding_mask: If specified, a mask of shape :math:`(N, S)` indicating which elements within ``key``
+                to ignore for the purpose of attention (i.e. treat as "padding"). For unbatched `query`, shape should be :math:`(S)`.
+                Binary and byte masks are supported.
+                For a binary mask, a ``True`` value indicates that the corresponding ``key`` value will be ignored for
+                the purpose of attention. For a float mask, it will be directly added to the corresponding ``key`` value.
+            need_weights: If specified, returns ``attn_output_weights`` in addition to ``attn_outputs``.
+                Default: ``True``.
+            attn_mask: If specified, a 2D or 3D mask preventing attention to certain positions. Must be of shape
+                :math:`(L, S)` or :math:`(N\cdot\text{num\_heads}, L, S)`, where :math:`N` is the batch size,
+                :math:`L` is the target sequence length, and :math:`S` is the source sequence length. A 2D mask will be
+                broadcasted across the batch while a 3D mask allows for a different mask for each entry in the batch.
+                Binary, byte, and float masks are supported. For a binary mask, a ``True`` value indicates that the
+                corresponding position is not allowed to attend. For a byte mask, a non-zero value indicates that the
+                corresponding position is not allowed to attend. For a float mask, the mask values will be added to
+                the attention weight.
+            average_attn_weights: If true, indicates that the returned ``attn_weights`` should be averaged across
+                heads. Otherwise, ``attn_weights`` are provided separately per head. Note that this flag only has an
+                effect when ``need_weights=True``. Default: ``True`` (i.e. average weights across heads)
+
+        Outputs:
+            - **attn_output** - Attention outputs of shape :math:`(L, E)` when input is unbatched,
+              :math:`(L, N, E)` when ``batch_first=False`` or :math:`(N, L, E)` when ``batch_first=True``,
+              where :math:`L` is the target sequence length, :math:`N` is the batch size, and :math:`E` is the
+              embedding dimension ``embed_dim``.
+            - **attn_output_weights** - Only returned when ``need_weights=True``. If ``average_attn_weights=True``,
+              returns attention weights averaged across heads of shape :math:`(L, S)` when input is unbatched or
+              :math:`(N, L, S)`, where :math:`N` is the batch size, :math:`L` is the target sequence length, and
+              :math:`S` is the source sequence length. If ``average_attn_weights=False``, returns attention weights per
+              head of shape :math:`(\text{num\_heads}, L, S)` when input is unbatched or :math:`(N, \text{num\_heads}, L, S)`.
+
+            .. note::
+                `batch_first` argument is ignored for unbatched inputs.
+        """
+        is_batched = query.dim() == 3
+        if key_padding_mask is not None:
+            _kpm_dtype = key_padding_mask.dtype
+            if _kpm_dtype != torch.bool and not torch.is_floating_point(
+                    key_padding_mask
+            ):
+                raise AssertionError(
+                    "only bool and floating types of key_padding_mask are supported"
+                )
+        why_not_fast_path = ""
+        if not is_batched:
+            why_not_fast_path = f"input not batched; expected query.dim() of 3 but got {query.dim()}"
+        elif query is not key or key is not value:
+            # When lifting this restriction, don't forget to either
+            # enforce that the dtypes all match or test cases where
+            # they don't!
+            why_not_fast_path = "non-self attention was used (query, key, and value are not the same Tensor)"
+        elif (
+                self.in_proj_bias is not None
+                and query.dtype != self.in_proj_bias.dtype
+        ):
+            why_not_fast_path = f"dtypes of query ({query.dtype}) and self.in_proj_bias ({self.in_proj_bias.dtype}) don't match"
+        elif (
+                self.in_proj_weight is not None
+                and query.dtype != self.in_proj_weight.dtype
+        ):
+            # this case will fail anyway, but at least they'll get a useful error message.
+            why_not_fast_path = f"dtypes of query ({query.dtype}) and self.in_proj_weight ({self.in_proj_weight.dtype}) don't match"
+        elif self.training:
+            why_not_fast_path = "training is enabled"
+        elif not self.batch_first:
+            why_not_fast_path = "batch_first was not True"
+        elif self.bias_k is not None:
+            why_not_fast_path = "self.bias_k was not None"
+        elif self.bias_v is not None:
+            why_not_fast_path = "self.bias_v was not None"
+        elif self.dropout:
+            why_not_fast_path = f"dropout was {self.dropout}, required zero"
+        elif self.add_zero_attn:
+            why_not_fast_path = "add_zero_attn was enabled"
+        elif not self._qkv_same_embed_dim:
+            why_not_fast_path = "_qkv_same_embed_dim was not True"
+        elif attn_mask is not None:
+            why_not_fast_path = "attn_mask was not None"
+        elif query.is_nested and key_padding_mask is not None:
+            why_not_fast_path = (
+                "key_padding_mask is not supported with NestedTensor input"
+            )
+        elif self.num_heads % 2 == 1:
+            why_not_fast_path = "num_heads is odd"
+        elif torch.is_autocast_enabled():
+            why_not_fast_path = "autocast is enabled"
+
+        if not why_not_fast_path:
+            tensor_args = (
+                query,
+                key,
+                value,
+                self.in_proj_weight,
+                self.in_proj_bias,
+                self.out_proj.weight,
+                self.out_proj.bias,
+            )
+            # We have to use list comprehensions below because TorchScript does not support
+            # generator expressions.
+            if torch.overrides.has_torch_function(tensor_args):
+                why_not_fast_path = "some Tensor argument has_torch_function"
+            elif not all(
+                    [
+                        (x is None or x.is_cuda or "cpu" in str(x.device))
+                        for x in tensor_args
+                    ]
+            ):
+                why_not_fast_path = (
+                    "some Tensor argument is neither CUDA nor CPU"
+                )
+            elif torch.is_grad_enabled() and any(
+                    [x is not None and x.requires_grad for x in tensor_args]
+            ):
+                why_not_fast_path = (
+                    "grad is enabled and at least one of query or the "
+                    "input/output projection weights or biases requires_grad"
+                )
+            if not why_not_fast_path:
+                return torch._native_multi_head_attention(
+                    query,
+                    key,
+                    value,
+                    self.embed_dim,
+                    self.num_heads,
+                    self.in_proj_weight,
+                    self.in_proj_bias,
+                    self.out_proj.weight,
+                    self.out_proj.bias,
+                    key_padding_mask
+                    if key_padding_mask is not None
+                    else attn_mask,
+                    need_weights,
+                    average_attn_weights,
+                    1
+                    if key_padding_mask is not None
+                    else 0
+                    if attn_mask is not None
+                    else None,
+                )
+
+        any_nested = query.is_nested or key.is_nested or value.is_nested
+        assert not any_nested, (
+                "MultiheadAttention does not support NestedTensor outside of its fast path. "
+                + f"The fast path was not hit because {why_not_fast_path}"
+        )
+
+        if self.batch_first and is_batched:
+            # make sure that the transpose op does not affect the "is" property
+            if key is value:
+                if query is key:
+                    query = key = value = query.transpose(1, 0)
+                else:
+                    query, key = [x.transpose(1, 0) for x in (query, key)]
+                    value = key
+            else:
+                query, key, value = [
+                    x.transpose(1, 0) for x in (query, key, value)
+                ]
+
+        if not self._qkv_same_embed_dim:
+            attn_output, attn_output_weights = F.multi_head_attention_forward(
+                query,
+                key,
+                value,
+                self.embed_dim,
+                self.num_heads,
+                self.in_proj_weight,
+                self.in_proj_bias,
+                self.bias_k,
+                self.bias_v,
+                self.add_zero_attn,
+                self.dropout,
+                self.out_proj.weight,
+                self.out_proj.bias,
+                training=self.training,
+                key_padding_mask=key_padding_mask,
+                need_weights=need_weights,
+                attn_mask=attn_mask,
+                use_separate_proj_weight=True,
+                q_proj_weight=self.q_proj_weight,
+                k_proj_weight=self.k_proj_weight,
+                v_proj_weight=self.v_proj_weight,
+                average_attn_weights=average_attn_weights,
+            )
+        else:
+            attn_output, attn_output_weights = F.multi_head_attention_forward(
+                query,
+                key,
+                value,
+                self.embed_dim,
+                self.num_heads,
+                self.in_proj_weight,
+                self.in_proj_bias,
+                self.bias_k,
+                self.bias_v,
+                self.add_zero_attn,
+                self.dropout,
+                self.out_proj.weight,
+                self.out_proj.bias,
+                training=self.training,
+                key_padding_mask=key_padding_mask,
+                need_weights=need_weights,
+                attn_mask=attn_mask,
+                average_attn_weights=average_attn_weights,
+            )
+        if self.batch_first and is_batched:
+            return attn_output.transpose(1, 0), attn_output_weights
+        else:
+            return attn_output, attn_output_weights
+
+    def infer(self,
+              x: Tensor,
+              key_padding_mask: Optional[Tensor] = None,
+              need_weights: bool = True,
+              attn_mask: Optional[Tensor] = None,
+              average_attn_weights: bool = True,
+              past_kv = None,
+              use_cache = False
+              ):
+        # x = x.transpose(1, 0)
+        y, kv = multi_head_attention_forward(
+                x=x,
+                ipw=self.in_proj_weight,
+                ipb=self.in_proj_bias,
+                opw=self.out_proj.weight,
+                opb=self.out_proj.bias,
+                n_head=self.num_heads,
+                attn_mask=attn_mask,
+                past_kv=past_kv,
+                use_cache=use_cache,
+        )
+        return (y, kv)

apps/audio_cloning/vallex/modules/embedding.py

@@ -0,0 +1,97 @@
+# Copyright    2023                             (authors: Feiteng Li)
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import math
+
+import torch
+import torch.nn as nn
+
+
+class TokenEmbedding(nn.Module):
+    def __init__(
+        self,
+        dim_model: int,
+        vocab_size: int,
+        dropout: float = 0.0,
+    ):
+        super().__init__()
+
+        self.vocab_size = vocab_size
+        self.dim_model = dim_model
+
+        self.dropout = torch.nn.Dropout(p=dropout)
+        self.word_embeddings = nn.Embedding(self.vocab_size, self.dim_model)
+
+    @property
+    def weight(self) -> torch.Tensor:
+        return self.word_embeddings.weight
+
+    def embedding(self, index: int) -> torch.Tensor:
+        return self.word_embeddings.weight[index : index + 1]
+
+    def forward(self, x: torch.Tensor):
+        X = self.word_embeddings(x)
+        X = self.dropout(X)
+
+        return X
+
+
+class SinePositionalEmbedding(nn.Module):
+    def __init__(
+        self,
+        dim_model: int,
+        dropout: float = 0.0,
+        scale: bool = False,
+        alpha: bool = False,
+    ):
+        super().__init__()
+        self.dim_model = dim_model
+        self.x_scale = math.sqrt(dim_model) if scale else 1.0
+        self.alpha = nn.Parameter(torch.ones(1), requires_grad=alpha)
+        self.dropout = torch.nn.Dropout(p=dropout)
+
+        self.reverse = False
+        self.pe = None
+        self.extend_pe(torch.tensor(0.0).expand(1, 4000))
+
+    def extend_pe(self, x):
+        """Reset the positional encodings."""
+        if self.pe is not None:
+            if self.pe.size(1) >= x.size(1):
+                if self.pe.dtype != x.dtype or self.pe.device != x.device:
+                    self.pe = self.pe.to(dtype=x.dtype, device=x.device)
+                return
+        pe = torch.zeros(x.size(1), self.dim_model)
+        if self.reverse:
+            position = torch.arange(
+                x.size(1) - 1, -1, -1.0, dtype=torch.float32
+            ).unsqueeze(1)
+        else:
+            position = torch.arange(
+                0, x.size(1), dtype=torch.float32
+            ).unsqueeze(1)
+        div_term = torch.exp(
+            torch.arange(0, self.dim_model, 2, dtype=torch.float32)
+            * -(math.log(10000.0) / self.dim_model)
+        )
+        pe[:, 0::2] = torch.sin(position * div_term)
+        pe[:, 1::2] = torch.cos(position * div_term)
+        pe = pe.unsqueeze(0)
+        self.pe = pe.to(device=x.device, dtype=x.dtype).detach()
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        self.extend_pe(x)
+        output = x.unsqueeze(-1) if x.ndim == 2 else x
+        output = output * self.x_scale + self.alpha * self.pe[:, : x.size(1)]
+        return self.dropout(output)

apps/audio_cloning/vallex/modules/optim.py

@@ -0,0 +1,1105 @@
+# Copyright      2022  Xiaomi Corp.        (authors: Daniel Povey)
+#
+# See ../LICENSE for clarification regarding multiple authors
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import contextlib
+import logging
+import random
+from collections import defaultdict
+from typing import List, Optional, Tuple, Union
+
+import torch
+from lhotse.utils import fix_random_seed
+from torch import Tensor
+from torch.optim import Optimizer
+
+
+class BatchedOptimizer(Optimizer):
+    """
+    This class adds to class Optimizer the capability to optimize parameters in batches:
+    it will stack the parameters and their grads for you so the optimizer can work
+    on tensors with an extra leading dimension.  This is intended for speed with GPUs,
+    as it reduces the number of kernels launched in the optimizer.
+
+    Args:
+      params:
+    """
+
+    def __init__(self, params, defaults):
+        super(BatchedOptimizer, self).__init__(params, defaults)
+
+    @contextlib.contextmanager
+    def batched_params(self, param_group, group_params_names):
+        """
+        This function returns (technically, yields) a list of
+          of tuples (p, state), where
+        p is a `fake` parameter that is stacked (over axis 0) from real parameters
+        that share the same shape, and its gradient is also stacked;
+        `state` is the state corresponding to this batch of parameters
+        (it will be physically located in the "state" for one of the real
+        parameters, the last one that has any particular shape and dtype).
+
+        This function is decorated as a context manager so that it can
+        write parameters back to their "real" locations.
+
+        The idea is, instead of doing:
+        <code>
+          for p in group["params"]:
+             state = self.state[p]
+             ...
+        </code>
+        you can do:
+        <code>
+          with self.batched_params(group["params"]) as batches:
+             for p, state, p_names in batches:
+                 ...
+        </code>
+
+        Args:
+          group: a parameter group, which is a list of parameters; should be
+                one of self.param_groups.
+          group_params_names: name for each parameter in group,
+                which is List[str].
+        """
+        batches = defaultdict(
+            list
+        )  # `batches` maps from tuple (dtype_as_str,*shape) to list of nn.Parameter
+        batches_names = defaultdict(
+            list
+        )  # `batches` maps from tuple (dtype_as_str,*shape) to list of str
+
+        assert len(param_group) == len(group_params_names)
+        for p, named_p in zip(param_group, group_params_names):
+            key = (str(p.dtype), *p.shape)
+            batches[key].append(p)
+            batches_names[key].append(named_p)
+
+        batches_names_keys = list(batches_names.keys())
+        sorted_idx = sorted(
+            range(len(batches_names)), key=lambda i: batches_names_keys[i]
+        )
+        batches_names = [
+            batches_names[batches_names_keys[idx]] for idx in sorted_idx
+        ]
+        batches = [batches[batches_names_keys[idx]] for idx in sorted_idx]
+
+        stacked_params_dict = dict()
+
+        # turn batches into a list, in deterministic order.
+        # tuples will contain tuples of (stacked_param, state, stacked_params_names),
+        # one for each batch in `batches`.
+        tuples = []
+
+        for batch, batch_names in zip(batches, batches_names):
+            p = batch[0]
+            # we arbitrarily store the state in the
+            # state corresponding to the 1st parameter in the
+            # group.  class Optimizer will take care of saving/loading state.
+            state = self.state[p]
+            p_stacked = torch.stack(batch)
+            grad = torch.stack(
+                [
+                    torch.zeros_like(p) if p.grad is None else p.grad
+                    for p in batch
+                ]
+            )
+            p_stacked.grad = grad
+            stacked_params_dict[key] = p_stacked
+            tuples.append((p_stacked, state, batch_names))
+
+        yield tuples  # <-- calling code will do the actual optimization here!
+
+        for ((stacked_params, _state, _names), batch) in zip(tuples, batches):
+            for i, p in enumerate(batch):  # batch is list of Parameter
+                p.copy_(stacked_params[i])
+
+
+class ScaledAdam(BatchedOptimizer):
+    """
+     Implements 'Scaled Adam', a variant of Adam where we scale each parameter's update
+     proportional to the norm of that parameter; and also learn the scale of the parameter,
+     in log space, subject to upper and lower limits (as if we had factored each parameter as
+     param = underlying_param * log_scale.exp())
+
+
+     Args:
+          params:  The parameters or param_groups to optimize (like other Optimizer subclasses)
+              lr:  The learning rate.  We will typically use a learning rate schedule that starts
+                   at 0.03 and decreases over time, i.e. much higher than other common
+                   optimizers.
+     clipping_scale: (e.g. 2.0)
+                   A scale for gradient-clipping: if specified, the normalized gradients
+                   over the whole model will be clipped to have 2-norm equal to
+                   `clipping_scale` times the median 2-norm over the most recent period
+                   of `clipping_update_period` minibatches.  By "normalized gradients",
+                   we mean after multiplying by the rms parameter value for this tensor
+                   [for non-scalars]; this is appropriate because our update is scaled
+                   by this quantity.
+            betas: beta1,beta2 are momentum constants for regular momentum, and moving sum-sq grad.
+                   Must satisfy 0 < beta <= beta2 < 1.
+     scalar_lr_scale: A scaling factor on the learning rate, that we use to update the
+                   scale of each parameter tensor and scalar parameters of the mode..
+                   If each parameter were decomposed
+                   as p * p_scale.exp(), where (p**2).mean().sqrt() == 1.0, scalar_lr_scale
+                   would be a the scaling factor on the learning rate of p_scale.
+              eps:  A general-purpose epsilon to prevent division by zero
+    param_min_rms: Minimum root-mean-square value of parameter tensor, for purposes of
+                   learning the scale on the parameters (we'll constrain the rms of each non-scalar
+                   parameter tensor to be >= this value)
+    param_max_rms: Maximum root-mean-square value of parameter tensor, for purposes of
+                   learning the scale on the parameters (we'll constrain the rms of each non-scalar
+                   parameter tensor to be <= this value)
+       scalar_max: Maximum absolute value for scalar parameters (applicable if your
+                   model has any parameters with numel() == 1).
+    size_update_period: The periodicity, in steps, with which we update the size (scale)
+                   of the parameter tensor.  This is provided to save a little time
+                   in the update.
+     clipping_update_period: if clipping_scale is specified, this is the period
+    """
+
+    def __init__(
+        self,
+        params,
+        lr=3e-02,
+        clipping_scale=None,
+        betas=(0.9, 0.98),
+        scalar_lr_scale=0.1,
+        eps=1.0e-08,
+        param_min_rms=1.0e-05,
+        param_max_rms=3.0,
+        scalar_max=10.0,
+        size_update_period=4,
+        clipping_update_period=100,
+        parameters_names=None,
+        show_dominant_parameters=True,
+    ):
+
+        assert parameters_names is not None, (
+            "Please prepare parameters_names,"
+            "which is a List[List[str]]. Each List[str] is for a group"
+            "and each str is for a parameter"
+        )
+        defaults = dict(
+            lr=lr,
+            clipping_scale=clipping_scale,
+            betas=betas,
+            scalar_lr_scale=scalar_lr_scale,
+            eps=eps,
+            param_min_rms=param_min_rms,
+            param_max_rms=param_max_rms,
+            scalar_max=scalar_max,
+            size_update_period=size_update_period,
+            clipping_update_period=clipping_update_period,
+        )
+
+        super(ScaledAdam, self).__init__(params, defaults)
+        assert len(self.param_groups) == len(parameters_names)
+        self.parameters_names = parameters_names
+        self.show_dominant_parameters = show_dominant_parameters
+
+    def __setstate__(self, state):
+        super(ScaledAdam, self).__setstate__(state)
+
+    @torch.no_grad()
+    def step(self, closure=None):
+        """Performs a single optimization step.
+
+        Arguments:
+            closure (callable, optional): A closure that reevaluates the model
+                and returns the loss.
+        """
+        loss = None
+        if closure is not None:
+            with torch.enable_grad():
+                loss = closure()
+
+        batch = True
+
+        for group, group_params_names in zip(
+            self.param_groups, self.parameters_names
+        ):
+
+            with self.batched_params(
+                group["params"], group_params_names
+            ) as batches:
+
+                # batches is list of pairs (stacked_param, state).  stacked_param is like
+                # a regular parameter, and will have a .grad, but the 1st dim corresponds to
+                # a stacking dim, it is not a real dim.
+
+                if (
+                    len(batches[0][1]) == 0
+                ):  # if len(first state) == 0: not yet initialized
+                    clipping_scale = 1
+                else:
+                    clipping_scale = self._get_clipping_scale(group, batches)
+
+                for p, state, _ in batches:
+                    # Perform optimization step.
+                    # grad is not going to be None, we handled that when creating the batches.
+                    grad = p.grad
+                    if grad.is_sparse:
+                        raise RuntimeError(
+                            "ScaledAdam optimizer does not support sparse gradients"
+                        )
+                    # State initialization
+                    if len(state) == 0:
+                        self._init_state(group, p, state)
+
+                    self._step_one_batch(group, p, state, clipping_scale)
+
+        return loss
+
+    def _init_state(self, group: dict, p: Tensor, state: dict):
+        """
+        Initializes state dict for parameter 'p'.  Assumes that dim 0 of tensor p
+        is actually the batch dimension, corresponding to batched-together
+        parameters of a given shape.
+
+
+        Args:
+           group:   Dict to look up configuration values.
+               p: The parameter that we are initializing the state for
+           state: Dict from string to whatever state we are initializing
+        """
+        size_update_period = group["size_update_period"]
+
+        state["step"] = 0
+
+        kwargs = {"device": p.device, "dtype": p.dtype}
+
+        # 'delta' implements conventional momentum.  There are
+        # several different kinds of update going on, so rather than
+        # compute "exp_avg" like in Adam, we store and decay a
+        # parameter-change "delta", which combines all forms of
+        # update.  this is equivalent to how it's done in Adam,
+        # except for the first few steps.
+        state["delta"] = torch.zeros_like(
+            p, memory_format=torch.preserve_format
+        )
+
+        batch_size = p.shape[0]
+        numel = p.numel() // batch_size
+        numel = p.numel()
+
+        if numel > 1:
+            # "param_rms" just periodically records the scalar root-mean-square value of
+            # the parameter tensor.
+            # it has a shape like (batch_size, 1, 1, 1, 1)
+            param_rms = (
+                (p ** 2).mean(dim=list(range(1, p.ndim)), keepdim=True).sqrt()
+            )
+            state["param_rms"] = param_rms
+
+            state["scale_exp_avg_sq"] = torch.zeros_like(param_rms)
+            state["scale_grads"] = torch.zeros(
+                size_update_period, *param_rms.shape, **kwargs
+            )
+
+        # exp_avg_sq is the weighted sum of scaled gradients. as in Adam.
+        state["exp_avg_sq"] = torch.zeros_like(
+            p, memory_format=torch.preserve_format
+        )
+
+    def _get_clipping_scale(
+        self, group: dict, tuples: List[Tuple[Tensor, dict, List[str]]]
+    ) -> float:
+        """
+        Returns a scalar factor <= 1.0 that dictates gradient clipping, i.e. we will scale the gradients
+        by this amount before applying the rest of the update.
+
+        Args:
+           group: the parameter group, an item in self.param_groups
+           tuples: a list of tuples of (param, state, param_names)
+                where param is a batched set of parameters,
+                with a .grad (1st dim is batch dim)
+                and state is the state-dict where optimization parameters are kept.
+                param_names is a List[str] while each str is name for a parameter
+                in batched set of parameters "param".
+        """
+        assert len(tuples) >= 1
+        clipping_scale = group["clipping_scale"]
+        (first_p, first_state, _) = tuples[0]
+        step = first_state["step"]
+        if clipping_scale is None or step == 0:
+            # no clipping.  return early on step == 0 because the other
+            # parameters' state won't have been initialized yet.
+            return 1.0
+        clipping_update_period = group["clipping_update_period"]
+
+        tot_sumsq = torch.tensor(0.0, device=first_p.device)
+        for (p, state, param_names) in tuples:
+            grad = p.grad
+            if grad.is_sparse:
+                raise RuntimeError(
+                    "ScaledAdam optimizer does not support sparse gradients"
+                )
+            if p.numel() == p.shape[0]:  # a batch of scalars
+                tot_sumsq += (
+                    grad ** 2
+                ).sum()  # sum() to change shape [1] to []
+            else:
+                tot_sumsq += ((grad * state["param_rms"]) ** 2).sum()
+
+        tot_norm = tot_sumsq.sqrt()
+        if "model_norms" not in first_state:
+            first_state["model_norms"] = torch.zeros(
+                clipping_update_period, device=p.device
+            )
+        first_state["model_norms"][step % clipping_update_period] = tot_norm
+
+        if step % clipping_update_period == 0:
+            # Print some stats.
+            # We don't reach here if step == 0 because we would have returned
+            # above.
+            sorted_norms = first_state["model_norms"].sort()[0].to("cpu")
+            quartiles = []
+            for n in range(0, 5):
+                index = min(
+                    clipping_update_period - 1,
+                    (clipping_update_period // 4) * n,
+                )
+                quartiles.append(sorted_norms[index].item())
+
+            median = quartiles[2]
+            threshold = clipping_scale * median
+            first_state["model_norm_threshold"] = threshold
+            percent_clipped = (
+                first_state["num_clipped"] * 100.0 / clipping_update_period
+                if "num_clipped" in first_state
+                else 0.0
+            )
+            first_state["num_clipped"] = 0
+            quartiles = " ".join(["%.3e" % x for x in quartiles])
+            logging.info(
+                f"Clipping_scale={clipping_scale}, grad-norm quartiles {quartiles}, "
+                f"threshold={threshold:.3e}, percent-clipped={percent_clipped:.1f}"
+            )
+
+        if step < clipping_update_period:
+            return 1.0  # We have not yet estimated a norm to clip to.
+        else:
+            try:
+                model_norm_threshold = first_state["model_norm_threshold"]
+            except KeyError:
+                logging.info(
+                    "Warning: model_norm_threshold not in state: possibly "
+                    "you changed config when restarting, adding clipping_scale option?"
+                )
+                return 1.0
+            ans = min(1.0, (model_norm_threshold / (tot_norm + 1.0e-20)).item())
+            if ans < 1.0:
+                first_state["num_clipped"] += 1
+            if ans < 0.1:
+                logging.warn(
+                    f"Scaling gradients by {ans}, model_norm_threshold={model_norm_threshold}"
+                )
+                if self.show_dominant_parameters:
+                    assert p.shape[0] == len(param_names)
+                    self._show_gradient_dominating_parameter(tuples, tot_sumsq)
+            return ans
+
+    def _show_gradient_dominating_parameter(
+        self, tuples: List[Tuple[Tensor, dict, List[str]]], tot_sumsq: Tensor
+    ):
+        """
+        Show information of parameter wihch dominanting tot_sumsq.
+
+        Args:
+           tuples: a list of tuples of (param, state, param_names)
+                where param is a batched set of parameters,
+                with a .grad (1st dim is batch dim)
+                and state is the state-dict where optimization parameters are kept.
+                param_names is a List[str] while each str is name for a parameter
+                in batched set of parameters "param".
+            tot_sumsq: sumsq of all parameters. Though it's could be calculated
+                from tuples, we still pass it to save some time.
+        """
+        all_sumsq_orig = {}
+        for (p, state, batch_param_names) in tuples:
+            # p is a stacked batch parameters.
+            batch_grad = p.grad
+            if p.numel() == p.shape[0]:  # a batch of scalars
+                batch_sumsq_orig = batch_grad ** 2
+                # Dummpy values used by following `zip` statement.
+                batch_rms_orig = torch.ones(p.shape[0])
+            else:
+                batch_rms_orig = state["param_rms"]
+                batch_sumsq_orig = ((batch_grad * batch_rms_orig) ** 2).sum(
+                    dim=list(range(1, batch_grad.ndim))
+                )
+
+            for name, sumsq_orig, rms, grad in zip(
+                batch_param_names, batch_sumsq_orig, batch_rms_orig, batch_grad
+            ):
+
+                proportion_orig = sumsq_orig / tot_sumsq
+                all_sumsq_orig[name] = (proportion_orig, sumsq_orig, rms, grad)
+
+        assert torch.isclose(
+            sum([value[0] for value in all_sumsq_orig.values()]).cpu(),
+            torch.tensor(1.0),
+        )
+        sorted_by_proportion = {
+            k: v
+            for k, v in sorted(
+                all_sumsq_orig.items(),
+                key=lambda item: item[1][0],
+                reverse=True,
+            )
+        }
+        dominant_param_name = next(iter(sorted_by_proportion))
+        (
+            dominant_proportion,
+            dominant_sumsq,
+            dominant_rms,
+            dominant_grad,
+        ) = sorted_by_proportion[dominant_param_name]
+        logging.info(
+            f"Parameter Dominanting tot_sumsq {dominant_param_name}"
+            f" with proportion {dominant_proportion:.2f},"
+            f" where dominant_sumsq=(grad_sumsq*orig_rms_sq)"
+            f"={dominant_sumsq:.3e},"
+            f" grad_sumsq = {(dominant_grad**2).sum():.3e},"
+            f" orig_rms_sq={(dominant_rms**2).item():.3e}"
+        )
+
+    def _step_one_batch(
+        self, group: dict, p: Tensor, state: dict, clipping_scale: float
+    ):
+        """
+        Do the step for one parameter, which is actually going to be a batch of
+        `real` parameters, with dim 0 as the batch dim.
+        Args:
+                  group:  dict to look up configuration values
+                    p: parameter to update (actually multiple parameters stacked together
+                       as a batch)
+                  state: state-dict for p, to look up the optimizer state
+        """
+        lr = group["lr"]
+        size_update_period = group["size_update_period"]
+        beta1 = group["betas"][0]
+
+        grad = p.grad
+        if clipping_scale != 1.0:
+            grad = grad * clipping_scale
+        step = state["step"]
+        delta = state["delta"]
+
+        delta.mul_(beta1)
+        batch_size = p.shape[0]
+        numel = p.numel() // batch_size
+        if numel > 1:
+            # Update the size/scale of p, and set param_rms
+            scale_grads = state["scale_grads"]
+            scale_grads[step % size_update_period] = (p * grad).sum(
+                dim=list(range(1, p.ndim)), keepdim=True
+            )
+            if step % size_update_period == size_update_period - 1:
+                param_rms = state["param_rms"]  # shape: (batch_size, 1, 1, ..)
+                param_rms.copy_(
+                    (p ** 2)
+                    .mean(dim=list(range(1, p.ndim)), keepdim=True)
+                    .sqrt()
+                )
+                if step > 0:
+                    # self._size_update() learns the overall scale on the
+                    # parameter, by shrinking or expanding it.
+                    self._size_update(group, scale_grads, p, state)
+
+        if numel == 1:
+            # For parameters with 1 element we just use regular Adam.
+            # Updates delta.
+            self._step_scalar(group, p, state)
+        else:
+            self._step(group, p, state)
+
+        state["step"] = step + 1
+
+    def _size_update(
+        self, group: dict, scale_grads: Tensor, p: Tensor, state: dict
+    ) -> None:
+        """
+               Called only where p.numel() > 1, this updates the scale of the parameter.
+               If we imagine: p =  underlying_param * scale.exp(), and we are doing
+               gradient descent on underlying param and on scale, this function does the update
+               on `scale`.
+
+               Args:
+              group: dict to look up configuration values
+        scale_grads: a tensor of shape (size_update_period, batch_size, 1, 1,...) containing
+                      grads w.r.t. the scales.
+                  p:  The parameter to update
+               state: The state-dict of p
+        """
+
+        param_rms = state["param_rms"]
+        beta1, beta2 = group["betas"]
+        size_lr = group["lr"] * group["scalar_lr_scale"]
+        param_min_rms = group["param_min_rms"]
+        param_max_rms = group["param_max_rms"]
+        eps = group["eps"]
+        step = state["step"]
+        batch_size = p.shape[0]
+
+        size_update_period = scale_grads.shape[0]
+        # correct beta2 for the size update period: we will have
+        # faster decay at this level.
+        beta2_corr = beta2 ** size_update_period
+
+        scale_exp_avg_sq = state[
+            "scale_exp_avg_sq"
+        ]  # shape: (batch_size, 1, 1, ..)
+        scale_exp_avg_sq.mul_(beta2_corr).add_(
+            (scale_grads ** 2).mean(
+                dim=0
+            ),  # mean over dim `size_update_period`
+            alpha=1 - beta2_corr,
+        )  # shape is (batch_size, 1, 1, ...)
+
+        # The 1st time we reach here is when size_step == 1.
+        size_step = (step + 1) // size_update_period
+        bias_correction2 = 1 - beta2_corr ** size_step
+        # we don't bother with bias_correction1; this will help prevent divergence
+        # at the start of training.
+
+        denom = scale_exp_avg_sq.sqrt() + eps
+
+        scale_step = (
+            -size_lr
+            * (bias_correction2 ** 0.5)
+            * scale_grads.sum(dim=0)
+            / denom
+        )
+
+        is_too_small = param_rms < param_min_rms
+        is_too_large = param_rms > param_max_rms
+
+        # when the param gets too small, just don't shrink it any further.
+        scale_step.masked_fill_(is_too_small, 0.0)
+        # when it gets too large, stop it from getting any larger.
+        scale_step.masked_fill_(is_too_large, -size_lr * size_update_period)
+        delta = state["delta"]
+        # the factor of (1-beta1) relates to momentum.
+        delta.add_(p * scale_step, alpha=(1 - beta1))
+
+    def _step(self, group: dict, p: Tensor, state: dict):
+        """
+        This function does the core update of self.step(), in the case where the members of
+        the batch have more than 1 element.
+
+        Args:
+            group: A dict which will be used to look up configuration values
+                p: The parameter to be updated
+             grad: The grad of p
+            state: The state-dict corresponding to parameter p
+
+        This function modifies p.
+        """
+        grad = p.grad
+        lr = group["lr"]
+        beta1, beta2 = group["betas"]
+        eps = group["eps"]
+        param_min_rms = group["param_min_rms"]
+        step = state["step"]
+
+        exp_avg_sq = state["exp_avg_sq"]
+        exp_avg_sq.mul_(beta2).addcmul_(grad, grad, value=(1 - beta2))
+
+        this_step = state["step"] - (
+            state["zero_step"] if "zero_step" in state else 0
+        )
+        bias_correction2 = 1 - beta2 ** (this_step + 1)
+        if bias_correction2 < 0.99:
+            # note: not in-place.
+            exp_avg_sq = exp_avg_sq * (1.0 / bias_correction2)
+
+        denom = exp_avg_sq.sqrt()
+        denom += eps
+        grad = grad / denom
+
+        alpha = -lr * (1 - beta1) * state["param_rms"].clamp(min=param_min_rms)
+
+        delta = state["delta"]
+        delta.add_(grad * alpha)
+        p.add_(delta)
+
+    def _step_scalar(self, group: dict, p: Tensor, state: dict):
+        """
+        A simplified form of the core update for scalar tensors, where we cannot get a good
+        estimate of the parameter rms.
+        """
+        beta1, beta2 = group["betas"]
+        scalar_max = group["scalar_max"]
+        eps = group["eps"]
+        lr = group["lr"] * group["scalar_lr_scale"]
+        grad = p.grad
+
+        exp_avg_sq = state["exp_avg_sq"]  # shape: (batch_size,)
+        exp_avg_sq.mul_(beta2).addcmul_(grad, grad, value=1 - beta2)
+
+        # bias_correction2 is like in Adam.  Don't bother with bias_correction1;
+        # slower update at the start will help stability anyway.
+        bias_correction2 = 1 - beta2 ** (state["step"] + 1)
+        denom = (exp_avg_sq / bias_correction2).sqrt() + eps
+
+        delta = state["delta"]
+        delta.add_(grad / denom, alpha=-lr * (1 - beta1))
+        p.clamp_(min=-scalar_max, max=scalar_max)
+        p.add_(delta)
+
+
+class LRScheduler(object):
+    """
+    Base-class for learning rate schedulers where the learning-rate depends on both the
+    batch and the epoch.
+    """
+
+    def __init__(self, optimizer: Optimizer, verbose: bool = False):
+        # Attach optimizer
+        if not isinstance(optimizer, Optimizer):
+            raise TypeError(
+                "{} is not an Optimizer".format(type(optimizer).__name__)
+            )
+        self.optimizer = optimizer
+        self.verbose = verbose
+
+        for group in optimizer.param_groups:
+            group.setdefault("base_lr", group["lr"])
+
+        self.base_lrs = [group["base_lr"] for group in optimizer.param_groups]
+
+        self.epoch = 0
+        self.batch = 0
+
+    def state_dict(self):
+        """Returns the state of the scheduler as a :class:`dict`.
+
+        It contains an entry for every variable in self.__dict__ which
+        is not the optimizer.
+        """
+        return {
+            "base_lrs": self.base_lrs,
+            "epoch": self.epoch,
+            "batch": self.batch,
+        }
+
+    def load_state_dict(self, state_dict):
+        """Loads the schedulers state.
+
+        Args:
+            state_dict (dict): scheduler state. Should be an object returned
+                from a call to :meth:`state_dict`.
+        """
+        self.__dict__.update(state_dict)
+
+    def get_last_lr(self) -> List[float]:
+        """Return last computed learning rate by current scheduler.  Will be a list of float."""
+        return self._last_lr
+
+    def get_lr(self):
+        # Compute list of learning rates from self.epoch and self.batch and
+        # self.base_lrs; this must be overloaded by the user.
+        # e.g. return [some_formula(self.batch, self.epoch, base_lr) for base_lr in self.base_lrs ]
+        raise NotImplementedError
+
+    def step_batch(self, batch: Optional[int] = None) -> None:
+        # Step the batch index, or just set it.  If `batch` is specified, it
+        # must be the batch index from the start of training, i.e. summed over
+        # all epochs.
+        # You can call this in any order; if you don't provide 'batch', it should
+        # of course be called once per batch.
+        if batch is not None:
+            self.batch = batch
+        else:
+            self.batch = self.batch + 1
+        self._set_lrs()
+
+    def step_epoch(self, epoch: Optional[int] = None):
+        # Step the epoch index, or just set it.  If you provide the 'epoch' arg,
+        # you should call this at the start of the epoch; if you don't provide the 'epoch'
+        # arg, you should call it at the end of the epoch.
+        if epoch is not None:
+            self.epoch = epoch
+        else:
+            self.epoch = self.epoch + 1
+        self._set_lrs()
+
+    def _set_lrs(self):
+        values = self.get_lr()
+        assert len(values) == len(self.optimizer.param_groups)
+
+        for i, data in enumerate(zip(self.optimizer.param_groups, values)):
+            param_group, lr = data
+            param_group["lr"] = lr
+            self.print_lr(self.verbose, i, lr)
+        self._last_lr = [group["lr"] for group in self.optimizer.param_groups]
+
+    def print_lr(self, is_verbose, group, lr):
+        """Display the current learning rate."""
+        if is_verbose:
+            logging.info(
+                f"Epoch={self.epoch}, batch={self.batch}: adjusting learning rate"
+                f" of group {group} to {lr:.4e}."
+            )
+
+
+class Eden(LRScheduler):
+    """
+    Eden scheduler.
+    The basic formula (before warmup) is:
+      lr = base_lr * (((batch**2 + lr_batches**2) / lr_batches**2) ** -0.25 *
+                     (((epoch**2 + lr_epochs**2) / lr_epochs**2) ** -0.25)) * warmup
+    where `warmup` increases from linearly 0.5 to 1 over `warmup_batches` batches
+    and then stays constant at 1.
+
+
+     E.g. suggest base_lr = 0.04 (passed to optimizer) if used with ScaledAdam
+
+    Args:
+        optimizer: the optimizer to change the learning rates on
+        lr_batches: the number of batches after which we start significantly
+              decreasing the learning rate, suggest 5000.
+        lr_epochs: the number of epochs after which we start significantly
+              decreasing the learning rate, suggest 6 if you plan to do e.g.
+              20 to 40 epochs, but may need smaller number if dataset is huge
+              and you will do few epochs.
+    """
+
+    def __init__(
+        self,
+        optimizer: Optimizer,
+        lr_batches: Union[int, float],
+        lr_epochs: Union[int, float],
+        warmup_batches: Union[int, float] = 500.0,
+        verbose: bool = False,
+    ):
+        super(Eden, self).__init__(optimizer, verbose)
+        self.lr_batches = lr_batches
+        self.lr_epochs = lr_epochs
+        self.warmup_batches = warmup_batches
+
+    def get_lr(self):
+        factor = (
+            (self.batch ** 2 + self.lr_batches ** 2) / self.lr_batches ** 2
+        ) ** -0.25 * (
+            ((self.epoch ** 2 + self.lr_epochs ** 2) / self.lr_epochs ** 2)
+            ** -0.25
+        )
+        warmup_factor = (
+            1.0
+            if self.batch >= self.warmup_batches
+            else 0.5 + 0.5 * (self.batch / self.warmup_batches)
+        )
+
+        return [x * factor * warmup_factor for x in self.base_lrs]
+
+
+def _test_eden():
+    m = torch.nn.Linear(100, 100)
+    optim = ScaledAdam(m.parameters(), lr=0.03)
+
+    scheduler = Eden(optim, lr_batches=100, lr_epochs=2, verbose=True)
+
+    for epoch in range(10):
+        scheduler.step_epoch(epoch)  # sets epoch to `epoch`
+
+        for step in range(20):
+            x = torch.randn(200, 100).detach()
+            x.requires_grad = True
+            y = m(x)
+            dy = torch.randn(200, 100).detach()
+            f = (y * dy).sum()
+            f.backward()
+
+            optim.step()
+            scheduler.step_batch()
+            optim.zero_grad()
+
+    logging.info(f"last lr = {scheduler.get_last_lr()}")
+    logging.info(f"state dict = {scheduler.state_dict()}")
+
+
+# This is included mostly as a baseline for ScaledAdam.
+class Eve(Optimizer):
+    """
+    Implements Eve algorithm.  This is a modified version of AdamW with a special
+    way of setting the weight-decay / shrinkage-factor, which is designed to make the
+    rms of the parameters approach a particular target_rms (default: 0.1).  This is
+    for use with networks with 'scaled' versions of modules (see scaling.py), which
+    will be close to invariant to the absolute scale on the parameter matrix.
+
+    The original Adam algorithm was proposed in `Adam: A Method for Stochastic Optimization`_.
+    The AdamW variant was proposed in `Decoupled Weight Decay Regularization`_.
+    Eve is unpublished so far.
+
+    Arguments:
+        params (iterable): iterable of parameters to optimize or dicts defining
+            parameter groups
+        lr (float, optional): learning rate (default: 1e-3)
+        betas (Tuple[float, float], optional): coefficients used for computing
+            running averages of gradient and its square (default: (0.9, 0.999))
+        eps (float, optional): term added to the denominator to improve
+            numerical stability (default: 1e-8)
+        weight_decay (float, optional): weight decay coefficient (default: 3e-4;
+            this value means that the weight would decay significantly after
+            about 3k minibatches.  Is not multiplied by learning rate, but
+            is conditional on RMS-value of parameter being > target_rms.
+        target_rms (float, optional): target root-mean-square value of
+           parameters, if they fall below this we will stop applying weight decay.
+
+
+    .. _Adam: A Method for Stochastic Optimization:
+        https://arxiv.org/abs/1412.6980
+    .. _Decoupled Weight Decay Regularization:
+        https://arxiv.org/abs/1711.05101
+    .. _On the Convergence of Adam and Beyond:
+        https://openreview.net/forum?id=ryQu7f-RZ
+    """
+
+    def __init__(
+        self,
+        params,
+        lr=1e-3,
+        betas=(0.9, 0.98),
+        eps=1e-8,
+        weight_decay=1e-3,
+        target_rms=0.1,
+    ):
+        if not 0.0 <= lr:
+            raise ValueError("Invalid learning rate: {}".format(lr))
+        if not 0.0 <= eps:
+            raise ValueError("Invalid epsilon value: {}".format(eps))
+        if not 0.0 <= betas[0] < 1.0:
+            raise ValueError(
+                "Invalid beta parameter at index 0: {}".format(betas[0])
+            )
+        if not 0.0 <= betas[1] < 1.0:
+            raise ValueError(
+                "Invalid beta parameter at index 1: {}".format(betas[1])
+            )
+        if not 0 <= weight_decay <= 0.1:
+            raise ValueError(
+                "Invalid weight_decay value: {}".format(weight_decay)
+            )
+        if not 0 < target_rms <= 10.0:
+            raise ValueError("Invalid target_rms value: {}".format(target_rms))
+        defaults = dict(
+            lr=lr,
+            betas=betas,
+            eps=eps,
+            weight_decay=weight_decay,
+            target_rms=target_rms,
+        )
+        super(Eve, self).__init__(params, defaults)
+
+    def __setstate__(self, state):
+        super(Eve, self).__setstate__(state)
+
+    @torch.no_grad()
+    def step(self, closure=None):
+        """Performs a single optimization step.
+
+        Arguments:
+            closure (callable, optional): A closure that reevaluates the model
+                and returns the loss.
+        """
+        loss = None
+        if closure is not None:
+            with torch.enable_grad():
+                loss = closure()
+
+        for group in self.param_groups:
+            for p in group["params"]:
+                if p.grad is None:
+                    continue
+
+                # Perform optimization step
+                grad = p.grad
+                if grad.is_sparse:
+                    raise RuntimeError(
+                        "AdamW does not support sparse gradients"
+                    )
+
+                state = self.state[p]
+
+                # State initialization
+                if len(state) == 0:
+                    state["step"] = 0
+                    # Exponential moving average of gradient values
+                    state["exp_avg"] = torch.zeros_like(
+                        p, memory_format=torch.preserve_format
+                    )
+                    # Exponential moving average of squared gradient values
+                    state["exp_avg_sq"] = torch.zeros_like(
+                        p, memory_format=torch.preserve_format
+                    )
+
+                exp_avg, exp_avg_sq = state["exp_avg"], state["exp_avg_sq"]
+
+                beta1, beta2 = group["betas"]
+
+                state["step"] += 1
+                bias_correction1 = 1 - beta1 ** state["step"]
+                bias_correction2 = 1 - beta2 ** state["step"]
+
+                # Decay the first and second moment running average coefficient
+                exp_avg.mul_(beta1).add_(grad, alpha=1 - beta1)
+                exp_avg_sq.mul_(beta2).addcmul_(grad, grad, value=1 - beta2)
+                denom = (exp_avg_sq.sqrt() * (bias_correction2 ** -0.5)).add_(
+                    group["eps"]
+                )
+
+                step_size = group["lr"] / bias_correction1
+                target_rms = group["target_rms"]
+                weight_decay = group["weight_decay"]
+
+                if p.numel() > 1:
+                    # avoid applying this weight-decay on "scaling factors"
+                    # (which are scalar).
+                    is_above_target_rms = p.norm() > (
+                        target_rms * (p.numel() ** 0.5)
+                    )
+                    p.mul_(1 - (weight_decay * is_above_target_rms))
+
+                p.addcdiv_(exp_avg, denom, value=-step_size)
+
+                # if random.random() < 0.0005:
+                #     step = (exp_avg / denom) * step_size
+                #     logging.info(
+                #         f"Delta rms = {(step**2).mean().item()}, shape = {step.shape}"
+                #     )
+
+        return loss
+
+
+def _test_scaled_adam(hidden_dim: int):
+    import timeit
+
+    from scaling import ScaledLinear
+
+    E = 100
+    B = 4
+    T = 2
+    logging.info("in test_eve_cain")
+    # device = torch.device('cuda')
+    device = torch.device("cpu")
+    dtype = torch.float32
+
+    fix_random_seed(42)
+    # these input_magnitudes and output_magnitudes are to test that
+    # Abel is working as we expect and is able to adjust scales of
+    # different dims differently.
+    input_magnitudes = (1.0 * torch.randn(E, dtype=dtype, device=device)).exp()
+    output_magnitudes = (1.0 * torch.randn(E, dtype=dtype, device=device)).exp()
+
+    for iter in [1, 0]:
+        fix_random_seed(42)
+        Linear = torch.nn.Linear if iter == 0 else ScaledLinear
+
+        m = torch.nn.Sequential(
+            Linear(E, hidden_dim),
+            torch.nn.PReLU(),
+            Linear(hidden_dim, hidden_dim),
+            torch.nn.PReLU(),
+            Linear(hidden_dim, E),
+        ).to(device)
+
+        train_pairs = [
+            (
+                100.0
+                * torch.randn(B, T, E, device=device, dtype=dtype)
+                * input_magnitudes,
+                torch.randn(B, T, E, device=device, dtype=dtype)
+                * output_magnitudes,
+            )
+            for _ in range(20)
+        ]
+
+        if iter == 0:
+            optim = Eve(m.parameters(), lr=0.003)
+        elif iter == 1:
+            optim = ScaledAdam(m.parameters(), lr=0.03, clipping_scale=2.0)
+        scheduler = Eden(optim, lr_batches=200, lr_epochs=5, verbose=False)
+
+        start = timeit.default_timer()
+        avg_loss = 0.0
+        for epoch in range(180):
+            scheduler.step_epoch()
+            # if epoch == 100 and iter in [2,3]:
+            #    optim.reset_speedup()  # check it doesn't crash.
+
+            # if epoch == 130:
+            #    opts = diagnostics.TensorDiagnosticOptions(
+            #        2 ** 22
+            #    )  # allow 4 megabytes per sub-module
+            #    diagnostic = diagnostics.attach_diagnostics(m, opts)
+
+            for n, (x, y) in enumerate(train_pairs):
+                y_out = m(x)
+                loss = ((y_out - y) ** 2).mean() * 100.0
+                if epoch == 0 and n == 0:
+                    avg_loss = loss.item()
+                else:
+                    avg_loss = 0.98 * avg_loss + 0.02 * loss.item()
+                if n == 0 and epoch % 5 == 0:
+                    # norm1 = '%.2e' % (m[0].weight**2).mean().sqrt().item()
+                    # norm1b = '%.2e' % (m[0].bias**2).mean().sqrt().item()
+                    # norm2 = '%.2e' % (m[2].weight**2).mean().sqrt().item()
+                    # norm2b = '%.2e' % (m[2].bias**2).mean().sqrt().item()
+                    # scale1 = '%.2e' % (m[0].weight_scale.exp().item())
+                    # scale1b = '%.2e' % (m[0].bias_scale.exp().item())
+                    # scale2 = '%.2e' % (m[2].weight_scale.exp().item())
+                    # scale2b = '%.2e' % (m[2].bias_scale.exp().item())
+                    lr = scheduler.get_last_lr()[0]
+                    logging.info(
+                        f"Iter {iter}, epoch {epoch}, batch {n}, avg_loss {avg_loss:.4g}, lr={lr:.4e}"
+                    )  # , norms={norm1,norm1b,norm2,norm2b}") # scales={scale1,scale1b,scale2,scale2b}
+                loss.log().backward()
+                optim.step()
+                optim.zero_grad()
+                scheduler.step_batch()
+
+        # diagnostic.print_diagnostics()
+
+        stop = timeit.default_timer()
+        logging.info(f"Iter={iter}, Time taken: {stop - start}")
+
+        logging.info(f"last lr = {scheduler.get_last_lr()}")
+        # logging.info("state dict = ", scheduler.state_dict())
+        # logging.info("optim state_dict = ", optim.state_dict())
+        logging.info(f"input_magnitudes = {input_magnitudes}")
+        logging.info(f"output_magnitudes = {output_magnitudes}")
+
+
+if __name__ == "__main__":
+    torch.set_num_threads(1)
+    torch.set_num_interop_threads(1)
+    logging.getLogger().setLevel(logging.INFO)
+    import subprocess
+
+    s = subprocess.check_output(
+        "git status -uno .; git log -1; git diff HEAD .", shell=True
+    )
+    logging.info(s)
+    import sys
+
+    if len(sys.argv) > 1:
+        hidden_dim = int(sys.argv[1])
+    else:
+        hidden_dim = 200
+
+    _test_scaled_adam(hidden_dim)
+    _test_eden()

apps/audio_cloning/vallex/modules/scaling.py

@@ -0,0 +1,1369 @@
+# Copyright    2022  Xiaomi Corp.        (authors: Daniel Povey)
+#
+# See ../../../../LICENSE for clarification regarding multiple authors
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+
+import logging
+import math
+import random
+from typing import Optional, Tuple, Union
+
+import torch
+import torch.nn as nn
+from torch import Tensor
+
+
+class Transpose(nn.Identity):
+    """(N, T, D) -> (N, D, T)"""
+
+    def forward(self, input: torch.Tensor) -> torch.Tensor:
+        return input.transpose(1, 2)
+
+
+class ActivationBalancerFunction(torch.autograd.Function):
+    @staticmethod
+    def forward(
+        ctx,
+        x: Tensor,
+        scale_factor: Tensor,
+        sign_factor: Optional[Tensor],
+        channel_dim: int,
+    ) -> Tensor:
+        if channel_dim < 0:
+            channel_dim += x.ndim
+        ctx.channel_dim = channel_dim
+        xgt0 = x > 0
+        if sign_factor is None:
+            ctx.save_for_backward(xgt0, scale_factor)
+        else:
+            ctx.save_for_backward(xgt0, scale_factor, sign_factor)
+        return x
+
+    @staticmethod
+    def backward(ctx, x_grad: Tensor) -> Tuple[Tensor, None, None, None]:
+        if len(ctx.saved_tensors) == 3:
+            xgt0, scale_factor, sign_factor = ctx.saved_tensors
+            for _ in range(ctx.channel_dim, x_grad.ndim - 1):
+                scale_factor = scale_factor.unsqueeze(-1)
+                sign_factor = sign_factor.unsqueeze(-1)
+            factor = sign_factor + scale_factor * (xgt0.to(x_grad.dtype) - 0.5)
+        else:
+            xgt0, scale_factor = ctx.saved_tensors
+            for _ in range(ctx.channel_dim, x_grad.ndim - 1):
+                scale_factor = scale_factor.unsqueeze(-1)
+            factor = scale_factor * (xgt0.to(x_grad.dtype) - 0.5)
+        neg_delta_grad = x_grad.abs() * factor
+        return (
+            x_grad - neg_delta_grad,
+            None,
+            None,
+            None,
+        )
+
+
+def _compute_scale_factor(
+    x: Tensor,
+    channel_dim: int,
+    min_abs: float,
+    max_abs: float,
+    gain_factor: float,
+    max_factor: float,
+) -> Tensor:
+    if channel_dim < 0:
+        channel_dim += x.ndim
+    sum_dims = [d for d in range(x.ndim) if d != channel_dim]
+    x_abs_mean = torch.mean(x.abs(), dim=sum_dims).to(torch.float32)
+
+    if min_abs == 0.0:
+        below_threshold = 0.0
+    else:
+        # below_threshold is 0 if x_abs_mean > min_abs, can be at most max_factor if
+        # x_abs)_mean , min_abs.
+        below_threshold = ((min_abs - x_abs_mean) * (gain_factor / min_abs)).clamp(
+            min=0, max=max_factor
+        )
+
+    above_threshold = ((x_abs_mean - max_abs) * (gain_factor / max_abs)).clamp(
+        min=0, max=max_factor
+    )
+
+    return below_threshold - above_threshold
+
+
+def _compute_sign_factor(
+    x: Tensor,
+    channel_dim: int,
+    min_positive: float,
+    max_positive: float,
+    gain_factor: float,
+    max_factor: float,
+) -> Tensor:
+    if channel_dim < 0:
+        channel_dim += x.ndim
+    sum_dims = [d for d in range(x.ndim) if d != channel_dim]
+    proportion_positive = torch.mean((x > 0).to(torch.float32), dim=sum_dims)
+    if min_positive == 0.0:
+        factor1 = 0.0
+    else:
+        # 0 if proportion_positive >= min_positive, else can be
+        # as large as max_factor.
+        factor1 = (
+            (min_positive - proportion_positive) * (gain_factor / min_positive)
+        ).clamp_(min=0, max=max_factor)
+
+    if max_positive == 1.0:
+        factor2 = 0.0
+    else:
+        # 0 if self.proportion_positive <= max_positive, else can be
+        # as large as -max_factor.
+        factor2 = (
+            (proportion_positive - max_positive) * (gain_factor / (1.0 - max_positive))
+        ).clamp_(min=0, max=max_factor)
+    sign_factor = factor1 - factor2
+    # require min_positive != 0 or max_positive != 1:
+    assert not isinstance(sign_factor, float)
+    return sign_factor
+
+
+class ActivationScaleBalancerFunction(torch.autograd.Function):
+    """
+    This object is used in class ActivationBalancer when the user specified
+    min_positive=0, max_positive=1, so there are no constraints on the signs
+    of the activations and only the absolute value has a constraint.
+    """
+
+    @staticmethod
+    def forward(
+        ctx,
+        x: Tensor,
+        sign_factor: Tensor,
+        scale_factor: Tensor,
+        channel_dim: int,
+    ) -> Tensor:
+        if channel_dim < 0:
+            channel_dim += x.ndim
+        ctx.channel_dim = channel_dim
+        xgt0 = x > 0
+        ctx.save_for_backward(xgt0, sign_factor, scale_factor)
+        return x
+
+    @staticmethod
+    def backward(ctx, x_grad: Tensor) -> Tuple[Tensor, None, None, None]:
+        xgt0, sign_factor, scale_factor = ctx.saved_tensors
+        for _ in range(ctx.channel_dim, x_grad.ndim - 1):
+            sign_factor = sign_factor.unsqueeze(-1)
+            scale_factor = scale_factor.unsqueeze(-1)
+
+        factor = sign_factor + scale_factor * (xgt0.to(x_grad.dtype) - 0.5)
+        neg_delta_grad = x_grad.abs() * factor
+        return (
+            x_grad - neg_delta_grad,
+            None,
+            None,
+            None,
+        )
+
+
+class RandomClampFunction(torch.autograd.Function):
+    @staticmethod
+    def forward(
+        ctx,
+        x: Tensor,
+        min: Optional[float],
+        max: Optional[float],
+        prob: float,
+        reflect: float,
+    ) -> Tensor:
+        x_clamped = torch.clamp(x, min=min, max=max)
+        mask = torch.rand_like(x) < prob
+        ans = torch.where(mask, x_clamped, x)
+        if x.requires_grad:
+            ctx.save_for_backward(ans == x)
+            ctx.reflect = reflect
+        if reflect != 0.0:
+            ans = ans * (1.0 + reflect) - (x * reflect)
+        return ans
+
+    @staticmethod
+    def backward(ctx, ans_grad: Tensor) -> Tuple[Tensor, None, None, None, None]:
+        (is_same,) = ctx.saved_tensors
+        x_grad = ans_grad * is_same.to(ans_grad.dtype)
+        reflect = ctx.reflect
+        if reflect != 0.0:
+            x_grad = x_grad * (1.0 + reflect) - (ans_grad * reflect)
+        return x_grad, None, None, None, None
+
+
+def random_clamp(
+    x: Tensor,
+    min: Optional[float] = None,
+    max: Optional[float] = None,
+    prob: float = 0.5,
+    reflect: float = 0.0,
+):
+    return RandomClampFunction.apply(x, min, max, prob, reflect)
+
+
+def random_cast_to_half(x: Tensor, min_abs: float = 5.0e-06) -> Tensor:
+    """
+    A randomized way of casting a floating point value to half precision.
+    """
+    if x.dtype == torch.float16:
+        return x
+    x_abs = x.abs()
+    is_too_small = x_abs < min_abs
+    # for elements where is_too_small is true, random_val will contain +-min_abs with
+    # probability (x.abs() / min_abs), and 0.0 otherwise.  [so this preserves expectations,
+    # for those elements].
+    random_val = min_abs * x.sign() * (torch.rand_like(x) * min_abs < x_abs)
+    return torch.where(is_too_small, random_val, x).to(torch.float16)
+
+
+class RandomGradFunction(torch.autograd.Function):
+    """
+    Does nothing in forward pass; in backward pass, gets rid of very small grads using
+    randomized approach that preserves expectations (intended to reduce roundoff).
+    """
+
+    @staticmethod
+    def forward(ctx, x: Tensor, min_abs: float) -> Tensor:
+        ctx.min_abs = min_abs
+        return x
+
+    @staticmethod
+    def backward(ctx, ans_grad: Tensor) -> Tuple[Tensor, None]:
+        if ans_grad.dtype == torch.float16:
+            return (
+                random_cast_to_half(ans_grad.to(torch.float32), min_abs=ctx.min_abs),
+                None,
+            )
+        else:
+            return ans_grad, None
+
+
+class RandomGrad(torch.nn.Module):
+    """
+    Gets rid of very small gradients using an expectation-preserving method, intended to increase
+    accuracy of training when using amp (automatic mixed precision)
+    """
+
+    def __init__(self, min_abs: float = 5.0e-06):
+        super(RandomGrad, self).__init__()
+        self.min_abs = min_abs
+
+    def forward(self, x: Tensor):
+        if torch.jit.is_scripting() or not self.training or torch.jit.is_tracing():
+            return x
+        else:
+            return RandomGradFunction.apply(x, self.min_abs)
+
+
+class SoftmaxFunction(torch.autograd.Function):
+    """
+    Tries to handle half-precision derivatives in a randomized way that should
+    be more accurate for training than the default behavior.
+    """
+
+    @staticmethod
+    def forward(ctx, x: Tensor, dim: int):
+        ans = x.softmax(dim=dim)
+        # if x dtype is float16, x.softmax() returns a float32 because
+        # (presumably) that op does not support float16, and autocast
+        # is enabled.
+        if torch.is_autocast_enabled():
+            ans = ans.to(torch.float16)
+        ctx.save_for_backward(ans)
+        ctx.x_dtype = x.dtype
+        ctx.dim = dim
+        return ans
+
+    @staticmethod
+    def backward(ctx, ans_grad: Tensor):
+        (ans,) = ctx.saved_tensors
+        with torch.cuda.amp.autocast(enabled=False):
+            ans_grad = ans_grad.to(torch.float32)
+            ans = ans.to(torch.float32)
+            x_grad = ans_grad * ans
+            x_grad = x_grad - ans * x_grad.sum(dim=ctx.dim, keepdim=True)
+            return x_grad, None
+
+
+def softmax(x: Tensor, dim: int):
+    if torch.jit.is_scripting() or torch.jit.is_tracing():
+        return x.softmax(dim)
+
+    return SoftmaxFunction.apply(x, dim)
+
+
+class MaxEigLimiterFunction(torch.autograd.Function):
+    @staticmethod
+    def forward(
+        ctx,
+        x: Tensor,
+        coeffs: Tensor,
+        direction: Tensor,
+        channel_dim: int,
+        grad_scale: float,
+    ) -> Tensor:
+        ctx.channel_dim = channel_dim
+        ctx.grad_scale = grad_scale
+        ctx.save_for_backward(x.detach(), coeffs.detach(), direction.detach())
+        return x
+
+    @staticmethod
+    def backward(ctx, x_grad, *args):
+        with torch.enable_grad():
+            (x_orig, coeffs, new_direction) = ctx.saved_tensors
+            x_orig.requires_grad = True
+            num_channels = x_orig.shape[ctx.channel_dim]
+            x = x_orig.transpose(ctx.channel_dim, -1).reshape(-1, num_channels)
+            new_direction.requires_grad = False
+            x = x - x.mean(dim=0)
+            x_var = (x**2).mean()
+            x_residual = x - coeffs * new_direction
+            x_residual_var = (x_residual**2).mean()
+            # `variance_proportion` is the proportion of the variance accounted for
+            # by the top eigen-direction.  This is to be minimized.
+            variance_proportion = (x_var - x_residual_var) / (x_var + 1.0e-20)
+            variance_proportion.backward()
+        x_orig_grad = x_orig.grad
+        x_extra_grad = (
+            x_orig.grad
+            * ctx.grad_scale
+            * x_grad.norm()
+            / (x_orig_grad.norm() + 1.0e-20)
+        )
+        return x_grad + x_extra_grad.detach(), None, None, None, None
+
+
+class BasicNorm(torch.nn.Module):
+    """
+    This is intended to be a simpler, and hopefully cheaper, replacement for
+    LayerNorm.  The observation this is based on, is that Transformer-type
+    networks, especially with pre-norm, sometimes seem to set one of the
+    feature dimensions to a large constant value (e.g. 50), which "defeats"
+    the LayerNorm because the output magnitude is then not strongly dependent
+    on the other (useful) features.  Presumably the weight and bias of the
+    LayerNorm are required to allow it to do this.
+
+    So the idea is to introduce this large constant value as an explicit
+    parameter, that takes the role of the "eps" in LayerNorm, so the network
+    doesn't have to do this trick.  We make the "eps" learnable.
+
+    Args:
+       num_channels: the number of channels, e.g. 512.
+      channel_dim: the axis/dimension corresponding to the channel,
+        interprted as an offset from the input's ndim if negative.
+        shis is NOT the num_channels; it should typically be one of
+        {-2, -1, 0, 1, 2, 3}.
+       eps: the initial "epsilon" that we add as ballast in:
+             scale = ((input_vec**2).mean() + epsilon)**-0.5
+          Note: our epsilon is actually large, but we keep the name
+          to indicate the connection with conventional LayerNorm.
+       learn_eps: if true, we learn epsilon; if false, we keep it
+         at the initial value.
+    eps_min: float
+    eps_max: float
+    """
+
+    def __init__(
+        self,
+        num_channels: int,
+        channel_dim: int = -1,  # CAUTION: see documentation.
+        eps: float = 0.25,
+        learn_eps: bool = True,
+        eps_min: float = -3.0,
+        eps_max: float = 3.0,
+    ) -> None:
+        super(BasicNorm, self).__init__()
+        self.num_channels = num_channels
+        self.channel_dim = channel_dim
+        if learn_eps:
+            self.eps = nn.Parameter(torch.tensor(eps).log().detach())
+        else:
+            self.register_buffer("eps", torch.tensor(eps).log().detach())
+        self.eps_min = eps_min
+        self.eps_max = eps_max
+
+    def forward(self, x: Tensor) -> Tensor:
+        assert x.shape[self.channel_dim] == self.num_channels
+        eps = self.eps
+        if self.training and random.random() < 0.25:
+            # with probability 0.25, in training mode, clamp eps between the min
+            # and max; this will encourage it to learn parameters within the
+            # allowed range by making parameters that are outside the allowed
+            # range noisy.
+
+            # gradients to allow the parameter to get back into the allowed
+            # region if it happens to exit it.
+            eps = eps.clamp(min=self.eps_min, max=self.eps_max)
+        scales = (
+            torch.mean(x**2, dim=self.channel_dim, keepdim=True) + eps.exp()
+        ) ** -0.5
+        return x * scales
+
+
+def ScaledLinear(*args, initial_scale: float = 1.0, **kwargs) -> nn.Linear:
+    """
+    Behaves like a constructor of a modified version of nn.Linear
+    that gives an easy way to set the default initial parameter scale.
+
+    Args:
+        Accepts the standard args and kwargs that nn.Linear accepts
+        e.g. in_features, out_features, bias=False.
+
+        initial_scale: you can override this if you want to increase
+           or decrease the initial magnitude of the module's output
+           (affects the initialization of weight_scale and bias_scale).
+           Another option, if you want to do something like this, is
+           to re-initialize the parameters.
+    """
+    ans = nn.Linear(*args, **kwargs)
+    with torch.no_grad():
+        ans.weight[:] *= initial_scale
+        if ans.bias is not None:
+            torch.nn.init.uniform_(ans.bias, -0.1 * initial_scale, 0.1 * initial_scale)
+    return ans
+
+
+def ScaledConv1d(
+    *args,
+    initial_scale: float = 1.0,
+    kernel_size: int = 3,
+    padding: str = "same",
+    **kwargs,
+) -> nn.Conv1d:
+    """
+    Behaves like a constructor of a modified version of nn.Conv1d
+    that gives an easy way to set the default initial parameter scale.
+
+    Args:
+        Accepts the standard args and kwargs that nn.Linear accepts
+        e.g. in_features, out_features, bias=False.
+
+        initial_scale: you can override this if you want to increase
+           or decrease the initial magnitude of the module's output
+           (affects the initialization of weight_scale and bias_scale).
+           Another option, if you want to do something like this, is
+           to re-initialize the parameters.
+    """
+    ans = nn.Conv1d(*args, kernel_size=kernel_size, padding=padding, **kwargs)
+    with torch.no_grad():
+        ans.weight[:] *= initial_scale
+        if ans.bias is not None:
+            torch.nn.init.uniform_(ans.bias, -0.1 * initial_scale, 0.1 * initial_scale)
+    return ans
+
+
+def TransposeScaledConv1d(
+    *args,
+    initial_scale: float = 1.0,
+    kernel_size: int = 3,
+    padding: str = "same",
+    **kwargs,
+) -> nn.Sequential:
+    """
+    Transpose -> ScaledConv1d
+    """
+    return nn.Sequential(
+        Transpose(),
+        ScaledConv1d(
+            *args,
+            initial_scale=initial_scale,
+            kernel_size=kernel_size,
+            padding=padding,
+            **kwargs,
+        ),
+    )
+
+
+def ScaledConv1dTranspose(
+    *args,
+    initial_scale: float = 1.0,
+    kernel_size: int = 3,
+    padding: str = "same",
+    **kwargs,
+) -> nn.Sequential:
+    """
+    Transpose -> ScaledConv1d
+    """
+    return nn.Sequential(
+        ScaledConv1d(
+            *args,
+            initial_scale=initial_scale,
+            kernel_size=kernel_size,
+            padding=padding,
+            **kwargs,
+        ),
+        Transpose(),
+    )
+
+
+def TransposeConv1d(
+    *args, kernel_size: int = 3, padding: str = "same", **kwargs
+) -> nn.Sequential:
+    """
+    Transpose -> Conv1d
+    """
+    return nn.Sequential(
+        Transpose(),
+        nn.Conv1d(*args, kernel_size=kernel_size, padding=padding, **kwargs),
+    )
+
+
+def Conv1dTranspose(
+    *args, kernel_size: int = 3, padding: str = "same", **kwargs
+) -> nn.Sequential:
+    """
+    ScaledConv1d -> Transpose
+    """
+    return nn.Sequential(
+        nn.Conv1d(*args, kernel_size=kernel_size, padding=padding, **kwargs),
+        Transpose(),
+    )
+
+
+class SRLinear(nn.Linear):
+    """https://arxiv.org/abs/2303.06296
+    Stabilizing Transformer Training by Preventing Attention Entropy Collapse
+    """
+
+    def __init__(self, in_features, out_features, bias=True, **kwargs):
+        super().__init__(in_features, out_features, bias=bias, **kwargs)
+        self.register_buffer(
+            "u", nn.functional.normalize(torch.randn(in_features), dim=0)
+        )
+        with torch.no_grad():
+            sigma = self.get_sigma()
+        self.register_buffer("spectral_norm", sigma)
+        self.sigma = nn.Parameter(torch.ones(1))
+
+    def get_sigma(self):
+        with torch.no_grad():
+            u = self.u
+            v = self.weight.mv(u)
+            v = nn.functional.normalize(v, dim=0)
+            u = self.weight.T.mv(v)
+            u = nn.functional.normalize(u, dim=0)
+            self.u.data.copy_(u)
+        return torch.einsum("c,cd,d->", v, self.weight, u)
+
+    def get_weight(self):
+        sigma = self.get_sigma()
+        if self.training:
+            self.spectral_norm.data.copy_(sigma)
+        weight = (self.sigma / sigma) * self.weight
+        return weight
+
+    def forward(self, x):
+        return nn.functional.linear(x, self.get_weight(), self.bias)
+
+
+class SRConv1d(SRLinear):
+    def __init__(
+        self,
+        in_features,
+        out_features,
+        kernel_size,
+        stride: int = 1,
+        padding: str = "same",
+        bias: bool = True,
+        **kwargs,
+    ):
+        in_features = in_features * kernel_size
+        super().__init__(in_features, out_features, bias=bias, **kwargs)
+        nn.init.kaiming_uniform_(self.weight, a=math.sqrt(5))
+        self.kernel_size = kernel_size
+        self.stride = stride
+        self.padding = padding
+
+    def forward(self, x):
+        in_features = self.in_features // self.kernel_size
+        weight = self.get_weight().view(
+            self.out_features, in_features, self.kernel_size
+        )
+        return nn.functional.conv1d(
+            x, weight, bias=self.bias, stride=self.stride, padding=self.padding
+        )
+
+
+def TransposeSRConv1d(
+    *args, kernel_size: int = 3, padding: str = "same", **kwargs
+) -> nn.Sequential:
+    """
+    Transpose -> SRConv1d
+    """
+    return nn.Sequential(
+        Transpose(),
+        SRConv1d(*args, kernel_size=kernel_size, padding=padding, **kwargs),
+    )
+
+
+def SRConv1dTranspose(
+    *args, kernel_size: int = 3, padding: str = "same", **kwargs
+) -> nn.Sequential:
+    """
+    SRConv1d -> Transpose
+    """
+    return nn.Sequential(
+        SRConv1d(*args, kernel_size=kernel_size, padding=padding, **kwargs),
+        Transpose(),
+    )
+
+
+class ActivationBalancer(torch.nn.Module):
+    """
+    Modifies the backpropped derivatives of a function to try to encourage, for
+    each channel, that it is positive at least a proportion `threshold` of the
+    time.  It does this by multiplying negative derivative values by up to
+    (1+max_factor), and positive derivative values by up to (1-max_factor),
+    interpolated from 1 at the threshold to those extremal values when none
+    of the inputs are positive.
+
+    Args:
+           num_channels: the number of channels
+           channel_dim: the dimension/axis corresponding to the channel, e.g.
+               -1, 0, 1, 2; will be interpreted as an offset from x.ndim if negative.
+           min_positive: the minimum, per channel, of the proportion of the time
+               that (x > 0), below which we start to modify the derivatives.
+           max_positive: the maximum, per channel, of the proportion of the time
+               that (x > 0), above which we start to modify the derivatives.
+           max_factor: the maximum factor by which we modify the derivatives for
+              either the sign constraint or the magnitude constraint;
+              e.g. with max_factor=0.02, the the derivatives would be multiplied by
+              values in the range [0.98..1.02].
+           sign_gain_factor: determines the 'gain' with which we increase the
+              change in gradient once the constraints on min_positive and max_positive
+              are violated.
+           scale_gain_factor: determines the 'gain' with which we increase the
+              change in gradient once the constraints on min_abs and max_abs
+              are violated.
+           min_abs:  the minimum average-absolute-value difference from the mean
+               value per channel, which we allow, before we start to modify
+               the derivatives to prevent this.
+           max_abs:  the maximum average-absolute-value difference from the mean
+               value per channel, which we allow, before we start to modify
+               the derivatives to prevent this.
+          min_prob: determines the minimum probability with which we modify the
+             gradients for the {min,max}_positive and {min,max}_abs constraints,
+             on each forward().  This is done randomly to prevent all layers
+             from doing it at the same time.  Early in training we may use
+             higher probabilities than this; it will decay to this value.
+    """
+
+    def __init__(
+        self,
+        num_channels: int,
+        channel_dim: int,
+        min_positive: float = 0.05,
+        max_positive: float = 0.95,
+        max_factor: float = 0.04,
+        sign_gain_factor: float = 0.01,
+        scale_gain_factor: float = 0.02,
+        min_abs: float = 0.2,
+        max_abs: float = 100.0,
+        min_prob: float = 0.1,
+    ):
+        super(ActivationBalancer, self).__init__()
+        self.num_channels = num_channels
+        self.channel_dim = channel_dim
+        self.min_positive = min_positive
+        self.max_positive = max_positive
+        self.max_factor = max_factor
+        self.min_abs = min_abs
+        self.max_abs = max_abs
+        self.min_prob = min_prob
+        self.sign_gain_factor = sign_gain_factor
+        self.scale_gain_factor = scale_gain_factor
+
+        # count measures how many times the forward() function has been called.
+        # We occasionally sync this to a tensor called `count`, that exists to
+        # make sure it is synced to disk when we load and save the model.
+        self.cpu_count = 0
+        self.register_buffer("count", torch.tensor(0, dtype=torch.int64))
+
+    def forward(self, x: Tensor) -> Tensor:
+        if torch.jit.is_scripting() or not x.requires_grad or torch.jit.is_tracing():
+            return _no_op(x)
+
+        count = self.cpu_count
+        self.cpu_count += 1
+
+        if random.random() < 0.01:
+            # Occasionally sync self.cpu_count with self.count.
+            # count affects the decay of 'prob'.  don't do this on every iter,
+            # because syncing with the GPU is slow.
+            self.cpu_count = max(self.cpu_count, self.count.item())
+            self.count.fill_(self.cpu_count)
+
+        # the prob of doing some work exponentially decreases from 0.5 till it hits
+        # a floor at min_prob (==0.1, by default)
+        prob = max(self.min_prob, 0.5 ** (1 + (count / 4000.0)))
+
+        if random.random() < prob:
+            sign_gain_factor = 0.5
+            if self.min_positive != 0.0 or self.max_positive != 1.0:
+                sign_factor = _compute_sign_factor(
+                    x,
+                    self.channel_dim,
+                    self.min_positive,
+                    self.max_positive,
+                    gain_factor=self.sign_gain_factor / prob,
+                    max_factor=self.max_factor,
+                )
+            else:
+                sign_factor = None
+
+            scale_factor = _compute_scale_factor(
+                x.detach(),
+                self.channel_dim,
+                min_abs=self.min_abs,
+                max_abs=self.max_abs,
+                gain_factor=self.scale_gain_factor / prob,
+                max_factor=self.max_factor,
+            )
+            return ActivationBalancerFunction.apply(
+                x,
+                scale_factor,
+                sign_factor,
+                self.channel_dim,
+            )
+        else:
+            return _no_op(x)
+
+
+def penalize_abs_values_gt(x: Tensor, limit: float, penalty: float) -> Tensor:
+    """
+    Returns x unmodified, but in backprop will put a penalty for the excess of
+    the absolute values of elements of x over the limit "limit".  E.g. if
+    limit == 10.0, then if x has any values over 10 it will get a penalty.
+
+    Caution: the value of this penalty will be affected by grad scaling used
+    in automatic mixed precision training.  For this reasons we use this,
+    it shouldn't really matter, or may even be helpful; we just use this
+    to disallow really implausible values of scores to be given to softmax.
+    """
+    x_sign = x.sign()
+    over_limit = (x.abs() - limit) > 0
+    # The following is a memory efficient way to penalize the absolute values of
+    # x that's over the limit.  (The memory efficiency comes when you think
+    # about which items torch needs to cache for the autograd, and which ones it
+    # can throw away).  The numerical value of aux_loss as computed here will
+    # actually be larger than it should be, by limit * over_limit.sum(), but it
+    # has the same derivative as the real aux_loss which is penalty * (x.abs() -
+    # limit).relu().
+    aux_loss = penalty * ((x_sign * over_limit).to(torch.int8) * x)
+    # note: we don't do sum() here on aux)_loss, but it's as if we had done
+    # sum() due to how with_loss() works.
+    x = with_loss(x, aux_loss)
+    # you must use x for something, or this will be ineffective.
+    return x
+
+
+def _diag(x: Tensor):  # like .diag(), but works for tensors with 3 dims.
+    if x.ndim == 2:
+        return x.diag()
+    else:
+        (batch, dim, dim) = x.shape
+        x = x.reshape(batch, dim * dim)
+        x = x[:, :: dim + 1]
+        assert x.shape == (batch, dim)
+        return x
+
+
+def _whitening_metric(x: Tensor, num_groups: int):
+    """
+    Computes the "whitening metric", a value which will be 1.0 if all the eigenvalues of
+    of the centered feature covariance are the same within each group's covariance matrix
+    and also between groups.
+    Args:
+        x: a Tensor of shape (*, num_channels)
+     num_groups:  the number of groups of channels, a number >=1 that divides num_channels
+    Returns:
+        Returns a scalar Tensor that will be 1.0 if the data is "perfectly white" and
+    greater than 1.0 otherwise.
+    """
+    assert x.dtype != torch.float16
+    x = x.reshape(-1, x.shape[-1])
+    (num_frames, num_channels) = x.shape
+    assert num_channels % num_groups == 0
+    channels_per_group = num_channels // num_groups
+    x = x.reshape(num_frames, num_groups, channels_per_group).transpose(0, 1)
+    # x now has shape (num_groups, num_frames, channels_per_group)
+    # subtract the mean so we use the centered, not uncentered, covariance.
+    # My experience has been that when we "mess with the gradients" like this,
+    # it's better not do anything that tries to move the mean around, because
+    # that can easily cause instability.
+    x = x - x.mean(dim=1, keepdim=True)
+    # x_covar: (num_groups, channels_per_group, channels_per_group)
+    x_covar = torch.matmul(x.transpose(1, 2), x)
+    x_covar_mean_diag = _diag(x_covar).mean()
+    # the following expression is what we'd get if we took the matrix product
+    # of each covariance and measured the mean of its trace, i.e.
+    # the same as _diag(torch.matmul(x_covar, x_covar)).mean().
+    x_covarsq_mean_diag = (x_covar**2).sum() / (num_groups * channels_per_group)
+    # this metric will be >= 1.0; the larger it is, the less 'white' the data was.
+    metric = x_covarsq_mean_diag / (x_covar_mean_diag**2 + 1.0e-20)
+    return metric
+
+
+class WhiteningPenaltyFunction(torch.autograd.Function):
+    @staticmethod
+    def forward(
+        ctx,
+        x: Tensor,
+        num_groups: int,
+        whitening_limit: float,
+        grad_scale: float,
+    ) -> Tensor:
+        ctx.save_for_backward(x)
+        ctx.num_groups = num_groups
+        ctx.whitening_limit = whitening_limit
+        ctx.grad_scale = grad_scale
+        return x
+
+    @staticmethod
+    def backward(ctx, x_grad: Tensor):
+        (x_orig,) = ctx.saved_tensors
+        with torch.enable_grad():
+            with torch.cuda.amp.autocast(enabled=False):
+                x_detached = x_orig.to(torch.float32).detach()
+                x_detached.requires_grad = True
+
+                metric = _whitening_metric(x_detached, ctx.num_groups)
+
+                if random.random() < 0.005 or __name__ == "__main__":
+                    logging.info(
+                        f"Whitening: num_groups={ctx.num_groups}, num_channels={x_orig.shape[-1]}, "
+                        f"metric={metric.item():.2f} vs. limit={ctx.whitening_limit}"
+                    )
+
+                (metric - ctx.whitening_limit).relu().backward()
+                penalty_grad = x_detached.grad
+                scale = ctx.grad_scale * (
+                    x_grad.to(torch.float32).norm() / (penalty_grad.norm() + 1.0e-20)
+                )
+                penalty_grad = penalty_grad * scale
+        return x_grad + penalty_grad.to(x_grad.dtype), None, None, None
+
+
+class Whiten(nn.Module):
+    def __init__(
+        self,
+        num_groups: int,
+        whitening_limit: float,
+        prob: Union[float, Tuple[float, float]],
+        grad_scale: float,
+    ):
+        """
+        Args:
+          num_groups: the number of groups to divide the channel dim into before
+            whitening.  We will attempt to make the feature covariance
+            within each group, after mean subtraction, as "white" as possible,
+            while having the same trace across all groups.
+         whitening_limit: a value greater than 1.0, that dictates how much
+           freedom we have to violate the constraints.  1.0 would mean perfectly
+           white, with exactly the same trace across groups; larger values
+           give more freedom.  E.g. 2.0.
+         prob: the probability with which we apply the gradient modification
+           (also affects the grad scale).  May be supplied as a float,
+           or as a pair (min_prob, max_prob)
+
+          grad_scale: determines the scale on the gradient term from this object,
+            relative to the rest of the gradient on the attention weights.
+            E.g. 0.02 (you may want to use smaller values than this if prob is large)
+        """
+        super(Whiten, self).__init__()
+        assert num_groups >= 1
+        assert whitening_limit >= 1
+        assert grad_scale >= 0
+        self.num_groups = num_groups
+        self.whitening_limit = whitening_limit
+        if isinstance(prob, float):
+            assert 0 < prob <= 1
+            self.prob = prob
+        else:
+            (self.min_prob, self.max_prob) = prob
+            assert 0 < self.min_prob < self.max_prob <= 1
+            self.prob = self.max_prob
+
+        self.grad_scale = grad_scale
+
+    def forward(self, x: Tensor) -> Tensor:
+        """
+        In the forward pass, this function just returns the input unmodified.
+        In the backward pass, it will modify the gradients to ensure that the
+        distribution in each group has close to (lambda times I) as the covariance
+        after mean subtraction, with the same lambda across groups.
+        For whitening_limit > 1, there will be more freedom to violate this
+        constraint.
+
+        Args:
+           x: the input of shape (*, num_channels)
+
+        Returns:
+            x, unmodified.   You should make sure
+        you use the returned value, or the graph will be freed
+        and nothing will happen in backprop.
+        """
+        if not x.requires_grad or random.random() > self.prob or self.grad_scale == 0:
+            return _no_op(x)
+        else:
+            if hasattr(self, "min_prob") and random.random() < 0.25:
+                # occasionally switch between min_prob and max_prob, based on whether
+                # we are above or below the threshold.
+                if (
+                    _whitening_metric(x.to(torch.float32), self.num_groups)
+                    > self.whitening_limit
+                ):
+                    # there would be a change to the grad.
+                    self.prob = self.max_prob
+                else:
+                    self.prob = self.min_prob
+
+            return WhiteningPenaltyFunction.apply(
+                x, self.num_groups, self.whitening_limit, self.grad_scale
+            )
+
+
+class WithLoss(torch.autograd.Function):
+    @staticmethod
+    def forward(ctx, x: Tensor, y: Tensor):
+        ctx.y_shape = y.shape
+        return x
+
+    @staticmethod
+    def backward(ctx, ans_grad: Tensor):
+        return ans_grad, torch.ones(
+            ctx.y_shape, dtype=ans_grad.dtype, device=ans_grad.device
+        )
+
+
+def with_loss(x, y):
+    if torch.jit.is_scripting() or torch.jit.is_tracing():
+        return x
+    # returns x but adds y.sum() to the loss function.
+    return WithLoss.apply(x, y)
+
+
+def _no_op(x: Tensor) -> Tensor:
+    if torch.jit.is_scripting() or torch.jit.is_tracing():
+        return x
+    else:
+        # a no-op function that will have a node in the autograd graph,
+        # to avoid certain bugs relating to backward hooks
+        return x.chunk(1, dim=-1)[0]
+
+
+class Identity(torch.nn.Module):
+    def __init__(self):
+        super(Identity, self).__init__()
+
+    def forward(self, x):
+        return _no_op(x)
+
+
+class MaxEig(torch.nn.Module):
+    """
+    Modifies the backpropped derivatives of a function to try to discourage
+    that any given direction in activation space accounts for more than
+    a specified proportion of the covariance (e.g. 0.2).
+
+
+    Args:
+           num_channels: the number of channels
+           channel_dim: the dimension/axis corresponding to the channel, e.g.
+               -1, 0, 1, 2; will be interpreted as an offset from x.ndim if negative.
+           max_var_per_eig:  the maximum proportion of the variance of the
+               features/channels, after mean subtraction, that can come from
+               any given eigenvalue.
+           min_prob: the minimum probability with which we apply this during any invocation
+               of forward(), assuming last time we applied the constraint it was
+               not active; supplied for speed.
+           scale: determines the scale with which we modify the gradients, relative
+               to the existing / unmodified gradients
+    """
+
+    def __init__(
+        self,
+        num_channels: int,
+        channel_dim: int,
+        max_var_per_eig: float = 0.2,
+        min_prob: float = 0.01,
+        scale: float = 0.01,
+    ):
+        super(MaxEig, self).__init__()
+        self.num_channels = num_channels
+        self.channel_dim = channel_dim
+        self.scale = scale
+        assert max_var_per_eig == 0.0 or max_var_per_eig > 1.0 / num_channels
+        self.max_var_per_eig = max_var_per_eig
+
+        # we figure out the dominant direction using the power method: starting with
+        # a random vector, keep multiplying by the covariance and renormalizing.
+        with torch.no_grad():
+            # arbitrary.. would use randn() but want to leave the rest of the model's
+            # random parameters unchanged for comparison
+            direction = torch.arange(num_channels).to(torch.float)
+            direction = direction / direction.norm()
+            self.register_buffer("max_eig_direction", direction)
+
+        self.min_prob = min_prob
+        # cur_prob is the current probability we'll use to apply the ActivationBalancer.
+        # We'll regress this towards prob, each time we try to apply it and it is not
+        # active.
+        self.cur_prob = 1.0
+
+    def forward(self, x: Tensor) -> Tensor:
+        if (
+            torch.jit.is_scripting()
+            or self.max_var_per_eig <= 0
+            or random.random() > self.cur_prob
+            or torch.jit.is_tracing()
+        ):
+            return _no_op(x)
+
+        with torch.cuda.amp.autocast(enabled=False):
+            eps = 1.0e-20
+            orig_x = x
+            x = x.to(torch.float32)
+            with torch.no_grad():
+                x = x.transpose(self.channel_dim, -1).reshape(-1, self.num_channels)
+                x = x - x.mean(dim=0)
+                new_direction, coeffs = self._find_direction_coeffs(
+                    x, self.max_eig_direction
+                )
+                x_var = (x**2).mean()
+                x_residual = x - coeffs * new_direction
+                x_residual_var = (x_residual**2).mean()
+
+                # `variance_proportion` is the proportion of the variance accounted for
+                # by the top eigen-direction.
+                variance_proportion = (x_var - x_residual_var) / (x_var + 1.0e-20)
+
+                # ensure new direction is nonzero even if x == 0, by including `direction`.
+                self._set_direction(0.1 * self.max_eig_direction + new_direction)
+
+            if random.random() < 0.01 or __name__ == "__main__":
+                logging.info(
+                    f"variance_proportion = {variance_proportion.item()}, shape={tuple(orig_x.shape)}, cur_prob={self.cur_prob}"
+                )
+
+            if variance_proportion >= self.max_var_per_eig:
+                # The constraint is active.  Note, we should quite rarely
+                # reach here, only near the beginning of training if we are
+                # starting to diverge, should this constraint be active.
+                cur_prob = self.cur_prob
+                self.cur_prob = 1.0  # next time, do the update with probability 1.0.
+                return MaxEigLimiterFunction.apply(
+                    orig_x, coeffs, new_direction, self.channel_dim, self.scale
+                )
+            else:
+                # let self.cur_prob exponentially approach self.min_prob, as
+                # long as the constraint is inactive.
+                self.cur_prob = 0.75 * self.cur_prob + 0.25 * self.min_prob
+                return orig_x
+
+    def _set_direction(self, direction: Tensor):
+        """
+        Sets self.max_eig_direction to a normalized version of `direction`
+        """
+        direction = direction.detach()
+        direction = direction / direction.norm()
+        direction_sum = direction.sum().item()
+        if direction_sum - direction_sum == 0:  # no inf/nan
+            self.max_eig_direction[:] = direction
+        else:
+            logging.info(
+                f"Warning: sum of direction in MaxEig is {direction_sum}, "
+                "num_channels={self.num_channels}, channel_dim={self.channel_dim}"
+            )
+
+    def _find_direction_coeffs(
+        self, x: Tensor, prev_direction: Tensor
+    ) -> Tuple[Tensor, Tensor, Tensor]:
+        """
+            Figure out (an approximation to) the proportion of the variance of a set of
+            feature vectors that can be attributed to the top eigen-direction.
+            Args:
+             x: a Tensor of shape (num_frames, num_channels), with num_frames > 1.
+          prev_direction:  a Tensor of shape (num_channels,), that is our previous estimate
+                   of the top eigen-direction, or a random direction if this is the first
+                   iteration.  Does not have to be normalized, but should be nonzero.
+
+        Returns: (cur_direction, coeffs), where:
+             cur_direction: a Tensor of shape (num_channels,) that is the current
+                estimate of the top eigen-direction.
+             coeffs: a Tensor of shape (num_frames, 1) that minimizes, or
+                approximately minimizes, (x - coeffs * cur_direction).norm()
+        """
+        (num_frames, num_channels) = x.shape
+        assert num_channels > 1 and num_frames > 1
+        assert prev_direction.shape == (num_channels,)
+        # `coeffs` are the coefficients of `prev_direction` in x.
+        # actually represent the coeffs up to a constant positive factor.
+        coeffs = (x * prev_direction).sum(dim=1, keepdim=True) + 1.0e-10
+        cur_direction = (x * coeffs).sum(dim=0) / ((coeffs**2).sum() + 1.0e-20)
+        return cur_direction, coeffs
+
+
+class DoubleSwishFunction(torch.autograd.Function):
+    """
+      double_swish(x) = x * torch.sigmoid(x-1)
+    This is a definition, originally motivated by its close numerical
+    similarity to swish(swish(x)), where swish(x) =  x * sigmoid(x).
+
+    Memory-efficient derivative computation:
+     double_swish(x) = x * s, where s(x) = torch.sigmoid(x-1)
+     double_swish'(x) = d/dx double_swish(x) =  x * s'(x) + x' * s(x) = x * s'(x) + s(x).
+     Now, s'(x) = s(x) * (1-s(x)).
+     double_swish'(x) =  x * s'(x) + s(x).
+                      =  x * s(x) * (1-s(x)) + s(x).
+                     = double_swish(x) * (1-s(x)) + s(x)
+     ... so we just need to remember s(x) but not x itself.
+    """
+
+    @staticmethod
+    def forward(ctx, x: Tensor) -> Tensor:
+        requires_grad = x.requires_grad
+        x_dtype = x.dtype
+        if x.dtype == torch.float16:
+            x = x.to(torch.float32)
+
+        s = torch.sigmoid(x - 1.0)
+        y = x * s
+
+        if requires_grad:
+            deriv = y * (1 - s) + s
+            # notes on derivative of x * sigmoid(x - 1):
+            # https://www.wolframalpha.com/input?i=d%2Fdx+%28x+*+sigmoid%28x-1%29%29
+            # min \simeq -0.043638.  Take floor as -0.043637 so it's a lower bund
+            # max \simeq 1.1990.   Take ceil to be 1.2 so it's an upper bound.
+            # the combination of "+ torch.rand_like(deriv)" and casting to torch.uint8 (which
+            # floors), should be expectation-preserving.
+            floor = -0.043637
+            ceil = 1.2
+            d_scaled = (deriv - floor) * (255.0 / (ceil - floor)) + torch.rand_like(
+                deriv
+            )
+            if __name__ == "__main__":
+                # for self-testing only.
+                assert d_scaled.min() >= 0.0
+                assert d_scaled.max() < 256.0
+            d_int = d_scaled.to(torch.uint8)
+            ctx.save_for_backward(d_int)
+        if x.dtype == torch.float16 or torch.is_autocast_enabled():
+            y = y.to(torch.float16)
+        return y
+
+    @staticmethod
+    def backward(ctx, y_grad: Tensor) -> Tensor:
+        (d,) = ctx.saved_tensors
+        # the same constants as used in forward pass.
+        floor = -0.043637
+        ceil = 1.2
+        d = d * ((ceil - floor) / 255.0) + floor
+        return y_grad * d
+
+
+class DoubleSwish(torch.nn.Module):
+    def forward(self, x: Tensor) -> Tensor:
+        """Return double-swish activation function which is an approximation to Swish(Swish(x)),
+        that we approximate closely with x * sigmoid(x-1).
+        """
+        if torch.jit.is_scripting() or torch.jit.is_tracing():
+            return x * torch.sigmoid(x - 1.0)
+        return DoubleSwishFunction.apply(x)
+
+
+def BalancedDoubleSwish(
+    d_model, channel_dim=-1, max_abs=10.0, min_prob=0.25
+) -> nn.Sequential:
+    """
+    ActivationBalancer -> DoubleSwish
+    """
+    balancer = ActivationBalancer(
+        d_model, channel_dim=channel_dim, max_abs=max_abs, min_prob=min_prob
+    )
+    return nn.Sequential(
+        balancer,
+        DoubleSwish(),
+    )
+
+
+def _test_max_eig():
+    for proportion in [0.1, 0.5, 10.0]:
+        logging.info(f"proportion = {proportion}")
+        x = torch.randn(100, 128)
+        direction = torch.randn(128)
+        coeffs = torch.randn(100, 1)
+        x += proportion * direction * coeffs
+
+        x.requires_grad = True
+
+        num_channels = 128
+        m = MaxEig(
+            num_channels,
+            1,
+            0.5,
+            scale=0.1,  # channel_dim  # max_var_per_eig
+        )  # grad_scale
+
+        for _ in range(4):
+            y = m(x)
+
+        y_grad = torch.randn_like(x)
+        y.backward(gradient=y_grad)
+
+        if proportion < 0.2:
+            assert torch.allclose(x.grad, y_grad, atol=1.0e-02)
+        elif proportion > 1.0:
+            assert not torch.allclose(x.grad, y_grad)
+
+
+def _test_whiten():
+    for proportion in [0.1, 0.5, 10.0]:
+        logging.info(f"_test_whiten(): proportion = {proportion}")
+        x = torch.randn(100, 128)
+        direction = torch.randn(128)
+        coeffs = torch.randn(100, 1)
+        x += proportion * direction * coeffs
+
+        x.requires_grad = True
+
+        num_channels = 128
+        m = Whiten(
+            1,
+            5.0,
+            prob=1.0,
+            grad_scale=0.1,  # num_groups  # whitening_limit,
+        )  # grad_scale
+
+        for _ in range(4):
+            y = m(x)
+
+        y_grad = torch.randn_like(x)
+        y.backward(gradient=y_grad)
+
+        if proportion < 0.2:
+            assert torch.allclose(x.grad, y_grad)
+        elif proportion > 1.0:
+            assert not torch.allclose(x.grad, y_grad)
+
+
+def _test_activation_balancer_sign():
+    probs = torch.arange(0, 1, 0.01)
+    N = 1000
+    x = 1.0 * ((2.0 * (torch.rand(probs.numel(), N) < probs.unsqueeze(-1))) - 1.0)
+    x = x.detach()
+    x.requires_grad = True
+    m = ActivationBalancer(
+        probs.numel(),
+        channel_dim=0,
+        min_positive=0.05,
+        max_positive=0.95,
+        max_factor=0.2,
+        min_abs=0.0,
+    )
+
+    y_grad = torch.sign(torch.randn(probs.numel(), N))
+
+    y = m(x)
+    y.backward(gradient=y_grad)
+    print("_test_activation_balancer_sign: x = ", x)
+    print("_test_activation_balancer_sign: y grad = ", y_grad)
+    print("_test_activation_balancer_sign: x grad = ", x.grad)
+
+
+def _test_activation_balancer_magnitude():
+    magnitudes = torch.arange(0, 1, 0.01)
+    N = 1000
+    x = torch.sign(torch.randn(magnitudes.numel(), N)) * magnitudes.unsqueeze(-1)
+    x = x.detach()
+    x.requires_grad = True
+    m = ActivationBalancer(
+        magnitudes.numel(),
+        channel_dim=0,
+        min_positive=0.0,
+        max_positive=1.0,
+        max_factor=0.2,
+        min_abs=0.2,
+        max_abs=0.8,
+        min_prob=1.0,
+    )
+
+    y_grad = torch.sign(torch.randn(magnitudes.numel(), N))
+
+    y = m(x)
+    y.backward(gradient=y_grad)
+    print("_test_activation_balancer_magnitude: x = ", x)
+    print("_test_activation_balancer_magnitude: y grad = ", y_grad)
+    print("_test_activation_balancer_magnitude: x grad = ", x.grad)
+
+
+def _test_basic_norm():
+    num_channels = 128
+    m = BasicNorm(num_channels=num_channels, channel_dim=1)
+
+    x = torch.randn(500, num_channels)
+
+    y = m(x)
+
+    assert y.shape == x.shape
+    x_rms = (x**2).mean().sqrt()
+    y_rms = (y**2).mean().sqrt()
+    print("x rms = ", x_rms)
+    print("y rms = ", y_rms)
+    assert y_rms < x_rms
+    assert y_rms > 0.5 * x_rms
+
+
+def _test_double_swish_deriv():
+    x = torch.randn(10, 12, dtype=torch.double) * 3.0
+    x.requires_grad = True
+    m = DoubleSwish()
+
+    tol = (1.2 - (-0.043637)) / 255.0
+    torch.autograd.gradcheck(m, x, atol=tol)
+
+    # for self-test.
+    x = torch.randn(1000, 1000, dtype=torch.double) * 3.0
+    x.requires_grad = True
+    y = m(x)
+
+
+def _test_softmax():
+    a = torch.randn(2, 10, dtype=torch.float64)
+    b = a.clone()
+    a.requires_grad = True
+    b.requires_grad = True
+    a.softmax(dim=1)[:, 0].sum().backward()
+    print("a grad = ", a.grad)
+    softmax(b, dim=1)[:, 0].sum().backward()
+    print("b grad = ", b.grad)
+    assert torch.allclose(a.grad, b.grad)
+
+
+if __name__ == "__main__":
+    logging.getLogger().setLevel(logging.INFO)
+    torch.set_num_threads(1)
+    torch.set_num_interop_threads(1)
+    _test_softmax()
+    _test_whiten()
+    _test_max_eig()
+    _test_activation_balancer_sign()
+    _test_activation_balancer_magnitude()
+    _test_basic_norm()
+    _test_double_swish_deriv()

apps/audio_cloning/vallex/modules/scheduler.py

@@ -0,0 +1,78 @@
+#!/usr/bin/env python3
+# Copyright    2023                           (authors: Feiteng Li)
+#
+# See ../../../../LICENSE for clarification regarding multiple authors
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+
+import torch
+
+from modules.optim import Eden
+
+
+def calc_lr(step, dim_embed, warmup_steps):
+    return dim_embed ** (-0.5) * min(
+        step ** (-0.5), step * warmup_steps ** (-1.5)
+    )
+
+
+class NoamScheduler(torch.optim.lr_scheduler._LRScheduler):
+    def __init__(
+        self,
+        base_lr: float,
+        optimizer: torch.optim.Optimizer,
+        dim_embed: int,
+        warmup_steps: int,
+        last_epoch: int = -1,
+        verbose: bool = False,
+    ) -> None:
+
+        self.dim_embed = dim_embed
+        self.base_lr = base_lr
+        self.warmup_steps = warmup_steps
+        self.num_param_groups = len(optimizer.param_groups)
+
+        super().__init__(optimizer, last_epoch, verbose)
+
+    def get_lr(self) -> float:
+        lr = self.base_lr * calc_lr(
+            self._step_count, self.dim_embed, self.warmup_steps
+        )
+        return [lr] * self.num_param_groups
+
+    def set_step(self, step: int):
+        self._step_count = step
+
+
+def get_scheduler(params, optimizer):
+    if params.scheduler_name.lower() == "eden":
+        scheduler = Eden(optimizer, 5000, 4, warmup_batches=params.warmup_steps)
+    elif params.scheduler_name.lower() == "noam":
+        scheduler = NoamScheduler(
+            params.base_lr,
+            optimizer,
+            params.decoder_dim,
+            warmup_steps=params.warmup_steps,
+        )
+        # scheduler.set_step(params.start_batch or params.batch_idx_train)
+    elif params.scheduler_name.lower() == "cosine":
+        scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(
+            params.warmup_steps,
+            optimizer,
+            eta_min=params.base_lr,
+        )
+    else:
+        raise NotImplementedError(f"{params.scheduler_name}")
+
+    return scheduler

apps/audio_cloning/vallex/modules/transformer.py

@@ -0,0 +1,683 @@
+import copy
+import numbers
+from functools import partial
+from typing import Any, Callable, List, Optional, Tuple, Union
+
+import torch
+from torch import Tensor, nn
+from torch.nn import functional as F
+
+from .activation import MultiheadAttention
+from .scaling import ActivationBalancer, BalancedDoubleSwish
+from .scaling import BasicNorm as _BasicNorm
+
+_shape_t = Union[int, List[int], torch.Size]
+
+
+class LayerNorm(nn.Module):
+    __constants__ = ["normalized_shape", "eps", "elementwise_affine"]
+    normalized_shape: Tuple[int, ...]
+    eps: float
+    elementwise_affine: bool
+
+    def __init__(
+        self,
+        normalized_shape: _shape_t,
+        eps: float = 1e-5,
+        elementwise_affine: bool = True,
+        device=None,
+        dtype=None,
+    ) -> None:
+        factory_kwargs = {"device": device, "dtype": dtype}
+        super(LayerNorm, self).__init__()
+        if isinstance(normalized_shape, numbers.Integral):
+            # mypy error: incompatible types in assignment
+            normalized_shape = (normalized_shape,)  # type: ignore[assignment]
+        self.normalized_shape = tuple(normalized_shape)  # type: ignore[arg-type]
+        self.eps = eps
+        self.elementwise_affine = elementwise_affine
+        if self.elementwise_affine:
+            self.weight = nn.Parameter(
+                torch.empty(self.normalized_shape, **factory_kwargs)
+            )
+            self.bias = nn.Parameter(
+                torch.empty(self.normalized_shape, **factory_kwargs)
+            )
+        else:
+            self.register_parameter("weight", None)
+            self.register_parameter("bias", None)
+
+        self.reset_parameters()
+
+    def reset_parameters(self) -> None:
+        if self.elementwise_affine:
+            nn.init.ones_(self.weight)
+            nn.init.zeros_(self.bias)
+
+    def forward(self, input: Tensor, embedding: Any = None) -> Tensor:
+        if isinstance(input, tuple):
+            input, embedding = input
+            return (
+                F.layer_norm(
+                    input,
+                    self.normalized_shape,
+                    self.weight,
+                    self.bias,
+                    self.eps,
+                ),
+                embedding,
+            )
+
+        assert embedding is None
+        return F.layer_norm(
+            input, self.normalized_shape, self.weight, self.bias, self.eps
+        )
+
+    def extra_repr(self) -> str:
+        return (
+            "{normalized_shape}, eps={eps}, "
+            "elementwise_affine={elementwise_affine}".format(**self.__dict__)
+        )
+
+
+class AdaptiveLayerNorm(nn.Module):
+    r"""Adaptive Layer Normalization"""
+
+    def __init__(self, d_model, norm) -> None:
+        super(AdaptiveLayerNorm, self).__init__()
+        self.project_layer = nn.Linear(d_model, 2 * d_model)
+        self.norm = norm
+        self.d_model = d_model
+        self.eps = self.norm.eps
+
+    def forward(self, input: Tensor, embedding: Tensor = None) -> Tensor:
+        if isinstance(input, tuple):
+            input, embedding = input
+            weight, bias = torch.split(
+                self.project_layer(embedding),
+                split_size_or_sections=self.d_model,
+                dim=-1,
+            )
+            return (weight * self.norm(input) + bias, embedding)
+
+        weight, bias = torch.split(
+            self.project_layer(embedding),
+            split_size_or_sections=self.d_model,
+            dim=-1,
+        )
+        return weight * self.norm(input) + bias
+
+
+class BasicNorm(_BasicNorm):
+    def __init__(
+        self,
+        d_model: int,
+        eps: float = 1e-5,
+        device=None,
+        dtype=None,
+    ):
+        super(BasicNorm, self).__init__(d_model, eps=eps)
+
+    def forward(self, input: Tensor, embedding: Any = None) -> Tensor:
+        if isinstance(input, tuple):
+            input, embedding = input
+            return (
+                super(BasicNorm, self).forward(input),
+                embedding,
+            )
+
+        assert embedding is None
+        return super(BasicNorm, self).forward(input)
+
+
+class BalancedBasicNorm(nn.Module):
+    def __init__(
+        self,
+        d_model: int,
+        eps: float = 1e-5,
+        device=None,
+        dtype=None,
+    ):
+        super(BalancedBasicNorm, self).__init__()
+        self.balancer = ActivationBalancer(
+            d_model,
+            channel_dim=-1,
+            min_positive=0.45,
+            max_positive=0.55,
+            max_abs=6.0,
+        )
+        self.norm = BasicNorm(d_model, eps, device=device, dtype=dtype)
+
+    def forward(self, input: Tensor, embedding: Any = None) -> Tensor:
+        if isinstance(input, tuple):
+            input, embedding = input
+            return self.norm((self.balancer(input), embedding))
+
+        assert embedding is None
+        return self.norm(self.balancer(input))
+
+
+class IdentityNorm(nn.Module):
+    def __init__(
+        self,
+        d_model: int,
+        eps: float = 1e-5,
+        device=None,
+        dtype=None,
+    ) -> None:
+        super(IdentityNorm, self).__init__()
+
+    def forward(self, input: Tensor, embedding: Any = None) -> Tensor:
+        if isinstance(input, tuple):
+            return input
+
+        assert embedding is None
+        return input
+
+
+class TransformerEncoderLayer(nn.Module):
+    __constants__ = ["batch_first", "norm_first"]
+
+    def __init__(
+        self,
+        d_model: int,
+        nhead: int,
+        dim_feedforward: int = 2048,
+        dropout: float = 0.1,
+        activation: Union[str, Callable[[Tensor], Tensor]] = F.relu,
+        batch_first: bool = False,
+        norm_first: bool = False,
+        device=None,
+        dtype=None,
+        linear1_self_attention_cls: nn.Module = nn.Linear,
+        linear2_self_attention_cls: nn.Module = nn.Linear,
+        linear1_feedforward_cls: nn.Module = nn.Linear,
+        linear2_feedforward_cls: nn.Module = nn.Linear,
+        layer_norm_cls: nn.Module = LayerNorm,
+        layer_norm_eps: float = 1e-5,
+        adaptive_layer_norm=False,
+    ) -> None:
+        factory_kwargs = {"device": device, "dtype": dtype}
+        super(TransformerEncoderLayer, self).__init__()
+        self.self_attn = MultiheadAttention(
+            d_model,
+            nhead,
+            dropout=dropout,
+            batch_first=batch_first,
+            linear1_cls=linear1_self_attention_cls,
+            linear2_cls=linear2_self_attention_cls,
+            **factory_kwargs,
+        )
+
+        # Implementation of Feedforward model
+        self.linear1 = linear1_feedforward_cls(
+            d_model, dim_feedforward, **factory_kwargs
+        )
+        self.dropout = nn.Dropout(dropout)
+        self.linear2 = linear2_feedforward_cls(
+            dim_feedforward, d_model, **factory_kwargs
+        )
+
+        self.norm_first = norm_first
+        self.dropout1 = nn.Dropout(dropout)
+        self.dropout2 = nn.Dropout(dropout)
+
+        # Legacy string support for activation function.
+        if isinstance(activation, str):
+            activation = _get_activation_fn(activation)
+        elif isinstance(activation, partial):
+            activation = activation(d_model)
+        elif activation == BalancedDoubleSwish:
+            activation = BalancedDoubleSwish(d_model)
+
+        # # We can't test self.activation in forward() in TorchScript,
+        # # so stash some information about it instead.
+        # if activation is F.relu or isinstance(activation, torch.nn.ReLU):
+        #     self.activation_relu_or_gelu = 1
+        # elif activation is F.gelu or isinstance(activation, torch.nn.GELU):
+        #     self.activation_relu_or_gelu = 2
+        # else:
+        #     self.activation_relu_or_gelu = 0
+        self.activation = activation
+
+        norm1 = layer_norm_cls(d_model, eps=layer_norm_eps, **factory_kwargs)
+        if layer_norm_cls == IdentityNorm:
+            norm2 = BalancedBasicNorm(
+                d_model, eps=layer_norm_eps, **factory_kwargs
+            )
+        else:
+            norm2 = layer_norm_cls(
+                d_model, eps=layer_norm_eps, **factory_kwargs
+            )
+
+        if adaptive_layer_norm:
+            self.norm1 = AdaptiveLayerNorm(d_model, norm1)
+            self.norm2 = AdaptiveLayerNorm(d_model, norm2)
+        else:
+            self.norm1 = norm1
+            self.norm2 = norm2
+
+    def __setstate__(self, state):
+        super(TransformerEncoderLayer, self).__setstate__(state)
+        if not hasattr(self, "activation"):
+            self.activation = F.relu
+
+    def forward(
+        self,
+        src: Tensor,
+        src_mask: Optional[Tensor] = None,
+        src_key_padding_mask: Optional[Tensor] = None,
+    ) -> Tensor:
+        r"""Pass the input through the encoder layer.
+
+        Args:
+            src: the sequence to the encoder layer (required).
+            src_mask: the mask for the src sequence (optional).
+            src_key_padding_mask: the mask for the src keys per batch (optional).
+
+        Shape:
+            see the docs in Transformer class.
+        """
+        x, stage_embedding = src, None
+        is_src_tuple = False
+        if isinstance(src, tuple):
+            x, stage_embedding = src
+            is_src_tuple = True
+
+        if src_key_padding_mask is not None:
+            _skpm_dtype = src_key_padding_mask.dtype
+            if _skpm_dtype != torch.bool and not torch.is_floating_point(
+                src_key_padding_mask
+            ):
+                raise AssertionError(
+                    "only bool and floating types of key_padding_mask are supported"
+                )
+
+        if self.norm_first:
+            x = x + self._sa_block(
+                self.norm1(x, stage_embedding),
+                src_mask,
+                src_key_padding_mask,
+            )
+            x = x + self._ff_block(self.norm2(x, stage_embedding))
+        else:
+            x = self.norm1(
+                x + self._sa_block(x, src_mask, src_key_padding_mask),
+                stage_embedding,
+            )
+            x = self.norm2(x + self._ff_block(x), stage_embedding)
+
+        if is_src_tuple:
+            return (x, stage_embedding)
+        return x
+
+    def infer(
+        self,
+        src: Tensor,
+        src_mask: Optional[Tensor] = None,
+        src_key_padding_mask: Optional[Tensor] = None,
+        past_kv: Optional[Tensor] = None,
+        use_cache: bool = False,
+    ):
+        x, stage_embedding = src, None
+        is_src_tuple = False
+        if isinstance(src, tuple):
+            x, stage_embedding = src
+            is_src_tuple = True
+
+        if src_key_padding_mask is not None:
+            _skpm_dtype = src_key_padding_mask.dtype
+            if _skpm_dtype != torch.bool and not torch.is_floating_point(
+                src_key_padding_mask
+            ):
+                raise AssertionError(
+                    "only bool and floating types of key_padding_mask are supported"
+                )
+
+        if self.norm_first:
+            x_attn_out, kv = self.self_attn.infer(
+                self.norm1(x, stage_embedding),
+                attn_mask=src_mask,
+                key_padding_mask=src_key_padding_mask,
+                need_weights=False,
+                past_kv=past_kv,
+                use_cache=use_cache,
+            )
+            x = x + x_attn_out
+            x = x + self._ff_block(self.norm2(x, stage_embedding))
+
+        if is_src_tuple:
+            return (x, stage_embedding)
+        return (x, kv)
+
+    # self-attention block
+    def _sa_block(
+        self,
+        x: Tensor,
+        attn_mask: Optional[Tensor],
+        key_padding_mask: Optional[Tensor],
+    ) -> Tensor:
+        x = self.self_attn(
+            x,
+            x,
+            x,
+            attn_mask=attn_mask,
+            key_padding_mask=key_padding_mask,
+            need_weights=False,
+        )[0]
+        return self.dropout1(x)
+
+    # feed forward block
+    def _ff_block(self, x: Tensor) -> Tensor:
+        x = self.linear2(self.dropout(self.activation(self.linear1(x))))
+        return self.dropout2(x)
+
+
+class TransformerEncoder(nn.Module):
+    r"""TransformerEncoder is a stack of N encoder layers. Users can build the
+    BERT(https://arxiv.org/abs/1810.04805) model with corresponding parameters.
+
+    Args:
+        encoder_layer: an instance of the TransformerEncoderLayer() class (required).
+        num_layers: the number of sub-encoder-layers in the encoder (required).
+        norm: the layer normalization component (optional).
+        enable_nested_tensor: if True, input will automatically convert to nested tensor
+            (and convert back on output). This will improve the overall performance of
+            TransformerEncoder when padding rate is high. Default: ``True`` (enabled).
+
+    Examples::
+        >>> encoder_layer = TransformerEncoderLayer(d_model=512, nhead=8)
+        >>> transformer_encoder = TransformerEncoder(encoder_layer, num_layers=6)
+        >>> src = torch.rand(10, 32, 512)
+        >>> out = transformer_encoder(src)
+    """
+    __constants__ = ["norm"]
+
+    def __init__(self, encoder_layer, num_layers, norm=None):
+        super(TransformerEncoder, self).__init__()
+        self.layers = _get_clones(encoder_layer, num_layers)
+        self.num_layers = num_layers
+        self.norm = norm
+
+    def forward(
+        self,
+        src: Tensor,
+        mask: Optional[Tensor] = None,
+        src_key_padding_mask: Optional[Tensor] = None,
+        return_layer_states: bool = False,
+    ) -> Tensor:
+        r"""Pass the input through the encoder layers in turn.
+
+        Args:
+            src: the sequence to the encoder (required).
+            mask: the mask for the src sequence (optional).
+            src_key_padding_mask: the mask for the src keys per batch (optional).
+            return_layer_states: return layers' state (optional).
+
+        Shape:
+            see the docs in Transformer class.
+        """
+        if return_layer_states:
+            layer_states = []  # layers' output
+            output = src
+            for mod in self.layers:
+                output = mod(
+                    output,
+                    src_mask=mask,
+                    src_key_padding_mask=src_key_padding_mask,
+                )
+                layer_states.append(output[0])
+
+            if self.norm is not None:
+                output = self.norm(output)
+
+            return layer_states, output
+
+        output = src
+        for mod in self.layers:
+            output = mod(
+                output, src_mask=mask, src_key_padding_mask=src_key_padding_mask
+            )
+
+        if self.norm is not None:
+            output = self.norm(output)
+
+        return output
+
+    def infer(
+        self,
+        src: Tensor,
+        mask: Optional[Tensor] = None,
+        src_key_padding_mask: Optional[Tensor] = None,
+        return_layer_states: bool = False,
+        past_kv: Optional[Tensor] = None,
+        use_cache: bool = False,
+    ):
+        if past_kv is None:
+            past_length = 0
+            past_kv = tuple([None] * self.num_layers)
+        else:
+            past_length = past_kv[0][0].size(-2)
+        new_kv = () if use_cache else None
+        output = src
+        for mod, past_layer_kv in zip(self.layers, past_kv):
+            output, kv = mod.infer(
+                output, src_mask=mask, src_key_padding_mask=src_key_padding_mask, past_kv=past_layer_kv, use_cache=use_cache
+            )
+            if use_cache:
+                new_kv = new_kv + (kv,)
+
+        if self.norm is not None:
+            output = self.norm(output)
+
+        return output, new_kv
+
+
+class TransformerDecoderLayer(nn.Module):
+    __constants__ = ["batch_first", "norm_first"]
+
+    def __init__(
+        self,
+        d_model: int,
+        nhead: int,
+        dim_feedforward: int = 2048,
+        dropout: float = 0.1,
+        activation: Union[str, Callable[[Tensor], Tensor]] = F.relu,
+        linear1_self_attention_cls: nn.Module = nn.Linear,
+        linear2_self_attention_cls: nn.Module = nn.Linear,
+        linear1_feedforward_cls: nn.Module = nn.Linear,
+        linear2_feedforward_cls: nn.Module = nn.Linear,
+        batch_first: bool = False,
+        norm_first: bool = False,
+        device=None,
+        dtype=None,
+        layer_norm_cls: nn.Module = LayerNorm,
+        layer_norm_eps: float = 1e-5,
+        adaptive_layer_norm=False,
+    ) -> None:
+        factory_kwargs = {"device": device, "dtype": dtype}
+        super(TransformerDecoderLayer, self).__init__()
+        self.self_attn = MultiheadAttention(
+            d_model,
+            nhead,
+            dropout=dropout,
+            batch_first=batch_first,
+            linear1_cls=linear1_self_attention_cls,
+            linear2_cls=linear2_self_attention_cls,
+            **factory_kwargs,
+        )
+        self.multihead_attn = MultiheadAttention(
+            d_model,
+            nhead,
+            dropout=dropout,
+            batch_first=batch_first,
+            linear1_cls=linear1_self_attention_cls,
+            linear2_cls=linear2_self_attention_cls,
+            **factory_kwargs,
+        )
+        # Implementation of Feedforward model
+        self.linear1 = linear1_feedforward_cls(
+            d_model, dim_feedforward, **factory_kwargs
+        )
+        self.dropout = nn.Dropout(dropout)
+        self.linear2 = linear2_feedforward_cls(
+            dim_feedforward, d_model, **factory_kwargs
+        )
+
+        self.norm_first = norm_first
+        self.dropout1 = nn.Dropout(dropout)
+        self.dropout2 = nn.Dropout(dropout)
+        self.dropout3 = nn.Dropout(dropout)
+
+        # Legacy string support for activation function.
+        if isinstance(activation, str):
+            self.activation = _get_activation_fn(activation)
+        elif isinstance(activation, partial):
+            self.activation = activation(d_model)
+        elif activation == BalancedDoubleSwish:
+            self.activation = BalancedDoubleSwish(d_model)
+        else:
+            self.activation = activation
+
+        if adaptive_layer_norm:
+            norm1 = layer_norm_cls(
+                d_model, eps=layer_norm_eps, **factory_kwargs
+            )
+            norm2 = layer_norm_cls(
+                d_model, eps=layer_norm_eps, **factory_kwargs
+            )
+            norm3 = layer_norm_cls(
+                d_model, eps=layer_norm_eps, **factory_kwargs
+            )
+
+            self.norm1 = AdaptiveLayerNorm(d_model, norm1)
+            self.norm2 = AdaptiveLayerNorm(d_model, norm2)
+            self.norm3 = AdaptiveLayerNorm(d_model, norm3)
+        else:
+            self.norm1 = layer_norm_cls(
+                d_model, eps=layer_norm_eps, **factory_kwargs
+            )
+            self.norm2 = layer_norm_cls(
+                d_model, eps=layer_norm_eps, **factory_kwargs
+            )
+            if layer_norm_cls == IdentityNorm:
+                self.norm3 = BalancedBasicNorm(
+                    d_model, eps=layer_norm_eps, **factory_kwargs
+                )
+            else:
+                self.norm3 = layer_norm_cls(
+                    d_model, eps=layer_norm_eps, **factory_kwargs
+                )
+
+    def forward(
+        self,
+        tgt: Tensor,
+        memory: Tensor,
+        tgt_mask: Optional[Tensor] = None,
+        memory_mask: Optional[Tensor] = None,
+        tgt_key_padding_mask: Optional[Tensor] = None,
+        memory_key_padding_mask: Optional[Tensor] = None,
+    ) -> Tensor:
+        r"""Pass the inputs (and mask) through the decoder layer.
+
+        Args:
+            tgt: the sequence to the decoder layer (required).
+            memory: the sequence from the last layer of the encoder (required).
+            tgt_mask: the mask for the tgt sequence (optional).
+            memory_mask: the mask for the memory sequence (optional).
+            tgt_key_padding_mask: the mask for the tgt keys per batch (optional).
+            memory_key_padding_mask: the mask for the memory keys per batch (optional).
+
+        Shape:
+            see the docs in Transformer class.
+        """
+        tgt_is_tuple = False
+        if isinstance(tgt, tuple):
+            x, stage_embedding = tgt
+            tgt_is_tuple = True
+        else:
+            x, stage_embedding = tgt, None
+
+        if self.norm_first:
+            x = x + self._sa_block(
+                self.norm1(x, stage_embedding), tgt_mask, tgt_key_padding_mask
+            )
+            x = x + self._mha_block(
+                self.norm2(x, stage_embedding),
+                memory,
+                memory_mask,
+                memory_key_padding_mask,
+            )
+            x = x + self._ff_block(self.norm3(x, stage_embedding))
+        else:
+            x = self.norm1(
+                x + self._sa_block(x, tgt_mask, tgt_key_padding_mask),
+                stage_embedding,
+            )
+            x = self.norm2(
+                x
+                + self._mha_block(
+                    x, memory, memory_mask, memory_key_padding_mask
+                ),
+                stage_embedding,
+            )
+            x = self.norm3(x + self._ff_block(x), stage_embedding)
+
+        if tgt_is_tuple:
+            return (x, stage_embedding)
+        return x
+
+    # self-attention block
+    def _sa_block(
+        self,
+        x: Tensor,
+        attn_mask: Optional[Tensor],
+        key_padding_mask: Optional[Tensor],
+    ) -> Tensor:
+        x = self.self_attn(
+            x,
+            x,
+            x,
+            attn_mask=attn_mask,
+            key_padding_mask=key_padding_mask,
+            need_weights=False,
+        )[0]
+        return self.dropout1(x)
+
+    # multihead attention block
+    def _mha_block(
+        self,
+        x: Tensor,
+        mem: Tensor,
+        attn_mask: Optional[Tensor],
+        key_padding_mask: Optional[Tensor],
+    ) -> Tensor:
+        x = self.multihead_attn(
+            x,
+            mem,
+            mem,
+            attn_mask=attn_mask,
+            key_padding_mask=key_padding_mask,
+            need_weights=False,
+        )[0]
+        return self.dropout2(x)
+
+    # feed forward block
+    def _ff_block(self, x: Tensor) -> Tensor:
+        x = self.linear2(self.dropout(self.activation(self.linear1(x))))
+        return self.dropout3(x)
+
+
+def _get_clones(module, N):
+    return nn.ModuleList([copy.deepcopy(module) for i in range(N)])
+
+
+def _get_activation_fn(activation: str) -> Callable[[Tensor], Tensor]:
+    if activation == "relu":
+        return F.relu
+    elif activation == "gelu":
+        return F.gelu
+
+    raise RuntimeError(
+        "activation should be relu/gelu, not {}".format(activation)
+    )

apps/audio_cloning/vallex/presets/acou_1.npz

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:470ce66fc24a2d14e162343381f7d93ef0a3af51edf5fd37240c21f492b4e769
+size 15650

apps/audio_cloning/vallex/presets/acou_2.npz

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:ec1c5328751cadeed5356d4264759799ad96d33ea8dd4f8a3d0a80dd8ddb0e74
+size 15426

apps/audio_cloning/vallex/presets/acou_3.npz

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:03f241b094a32b3f542e74374183c6d15e8b70ae73ceeafb11bfd4ee6b8b4a3a
+size 15410

apps/audio_cloning/vallex/presets/acou_4.npz

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:52b96f32863f13f84cf7ac4a27d2bc95cea70c350a037f4d1890b20b8da9501e
+size 15506

apps/audio_cloning/vallex/presets/alan.npz

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:28838c3f0b2f9f315b34e9b940f30641306f0cadc5c527857cd1cc408547ed1c
+size 50002

apps/audio_cloning/vallex/presets/amused.npz

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:df3e882f3a62805b9aaf300d81822cd4eddeafee480503b7b78e32be2085fb11
+size 20882

apps/audio_cloning/vallex/presets/anger.npz

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:959cec6dc0b30219db0d70cdd165fe00bbdc098165cf9d67ccdd1ecf7a5da5be
+size 22090

apps/audio_cloning/vallex/presets/babara.npz

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:8106b2a98c3f70587f23ab46ed5bf73b1c9a770481c3620ab140bd3256010376
+size 11526

apps/audio_cloning/vallex/presets/bronya.npz

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:02eaada2c3d58866c813887ed9f871587ef5a7e976abc23382ce46a17b208001
+size 18106

apps/audio_cloning/vallex/presets/cafe.npz

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:d78d96f5829da8f69c327ff25958da5b451305fdc9c308f7e67f13cf8d640fea
+size 22442

apps/audio_cloning/vallex/presets/dingzhen.npz

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:4d19167c65eefef5e42dfaa1919ff5149ca0a93cb052396a47d1f42f9865f5f8
+size 18154

apps/audio_cloning/vallex/presets/disgust.npz

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:4443f0a395072700f2ec6101dbf2ad9d28968aa3e5809e384ea131832f894d7f
+size 39386