torchaudio 2.9.0__cp314-cp314-macosx_11_0_arm64.whl

torchaudio/pipelines/_tts/interface.py

@@ -0,0 +1,255 @@
+from abc import ABC, abstractmethod
+from typing import List, Optional, Tuple, Union
+
+from torch import Tensor
+from torchaudio.models import Tacotron2
+
+
+class _TextProcessor(ABC):
+    @property
+    @abstractmethod
+    def tokens(self):
+        """The tokens that the each value in the processed tensor represent.
+
+        :type: List[str]
+        """
+
+    @abstractmethod
+    def __call__(self, texts: Union[str, List[str]]) -> Tuple[Tensor, Tensor]:
+        """Encode the given (batch of) texts into numerical tensors
+
+        Args:
+            text (str or list of str): The input texts.
+
+        Returns:
+            (Tensor, Tensor):
+            Tensor:
+                The encoded texts. Shape: `(batch, max length)`
+            Tensor:
+                The valid length of each sample in the batch. Shape: `(batch, )`.
+        """
+
+
+class _Vocoder(ABC):
+    @property
+    @abstractmethod
+    def sample_rate(self):
+        """The sample rate of the resulting waveform
+
+        :type: float
+        """
+
+    @abstractmethod
+    def __call__(self, specgrams: Tensor, lengths: Optional[Tensor] = None) -> Tuple[Tensor, Optional[Tensor]]:
+        """Generate waveform from the given input, such as spectrogram
+
+        Args:
+            specgrams (Tensor):
+                The input spectrogram. Shape: `(batch, frequency bins, time)`.
+                The expected shape depends on the implementation.
+            lengths (Tensor, or None, optional):
+                The valid length of each sample in the batch. Shape: `(batch, )`.
+                (Default: `None`)
+
+        Returns:
+            (Tensor, Optional[Tensor]):
+            Tensor:
+                The generated waveform. Shape: `(batch, max length)`
+            Tensor or None:
+                The valid length of each sample in the batch. Shape: `(batch, )`.
+        """
+
+
+class Tacotron2TTSBundle(ABC):
+    """Data class that bundles associated information to use pretrained Tacotron2 and vocoder.
+
+    This class provides interfaces for instantiating the pretrained model along with
+    the information necessary to retrieve pretrained weights and additional data
+    to be used with the model.
+
+    Torchaudio library instantiates objects of this class, each of which represents
+    a different pretrained model. Client code should access pretrained models via these
+    instances.
+
+    Please see below for the usage and the available values.
+
+    Example - Character-based TTS pipeline with Tacotron2 and WaveRNN
+        >>> import torchaudio
+        >>>
+        >>> text = "Hello, T T S !"
+        >>> bundle = torchaudio.pipelines.TACOTRON2_WAVERNN_CHAR_LJSPEECH
+        >>>
+        >>> # Build processor, Tacotron2 and WaveRNN model
+        >>> processor = bundle.get_text_processor()
+        >>> tacotron2 = bundle.get_tacotron2()
+        Downloading:
+        100%|███████████████████████████████| 107M/107M [00:01<00:00, 87.9MB/s]
+        >>> vocoder = bundle.get_vocoder()
+        Downloading:
+        100%|███████████████████████████████| 16.7M/16.7M [00:00<00:00, 78.1MB/s]
+        >>>
+        >>> # Encode text
+        >>> input, lengths = processor(text)
+        >>>
+        >>> # Generate (mel-scale) spectrogram
+        >>> specgram, lengths, _ = tacotron2.infer(input, lengths)
+        >>>
+        >>> # Convert spectrogram to waveform
+        >>> waveforms, lengths = vocoder(specgram, lengths)
+        >>>
+        >>> torchaudio.save('hello-tts.wav', waveforms, vocoder.sample_rate)
+
+    Example - Phoneme-based TTS pipeline with Tacotron2 and WaveRNN
+        >>>
+        >>> # Note:
+        >>> #     This bundle uses pre-trained DeepPhonemizer as
+        >>> #     the text pre-processor.
+        >>> #     Please install deep-phonemizer.
+        >>> #     See https://github.com/as-ideas/DeepPhonemizer
+        >>> #     The pretrained weight is automatically downloaded.
+        >>>
+        >>> import torchaudio
+        >>>
+        >>> text = "Hello, TTS!"
+        >>> bundle = torchaudio.pipelines.TACOTRON2_WAVERNN_PHONE_LJSPEECH
+        >>>
+        >>> # Build processor, Tacotron2 and WaveRNN model
+        >>> processor = bundle.get_text_processor()
+        Downloading:
+        100%|███████████████████████████████| 63.6M/63.6M [00:04<00:00, 15.3MB/s]
+        >>> tacotron2 = bundle.get_tacotron2()
+        Downloading:
+        100%|███████████████████████████████| 107M/107M [00:01<00:00, 87.9MB/s]
+        >>> vocoder = bundle.get_vocoder()
+        Downloading:
+        100%|███████████████████████████████| 16.7M/16.7M [00:00<00:00, 78.1MB/s]
+        >>>
+        >>> # Encode text
+        >>> input, lengths = processor(text)
+        >>>
+        >>> # Generate (mel-scale) spectrogram
+        >>> specgram, lengths, _ = tacotron2.infer(input, lengths)
+        >>>
+        >>> # Convert spectrogram to waveform
+        >>> waveforms, lengths = vocoder(specgram, lengths)
+        >>>
+        >>> torchaudio.save('hello-tts.wav', waveforms, vocoder.sample_rate)
+    """
+
+    # Using the inner class so that these interfaces are not directly exposed on
+    # `torchaudio.pipelines`, but still listed in documentation.
+    # The thing is, text processing and vocoder are generic and we do not know what kind of
+    # new text processing and vocoder will be added in the future, so we want to make these
+    # interfaces specific to this Tacotron2TTS pipeline.
+
+    class TextProcessor(_TextProcessor):
+        """Interface of the text processing part of Tacotron2TTS pipeline
+
+        See :func:`torchaudio.pipelines.Tacotron2TTSBundle.get_text_processor` for the usage.
+        """
+
+    class Vocoder(_Vocoder):
+        """Interface of the vocoder part of Tacotron2TTS pipeline
+
+        See :func:`torchaudio.pipelines.Tacotron2TTSBundle.get_vocoder` for the usage.
+        """
+
+    @abstractmethod
+    def get_text_processor(self, *, dl_kwargs=None) -> TextProcessor:
+        """Create a text processor
+
+        For character-based pipeline, this processor splits the input text by character.
+        For phoneme-based pipeline, this processor converts the input text (grapheme) to
+        phonemes.
+
+        If a pre-trained weight file is necessary,
+        :func:`torch.hub.download_url_to_file` is used to downloaded it.
+
+        Args:
+            dl_kwargs (dictionary of keyword arguments,):
+                Passed to :func:`torch.hub.download_url_to_file`.
+
+        Returns:
+            TextProcessor:
+                A callable which takes a string or a list of strings as input and
+                returns Tensor of encoded texts and Tensor of valid lengths.
+                The object also has ``tokens`` property, which allows to recover the
+                tokenized form.
+
+        Example - Character-based
+            >>> text = [
+            >>>     "Hello World!",
+            >>>     "Text-to-speech!",
+            >>> ]
+            >>> bundle = torchaudio.pipelines.TACOTRON2_WAVERNN_CHAR_LJSPEECH
+            >>> processor = bundle.get_text_processor()
+            >>> input, lengths = processor(text)
+            >>>
+            >>> print(input)
+            tensor([[19, 16, 23, 23, 26, 11, 34, 26, 29, 23, 15,  2,  0,  0,  0],
+                    [31, 16, 35, 31,  1, 31, 26,  1, 30, 27, 16, 16, 14, 19,  2]],
+                   dtype=torch.int32)
+            >>>
+            >>> print(lengths)
+            tensor([12, 15], dtype=torch.int32)
+            >>>
+            >>> print([processor.tokens[i] for i in input[0, :lengths[0]]])
+            ['h', 'e', 'l', 'l', 'o', ' ', 'w', 'o', 'r', 'l', 'd', '!']
+            >>> print([processor.tokens[i] for i in input[1, :lengths[1]]])
+            ['t', 'e', 'x', 't', '-', 't', 'o', '-', 's', 'p', 'e', 'e', 'c', 'h', '!']
+
+        Example - Phoneme-based
+            >>> text = [
+            >>>     "Hello, T T S !",
+            >>>     "Text-to-speech!",
+            >>> ]
+            >>> bundle = torchaudio.pipelines.TACOTRON2_WAVERNN_PHONE_LJSPEECH
+            >>> processor = bundle.get_text_processor()
+            Downloading:
+            100%|███████████████████████████████| 63.6M/63.6M [00:04<00:00, 15.3MB/s]
+            >>> input, lengths = processor(text)
+            >>>
+            >>> print(input)
+            tensor([[54, 20, 65, 69, 11, 92, 44, 65, 38,  2,  0,  0,  0,  0],
+                    [81, 40, 64, 79, 81,  1, 81, 20,  1, 79, 77, 59, 37,  2]],
+                   dtype=torch.int32)
+            >>>
+            >>> print(lengths)
+            tensor([10, 14], dtype=torch.int32)
+            >>>
+            >>> print([processor.tokens[i] for i in input[0]])
+            ['HH', 'AH', 'L', 'OW', ' ', 'W', 'ER', 'L', 'D', '!', '_', '_', '_', '_']
+            >>> print([processor.tokens[i] for i in input[1]])
+            ['T', 'EH', 'K', 'S', 'T', '-', 'T', 'AH', '-', 'S', 'P', 'IY', 'CH', '!']
+        """
+
+    @abstractmethod
+    def get_vocoder(self, *, dl_kwargs=None) -> Vocoder:
+        """Create a vocoder module, based off of either WaveRNN or GriffinLim.
+
+        If a pre-trained weight file is necessary,
+        :func:`torch.hub.load_state_dict_from_url` is used to downloaded it.
+
+        Args:
+            dl_kwargs (dictionary of keyword arguments):
+                Passed to :func:`torch.hub.load_state_dict_from_url`.
+
+        Returns:
+            Vocoder:
+                A vocoder module, which takes spectrogram Tensor and an optional
+                length Tensor, then returns resulting waveform Tensor and an optional
+                length Tensor.
+        """
+
+    @abstractmethod
+    def get_tacotron2(self, *, dl_kwargs=None) -> Tacotron2:
+        """Create a Tacotron2 model with pre-trained weight.
+
+        Args:
+            dl_kwargs (dictionary of keyword arguments):
+                Passed to :func:`torch.hub.load_state_dict_from_url`.
+
+        Returns:
+            Tacotron2:
+                The resulting model.
+        """

torchaudio/pipelines/_tts/utils.py

@@ -0,0 +1,230 @@
+import logging
+import os
+
+import torch
+from torchaudio._internal import download_url_to_file, module_utils as _mod_utils
+
+
+def _get_chars():
+    return (
+        "_",
+        "-",
+        "!",
+        "'",
+        "(",
+        ")",
+        ",",
+        ".",
+        ":",
+        ";",
+        "?",
+        " ",
+        "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",
+    )
+
+
+def _get_phones():
+    return (
+        "_",
+        "-",
+        "!",
+        "'",
+        "(",
+        ")",
+        ",",
+        ".",
+        ":",
+        ";",
+        "?",
+        " ",
+        "AA",
+        "AA0",
+        "AA1",
+        "AA2",
+        "AE",
+        "AE0",
+        "AE1",
+        "AE2",
+        "AH",
+        "AH0",
+        "AH1",
+        "AH2",
+        "AO",
+        "AO0",
+        "AO1",
+        "AO2",
+        "AW",
+        "AW0",
+        "AW1",
+        "AW2",
+        "AY",
+        "AY0",
+        "AY1",
+        "AY2",
+        "B",
+        "CH",
+        "D",
+        "DH",
+        "EH",
+        "EH0",
+        "EH1",
+        "EH2",
+        "ER",
+        "ER0",
+        "ER1",
+        "ER2",
+        "EY",
+        "EY0",
+        "EY1",
+        "EY2",
+        "F",
+        "G",
+        "HH",
+        "IH",
+        "IH0",
+        "IH1",
+        "IH2",
+        "IY",
+        "IY0",
+        "IY1",
+        "IY2",
+        "JH",
+        "K",
+        "L",
+        "M",
+        "N",
+        "NG",
+        "OW",
+        "OW0",
+        "OW1",
+        "OW2",
+        "OY",
+        "OY0",
+        "OY1",
+        "OY2",
+        "P",
+        "R",
+        "S",
+        "SH",
+        "T",
+        "TH",
+        "UH",
+        "UH0",
+        "UH1",
+        "UH2",
+        "UW",
+        "UW0",
+        "UW1",
+        "UW2",
+        "V",
+        "W",
+        "Y",
+        "Z",
+        "ZH",
+    )
+
+
+def _to_tensor(indices):
+    lengths = torch.tensor([len(i) for i in indices], dtype=torch.int32)
+    values = [torch.tensor(i) for i in indices]
+    values = torch.nn.utils.rnn.pad_sequence(values, batch_first=True)
+    return values, lengths
+
+
+def _load_phonemizer(file, dl_kwargs):
+    if not _mod_utils.is_module_available("dp"):
+        raise RuntimeError("DeepPhonemizer is not installed. Please install it.")
+
+    from dp.phonemizer import Phonemizer
+    from dp.preprocessing.text import LanguageTokenizer, Preprocessor, SequenceTokenizer
+
+    # By default, dp issues DEBUG level log.
+    logger = logging.getLogger("dp")
+    orig_level = logger.level
+    logger.setLevel(logging.INFO)
+    try:
+        url = f"https://public-asai-dl-models.s3.eu-central-1.amazonaws.com/DeepPhonemizer/{file}"
+        directory = os.path.join(torch.hub.get_dir(), "checkpoints")
+        os.makedirs(directory, exist_ok=True)
+        path = os.path.join(directory, file)
+        if not os.path.exists(path):
+            dl_kwargs = {} if dl_kwargs is None else dl_kwargs
+            download_url_to_file(url, path, **dl_kwargs)
+        with torch.serialization.safe_globals([Preprocessor, LanguageTokenizer, SequenceTokenizer]):
+            return Phonemizer.from_checkpoint(path)
+    finally:
+        logger.setLevel(orig_level)
+
+
+def _unnormalize_waveform(waveform: torch.Tensor, bits: int) -> torch.Tensor:
+    r"""Transform waveform [-1, 1] to label [0, 2 ** bits - 1]"""
+    waveform = torch.clamp(waveform, -1, 1)
+    waveform = (waveform + 1.0) * (2**bits - 1) / 2
+    return torch.clamp(waveform, 0, 2**bits - 1).int()
+
+
+def _get_taco_params(n_symbols):
+    return {
+        "mask_padding": False,
+        "n_mels": 80,
+        "n_frames_per_step": 1,
+        "symbol_embedding_dim": 512,
+        "encoder_embedding_dim": 512,
+        "encoder_n_convolution": 3,
+        "encoder_kernel_size": 5,
+        "decoder_rnn_dim": 1024,
+        "decoder_max_step": 2000,
+        "decoder_dropout": 0.1,
+        "decoder_early_stopping": True,
+        "attention_rnn_dim": 1024,
+        "attention_hidden_dim": 128,
+        "attention_location_n_filter": 32,
+        "attention_location_kernel_size": 31,
+        "attention_dropout": 0.1,
+        "prenet_dim": 256,
+        "postnet_n_convolution": 5,
+        "postnet_kernel_size": 5,
+        "postnet_embedding_dim": 512,
+        "gate_threshold": 0.5,
+        "n_symbol": n_symbols,
+    }
+
+
+def _get_wrnn_params():
+    return {
+        "upsample_scales": [5, 5, 11],
+        "n_classes": 2**8,  # n_bits = 8
+        "hop_length": 275,
+        "n_res_block": 10,
+        "n_rnn": 512,
+        "n_fc": 512,
+        "kernel_size": 5,
+        "n_freq": 80,
+        "n_hidden": 128,
+        "n_output": 128,
+    }

torchaudio/pipelines/_wav2vec2/aligner.py

@@ -0,0 +1,87 @@
+from abc import ABC, abstractmethod
+from typing import Dict, List
+
+import torch
+import torchaudio.functional as F
+from torch import Tensor
+from torchaudio.functional import TokenSpan
+
+
+class ITokenizer(ABC):
+    @abstractmethod
+    def __call__(self, transcript: List[str]) -> List[List[str]]:
+        """Tokenize the given transcript (list of word)
+
+        .. note::
+
+           The toranscript must be normalized.
+
+        Args:
+            transcript (list of str): Transcript (list of word).
+
+        Returns:
+            (list of int): List of token sequences
+        """
+
+
+class Tokenizer(ITokenizer):
+    def __init__(self, dictionary: Dict[str, int]):
+        self.dictionary = dictionary
+
+    def __call__(self, transcript: List[str]) -> List[List[int]]:
+        return [[self.dictionary[c] for c in word] for word in transcript]
+
+
+def _align_emission_and_tokens(emission: Tensor, tokens: List[int], blank: int = 0):
+    device = emission.device
+    emission = emission.unsqueeze(0)
+    targets = torch.tensor([tokens], dtype=torch.int32, device=device)
+
+    aligned_tokens, scores = F.forced_align(emission, targets, blank=blank)
+
+    scores = scores.exp()  # convert back to probability
+    aligned_tokens, scores = aligned_tokens[0], scores[0]  # remove batch dimension
+    return aligned_tokens, scores
+
+
+class IAligner(ABC):
+    @abstractmethod
+    def __call__(self, emission: Tensor, tokens: List[List[int]]) -> List[List[TokenSpan]]:
+        """Generate list of time-stamped token sequences
+
+        Args:
+            emission (Tensor): Sequence of token probability distributions in log-domain.
+                Shape: `(time, tokens)`.
+            tokens (list of integer sequence): Tokenized transcript.
+                Output from :py:class:`torchaudio.pipelines.Wav2Vec2FABundle.Tokenizer`.
+
+        Returns:
+            (list of TokenSpan sequence): Tokens with time stamps and scores.
+        """
+
+
+def _unflatten(list_, lengths):
+    assert len(list_) == sum(lengths)
+    i = 0
+    ret = []
+    for l in lengths:
+        ret.append(list_[i : i + l])
+        i += l
+    return ret
+
+
+def _flatten(nested_list):
+    return [item for list_ in nested_list for item in list_]
+
+
+class Aligner(IAligner):
+    def __init__(self, blank):
+        self.blank = blank
+
+    def __call__(self, emission: Tensor, tokens: List[List[int]]) -> List[List[TokenSpan]]:
+        if emission.ndim != 2:
+            raise ValueError(f"The input emission must be 2D. Found: {emission.shape}")
+
+        aligned_tokens, scores = _align_emission_and_tokens(emission, _flatten(tokens), self.blank)
+        spans = F.merge_tokens(aligned_tokens, scores)
+        return _unflatten(spans, [len(ts) for ts in tokens])