diff --git a/CITATION.cff b/CITATION.cff new file mode 100644 index 0000000..fedba81 --- /dev/null +++ b/CITATION.cff @@ -0,0 +1,9 @@ +@techreport{kyutai2024moshi, + author = {Alexandre D\'efossez and Laurent Mazar\'e and Manu Orsini and Am\'elie Royer and + Patrick P\'erez and Herv\'e J\'egou and Edouard Grave and Neil Zeghidour}, + title = {Moshi: a speech-text foundation model for real-time dialogue}, + institution = {Kyutai}, + year={2024}, + month={September}, + url={http://kyutai.org/Moshi.pdf}, +} \ No newline at end of file diff --git a/LICENSE b/LICENSE new file mode 100644 index 0000000..31aa793 --- /dev/null +++ b/LICENSE @@ -0,0 +1,23 @@ +Permission is hereby granted, free of charge, to any +person obtaining a copy of this software and associated +documentation files (the "Software"), to deal in the +Software without restriction, including without +limitation the rights to use, copy, modify, merge, +publish, distribute, sublicense, and/or sell copies of +the Software, and to permit persons to whom the Software +is furnished to do so, subject to the following +conditions: + +The above copyright notice and this permission notice +shall be included in all copies or substantial portions +of the Software. + +THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF +ANY KIND, EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED +TO THE WARRANTIES OF MERCHANTABILITY, FITNESS FOR A +PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT +SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY +CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION +OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR +IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER +DEALINGS IN THE SOFTWARE. diff --git a/NOTICE b/NOTICE new file mode 100644 index 0000000..d186c16 --- /dev/null +++ b/NOTICE @@ -0,0 +1,30 @@ +This software contains code from the 'moshi' project, developed by Kyutai. + +The original source code can be found at: +https://github.com/kyutai-labs/moshi/tree/main/moshi + +The 'moshi' project is licensed under the MIT License, the text of which is included below. + +-------------------------------------------------------------------------------- + +MIT License + +Copyright (c) 2024 Kyutai + +Permission is hereby granted, free of charge, to any person obtaining a copy +of this software and associated documentation files (the "Software"), to deal +in the Software without restriction, including without limitation the rights +to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +copies of the Software, and to permit persons to whom the Software is +furnished to do so, subject to the following conditions: + +The above copyright notice and this permission notice shall be included in all +copies or substantial portions of the Software. + +THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE +SOFTWARE. diff --git a/README.md b/README.md new file mode 100644 index 0000000..708f4a3 --- /dev/null +++ b/README.md @@ -0,0 +1,81 @@ +# ComfyUI-KyutaiTTS + + +A custom node for ComfyUI that allows TTS generation with the [Kyutai TTS 1.6b en_fr model](https://huggingface.co/kyutai/tts-1.6b-en_fr) using [Kyutai offered voice models.](https://huggingface.co/kyutai/tts-voices) +The model's intended use is [https://github.com/kyutai-labs/delayed-streams-modeling](https://github.com/kyutai-labs/delayed-streams-modeling)which is not implemented here. +I made this version as it can generate large ammounts quickly and at acceptable quality for my use cases. +The model outputs at 24000Hz, some post processing can improve it if needed. + +## Features + +* **Text-to-Speech Generation:** Convert input text into spoken audio. +* **Model Selection:** Enter path to the Kyutai TTS model directory. +* **Voice Model Support:** Utilize various voice models (checkpoints) to customize the generated voice. +* **Device Selection:** Choose between CPU and CUDA (GPU) for processing, leveraging GPU acceleration for faster generation. +* **Adjustable Parameters:** Fine-tune speech generation with parameters such as: + * `n_q`: Number of quantization levels. + * `temp`: Sampling temperature for speech variability. + * `cfg_coef`: Classifier-free guidance coefficient for adherence to input. + * `padding_between`: Control silence between speech segments. + * `seed`: For reproducible audio generation. + +## Installation + +1. **Clone this Repository and install requirements:** + ```bash + cd ComfyUI/custom_nodes + git clone https://github.com/BobRandomNumber/ComfyUI-KyutaiTTS.git + pip install -r requirements.txt + ``` + +2. **Download Model Files:** + + * **Kyutai TTS Model (1.6B en_fr):** + Download all files from the `main` branch of the Hugging Face repository: + [https://huggingface.co/kyutai/tts-1.6b-en_fr/tree/main](https://huggingface.co/kyutai/tts-1.6b-en_fr/tree/main) + Create a dedicated folder for these files within your ComfyUI setup (e.g., `ComfyUI/models/checkpoints/KyutaiTTS`). The node expects the following files within this folder: + * `dsm_tts_1e68beda@240.safetensors` (Moshi weights) + * `tokenizer-e351c8d8-checkpoint125.safetensors` (Mimi weights) + * `tokenizer_spm_8k_en_fr_audio.model` (Tokenizer model) + * `config.json` (Model configuration) + + * **Kyutai TTS Voice Models:** + Download your desired voice models from the Hugging Face repository: + [https://huggingface.co/kyutai/tts-voices/tree/main](https://huggingface.co/kyutai/tts-voices/tree/main) + Place these voice models into your `ComfyUI/models/loras` directory or a subdirectory in loras. + +## Usage + +1. **Start ComfyUI:** +2. **Add the Node:** In your ComfyUI workflow, right-click and navigate to `Add Node` -> `Kyutai` -> `KyutaiTTS`. +3. **Set parameters:** + * **`text`**: Input the text you want to convert to speech. + * **`model_path`**: Input the path to the directory where you placed your Kyutai TTS model files (e.g., `C:\ComfyUI\models\checkpoints\KyutaiTTS`). + * **`voice_model`**: Select your desired voice model from the dropdown list. + * **`device`**: Choose `cuda` for GPU acceleration (recommended if available) or `cpu`. + * Adjust other parameters (`n_q`, `temp`, `cfg_coef`, `padding_between`, `seed`) as needed. +4. **Connect Output:** Connect the `AUDIO` output of the `KyutaiTTS` node to an audio playback or save node. +5. **Queue Prompt:** Queue your prompt to generate the audio. + +## Example Workflow + +An example workflow is provided in the `example_workflows` directory. + +![ComfyUI-KyutaiTTS Workflow Example](example_workflows/KyutaiTTS.png) + +## Troubleshooting + +* **`FileNotFoundError`**: Ensure that the `model_path` selected in the node points to the correct directory containing all required model files. Also, verify that your voice models are correctly placed and accessible by ComfyUI. +* **General Issues**: Always restart ComfyUI after making changes to custom nodes or model paths. Check the ComfyUI console for any error messages or warnings. + +## License + +This node pack integrates the Kyutai TTS model. Please refer to the original Kyutai TTS project's licensing information for details regarding the model and its components. + +## Attribution + +This custom node utilizes the Kyutai TTS model. + +* **Original Moshi Source:** [https://github.com/kyutai-labs/moshi/tree/main/moshi](https://github.com/kyutai-labs/moshi/tree/main/moshi) +* **Kyutai TTS Model (1.6B en_fr):** [https://huggingface.co/kyutai/tts-1.6b-en_fr](https://huggingface.co/kyutai/tts-1.6b-en_fr) +* **Kyutai TTS Voice Models:** [https://huggingface.co/kyutai/tts-voices](https://huggingface.co/kyutai/tts-voices) \ No newline at end of file diff --git a/__init__.py b/__init__.py new file mode 100644 index 0000000..d721463 --- /dev/null +++ b/__init__.py @@ -0,0 +1,3 @@ +from .nodes import NODE_CLASS_MAPPINGS, NODE_DISPLAY_NAME_MAPPINGS + +__all__ = ['NODE_CLASS_MAPPINGS', 'NODE_DISPLAY_NAME_MAPPINGS'] \ No newline at end of file diff --git a/example_workflows/KyutaiTTS.json b/example_workflows/KyutaiTTS.json new file mode 100644 index 0000000..afd39d4 --- /dev/null +++ b/example_workflows/KyutaiTTS.json @@ -0,0 +1,101 @@ +{ + "id": "00000000-0000-0000-0000-000000000000", + "revision": 0, + "last_node_id": 3, + "last_link_id": 1, + "nodes": [ + { + "id": 3, + "type": "SaveAudio", + "pos": [ + 1370, + 370 + ], + "size": [ + 320, + 112 + ], + "flags": {}, + "order": 1, + "mode": 0, + "inputs": [ + { + "name": "audio", + "type": "AUDIO", + "link": 1 + } + ], + "outputs": [], + "properties": { + "cnr_id": "comfy-core", + "ver": "0.3.43" + }, + "widgets_values": [ + "audio/KyutaiTTS" + ] + }, + { + "id": 2, + "type": "KyutaiTTS", + "pos": [ + 670, + 370 + ], + "size": [ + 670, + 304 + ], + "flags": {}, + "order": 0, + "mode": 0, + "inputs": [], + "outputs": [ + { + "name": "AUDIO", + "type": "AUDIO", + "links": [ + 1 + ] + } + ], + "properties": { + "Node name for S&R": "KyutaiTTS" + }, + "widgets_values": [ + "Hey there! How are you?", + "C:\\AI\\ComfyUI\\models\\checkpoints\\KyutaiTTS", + "tts-voices\\expresso\\ex03-ex01_laughing_002_channel2_232s.wav.1e68beda@240.safetensors", + "cuda", + 32, + 0.6, + 2, + 1, + 1650465257, + "randomize" + ] + } + ], + "links": [ + [ + 1, + 2, + 0, + 3, + 0, + "AUDIO" + ] + ], + "groups": [], + "config": {}, + "extra": { + "ds": { + "scale": 1.2100000000000002, + "offset": [ + -525.0992340844525, + -107.91739572572313 + ] + }, + "frontendVersion": "1.23.4" + }, + "version": 0.4 +} \ No newline at end of file diff --git a/example_workflows/KyutaiTTS.png b/example_workflows/KyutaiTTS.png new file mode 100644 index 0000000..03aa4c6 Binary files /dev/null and b/example_workflows/KyutaiTTS.png differ diff --git a/moshi_src/LICENSE b/moshi_src/LICENSE new file mode 100644 index 0000000..31aa793 --- /dev/null +++ b/moshi_src/LICENSE @@ -0,0 +1,23 @@ +Permission is hereby granted, free of charge, to any +person obtaining a copy of this software and associated +documentation files (the "Software"), to deal in the +Software without restriction, including without +limitation the rights to use, copy, modify, merge, +publish, distribute, sublicense, and/or sell copies of +the Software, and to permit persons to whom the Software +is furnished to do so, subject to the following +conditions: + +The above copyright notice and this permission notice +shall be included in all copies or substantial portions +of the Software. + +THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF +ANY KIND, EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED +TO THE WARRANTIES OF MERCHANTABILITY, FITNESS FOR A +PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT +SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY +CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION +OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR +IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER +DEALINGS IN THE SOFTWARE. diff --git a/moshi_src/LICENSE.audiocraft b/moshi_src/LICENSE.audiocraft new file mode 100644 index 0000000..b93be90 --- /dev/null +++ b/moshi_src/LICENSE.audiocraft @@ -0,0 +1,21 @@ +MIT License + +Copyright (c) Meta Platforms, Inc. and affiliates. + +Permission is hereby granted, free of charge, to any person obtaining a copy +of this software and associated documentation files (the "Software"), to deal +in the Software without restriction, including without limitation the rights +to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +copies of the Software, and to permit persons to whom the Software is +furnished to do so, subject to the following conditions: + +The above copyright notice and this permission notice shall be included in all +copies or substantial portions of the Software. + +THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE +SOFTWARE. diff --git a/moshi_src/moshi/__init__.py b/moshi_src/moshi/__init__.py new file mode 100644 index 0000000..58c3afa --- /dev/null +++ b/moshi_src/moshi/__init__.py @@ -0,0 +1,19 @@ +# Copyright (c) Kyutai, all rights reserved. +# This source code is licensed under the license found in the +# LICENSE file in the root directory of this source tree. + +""" +moshi is the inference codebase for Kyutai audio generation models. + +The code has been adapted from Audiocraft, see LICENSE.audiocraft + Copyright (c) Meta Platforms, Inc. and affiliates. +""" + +# flake8: noqa +from . import conditioners +from . import models +from . import modules +from . import quantization +from . import utils + +__version__ = "0.2.9a1" diff --git a/moshi_src/moshi/conditioners/__init__.py b/moshi_src/moshi/conditioners/__init__.py new file mode 100644 index 0000000..068cd41 --- /dev/null +++ b/moshi_src/moshi/conditioners/__init__.py @@ -0,0 +1,10 @@ +# flake8: noqa +# Copyright (c) Kyutai, all rights reserved. +# This source code is licensed under the license found in the +# LICENSE file in the root directory of this source tree. +""" +Modules to help doing generations under some fixed conditions. +""" + +from .base import (ConditionType, ConditionAttributes, ConditionFuser, ConditionProvider, + BaseConditioner, TensorCondition, ConditionTensors, dropout_all_conditions) diff --git a/moshi_src/moshi/conditioners/base.py b/moshi_src/moshi/conditioners/base.py new file mode 100644 index 0000000..7db3cf1 --- /dev/null +++ b/moshi_src/moshi/conditioners/base.py @@ -0,0 +1,432 @@ +# Copyright (c) Kyutai, all rights reserved. +# This source code is licensed under the license found in the +# LICENSE file in the root directory of this source tree. +# +# Adapted from +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. + +from collections import defaultdict +from dataclasses import dataclass, field +from itertools import chain +import logging +import typing as tp + +import torch +from torch import nn + +from ..modules.transformer import create_sin_embedding + + +logger = logging.getLogger(__name__) +TextCondition = tp.Optional[str] # a text condition can be a string or None (if doesn't exist) +ConditionTensors = dict[str, 'ConditionType'] + + +class ConditionType(tp.NamedTuple): + """Return type for a conditioner: both a condition tensor, and a mask indicating valid positions. + """ + condition: torch.Tensor + mask: torch.Tensor + + +@dataclass(frozen=True) +class TensorCondition: + """Looks quite similar to ConditionType, but represents the input to TensorConditioners. + `tensor` should be [B | 1, T, D], and `mask` should be `[B | 1, T]`. + """ + tensor: torch.Tensor + mask: torch.Tensor + + @staticmethod + def from_tensor(tensor: torch.Tensor): + B, T, _ = tensor.shape + mask = torch.ones(B, T, dtype=torch.bool, device=tensor.device) + return TensorCondition(tensor, mask) + + @staticmethod + def cat(conditions: tp.Sequence['TensorCondition']) -> 'TensorCondition': + assert conditions, "Cannot cat empty list." + ref_tensor = conditions[0].tensor + B, _, D = ref_tensor.shape + assert B == 1 + B = len(conditions) + T = max(condition.tensor.shape[1] for condition in conditions) + mask = torch.zeros(B, T, dtype=torch.bool, device=ref_tensor.device) + tensor = torch.zeros(B, T, D, dtype=ref_tensor.dtype, device=ref_tensor.device) + for b, condition in enumerate(conditions): + tensor[b, :condition.tensor.shape[1], :] = condition.tensor[0] + mask[b, :condition.mask.shape[1]] = condition.mask[0] + return TensorCondition(tensor, mask) + + +@dataclass +class ConditionAttributes: + """Standard class for representing the set of potential inputs to the conditioners. + Typically, `audiocraft.data.audio_dataset.SegmentInfo` will convert + to this class to make conditioning agnostic to the type of dataset. + + There are two kinds of conditionings: text (or None), or raw torch tensors (with a mask). + + """ + text: tp.Dict[str, tp.Optional[str]] = field(default_factory=dict) + tensor: tp.Dict[str, TensorCondition] = field(default_factory=dict) + + @property + def text_attributes(self) -> tp.Iterable[str]: + return self.text.keys() + + @property + def tensor_attributes(self) -> tp.Iterable[str]: + return self.text.keys() + + @staticmethod + def condition_types() -> tp.FrozenSet[str]: + return frozenset(["text", "tensor"]) + + def copy(self) -> 'ConditionAttributes': + return ConditionAttributes(dict(self.text), dict(self.tensor)) + + +Prepared = tp.TypeVar('Prepared') # represents the prepared condition input type. + + +class BaseConditioner(nn.Module, tp.Generic[Prepared]): + """Base model for all conditioner modules. + + Args: + dim (int): internal dim of the model. + output_dim (int): Output dim of the conditioner. + force_linear (bool, optional): Force linear projection even when `dim == output_dim`. + pad_empty (bool): if True, conditionings of 0 length will be padded to have length 1. + output_bias (bool): if True, the output projection will have a bias. + learn_padding (bool): if True, the padding value will be learnt, zero otherwise. + """ + + def __init__(self, dim: int, + output_dim: int, + device: tp.Union[torch.device, str], + force_linear: bool = True, + pad_empty: bool = True, + output_bias: bool = False, + learn_padding: bool = True): + super().__init__() + self.dim = dim + self.output_dim = output_dim + self.pad_empty = pad_empty + self.device = device + self.output_proj: nn.Module + if force_linear or dim != output_dim: + self.output_proj = nn.Linear(dim, output_dim, bias=output_bias, device=device) + assert not output_bias + else: + self.output_proj = nn.Identity() + self.learnt_padding: tp.Optional[torch.Tensor] + if learn_padding: + self.learnt_padding = nn.Parameter( + torch.randn(1, 1, output_dim, device=device), requires_grad=True) + self.learnt_padding.data *= 0.2 + else: + self.learnt_padding = None + + def prepare(self, *args, **kwargs) -> Prepared: + """Should be any part of the processing that will lead to a synchronization + point, e.g. BPE tokenization with transfer to the GPU. + + The returned value will be saved and return later when calling forward(). + """ + raise NotImplementedError() + + def _get_condition(self, inputs: Prepared) -> ConditionType: + """Gets input that should be used as conditioning (e.g, genre, description or a waveform). + Outputs a ConditionType, after the input data was embedded as a dense vector. + + Returns: + ConditionType: + - A tensor of size [B, T, dim] where B is the batch size, T is the length of the + output embedding and `dim` is the internal dimension of the embedding. + - And a mask indicating where the padding tokens, of shape `[B, T]`. + """ + raise NotImplementedError() + + def forward(self, inputs: Prepared) -> ConditionType: + cond, mask = self._get_condition(inputs) + B, T, C = cond.shape + if T == 0 and self.pad_empty: + cond = torch.zeros(B, T, C, device=cond.device, dtype=cond.dtype) + mask = torch.zeros_like(cond[..., 0], dtype=torch.bool) + + cond = self.output_proj(cond) + + maskf = mask.float()[..., None] + if self.learnt_padding is not None: + cond = cond * maskf + self.learnt_padding * (1 - maskf) + else: + cond = cond * maskf + return ConditionType(cond, mask) + + +class _BaseTextConditioner(BaseConditioner[Prepared]): + pass + + +class _BaseTensorConditioner(BaseConditioner[Prepared]): + pass + + +def dropout_tensor(condition: TensorCondition) -> TensorCondition: + """Utility function for nullifying a WavCondition object. + """ + return TensorCondition( + tensor=torch.zeros_like(condition.tensor), + mask=torch.zeros_like(condition.mask)) + + +def dropout_condition_(sample: ConditionAttributes, condition_type: str, condition: str) -> None: + """Utility function for nullifying an attribute inside a ConditionAttributes object. + Works in-place. + """ + valid_conditions = ConditionAttributes.condition_types() + if condition_type not in valid_conditions: + raise ValueError( + "dropout_condition got an unexpected condition type!" + f" expected one of {valid_conditions} but got '{condition_type}'") + + if condition not in getattr(sample, condition_type): + raise ValueError( + "dropout_condition received an unexpected condition!" + f" expected tensor={sample.tensor.keys()} and text={sample.text.keys()}" + f" but got '{condition}' of type '{condition_type}'!" + ) + + if condition_type == 'tensor': + tensor_condition = sample.tensor[condition] + sample.tensor[condition] = dropout_tensor(tensor_condition) + elif condition_type == 'text': + sample.text[condition] = None + else: + assert False + + +def dropout_all_conditions(attributes: tp.Sequence[ConditionAttributes]) -> list[ConditionAttributes]: + """ + Args: + attributes (list[ConditionAttributes]): All conditions attributes. + Returns: + list[ConditionAttributes]: Same with all conditions dropped. + """ + attributes = [attribute.copy() for attribute in attributes] + for condition_type in ConditionAttributes.condition_types(): + for attribute in attributes: + for condition in getattr(attribute, condition_type): + dropout_condition_(attribute, condition_type, condition) + return attributes + + +class ConditionProvider(nn.Module): + """Prepare and provide conditions given all the supported conditioners. + + Args: + conditioners (dict): Dictionary of conditioners. + device (torch.device or str, optional): Device for conditioners and output condition types. + """ + + def __init__(self, conditioners: tp.Dict[str, BaseConditioner], device: tp.Union[torch.device, str] = "cpu"): + super().__init__() + self.device = device + self.conditioners = nn.ModuleDict(conditioners).to(device) + + @property + def text_conditions(self): + return [k for k, v in self.conditioners.items() if isinstance(v, _BaseTextConditioner)] + + @property + def tensor_conditions(self): + return [k for k, v in self.conditioners.items() if isinstance(v, _BaseTensorConditioner)] + + def _collate_text(self, samples: tp.Sequence[ConditionAttributes]) -> tp.Dict[str, tp.List[tp.Optional[str]]]: + """Given a list of ConditionAttributes objects, compile a dictionary where the keys + are the attributes and the values are the aggregated input per attribute. + For example: + Input: + [ + ConditionAttributes(text={"genre": "Rock", "description": "A rock song with a guitar solo"}, wav=...), + ConditionAttributes(text={"genre": "Hip-hop", "description": "A hip-hop verse"}, wav=...), + ] + Output: + { + "genre": ["Rock", "Hip-hop"], + "description": ["A rock song with a guitar solo", "A hip-hop verse"] + } + + Args: + samples (list of ConditionAttributes): List of ConditionAttributes samples. + Returns: + dict[str, list[str, optional]]: A dictionary mapping an attribute name to text batch. + """ + out: tp.Dict[str, tp.List[tp.Optional[str]]] = defaultdict(list) + texts = [x.text for x in samples] + for text in texts: + for condition in self.text_conditions: + out[condition].append(text[condition]) + return out + + def _collate_tensors(self, samples: tp.Sequence[ConditionAttributes]) -> tp.Dict[str, TensorCondition]: + """For each tensor attribute, collate the tensor from individual batch items. + + Args: + samples (list of ConditionAttributes): List of ConditionAttributes samples. + Returns: + dict[str, TensorCondition]: A dictionary mapping an attribute name to tensor. + """ + per_attribute = defaultdict(list) + out: tp.Dict[str, TensorCondition] = {} + for sample in samples: + for attribute in self.tensor_conditions: + per_attribute[attribute].append(sample.tensor[attribute]) + + # stack all tensors to a single tensor + for attribute in self.tensor_conditions: + out[attribute] = TensorCondition.cat(per_attribute[attribute]) + + return out + + def prepare(self, inputs: tp.Sequence[ConditionAttributes]) -> tp.Dict[str, tp.Any]: + """Match attributes/tensors with existing conditioners in self, and call `prepare` for each one. + This should be called before starting any real GPU work to avoid synchronization points. + This will return a dict matching conditioner names to their arbitrary prepared representations. + + Args: + inputs (list[ConditionAttributes]): List of ConditionAttributes objects containing + text and tensors conditions. + """ + assert all([isinstance(x, ConditionAttributes) for x in inputs]), ( + "Got unexpected types input for conditioner! should be tp.List[ConditionAttributes]", + f" but types were {set([type(x) for x in inputs])}" + ) + + output = {} + text = self._collate_text(inputs) + tensors = self._collate_tensors(inputs) + + assert set(text.keys() | tensors.keys()).issubset(set(self.conditioners.keys())), ( + f"Got an unexpected attribute! Expected {self.conditioners.keys()}, ", + f"got {text.keys(), tensors.keys()}" + ) + + missing_inputs = set(self.conditioners.keys()) - (set(text.keys()) | set(tensors.keys())) + if missing_inputs: + raise RuntimeError(f"Some conditioners did not receive an input: {missing_inputs}") + for attribute, batch in chain(text.items(), tensors.items()): + conditioner = self.conditioners[attribute] + assert isinstance(conditioner, BaseConditioner) + output[attribute] = conditioner.prepare(batch) + return output + + def forward(self, prepared: tp.Dict[str, tp.Any]) -> tp.Dict[str, ConditionType]: + """Compute pairs of `(embedding, mask)` using the configured conditioners and the prepared representations. + The output is for example: + { + "genre": (torch.Tensor([B, 1, D_genre]), torch.Tensor([B, 1])), + "description": (torch.Tensor([B, T_desc, D_desc]), torch.Tensor([B, T_desc])), + ... + } + + Args: + prepared (dict): Dict of prepared representations as returned by `prepare()`. + """ + output = {} + for name, inputs in prepared.items(): + condition, mask = self.conditioners[name](inputs) + output[name] = ConditionType(condition, mask) + return output + + +class ConditionFuser(nn.Module): + """Condition fuser handles the logic to combine the different conditions + to the actual model input. + + Args: + fuse2cond (tp.Dict[str, str]): A dictionary that says how to fuse + each condition. For example: + { + "prepend": ["description"], + "sum": ["genre", "bpm"], + "cross": ["description"], + } + cross_attention_pos_emb (bool, optional): Use positional embeddings in cross attention. + cross_attention_pos_emb_scale (int): Scale for positional embeddings in cross attention if used. + """ + FUSING_METHODS = ["sum", "prepend", "cross"] + + def __init__(self, fuse2cond: tp.Dict[str, tp.List[str]], cross_attention_pos_emb: bool = False, + cross_attention_pos_emb_scale: float = 1.0): + super().__init__() + assert all( + [k in self.FUSING_METHODS for k in fuse2cond.keys()] + ), f"Got invalid fuse method, allowed methods: {self.FUSING_METHODS}" + self.cross_attention_pos_emb = cross_attention_pos_emb + self.cross_attention_pos_emb_scale = cross_attention_pos_emb_scale + self.fuse2cond: tp.Dict[str, tp.List[str]] = fuse2cond + self.cond2fuse: tp.Dict[str, str] = {} + for fuse_method, conditions in fuse2cond.items(): + for condition in conditions: + self.cond2fuse[condition] = fuse_method + if fuse_method not in ['cross', 'sum']: + raise RuntimeError("only `sum` and `cross` conditionings are supported " + f"for now, got {fuse_method}.") + + @property + def has_conditions(self) -> bool: + return bool(self.cond2fuse) + + @property + def has_prepend(self) -> bool: + """Is there a conditioning that needs to be prepending to the Transformer sequence.""" + return bool(self.fuse2cond['prepend']) + + def get_cross(self, conditions: ConditionTensors) -> torch.Tensor | None: + """Return the tensor to be provided for the cross attention.""" + cross = None + for name in self.fuse2cond['cross']: + cond, _ = conditions[name] + if cross is None: + cross = cond + else: + cross = torch.cat([cross, cond], dim=1) + + if self.cross_attention_pos_emb and cross is not None: + positions = torch.arange( + cross.shape[1], + device=cross.device + ).view(1, -1, 1) + pos_emb = create_sin_embedding(positions, cross.shape[-1]).to(cross.dtype) + cross = cross + self.cross_attention_pos_emb_scale * pos_emb + return cross + + def get_sum(self, conditions: ConditionTensors) -> torch.Tensor | None: + """Return the tensor to be provided as an extra sum offset shared for each step.""" + sum = None + for name in self.fuse2cond['sum']: + cond, _ = conditions[name] + assert cond.shape[1] == 1, cond.shape + if sum is None: + sum = cond + else: + sum = sum + cond + return sum + + def get_prepend(self, conditions: ConditionTensors) -> torch.Tensor | None: + """Return the tensor to be prepended to the transformer.""" + prepend = None + for name in self.fuse2cond['prepend']: + cond, _ = conditions[name] + if prepend is None: + prepend = cond + else: + prepend = torch.cat([cond, prepend], dim=1) + if prepend is not None: + sum = self.get_sum(conditions) + if sum is not None: + prepend = prepend + sum + return prepend diff --git a/moshi_src/moshi/conditioners/tensors.py b/moshi_src/moshi/conditioners/tensors.py new file mode 100644 index 0000000..40c0cc0 --- /dev/null +++ b/moshi_src/moshi/conditioners/tensors.py @@ -0,0 +1,16 @@ +# Copyright (c) Kyutai, all rights reserved. +# This source code is licensed under the license found in the +# LICENSE file in the root directory of this source tree. +from .base import _BaseTensorConditioner, TensorCondition, ConditionType + + +class TensorConditioner(_BaseTensorConditioner[TensorCondition]): + """Does basically nothing. + """ + + def prepare(self, tensor: TensorCondition) -> TensorCondition: + device = next(iter(self.parameters())).device + return TensorCondition(tensor.tensor.to(device=device), tensor.mask.to(device=device)) + + def _get_condition(self, inputs: TensorCondition) -> ConditionType: + return ConditionType(inputs.tensor, inputs.mask) diff --git a/moshi_src/moshi/conditioners/text.py b/moshi_src/moshi/conditioners/text.py new file mode 100644 index 0000000..10dc5ab --- /dev/null +++ b/moshi_src/moshi/conditioners/text.py @@ -0,0 +1,134 @@ +# Copyright (c) Kyutai, all rights reserved. +# This source code is licensed under the license found in the +# LICENSE file in the root directory of this source tree. +import hashlib +import logging +import typing as tp + +import torch +from torch import nn + + +from .base import _BaseTextConditioner, ConditionType + + +logger = logging.getLogger(__name__) + + +def length_to_mask(lengths: torch.Tensor, max_len: tp.Optional[int] = None) -> torch.Tensor: + """Utility function to convert a tensor of sequence lengths to a mask (useful when working on padded sequences). + For example: [3, 5] => [[1, 1, 1, 0, 0], [1, 1, 1, 1, 1]] + + Args: + lengths (torch.Tensor): tensor with lengths + max_len (int): can set the max length manually. Defaults to None. + Returns: + torch.Tensor: mask with 0s where there is pad tokens else 1s + """ + assert len(lengths.shape) == 1, "Length shape should be 1 dimensional." + final_length = lengths.max().item() if not max_len else max_len + final_length = max(final_length, 1) # if all seqs are of len zero we don't want a zero-size tensor + return torch.arange(final_length, device=lengths.device)[None, :] < lengths[:, None] + + +def hash_trick(word: str, vocab_size: int) -> int: + """Hash trick to pair each word with an index + + Args: + word (str): word we wish to convert to an index + vocab_size (int): size of the vocabulary + Returns: + int: index of the word in the embedding LUT + """ + hash = int(hashlib.sha256(word.encode("utf-8")).hexdigest(), 16) + return hash % vocab_size + + +class TokenizedText(tp.NamedTuple): + tokens: torch.Tensor # should be long tensor. + mask: torch.Tensor # should be bool tensor. + + +class TextConditioner(_BaseTextConditioner[TokenizedText]): + ... + + +class Tokenizer: + """Base tokenizer implementation + """ + def __call__(self, texts: tp.List[tp.Optional[str]]) -> TokenizedText: + raise NotImplementedError() + + +class NoopTokenizer(Tokenizer): + """This tokenizer should be used for global conditioners such as: artist, genre, key, etc. + The difference between this and WhiteSpaceTokenizer is that NoopTokenizer does not split + strings, so "Jeff Buckley" will get it's own index. Whereas WhiteSpaceTokenizer will + split it to ["Jeff", "Buckley"] and return an index per word. + + For example: + ["Queen", "ABBA", "Jeff Buckley"] => [43, 55, 101] + ["Metal", "Rock", "Classical"] => [0, 223, 51] + + When all possible values are known, one can use `possible_values` to provide the list + of possible tokens. If a token doesn't exist, `pad_idx` will be used instead. + """ + def __init__(self, n_bins: int, possible_values: list[str] | None = None): + self.n_bins = n_bins + self.pad_idx = n_bins + if possible_values is None: + self.possible_values = None + else: + self.possible_values = {value: idx for idx, value in enumerate(possible_values)} + assert n_bins >= len(possible_values) + + def __call__(self, texts: tp.List[tp.Optional[str]]) -> TokenizedText: + output, lengths = [], [] + for text in texts: + # if current sample doesn't have a certain attribute, replace with pad token + if text is None: + output.append(self.pad_idx) + lengths.append(0) + else: + if self.possible_values is None: + output.append(hash_trick(text, self.n_bins)) + else: + if text not in self.possible_values: + raise ValueError(f"'{text}' is not in possible_values {self.possible_values}") + output.append(self.possible_values[text]) + lengths.append(1) + + tokens = torch.tensor(output).int()[:, None] + mask = length_to_mask(torch.tensor(lengths)) + return TokenizedText(tokens, mask) + + +class LUTConditioner(TextConditioner): + """Lookup table TextConditioner. + + Args: + n_bins (int): Number of bins. + dim (int): Hidden dim of the model (text-encoder/LUT). + output_dim (int): Output dim of the conditioner. + pad_idx (int, optional): Index for padding token. Defaults to 0. + """ + def __init__(self, n_bins: int, tokenizer: str, possible_values: list[str] | None = None, + init_scale: float = 1., **kwargs): + super().__init__(**kwargs) + self.embed = nn.Embedding(n_bins + 1, self.dim) # n_bins + 1 for padding. + self.embed.weight.data *= init_scale + if tokenizer == 'noop': + self.tokenizer = NoopTokenizer(n_bins, possible_values) + else: + raise ValueError(f"unrecognized tokenizer `{tokenizer}`.") + + def prepare(self, x: tp.List[tp.Optional[str]]) -> TokenizedText: + device = self.embed.weight.device + tokens, mask = self.tokenizer(x) + tokens, mask = tokens.to(device), mask.to(device) + return TokenizedText(tokens.to(device), mask.to(device)) + + def _get_condition(self, inputs: TokenizedText) -> ConditionType: + tokens, mask = inputs + embeds = self.embed(tokens) + return ConditionType(embeds, mask) diff --git a/moshi_src/moshi/models/__init__.py b/moshi_src/moshi/models/__init__.py new file mode 100644 index 0000000..85b1fb2 --- /dev/null +++ b/moshi_src/moshi/models/__init__.py @@ -0,0 +1,14 @@ +# Copyright (c) Kyutai, all rights reserved. +# This source code is licensed under the license found in the +# LICENSE file in the root directory of this source tree. +""" +Models for the compression model Moshi, +""" + +# flake8: noqa +from .compression import ( + CompressionModel, + MimiModel, +) +from .lm import LMModel, LMGen +from .loaders import get_mimi, get_moshi_lm diff --git a/moshi_src/moshi/models/compression.py b/moshi_src/moshi/models/compression.py new file mode 100644 index 0000000..c1ae0d3 --- /dev/null +++ b/moshi_src/moshi/models/compression.py @@ -0,0 +1,488 @@ +# Copyright (c) Kyutai, all rights reserved. +# This source code is licensed under the license found in the +# LICENSE file in the root directory of this source tree. + +# Part of this file is adapted from encodec.py in https://github.com/facebookresearch/audiocraft +# released under the following license. +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the license found in the +# LICENSE file in the root directory of this source tree. +"""Compression models or wrapper around existing models. In particular, provides the implementation +for Mimi. Also defines the main interface that a model must follow to be usable as an audio tokenizer. +""" + +from abc import abstractmethod +from dataclasses import dataclass +import logging +import typing as tp + +import torch +from torch import nn + + +from ..quantization import ( + QuantizedResult, + BaseQuantizer, + SplitResidualVectorQuantizer, + ResidualVectorQuantizer, +) +from ..modules.conv import pad_for_conv1d +from ..modules.resample import ConvDownsample1d, ConvTrUpsample1d +from ..modules.streaming import StreamingModule, State, StateT +from ..utils.compile import CUDAGraphed + + +logger = logging.getLogger() + + +class CompressionModel(StreamingModule[StateT]): + """Base API for all compression model that aim at being used as audio tokenizers + with a language model. + """ + + @abstractmethod + def forward(self, x: torch.Tensor) -> QuantizedResult: ... + + @abstractmethod + def encode(self, x: torch.Tensor) -> torch.Tensor: + """See `MimiModel.encode`.""" + ... + + @abstractmethod + def decode(self, codes: torch.Tensor) -> torch.Tensor: + """See `MimiModel.decode`.""" + ... + + @abstractmethod + def decode_latent(self, codes: torch.Tensor) -> torch.Tensor: + """Decode from the discrete codes to continuous latent space.""" + ... + + @property + @abstractmethod + def channels(self) -> int: ... + + @property + @abstractmethod + def frame_size(self) -> int: ... + + @property + @abstractmethod + def frame_rate(self) -> float: ... + + @property + @abstractmethod + def sample_rate(self) -> int: ... + + @property + @abstractmethod + def cardinality(self) -> int: ... + + @property + @abstractmethod + def num_codebooks(self) -> int: ... + + @property + @abstractmethod + def total_codebooks(self) -> int: ... + + @abstractmethod + def set_num_codebooks(self, n: int): + """Set the active number of codebooks used by the quantizer.""" + ... + + +@dataclass +class _MimiState(State): + graphed_tr_enc: CUDAGraphed | None + graphed_tr_dec: CUDAGraphed | None + graphed_encoder: CUDAGraphed + graphed_decoder: CUDAGraphed + + +class MimiModel(CompressionModel[_MimiState]): + """Mimi model operating on the raw waveform. + + Args: + encoder (nn.Module): Encoder network. + decoder (nn.Module): Decoder network. + quantizer (qt.BaseQuantizer): Quantizer network. + frame_rate (float): Final frame rate of the quantized representatiopn. + encoder_frame_rate (float): frame rate of the encoder model. Note that if `frame_rate != encopder_frame_rate`, + the latent will be resampled linearly to match the desired `frame_rate` before and after quantization. + sample_rate (int): Audio sample rate. + channels (int): Number of audio channels. + causal (bool): Whether to use a causal version of the model. + encoder_transformer (nn.Module or None): optional transformer for the encoder. + decoder_transformer (nn.Module or None): optional transformer for the decoder. + resample_method (str): method to use for resampling the latent space before the quantizer. + upsample_channel_wise_bug (bool): controls whether the upsampling is channel wise. + Defaults to true to reproduce bug in original implementation. + freeze_encoder: whether to freeze the encoder weights. + freeze_quantizer: whether to freeze the quantizer weights. + freeze_quantizer_level: If positive, freeze the quantizer up to this level. + """ + + def __init__( + self, + encoder: nn.Module, + decoder: nn.Module, + quantizer: BaseQuantizer, + frame_rate: float, + encoder_frame_rate: float, + sample_rate: int, + channels: int, + causal: bool = False, + encoder_transformer: tp.Optional[nn.Module] = None, + decoder_transformer: tp.Optional[nn.Module] = None, + resample_method: str = "interpolate", + upsample_channel_wise_bug: bool = True, + freeze_encoder: bool = False, + freeze_quantizer: bool = False, + freeze_quantizer_level: int = -1, + ): + super().__init__() + self.encoder = encoder + self.decoder = decoder + self.encoder_transformer = encoder_transformer + self.decoder_transformer = decoder_transformer + self.quantizer = quantizer + self._frame_rate = frame_rate + self._sample_rate = sample_rate + self._channels = channels + self.encoder_frame_rate = encoder_frame_rate + + if freeze_encoder: + for p in self.encoder.parameters(): + p.requires_grad = False + if self.encoder_transformer is not None: + for p in self.encoder_transformer.parameters(): + p.requires_grad = False + for name, p in self.quantizer.named_parameters(): + if name.endswith("input_proj.weight"): + p.requires_grad = False + if freeze_quantizer: + self.quantizer.ema_frozen_(True) + self.freeze_quantizer = freeze_quantizer + self.freeze_quantizer_level = ( + freeze_quantizer_level + if freeze_quantizer_level > 0 + else self.quantizer.num_codebooks + ) + + # We will need the dimension for the resampling. In general the encoder will be a SeanetEncoder + # which exposes a `dimension` attribute. + dimension = encoder.dimension + assert isinstance( + dimension, int + ), f"Dimension should be int, got {dimension} of type {type(dimension)}." + self.dimension = dimension + + assert resample_method in [ + "interpolate", + "conv", + "avg_pool", + ], f"Invalid resample_method {resample_method}" + self.resample_method = resample_method + if encoder_frame_rate != frame_rate: + assert not ( + causal and resample_method == "interpolate" + ), "Cannot interpolate with causal model." + if resample_method in ["conv", "avg_pool"]: + assert ( + self.encoder_frame_rate > self.frame_rate + ), "Cannot upsample with conv." + downsample_stride = self.encoder_frame_rate / self.frame_rate + assert downsample_stride == int( + downsample_stride + ), f"Only integer strides are supported, got {downsample_stride}" + learnt = resample_method == "conv" + self.downsample = ConvDownsample1d( + int(downsample_stride), + dimension=dimension, + learnt=learnt, + causal=causal, + ) + if freeze_encoder: + for p in self.downsample.parameters(): + p.requires_grad = False + self.upsample = ConvTrUpsample1d( + int(downsample_stride), + dimension=dimension, + learnt=learnt, + causal=causal, + channel_wise=upsample_channel_wise_bug, + ) + + def _init_streaming_state(self, batch_size: int) -> _MimiState: + device = next(self.parameters()).device + disable = device.type != 'cuda' + graphed_tr_dec = None + graphed_tr_enc = None + if self.encoder_transformer is not None: + graphed_tr_enc = CUDAGraphed(self.encoder_transformer, disable=disable) + if self.decoder_transformer is not None: + graphed_tr_dec = CUDAGraphed(self.decoder_transformer, disable=disable) + graphed_encoder = CUDAGraphed(self.encoder, disable=disable) + graphed_decoder = CUDAGraphed(self.decoder, disable=disable) + return _MimiState(batch_size, device, graphed_tr_enc, graphed_tr_dec, graphed_encoder, graphed_decoder) + + @property + def channels(self) -> int: + return self._channels + + @property + def frame_rate(self) -> float: + return self._frame_rate + + @property + def sample_rate(self) -> int: + return self._sample_rate + + @property + def frame_size(self) -> int: + return int(self.sample_rate / self.frame_rate) + + @property + def total_codebooks(self): + """Total number of quantizer codebooks available.""" + return self.quantizer.total_codebooks + + @property + def num_codebooks(self): + """Active number of codebooks used by the quantizer.""" + return self.quantizer.num_codebooks + + def set_num_codebooks(self, n: int): + """Set the active number of codebooks used by the quantizer.""" + self.quantizer.set_num_codebooks(n) + + @property + def cardinality(self): + """Cardinality of each codebook.""" + return self.quantizer.cardinality + + def _to_framerate(self, x: torch.Tensor): + # Convert from the encoder frame rate to the overall framerate. + _, _, length = x.shape + frame_rate = self.encoder_frame_rate + new_frame_rate = self.frame_rate + if frame_rate == new_frame_rate: + return x + if self.resample_method == "interpolate": + target_length = int(length * new_frame_rate / frame_rate) + return nn.functional.interpolate(x, size=target_length, mode="linear") + else: + return self.downsample(x) + + def _to_encoder_framerate(self, x: torch.Tensor): + # Convert from overall framerate to the encoder frame rate. + _, _, length = x.shape + frame_rate = self.encoder_frame_rate + new_frame_rate = self.frame_rate + if frame_rate == new_frame_rate: + return x + if self.resample_method == "interpolate": + target_length = int(length * new_frame_rate / frame_rate) + return nn.functional.interpolate(x, size=target_length, mode="linear") + else: + return self.upsample(x) + + def forward(self, x: torch.Tensor) -> QuantizedResult: + assert x.dim() == 3 + length = x.shape[-1] + extra_metrics: tp.Dict[str, torch.Tensor] = {} + + if self.freeze_quantizer: + if isinstance(self.quantizer, SplitResidualVectorQuantizer): + self.quantizer.rvq_first.eval() + for i in range( + self.freeze_quantizer_level - self.quantizer.n_q_semantic + ): + self.quantizer.rvq_rest.vq.layers[i].eval() + elif isinstance(self.quantizer, ResidualVectorQuantizer): + for i in range(self.freeze_quantizer_level): + self.quantizer.vq.layers[i].eval() + else: + raise ValueError(f"Unsupported quantizer type {type(self.quantizer)}") + + emb = self.encoder(x) + if self.encoder_transformer is not None: + (emb,) = self.encoder_transformer(emb) + emb = self._to_framerate(emb) + expected_length = self.frame_rate * length / self.sample_rate + # Checking that we have the proper length given the advertised frame rate. + assert abs(emb.shape[-1] - expected_length) < 1, ( + emb.shape[-1], + expected_length, + ) + + q_res = self.quantizer(emb, self.frame_rate) + emb = q_res.x + emb = self._to_encoder_framerate(emb) + if self.decoder_transformer is not None: + (emb,) = self.decoder_transformer(emb) + + out = self.decoder(emb) + + # remove extra padding added by the encoder and decoder + assert out.shape[-1] >= length, (out.shape[-1], length) + out = out[..., :length] + + q_res.x = out + q_res.metrics.update(extra_metrics) + return q_res + + def _encode_to_unquantized_latent(self, x: torch.Tensor) -> torch.Tensor: + """Projects a batch of waveforms to unquantized latent space. + + Args: + x (torch.Tensor): Float tensor of shape [B, C, T]. + + Returns: + Unquantized embeddings. + """ + assert ( + x.dim() == 3 + ), f"CompressionModel._encode_to_unquantized_latent expects audio of shape [B, C, T] but got {x.shape}" + + state = self._streaming_state + frame_size = self.frame_size + + if state is None: + # The underlying convolutions no longer accept partial inputs, + # `x` needs to be exactly a multiple of the frame size, + # reproducing the previous padding behavior here. + x = pad_for_conv1d(x, frame_size, frame_size) + emb = self.encoder(x) + else: + if x.shape[-1] % frame_size != 0 or x.shape[-1] == 0: + raise RuntimeError( + f"Invalid input x of length {x.shape[-1]}. The length must be " + f"a positive multiple of the frame size {frame_size}. " + "You are responsible for buffering accordingly before feeding audio to Mimi.") + emb = state.graphed_encoder(x).clone() + if self.encoder_transformer is not None: + if state is None: + (emb,) = self.encoder_transformer(emb) + else: + assert state.graphed_tr_enc is not None + (emb,) = state.graphed_tr_enc(emb) + emb = self._to_framerate(emb) + return emb + + def encode(self, x: torch.Tensor) -> torch.Tensor: + """Encode the given input tensor to quantized representation. + + Args: + x (torch.Tensor): Float tensor of shape [B, C, T] + + Returns: + codes (torch.Tensor): an int tensor of shape [B, K, T] + with K the number of codebooks used and T the timestep. + """ + emb = self._encode_to_unquantized_latent(x) + codes = self.quantizer.encode(emb) + return codes + + def encode_to_latent(self, x: torch.Tensor, quantize: bool = True) -> torch.Tensor: + """Projects a batch of waveforms to latent space. + + Args: + x (torch.Tensor): Float tensor of shape [B, C, T]. + + Returns: + Embeddings, either quantized or not. + """ + emb = self._encode_to_unquantized_latent(x) + if not quantize: + return emb + else: + codes = self.quantizer.encode(emb) + return self.decode_latent(codes) + + def decode(self, codes: torch.Tensor): + """Decode the given codes to a reconstructed representation. + + Args: + codes (torch.Tensor): Int tensor of shape [B, K, T] + + Returns: + out (torch.Tensor): Float tensor of shape [B, C, T], the reconstructed audio. + """ + state = self._streaming_state + emb = self.decode_latent(codes) + emb = self._to_encoder_framerate(emb) + if self.decoder_transformer is not None: + if state is None: + (emb,) = self.decoder_transformer(emb) + else: + assert state.graphed_tr_dec is not None + (emb,) = state.graphed_tr_dec(emb) + if state is None: + out = self.decoder(emb) + else: + out = state.graphed_decoder(emb).clone() + # out contains extra padding added by the encoder and decoder + return out + + def decode_latent(self, codes: torch.Tensor) -> torch.Tensor: + """Decode from the discrete codes to continuous latent space.""" + return self.quantizer.decode(codes) + + +class WrapperCompressionModel(CompressionModel[State]): + """Base API for CompressionModel wrappers that do not depend on external frameworks.""" + + def __init__(self, model: CompressionModel): + super().__init__() + self.model = model + + def forward(self, x: torch.Tensor) -> QuantizedResult: + return self.model.forward(x) + + def encode(self, x: torch.Tensor) -> torch.Tensor: + return self.model.encode(x) + + def decode(self, codes: torch.Tensor) -> torch.Tensor: + return self.model.decode(codes) + + def decode_latent(self, codes: torch.Tensor) -> torch.Tensor: + return self.model.decode_latent(codes) + + def set_num_codebooks(self, n: int): + self.model.set_num_codebooks(n) + + @property + def quantizer(self): + return self.model.quantizer + + @property + def channels(self) -> int: + return self.model.channels + + @property + def frame_rate(self) -> float: + return self.model.frame_rate + + @property + def sample_rate(self) -> int: + return self.model.sample_rate + + @property + def frame_size(self) -> int: + return self.model.frame_size + + @property + def cardinality(self) -> int: + return self.model.cardinality + + @property + def num_codebooks(self) -> int: + return self.model.num_codebooks + + @property + def total_codebooks(self) -> int: + return self.model.total_codebooks diff --git a/moshi_src/moshi/models/lm.py b/moshi_src/moshi/models/lm.py new file mode 100644 index 0000000..2fc6f1c --- /dev/null +++ b/moshi_src/moshi/models/lm.py @@ -0,0 +1,837 @@ +# Copyright (c) Kyutai, all rights reserved. +# This source code is licensed under the license found in the +# LICENSE file in the root directory of this source tree. + +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the license found in the +# LICENSE file in the root directory of this source tree. + +from contextlib import ExitStack +from dataclasses import dataclass, field +from functools import partial +import logging +import typing as tp +import torch +from torch import nn +from ..conditioners import ConditionProvider, ConditionFuser, ConditionTensors +from ..utils.sampling import sample_token +from ..utils.compile import CUDAGraphed +from ..utils.quantize import replace_linear_with_qlinear +from ..modules.streaming import StreamingContainer, StreamingModule, State +from ..modules.transformer import StreamingTransformer, create_norm_fn +from .lm_utils import (_delay_sequence, + _undelay_sequence, + _init_layer, + ScaledEmbedding) + + +logger = logging.getLogger(__name__) + + +def scatter_with_mask_(tensor: torch.Tensor, dim: int, + index: torch.Tensor, value: torch.Tensor, mask: torch.Tensor) -> None: + """Scatter but skipping the updates that are masked.""" + old_value = tensor.gather(dim, index) + value = torch.where(mask, value, old_value) + tensor.scatter_(dim, index, value) + + +@dataclass +class LMOutput: + # The logits are already re-aligned with the input codes + # hence no extra shift is required, e.g. when computing CE + logits: torch.Tensor # [B, K, T, card] + mask: torch.Tensor # [B, K, T] + text_logits: torch.Tensor # [B, 1, T, text_card] + text_mask: torch.Tensor # [B, 1, T] + + +class LMModel(StreamingContainer): + """Transformer-based language model on multiple streams of codes. + + Args: + n_q (int): Number of parallel streams to model as input. + dep_q (int): Number of parallel streams to model in the depformer. + card (int): Cardinality, vocabulary size. + text_card (int): Cardinality of the text vocabulary. + text_card_out (int or None): Cardinality of output text, if different from the input. + demux_second_text_stream: (bool): Whether two text streams are muxed together with a cartesian product. + dim (int): Dimension of the transformer encoder. + num_heads (int): Number of heads for the transformer encoder. + hidden_scale (int): Scale for hidden feed forward dimension of the transformer encoder. + norm (str): Normalization method. + norm_emb (bool): Whether to normalize embeddings. + bias_proj (bool): Use bias for output projections. + depformer_*: params used for the Depformer Transformer, all the other will be shared. + depformer_multi_linear (bool): if True, uses one linear layer per codebook to project the + output of the main transformer to the Depformer latent space. + depformer_dim_feedforward (int| list[int]| None): If None, defaults to hidden_scale * depformer_dim. + depformer_weights_per_step_schedule (list[int] | None): mapping `CODEBOOK_INDEX -> WEIGHT_INDEX`, allowing + depformer_low_rank_embeddings (int | None): if provided, uses low rank embeddings, with a linear + existing_text_padding_id (int): token to use for the padding. + same_initial (bool): if True, uses the same initial tokens for both text and audio mode. + **kwargs: Additional parameters for the transformer encoder. + """ + + def __init__( + self, + delays: tp.List[int] = [0], + n_q: int = 8, + dep_q: int = 8, + card: int = 1024, + text_card: int = 32000, + text_card_out: int | None = None, + demux_second_text_stream: bool = False, + dim: int = 128, + num_heads: int = 8, + hidden_scale: int = 4, + norm: str = "layer_norm", + norm_emb: bool = False, + bias_proj: bool = False, + depformer_dim: int = 256, + depformer_dim_feedforward: int | list[int] | None = None, + depformer_multi_linear: bool = False, + depformer_weights_per_step: bool = False, + depformer_weights_per_step_schedule: list[int] | None = None, + depformer_low_rank_embeddings: int | None = None, + depformer_pos_emb: str = "sin", + existing_text_padding_id: int = 3, + existing_text_end_padding_id: int = 0, + extra_heads_num_heads: int = 0, + extra_heads_dim: int = 6, + context: tp.Optional[int] = None, + causal: bool = True, + condition_provider: tp.Optional[ConditionProvider] = None, + fuser: tp.Optional[ConditionFuser] = None, + quantize: bool = False, + device=None, + dtype=None, + gradient_checkpointing: bool = False, + **kwargs, + ): + super().__init__() + self.n_q = n_q + self.dep_q = dep_q + self.card = card + self.text_card = text_card + text_card_out = text_card if text_card_out is None else text_card_out + assert len(delays) == self.num_codebooks, f"expected {self.num_codebooks} delays, got {len(delays)}." + self.delays = delays + self.dim = dim + self.existing_text_padding_id = existing_text_padding_id + self.existing_text_end_padding_id = existing_text_end_padding_id + self.context = context + self.depformer_weights_per_step_schedule = depformer_weights_per_step_schedule + if depformer_weights_per_step_schedule is not None: + assert len(depformer_weights_per_step_schedule) == dep_q + EmbeddingFactory = partial( + ScaledEmbedding, + norm=norm_emb, + device=device, + dtype=dtype, + zero_idx=self.zero_token_id, + ) + self.emb = nn.ModuleList( + [EmbeddingFactory(self.card + 1, dim) for _ in range(n_q)] + ) + # Unlike for audio, here we authorize the model to output the special token. + self.text_emb = EmbeddingFactory(text_card + 1, dim, demux_second_stream=demux_second_text_stream) + + self.text_linear = nn.Linear(dim, text_card_out, bias=bias_proj) + depformer_prefix = "depformer_" + main_kwargs = { + k: v for k, v in kwargs.items() if not k.startswith(depformer_prefix) + } + self.transformer = StreamingTransformer( + d_model=dim, + num_heads=num_heads, + dim_feedforward=int(hidden_scale * dim), + norm=norm, + device=device, + dtype=dtype, + quantize=quantize, + context=context, + causal=causal, + checkpointing=gradient_checkpointing, + **main_kwargs, + ) + self.out_norm = create_norm_fn(norm, dim) + self.depformer_multi_linear = depformer_multi_linear + kwargs_dep = main_kwargs.copy() + kwargs_dep.update( + { + k.removeprefix(depformer_prefix): v + for k, v in kwargs.items() + if k.startswith(depformer_prefix) + } + ) + kwargs_dep["positional_embedding"] = depformer_pos_emb + kwargs_dep["context"] = None + kwargs_dep["cross_attention"] = False + if depformer_weights_per_step: + kwargs_dep["weights_per_step"] = dep_q + if depformer_multi_linear: + # One linear layer per codebook to project different informations from the main model. + num_in = dep_q + if depformer_weights_per_step_schedule: + num_in = max(depformer_weights_per_step_schedule) + 1 + self.depformer_in = nn.ModuleList( + [nn.Linear(dim, depformer_dim, bias=False) for _ in range(num_in)] + ) + else: + self.depformer_in = nn.ModuleList( + [nn.Linear(dim, depformer_dim, bias=False)] + ) + EmbeddingFactory = partial(EmbeddingFactory, low_rank=depformer_low_rank_embeddings) + if dep_q > 0: + # Only using up to dep_q - 1 because the last codebook is never an input to Depformer. + self.depformer_emb = nn.ModuleList( + [EmbeddingFactory(self.card + 1, depformer_dim) for _ in range(dep_q - 1)] + ) + self.depformer_text_emb = EmbeddingFactory( + text_card + 1, + depformer_dim, + demux_second_stream=demux_second_text_stream, + ) + if depformer_dim_feedforward is None: + depformer_dim_feedforward = int(hidden_scale * depformer_dim) + self.depformer = StreamingTransformer( + d_model=depformer_dim, + dim_feedforward=depformer_dim_feedforward, + norm=norm, + weights_per_step_schedule=depformer_weights_per_step_schedule, + causal=causal, + quantize=quantize, + checkpointing=gradient_checkpointing, + device=device, + dtype=dtype, + **kwargs_dep, + ) + # Depformer follow its own cycle of streaming entirely contained in one time step + # and should not follow the streaming of the steps dimensions. + self.depformer.set_streaming_detached(True) + else: # No-Depformer --- e.g., an ASR model + self.depformer_emb = None + self.depformer_text_emb = None + self.depformer = None + + self.extra_heads = nn.ModuleList( + [nn.Linear(dim, extra_heads_dim, bias=False) for _ in range(extra_heads_num_heads)] + ) + + dim = depformer_dim # we will directly apply the next linears to the output of the Depformer. + + self.linears = nn.ModuleList( + [nn.Linear(dim, self.card, bias=bias_proj) for _ in range(dep_q)] + ) + self.to(device=device, dtype=dtype) + # We always keep the condition provider as float32. + self.condition_provider = condition_provider + self.fuser = fuser + if self.condition_provider is not None: + self.condition_provider.to(device=device) + if self.fuser is not None: + self.fuser.to(device=device) + self._init_weights() + if quantize: + replace_linear_with_qlinear(self) + + @property + def initial_token_id(self) -> int: + """Token id for the start of sequence (audio).""" + return self.card + + @property + def text_initial_token_id(self) -> int: + """Token id for the start of sequence (text).""" + return self.text_card + + @property + def text_padding_token_id(self) -> int: + """Token id for text padding.""" + return self.existing_text_padding_id + + @property + def end_of_text_padding_id(self) -> int: + """Token id for optionally marking the last padding step for a word.""" + return self.existing_text_end_padding_id + + @property + def zero_token_id(self) -> int: + """Special value in the input tokens, indicating that no sampling should + happen for that value, and no input should be given to the model.""" + return -1 + + @property + def ungenerated_token_id(self) -> int: + """Special value that can be provided in the prompt to indicate that this specific + value should be predicted and sampled. This allows for partial teacher forcing, by generating + one modality, with the other one fixed. + """ + return -2 + + @property + def device(self) -> torch.device: + first_param = next(iter(self.parameters())) + return first_param.device + + @property + def dtype(self) -> torch.dtype: + first_param = next(iter(self.text_emb.parameters())) + return first_param.dtype + + @property + def num_codebooks(self) -> int: + return self.n_q + 1 + + @property + def num_audio_codebooks(self) -> int: + return self.n_q + + @property + def audio_offset(self) -> int: + return 1 + + def _get_initial_token(self) -> torch.Tensor: + # Returns the initial token that will be fed to the model to predict the very first timestep. + # The output shape will be [B, K, 1]. + device = next(iter(self.parameters())).device + zero = torch.full( + [1, 1, 1], self.zero_token_id, device=device, dtype=torch.long + ) + special = torch.full_like(zero, self.initial_token_id) + + text_special = torch.full_like(zero, self.text_initial_token_id) + audio_token = special + text_token = text_special + audio_token = audio_token.expand(-1, self.num_audio_codebooks, -1) + token = torch.cat([text_token, audio_token], dim=1) + return token + + def forward( + self, codes: torch.Tensor, + condition_tensors: tp.Optional[ConditionTensors] = None) -> LMOutput: + """Given an input tensor of codes [B, K, T] and list of conditions, returns the logits + along with masks indicating the valid positions at which to compute the loss. + The logits time steps are aligned with those in the input `code`. + Should only be used for training, not inference (use `LMGen` for that). + + Args: + codes (torch.Tensor): Input codes of shape [B, K, T] with B the batch size, + K the number of codebooks and T the number of timesteps. When text is supported, + the first 'codebook' corresponds to the text, and the remaining codebooks are for the audio. + condition_tensors (dict[str, ConditionType], optional): pre-computed conditioning tensors. + Returns: + LMOutput: Language model outputs, containing either text or audio logits, or both. + logits (torch.Tensor, or None) of shape [B, K, T, card] corresponding to the provided codes, + i.e. the first item corresponds to logits to predict the first code, meaning that + no additional shifting of codes and logits is required. + mask (torch.Tensor, or None) of shape [B, K, T], mask over valid and invalid positions. + Given the specified interleaving strategies, parts of the logits and codes should + not be considered as valid predictions because of invalid context. + text_logits (torch.Tensor, or None) of shape [B, 1, T, text_card]. + text_mask (torch.Tensor, or None) of shape [B, 1, T], mask over the valid positions for the text. + """ + B, K, T = codes.shape + assert K == self.num_codebooks, (K, self.num_codebooks) + # Delaying codes and removing the last time step that will never be an input. + initial = self._get_initial_token().expand(B, -1, -1) + delayed_codes = _delay_sequence(self.delays, codes, initial) + # Inserting the empty tokens for the first time step. + delayed_codes = torch.cat([initial, delayed_codes], dim=2) + + sum_condition: torch.Tensor | None = None + cross_attention_src: torch.Tensor | None = None + if condition_tensors is None: + assert self.fuser is None + else: + assert self.fuser is not None + sum_condition = self.fuser.get_sum(condition_tensors) + cross_attention_src = self.fuser.get_cross(condition_tensors) + + transformer_out, text_logits = self.forward_text(delayed_codes[:, :, :-1], sum_condition, cross_attention_src) + assert transformer_out.shape[0] == delayed_codes.shape[0] + assert transformer_out.shape[1] == delayed_codes.shape[2] - 1 + logits = self.forward_depformer_training(delayed_codes[:, :, 1:], transformer_out) + + # map back the logits on pattern sequence to logits on original codes: [B, K, S, card] -> [B, K, T, card] + # and provide the corresponding mask over invalid positions of tokens. We will with NaN values invalid positions + # to ensure they properly handled. + logits, logits_mask = _undelay_sequence( + self.delays[self.audio_offset:self.audio_offset + self.dep_q], + logits, fill_value=float('NaN')) + logits_mask &= (codes[:, self.audio_offset: self.audio_offset + self.dep_q] != self.zero_token_id) + text_logits, text_logits_mask = _undelay_sequence(self.delays[:1], text_logits, fill_value=float('NaN')) + text_logits_mask &= (codes[:, :1] != self.zero_token_id) + return LMOutput(logits, logits_mask, text_logits, text_logits_mask) + + def forward_text( + self, + sequence: torch.Tensor, sum_condition: torch.Tensor | None = None, + cross_attention_src: torch.Tensor | None = None, + ) -> tuple[torch.Tensor, torch.Tensor]: + B, K, S = sequence.shape + assert ( + K == self.num_codebooks + ), f"Sequence shape {sequence.shape} must match the number of codebooks." + input_sequence = sequence + input_ = None + for cb_index in range(self.num_audio_codebooks): + audio_emb = self.emb[cb_index]( + input_sequence[:, cb_index + self.audio_offset] + ) + input_ = audio_emb if input_ is None else input_ + audio_emb + text_emb = self.text_emb(input_sequence[:, 0]) + + input_ = text_emb if input_ is None else input_ + text_emb + if sum_condition is not None: + input_ = input_ + sum_condition.to(input_) + if cross_attention_src is not None: + cross_attention_src = cross_attention_src.to(input_) + transformer_out = self.transformer(input_, cross_attention_src=cross_attention_src) + if self.out_norm: + transformer_out = self.out_norm(transformer_out) + assert isinstance(transformer_out, torch.Tensor) + text_logits = self.text_linear(transformer_out) + text_logits = text_logits[:, None] + return transformer_out, text_logits + + def forward_depformer_training( + self, + sequence: torch.Tensor, + transformer_out: torch.Tensor, + ) -> torch.Tensor: + assert self.depformer_text_emb + assert self.depformer_emb + assert self.depformer + + B, K, T = sequence.shape + Ka = self.dep_q + assert ( + K == self.num_codebooks + ), f"Codebooks for Depformer training should be passed all at once, got {K}." + depformer_inputs = [] + for cb_index in range(Ka): + if self.depformer_multi_linear: + linear_index = cb_index + if self.depformer_weights_per_step_schedule is not None: + linear_index = self.depformer_weights_per_step_schedule[cb_index] + transformer_in = self.depformer_in[linear_index](transformer_out) + else: + transformer_in = self.depformer_in[0](transformer_out) + if cb_index == 0: + token_in = self.depformer_text_emb(sequence[:, 0]) + else: + token_in = self.depformer_emb[cb_index - 1](sequence[:, cb_index + self.audio_offset - 1]) + depformer_inputs.append(token_in + transformer_in) + depformer_input = torch.stack(depformer_inputs, 2) + # depformer_input is [B, T, K, depformer_dim], reshaping to [B * T, K, D] + depformer_input = depformer_input.view(B * T, Ka, -1) + depformer_output = self.depformer(depformer_input) + all_logits = [] + for cb_index in range(Ka): + logits = self.linears[cb_index](depformer_output[:, cb_index]) + all_logits.append(logits.view(B, T, -1)) + logits = torch.stack(all_logits, 1) + assert logits.dim() == 4, logits.shape # [B, Ka, T, card] + return logits + + def forward_depformer( + self, + depformer_cb_index: int, + sequence: torch.Tensor, + transformer_out: torch.Tensor, + ) -> torch.Tensor: + assert self.depformer_text_emb is not None + assert self.depformer_emb is not None + assert self.depformer is not None + B, K, S = sequence.shape + assert ( + K == 1 + ), f"Codebooks for Depformer streaming should be passed 1 by 1, got {K}." + assert ( + S == 1 + ), f"Steps for Depformer streaming should be passed 1 by 1, got {S}." + assert ( + transformer_out.shape[1] == 1 + ), "Transformer out should be a for a single step." + last_token_input: tp.Optional[torch.Tensor] = None + depformer_input = transformer_out + if self.depformer_multi_linear: + in_index = depformer_cb_index + if self.depformer_weights_per_step_schedule is not None: + in_index = self.depformer_weights_per_step_schedule[in_index] + depformer_input = self.depformer_in[in_index](depformer_input) + else: + depformer_input = self.depformer_in[0](depformer_input) + if depformer_cb_index == 0: + last_token_input = self.depformer_text_emb(sequence[:, 0]) + else: + last_token_input = self.depformer_emb[depformer_cb_index - 1]( + sequence[:, 0] + ) + assert last_token_input is not None + depformer_input = depformer_input + last_token_input + assert depformer_input.shape[1] == 1 + # depformer_input is [B, 1, depformer_dim]. + # The streaming state of the depformer ensures that the proper layer is run. + dep_output = self.depformer(depformer_input) + logits = self.linears[depformer_cb_index](dep_output) + logits = logits[:, None] + assert logits.dim() == 4, logits.shape # [B, Ka, S, card] + return logits + + def _init_weights(self): + """Initialization of the transformer module weights. + Mostly truncated gaussian, with `std = 1 / sqrt(dim_in)`. + Embeddings are also initialized with `1 / sqrt(dim)` rather than `1`. + Some layers are not going to be properly initialized: + - in_proj in MHA. + - depth transformer layers. + This is to match how our models were trained so far. + """ + + for emb_layer in self.emb: + _init_layer(emb_layer) + if self.depformer_emb is not None: + for emb_layer in self.depformer_emb: + _init_layer(emb_layer) + _init_layer(self.text_emb) + if self.depformer_text_emb is not None: + _init_layer(self.depformer_text_emb) + _init_layer(self.text_linear) + + for tr_layer in self.transformer.layers: + tr_layer.apply(_init_layer) + + for linear in self.linears: + _init_layer(linear) + + +@dataclass +class _LMGenState(State): + cache: torch.Tensor + initial: torch.Tensor + graphed_main: CUDAGraphed + graphed_depth: CUDAGraphed | None + offsets: torch.Tensor + offset_cpu: int = 0 + condition_sum: torch.Tensor | None = None + condition_cross: torch.Tensor | None = None + cfg_is_masked_until: torch.Tensor | None = None + exit_stack: ExitStack = field(default_factory=ExitStack) + reset_callback: tp.Callable[[torch.Tensor], None] | None = None + set_exec_mask_callback: tp.Callable[[torch.Tensor], None] | None = None + + def reset(self, reset_mask: torch.Tensor) -> None: + super().reset(reset_mask) + self.offsets[:] = torch.where(reset_mask, torch.zeros_like(self.offsets), self.offsets) + self.offset_cpu = 0 + if self.reset_callback is not None: + self.reset_callback(reset_mask) + + def set_exec_mask(self, exec_mask: torch.Tensor): + super().set_exec_mask(exec_mask) + if self.set_exec_mask_callback is not None: + self.set_exec_mask_callback(exec_mask) + + def __enter__(self): + self.exit_stack.__enter__() + + def __exit__(self, exc_type, exc_value, traceback): + self.exit_stack.__exit__(exc_type, exc_value, traceback) + + +class LMGen(StreamingModule[_LMGenState]): + def __init__( + self, + lm_model: LMModel, + use_sampling: bool = True, + temp: float = 0.8, + temp_text: float = 0.7, + top_k: int = 250, + top_k_text: int = 25, + cfg_coef: float = 1., + check: bool = False, + condition_tensors: ConditionTensors | None = None, + on_text_hook: tp.Optional[tp.Callable[[torch.Tensor], None]] = None, + on_text_logits_hook: tp.Optional[tp.Callable[[torch.Tensor], None]] = None, + on_audio_hook: tp.Optional[tp.Callable[[torch.Tensor], None]] = None, + support_out_of_sync: bool = False, + cfg_is_masked_until: list[int] | None = None, + cfg_is_no_text: bool = False, + ): + assert not lm_model.training, "generation shouldn't be used in training mode." + super().__init__() + + self.lm_model = lm_model + self.lm_model.set_streaming_detached(True) + self.use_sampling = use_sampling + self.temp = temp + self.temp_text = temp_text + self.top_k = top_k + self.top_k_text = top_k_text + self.cfg_coef = cfg_coef + self.check = check + self.max_delay = max( + lm_model.delays + ) # with delays, we need to generate a few more time steps. + self.delays_cuda = torch.tensor( + lm_model.delays, device=lm_model.device, dtype=torch.long + ) + self.condition_tensors = condition_tensors + self.on_text_hook = on_text_hook + self.on_text_logits_hook = on_text_logits_hook + self.on_audio_hook = on_audio_hook + self.support_out_of_sync = support_out_of_sync + self.cfg_is_masked_until = cfg_is_masked_until + self.cfg_is_no_text = cfg_is_no_text + if self.cfg_coef != 1.: + if not self.cfg_is_no_text and not self.cfg_is_masked_until: + assert self.lm_model.fuser is not None, "Model has no fuser, cannot do CFG." + assert self.condition_tensors, "Missing condition tensors for CFG." + + def _init_streaming_state(self, batch_size: int) -> _LMGenState: + lm_model = self.lm_model + initial = lm_model._get_initial_token() + cache = torch.full( + (batch_size, self.lm_model.num_codebooks, self.max_delay + 2), + lm_model.ungenerated_token_id, + device=lm_model.device, + dtype=torch.long, + ) + offsets = torch.zeros(batch_size, device=lm_model.device, dtype=torch.long) + + if self.lm_model.fuser is None: + assert not self.condition_tensors + condition_sum = None + condition_cross = None + else: + assert self.condition_tensors is not None + condition_sum = self.lm_model.fuser.get_sum(self.condition_tensors) + condition_cross = self.lm_model.fuser.get_cross(self.condition_tensors) + if condition_sum is not None: + condition_sum = condition_sum.to(self.lm_model.dtype) + if condition_cross is not None: + condition_cross = condition_cross.to(self.lm_model.dtype) + + disable = lm_model.device.type != 'cuda' + graphed_main = CUDAGraphed(lm_model.forward_text, disable=disable) + if lm_model.depformer is not None: + graphed_depth = CUDAGraphed(self.depformer_step, disable=disable) + else: + graphed_depth = None + + if self.cfg_is_masked_until is None: + cfg_is_masked_until = None + else: + cfg_is_masked_until = torch.tensor(self.cfg_is_masked_until, dtype=torch.long, device=lm_model.device) + + state = _LMGenState( + batch_size, lm_model.device, cache, initial, graphed_main, graphed_depth, + offsets, condition_sum=condition_sum, condition_cross=condition_cross, + cfg_is_masked_until=cfg_is_masked_until) + + if self.cfg_coef != 1.: + batch_size *= 2 + if state.condition_sum is not None: + assert state.condition_sum.shape[0] == batch_size, "cfg requires 2x more conditions." + if state.condition_cross is not None: + assert state.condition_cross.shape[0] == batch_size, "cfg requires 2x more conditions." + state.exit_stack.enter_context(self.lm_model.streaming(batch_size)) + + def _reset_callback(reset_mask: torch.Tensor) -> None: + if self.cfg_coef != 1.: + reset_mask = reset_mask.repeat(2) + self.lm_model.reset_streaming(reset_mask) + + def _set_exec_mask_callback(exec_mask: torch.Tensor) -> None: + if self.cfg_coef != 1.: + exec_mask = exec_mask.repeat(2) + self.lm_model.set_exec_mask(exec_mask) + + state.reset_callback = _reset_callback + state.set_exec_mask_callback = _set_exec_mask_callback + return state + + @torch.no_grad() + def _step(self, input_tokens: torch.Tensor, + depformer_replace_tokens: torch.Tensor | None = None + ) -> tuple[torch.Tensor, torch.Tensor] | None: + state = self._streaming_state + if state is None: + raise RuntimeError( + "You should wrap those calls with a `with lm_gen.streaming(): ...`." + ) + lm_model = self.lm_model + + assert input_tokens.dim() == 3, "Shape should be [B, K, T]." + B, Ki, S = input_tokens.shape + assert B == state.batch_size, f"Got a batch size {B}, expected {state.batch_size}" + assert S == 1, "Only support being given steps one by one." + needed_tokens = lm_model.num_codebooks - lm_model.dep_q - 1 + assert ( + Ki >= needed_tokens + ), f"We expect {needed_tokens} tokens from the user stream, got {Ki}." + + if Ki > needed_tokens: + input_tokens = input_tokens[:, :needed_tokens, :] + + CT = state.cache.shape[2] + + delays = self.delays_cuda[lm_model.dep_q + 1:] + write_positions = (state.offsets[:, None, None] + delays[:, None]) % CT + scatter_with_mask_(state.cache[:, lm_model.dep_q + 1:], -1, write_positions, input_tokens, + state.exec_mask[:, None, None]) + + is_init = state.offsets[:, None, None] <= self.delays_cuda[:, None] + is_init |= ~state.exec_mask[:, None, None] # we also give init tokens if not executing to avoid crashing. + positions = (state.offsets % CT)[:, None, None].expand_as(is_init) + input_ = state.cache.gather(dim=2, index=positions) + input_ = torch.where(is_init, state.initial, input_) + + if self.check: + # Check that we are not feeding in any value that is not generated yet. + assert not (input_ == lm_model.ungenerated_token_id).any(), ( + state.offsets, + input_, + ) + assert (input_[:, lm_model.audio_offset :] <= lm_model.card).all(), input_ + assert (input_[:, :1] <= lm_model.text_card).all() + + zero = torch.full((1,), self.lm_model.zero_token_id, dtype=torch.long, device=input_.device) + if self.cfg_coef != 1.: + if state.cfg_is_masked_until is not None: + limit = self.delays_cuda[:, None] + state.cfg_is_masked_until.view(-1, 1, 1) + is_zeroed = state.offsets[:, None, None] <= limit + + masked = torch.where(is_zeroed & ~is_init, zero, input_) + input_ = torch.cat([input_, masked], dim=0) + else: + input_ = input_.repeat(2, 1, 1) + if self.cfg_is_no_text: + input_[B:, :1] = torch.where(~is_init[:, :1], zero, input_[B:, :1]) + + transformer_out, text_logits = state.graphed_main(input_, state.condition_sum, state.condition_cross) + if self.cfg_coef != 1.: + logits, logits_null = text_logits.chunk(2) + if self.cfg_is_no_text: + text_logits = logits + else: + text_logits = logits_null + (logits - logits_null) * self.cfg_coef + # Shape of text_logits should be [B, K_text=1, T=1, Card_text] + if self.on_text_logits_hook: + self.on_text_logits_hook(text_logits) + text_token = sample_token( + text_logits.float(), + self.use_sampling, + self.temp_text, + self.top_k_text, + ) + assert text_token.dim() == 3, text_token.shape + assert text_token.shape[2] == 1 + assert text_token.shape[1] == 1, "Only one text stream supported." + text_token = text_token[:, 0, 0] # shape is [B] + if self.on_text_hook is not None: + self.on_text_hook(text_token) + if state.graphed_depth is None: + audio_tokens = None + elif depformer_replace_tokens is None: + audio_tokens = state.graphed_depth(text_token, transformer_out) + if self.on_audio_hook is not None: + self.on_audio_hook(audio_tokens) + else: + assert depformer_replace_tokens.dim() == 3 + audio_tokens = depformer_replace_tokens.squeeze(-1) + + state.offsets = torch.where(state.exec_mask, state.offsets + 1, state.offsets) + state.offset_cpu += 1 + positions = (state.offsets % CT)[:, None, None] + scatter_with_mask_(state.cache[:, :1], -1, positions, + text_token[:, None, None], state.exec_mask[:, None, None]) + if audio_tokens is not None: + audio_tokens = audio_tokens[:, :, None] + scatter_with_mask_( + state.cache[:, 1 : lm_model.dep_q + 1, :], + -1, + positions.expand_as(audio_tokens), + audio_tokens, + state.exec_mask[:, None, None], + ) + + if not self.support_out_of_sync and state.offset_cpu <= self.max_delay: + # When using out of sync exec, should not rely on this being None. + return None + B = state.cache.shape[0] + gen_delays_cuda = self.delays_cuda[: lm_model.dep_q + 1] + index = (state.offsets[:, None, None] - self.max_delay + gen_delays_cuda[:, None]) % CT + out = state.cache.gather(dim=2, index=index) + mask = (state.offsets <= self.max_delay) | ~state.exec_mask + out[mask, :, :] = lm_model.ungenerated_token_id + return out, transformer_out + + @torch.no_grad() + def step(self, input_tokens: torch.Tensor, + depformer_replace_tokens: torch.Tensor | None = None) -> torch.Tensor | None: + out = self._step(input_tokens, depformer_replace_tokens) + if out is None: + return None + return out[0] + + @torch.no_grad() + def step_with_extra_heads( + self, + input_tokens: torch.Tensor, + depformer_replace_tokens: torch.Tensor | None = None, + ) -> tuple[torch.Tensor, list[torch.Tensor]] | None: + out = self._step(input_tokens, depformer_replace_tokens) + if out is None: + return None + out, transformer_out = out + extra_heads = [extra_head(transformer_out) for extra_head in self.lm_model.extra_heads] + return out, extra_heads + + def depformer_step( + self, + text_token: torch.Tensor, + transformer_out: torch.Tensor, + ) -> torch.Tensor: + B, = text_token.shape + B_cfg = B + if self.cfg_coef != 1.: + B_cfg = 2 * B + prev_token = text_token + lm_model = self.lm_model + depformer_tokens: list[torch.Tensor] = [] + assert lm_model.depformer + assert not lm_model.depformer.is_streaming + with lm_model.depformer.streaming(B_cfg): + assert lm_model.depformer.is_streaming + for cb_index in range(lm_model.dep_q): + input_ = prev_token[:, None, None] + if self.cfg_coef != 1.: + input_ = input_.repeat(2, 1, 1) + logits = lm_model.forward_depformer(cb_index, input_, transformer_out) + if self.cfg_coef != 1.: + logits, logits_null = logits.chunk(2) + logits = logits_null + (logits - logits_null) * self.cfg_coef + next_token = sample_token( + logits.float(), + self.use_sampling, + self.temp, + self.top_k, + ) + assert next_token.shape == (B, 1, 1) + next_token = next_token[:, 0, 0] # shape is B + depformer_tokens.append(next_token) + prev_token = next_token + + assert len(depformer_tokens) == lm_model.dep_q, ( + len(depformer_tokens), + lm_model.dep_q, + ) + out = torch.stack(depformer_tokens, dim=1) + assert out.shape == (B, lm_model.dep_q), out.shape + return out diff --git a/moshi_src/moshi/models/lm_utils.py b/moshi_src/moshi/models/lm_utils.py new file mode 100644 index 0000000..de96ebf --- /dev/null +++ b/moshi_src/moshi/models/lm_utils.py @@ -0,0 +1,124 @@ +import math +import typing as tp +import torch +from torch import nn + +from ..modules.transformer import create_norm_fn + + +def _delay_sequence(delays: tp.List[int], tensor: torch.Tensor, padding: torch.Tensor) -> torch.Tensor: + B, K, T = tensor.shape + assert len(delays) == K, (len(delays), K) + outs = [] + + for k, delay in enumerate(delays): + assert delay >= 0 + line = tensor[:, k].roll(delay, dims=1) + if delay > 0: + line[:, :delay] = padding[:, k] + outs.append(line) + return torch.stack(outs, dim=1) + + +def _undelay_sequence(delays: tp.List[int], tensor: torch.Tensor, + fill_value: tp.Union[int, float] = float('NaN')) -> tp.Tuple[torch.Tensor, torch.Tensor]: + B, K, T, *_ = tensor.shape + assert len(delays) == K + mask = torch.ones(B, K, T, dtype=torch.bool, device=tensor.device) + outs = [] + if all([delay == 0 for delay in delays]): + return tensor, mask + for k, delay in enumerate(delays): + assert delay >= 0 + line = tensor[:, k].roll(-delay, dims=1) + if delay > 0: + line[:, -delay:] = fill_value + mask[:, k, -delay:] = 0 + outs.append(line) + return torch.stack(outs, dim=1), mask + + +def _get_init_fn(input_dim: int) -> tp.Callable[[torch.Tensor], None]: + def _init(x: torch.Tensor) -> None: + std = 1 / math.sqrt(input_dim) + x_orig = x + if x.device.type == 'cpu' and x.dtype in [torch.float16, torch.bfloat16]: + x = x.float() + + torch.nn.init.trunc_normal_(x, mean=0.0, std=std, a=-3 * std, b=3 * std) + if x_orig is not x: + x_orig.data[:] = x.to(x_orig) + return _init + + +def _init_layer(m: nn.Module, + zero_bias_init: bool = True): + if isinstance(m, nn.Linear): + init_fn = _get_init_fn(m.in_features) + init_fn(m.weight) + if zero_bias_init and m.bias is not None: + nn.init.constant_(m.bias, 0) + elif isinstance(m, nn.Embedding): + init_fn = _get_init_fn(m.embedding_dim) + init_fn(m.weight) + + +class ScaledEmbedding(nn.Embedding): + """Boost learning rate for embeddings (with `scale`). + + Args: + norm (bool): if True, uses a layer norm after the embedding. + zero_idx (int): special value indicating that the output should be exactly 0. + low_rank (int | None): if provided, uses low rank embedding with a linear layer to reach + the desired dimension. Quite efficient for reducing the number of weights for very large vocabs. + lr (float or None): learning rate to use, only valid if the `make_optim_group()` method is used. + demux_second_stream (bool): input tokens can be the cartesian product of the vocab size, + and they will be demuxed, e.g. `(tok2 * card + tok1)`. In that case the same embedding + is used with different linear matrices. + """ + + def __init__(self, num_embeddings: int, embedding_dim: int, + *args, norm: bool = False, zero_idx: int = -1, + low_rank: int | None = None, lr: float | None = None, + demux_second_stream: bool = False, **kwargs): + super().__init__(num_embeddings, low_rank or embedding_dim, *args, **kwargs) + self.norm = None + if norm: + self.norm = create_norm_fn("layer_norm", self.embedding_dim) + assert zero_idx < 0, "Please use negative values for the zero_idx." + self.zero_idx = zero_idx + self.lr = lr + self.low_rank = None + if low_rank is not None: + self.low_rank = nn.Linear(low_rank, embedding_dim, bias=False) + + self.demux_second_stream = demux_second_stream + assert self.zero_idx == -1, "When demuxing a second stream, zero_idx must be -1." + if self.demux_second_stream: + assert not norm + self.out1 = nn.Linear(low_rank or embedding_dim, embedding_dim, bias=False) + self.out2 = nn.Linear(low_rank or embedding_dim, embedding_dim, bias=False) + + def forward(self, input, *args, **kwargs): + is_zero = input == self.zero_idx + zero = torch.zeros(1, dtype=input.dtype, device=input.device) + input = input.clamp(min=0) + if self.demux_second_stream: + left = input % self.num_embeddings + right = input // self.num_embeddings + # Right is itself between [-1, ..., card - 1], with -1 being the zero value. + right = right - 1 + left = super().forward(left, *args, **kwargs) + right_zero = (right < 0)[..., None] + right.clamp_(min=0) + right = super().forward(right, *args, **kwargs) + y = self.out1(left) + torch.where(right_zero, zero, self.out2(right)) + y = torch.where(is_zero[..., None], zero, y) + else: + y = super().forward(input, *args, **kwargs) + if self.norm is not None: + y = self.norm(y) + y = torch.where(is_zero[..., None], zero, y) + if self.low_rank is not None: + y = self.low_rank(y) + return y diff --git a/moshi_src/moshi/models/loaders.py b/moshi_src/moshi/models/loaders.py new file mode 100644 index 0000000..7aa7d31 --- /dev/null +++ b/moshi_src/moshi/models/loaders.py @@ -0,0 +1,481 @@ +# Copyright (c) Kyutai, all rights reserved. +# This source code is licensed under the license found in the +# LICENSE file in the root directory of this source tree. +"""Retrieves the pretrained models for Moshi and Mimi.""" + +from dataclasses import dataclass, field +import json +from pathlib import Path +import warnings +from huggingface_hub import hf_hub_download + +try: + from huggingface_hub.errors import EntryNotFoundError +except ImportError: + from huggingface_hub.utils import EntryNotFoundError # pyright: ignore +from safetensors.torch import load_model, load_file +import sentencepiece +import torch +import typing as tp +from .compression import MimiModel +from ..conditioners import BaseConditioner, ConditionProvider, ConditionFuser +from .lm import LMModel +from ..modules import SEANetEncoder, SEANetDecoder, transformer +from ..quantization import SplitResidualVectorQuantizer +from ..modules.lora import replace_all_linear_with_lora, replace_lora_with_linear + + +SAMPLE_RATE = 24000 +FRAME_RATE = 12.5 + +TEXT_TOKENIZER_NAME = "tokenizer_spm_32k_3.model" +MOSHI_NAME = "model.safetensors" +MOSHI_Q8_NAME = "model.q8.safetensors" +MIMI_NAME = "tokenizer-e351c8d8-checkpoint125.safetensors" +DEFAULT_REPO = "kyutai/moshiko-pytorch-bf16" + + +_seanet_kwargs = { + "channels": 1, + "dimension": 512, + "causal": True, + "n_filters": 64, + "n_residual_layers": 1, + "activation": "ELU", + "compress": 2, + "dilation_base": 2, + "disable_norm_outer_blocks": 0, + "kernel_size": 7, + "residual_kernel_size": 3, + "last_kernel_size": 3, + # We train using weight_norm but then the weights are pre-processed for inference so + # that we can use a normal convolution. + "norm": "none", + "pad_mode": "constant", + "ratios": [8, 6, 5, 4], + "true_skip": True, +} +_quantizer_kwargs = { + "dimension": 256, + "n_q": 32, + "bins": 2048, + "input_dimension": _seanet_kwargs["dimension"], + "output_dimension": _seanet_kwargs["dimension"], +} +_transformer_kwargs = { + "d_model": _seanet_kwargs["dimension"], + "num_heads": 8, + "num_layers": 8, + "causal": True, + "layer_scale": 0.01, + "context": 250, + "conv_layout": True, + "max_period": 10000, + "gating": "none", + "norm": "layer_norm", + "positional_embedding": "rope", + "dim_feedforward": 2048, + "input_dimension": _seanet_kwargs["dimension"], + "output_dimensions": [_seanet_kwargs["dimension"]], +} + +_lm_kwargs = { + "dim": 4096, + "text_card": 32000, + "existing_text_padding_id": 3, + "n_q": 16, + "dep_q": 8, + "card": _quantizer_kwargs["bins"], + "num_heads": 32, + "num_layers": 32, + "hidden_scale": 4.125, + "causal": True, + "layer_scale": None, + "context": 3000, + "max_period": 10000, + "gating": "silu", + "norm": "rms_norm_f32", + "positional_embedding": "rope", + "depformer_dim": 1024, + "depformer_dim_feedforward": int(4.125 * 1024), + "depformer_num_heads": 16, + "depformer_num_layers": 6, + "depformer_layer_scale": None, + "depformer_multi_linear": True, + "depformer_context": 8, + "depformer_max_period": 10000, + "depformer_gating": "silu", + "depformer_pos_emb": "none", + "depformer_weights_per_step": True, + "delays": [0, 0, 1, 1, 1, 1, 1, 1, 1, 0, 1, 1, 1, 1, 1, 1, 1], +} + + +def hf_get(filename: str | Path, hf_repo: str | None = None, + check_local_file_exists: bool = False) -> Path: + if isinstance(filename, Path): + return filename + if filename.startswith("hf://"): + parts = filename.removeprefix("hf://").split("/") + repo_name = parts[0] + "/" + parts[1] + filename = "/".join(parts[2:]) + return Path(hf_hub_download(repo_name, filename)) + elif filename.startswith("file://"): + # Provide a way to force the read of a local file. + filename = filename.removeprefix("file://") + return Path(filename) + elif hf_repo is not None: + if check_local_file_exists: + if Path(filename).exists(): + return Path(filename) + return Path(hf_hub_download(hf_repo, filename)) + else: + return Path(filename) + + +@dataclass +class CheckpointInfo: + """ + Contains the paths to each sub model, along with some extra configuration. + + Args: + moshi_weights: path to the checkpoint for the Moshi LM. + mimi_weights: path to the checkpoint for the Mimi audio tokenizer. + tokenizer: path to the text tokenizer. + lm_config: config for instantiating the LM model. + Can be None if the original Moshi 7B config should be used. + raw_config: raw config, including original keys not intended for the LM. + model_type: indicate the intended use, should be `moshi` or `hibiki`. + lora_weights: path to an optional checkpoint with lora weights. + lm_gen_config: optional default params to use for generation with this model. + tts_config: optional TTS specific configuration. + stt_config: optional STT specific configuration. + model_id: optional dict containing tracability information on the model origin, in particular + its signature and epoch. + """ + + moshi_weights: Path + mimi_weights: Path + tokenizer: Path + lm_config: dict | None = None + raw_config: dict | None = None + model_type: str = "moshi" + lora_weights: Path | None = None + lm_gen_config: dict = field(default_factory=dict) + tts_config: dict = field(default_factory=dict) + stt_config: dict = field(default_factory=dict) + model_id: dict = field(default_factory=dict) + + @staticmethod + def from_hf_repo( + hf_repo: str, + moshi_weights: Path | str | None = None, + mimi_weights: Path | str | None = None, + tokenizer: Path | str | None = None, + config_path: Path | str | None = None, + lora_weights: Path | str | None = None, + ) -> "CheckpointInfo": + """Downloads the checkpoints from the given repo, along with its config. + + Extra overrides are possible for each of Moshi, Mimi, or the text tokenizer, + which should be either a Path to a local file or a string representing a path + to a local file or starting with `hf://` for pointing to a file in another repo. + + Finally, a `config_path` can be provided to override the config from the repository. + """ + if config_path is None: + try: + config_path = hf_hub_download(hf_repo, "config.json") + except EntryNotFoundError: + # No config.json, which might indicate legacy repository. + warnings.warn( + f"Repository {hf_repo} contains no config.json. " + "Assuming this is a Moshi 7B. Support for such repository " + "might be removed in the future." + ) + if config_path is None: + moshi_name = MOSHI_NAME + mimi_name = MIMI_NAME + tokenizer_name = TEXT_TOKENIZER_NAME + lm_config = None + raw_config = None + model_type = "moshi" + lm_gen_config = {} + tts_config = {} + stt_config = {} + model_id = {} + lora_name = None + else: + raw_config = json.loads(Path(config_path).read_text()) + lm_config = dict(raw_config) + moshi_name = lm_config.pop("moshi_name", MOSHI_NAME) + mimi_name = lm_config.pop("mimi_name", MIMI_NAME) + tokenizer_name = lm_config.pop("tokenizer_name", TEXT_TOKENIZER_NAME) + lora_name = lm_config.pop("lora_name", None) + model_type = lm_config.pop("model_type", "moshi") + lm_gen_config = lm_config.pop("lm_gen_config", {}) + tts_config = lm_config.pop("tts_config", {}) + stt_config = lm_config.pop("stt_config", {}) + model_id = lm_config.pop("model_id", {}) + + if moshi_weights is None: + moshi_weights_final = hf_get(moshi_name, hf_repo) + else: + moshi_weights_final = hf_get(moshi_weights) + + if mimi_weights is None: + mimi_weights_final = hf_get(mimi_name, hf_repo) + else: + mimi_weights_final = hf_get(mimi_weights) + + if tokenizer is None: + tokenizer_final = hf_get(tokenizer_name, hf_repo) + else: + tokenizer_final = hf_get(tokenizer) + + if lora_weights is None and lora_name: + lora_weights_final = hf_get(lora_name, hf_repo) + elif lora_weights is not None: + lora_weights_final = hf_get(lora_weights) + else: + lora_weights_final = None + + return CheckpointInfo( + moshi_weights_final, + mimi_weights_final, + tokenizer_final, + lm_config, + raw_config, + model_type, + lora_weights_final, + lm_gen_config=lm_gen_config, + tts_config=tts_config, + stt_config=stt_config, + model_id=model_id, + ) + + def get_mimi(self, device: torch.device | str = "cpu") -> MimiModel: + if self.lm_config is None: + num_codebooks = 8 + else: + num_codebooks = max(self.lm_config["dep_q"], self.lm_config["n_q"] - self.lm_config["dep_q"]) + return get_mimi(self.mimi_weights, num_codebooks=num_codebooks, device=device) + + def get_moshi( + self, + device: torch.device | str = "cpu", + dtype: torch.dtype = torch.bfloat16, + load_weight: bool = True, + **kwargs, + ) -> LMModel: + model = get_moshi_lm( + self.moshi_weights if load_weight else None, + lm_kwargs=self.lm_config, + device=device, + dtype=dtype, + lora_weights=self.lora_weights, + **kwargs, + ) + if self.model_type == "hibiki": + # Sometime the model samples the EOS (2) too early, which we want to ignore. + # We keep generating if the input file is not finished, and this is a way + # to implicitely replace early EOS with PAD. + model.text_emb.weight.data[2] = model.text_emb.weight.data[3] + return model + + def get_text_tokenizer(self) -> sentencepiece.SentencePieceProcessor: + return sentencepiece.SentencePieceProcessor(str(self.tokenizer)) # type: ignore + + +def _is_safetensors(path: Path | str) -> bool: + return Path(path).suffix in (".safetensors", ".sft", ".sfts") + + +def get_mimi( + filename: str | Path | None, device: torch.device | str = "cpu", num_codebooks: int = 8 +) -> MimiModel: + """Return a pretrained Mimi model, or unintialized if `filename` is None.""" + encoder = SEANetEncoder(**_seanet_kwargs) + decoder = SEANetDecoder(**_seanet_kwargs) + encoder_transformer = transformer.ProjectedTransformer( + device=device, **_transformer_kwargs + ) + decoder_transformer = transformer.ProjectedTransformer( + device=device, **_transformer_kwargs + ) + quantizer = SplitResidualVectorQuantizer( + **_quantizer_kwargs, + ) + model = MimiModel( + encoder, + decoder, + quantizer, + channels=1, + sample_rate=SAMPLE_RATE, + frame_rate=FRAME_RATE, + encoder_frame_rate=SAMPLE_RATE / encoder.hop_length, + causal=True, + resample_method="conv", + encoder_transformer=encoder_transformer, + decoder_transformer=decoder_transformer, + ).to(device=device) + model.eval() + if filename is not None: + if _is_safetensors(filename): + load_model(model, filename, device=str(device)) + else: + pkg = torch.load(filename, "cpu") + model.load_state_dict(pkg["model"]) + model.set_num_codebooks(num_codebooks) + return model + + +def get_moshi_lm( + filename: str | Path | None, + lm_kwargs: tp.Optional[tp.Dict[str, tp.Any]] = None, + device: torch.device | str = "cpu", + dtype: torch.dtype = torch.bfloat16, + lora_weights: str | Path | None = None, + fuse_lora: bool = False, + lm_kwargs_overrides={}, +) -> LMModel: + if lm_kwargs is None: + lm_kwargs = _lm_kwargs + lm_kwargs = dict(lm_kwargs) + assert lm_kwargs is not None + + if "conditioners" in lm_kwargs: + lm_kwargs["condition_provider"] = get_conditioner_provider( + lm_kwargs["dim"], device, lm_kwargs + ) + del lm_kwargs["conditioners"] + if "fuser" in lm_kwargs: + lm_kwargs["fuser"] = get_condition_fuser(lm_kwargs) + + lm_kwargs = lm_kwargs | lm_kwargs_overrides + assert lm_kwargs is not None + + # deprecated params. + lm_kwargs.pop("depformer_causal", None) + + # moved params + if 'demux_second_stream' in lm_kwargs: + lm_kwargs['demux_second_text_stream'] = lm_kwargs.pop('demux_second_stream') + + # lora params. + lora = lm_kwargs.pop("lora", False) + lora_rank = lm_kwargs.pop("lora_rank", 128) + lora_scaling = lm_kwargs.pop("lora_scaling", 2.0) + + init_device = device + if filename is not None: + init_device = torch.device('meta') + + model = LMModel( + device=init_device, + dtype=dtype, + **lm_kwargs) + + if filename is not None: + if _is_safetensors(filename): + state = load_file(filename, device=str(device)) + for key, value in state.items(): + if value.dtype.is_floating_point: + if key.startswith('condition_provider.') or key.startswith('fuser.'): + value = value.float() + else: + value = value.to(dtype) + state[key] = value + model.load_state_dict(state, assign=True) + + else: + pkg = torch.load(filename, "cpu",) + model.load_state_dict(pkg["fsdp_best_state"]["model"], assign=True) + + if lora: + assert not lm_kwargs.get("quantize"), ( + "LoRA and quantization are incompatible for now." + ) + model = get_lora_moshi( + model=model, + lora_rank=lora_rank, + lora_scaling=lora_scaling, + lora_weights=lora_weights, + device=device, + dtype=dtype, + fuse_lora=fuse_lora, + ) + else: + assert lora_weights is None, ( + "`lora` is False, but received some lora_weights to load." + ) + model.eval() + return model + + +def get_conditioner( + output_dim: int, device: torch.device | str, conditioner_cfg: dict +) -> BaseConditioner: + conditioner_type = conditioner_cfg["type"] + conditioner_kwargs = conditioner_cfg[conditioner_type] + conditioner_kwargs.update({"output_dim": output_dim, "device": device}) + if conditioner_type == "lut": + from ..conditioners.text import LUTConditioner + return LUTConditioner(**conditioner_kwargs) + elif conditioner_type == "tensor": + from ..conditioners.tensors import TensorConditioner + return TensorConditioner(**conditioner_kwargs) + else: + raise RuntimeError(f"Unknow conditioner type {conditioner_type}.") + + +def get_conditioner_provider( + output_dim: int, device: torch.device | str, cfg: dict +) -> ConditionProvider: + """Instantiate a conditioning model.""" + conditioners: tp.Dict[str, BaseConditioner] = {} + for cond, cond_cfg in cfg["conditioners"].items(): + conditioners[cond] = get_conditioner(output_dim, device, cond_cfg) + conditioner = ConditionProvider(conditioners, device=device) + return conditioner + + +def get_condition_fuser(cfg: dict) -> ConditionFuser: + """Instantiate a condition fuser object.""" + fuser_cfg = cfg["fuser"] + fuser_methods = ["sum", "cross", "prepend"] + fuse2cond = {k: fuser_cfg.get(k, []) for k in fuser_methods} + kwargs = {k: v for k, v in fuser_cfg.items() if k not in fuser_methods} + fuser = ConditionFuser(fuse2cond=fuse2cond, **kwargs) + return fuser + + +def get_lora_moshi( + model: LMModel, + lora_weights: str | Path | None, + lora_rank: int, + lora_scaling: float, + dtype: torch.dtype = torch.bfloat16, + device: torch.device | str = "cpu", + fuse_lora: bool = True, +) -> LMModel: + init_device = device + if lora_weights is not None: + init_device = torch.device('meta') + replace_all_linear_with_lora(model, lora_rank, lora_scaling, device=init_device) + if lora_weights is not None: + assert _is_safetensors(lora_weights), "LoRA weights must be a safetensors file." + lora_state_dict = load_file(lora_weights, device=str(device)) + for key, value in lora_state_dict.items(): + if value.dtype.is_floating_point: + value = value.to(dtype=dtype) + lora_state_dict[key] = value + res = model.load_state_dict(lora_state_dict, strict=False, assign=True) + if res.unexpected_keys: + raise RuntimeError( + f"unexpected_keys in the lora weights: {res.unexpected_keys}" + ) + model = model.to(dtype=dtype, device=device) + if fuse_lora: + replace_lora_with_linear(model) + return model diff --git a/moshi_src/moshi/models/tts.py b/moshi_src/moshi/models/tts.py new file mode 100644 index 0000000..124291e --- /dev/null +++ b/moshi_src/moshi/models/tts.py @@ -0,0 +1,621 @@ +# Copyright (c) Kyutai, all rights reserved. +# This source code is licensed under the license found in the +# LICENSE file in the root directory of this source tree. +""" +Implements the logic for the state machine around the Delayed Streams Modeling (DSM) based TTS model. +Things are more complex than for STT models, where we can simply force feed the audio tokens and +sample the text ones. For TTS we start from pure text, not text properly padded with time alignment. +We co-generate the padded text sequence along with the audio output from the original text by +having the model signal us when it thinks the next step will be the start of a word. We then pop a word +to feed and feed it the token representation of the word over the next few steps. +""" + +from collections import deque +from dataclasses import dataclass, field +from functools import cached_property +import re +from pathlib import Path +import typing as tp + +from safetensors.torch import load_file +from sentencepiece import SentencePieceProcessor +import sphn +import torch + +from ..conditioners import ConditionAttributes, dropout_all_conditions, TensorCondition +from ..conditioners.text import LUTConditioner +from . import loaders, MimiModel, LMModel, LMGen + + +DEFAULT_DSM_TTS_REPO = 'kyutai/tts-1.6b-en_fr' +DEFAULT_DSM_TTS_VOICE_REPO = 'kyutai/tts-voices' + + +@dataclass +class TokenIds: + """ + The token ids for special tokens: + - card: text cardinality, including the initial token (1 + tokenizer cardinality). + This is used for multiplexing multiple input tokens into the text stream. + - new_word: a new word is starting. + - pad: padding, nothing happens. + - main: indicates the start of turn of the main speaker. + - other: indicates the start of turn of the other speaker. + - zero: special value that is embedded to exactly 0. + - ungenerated: indicate that a value is not yet generated but should be + + """ + card: int + new_word: int = 0 + pad: int = 3 + main: int = 1 + other: int = 2 + zero: int = -1 + ungenerated: int = -2 + + +@dataclass +class Entry: + """One word to generate. + + Args: + tokens: list of tokens for this word. + text: word as string. + padding: if > 0, we will prevent the model from sampling a new word for that + many steps after the current word. Note that even for `padding=0`, the model + will be forbidden to sample a new word until all the tokens for the current word are consumed. + audio_tokens: is used when some audio should be used as a prefix in the model.""" + tokens: list[int] + text: str + padding: int = 0 + audio_tokens: torch.Tensor | None = None + + +@dataclass +class State: + """State of the TTS Machine. + + Args: + entries: queue containing the entries to generate. + remaining_padding: how many times the model can still sample a pad. + forced_padding: how many times the model is still forced to sample a pad. + queued: queue containing the main stream text tokens to feed. + lookahead_queued: queue containing the lookahead text tokens to feed. + end_step: once we reach the end of the generation, this is set to the current step. + The end of the generation is once the model samples a `word` but `entries` is empty. + consumption_times: list of steps at which each entry in `entries` was consumed. + transcript: list of tuples `(word, step)`, at which each word was consumed. + audio_tokens_remaining: when using some audio as prefix, this would contain the remaining + audio tokens to force into the model. + zero_text_remaining: when using an audio only prefix, for how long we should still force + the text to be `zero`. + """ + entries: deque[Entry] + remaining_padding: int + forced_padding: int + queued: deque[int] = field(default_factory=deque) + lookahead_queued: deque[int] = field(default_factory=deque) + end_step: int | None = None + consumption_times: list[int] = field(default_factory=list) + transcript: list[tuple[str, int]] = field(default_factory=list) + + def get_tokens_ahead(self, lookahead: int) -> list[int]: + assert lookahead > 0 + for entry in self.entries: + if entry.tokens: + lookahead -= 1 + if lookahead == 0: + return entry.tokens + return [] + + +def _delayed(codes: torch.Tensor, delays: list[int], fill_value: int) -> torch.Tensor: + # Apply the acoustic delay on the provided audio tokens. + K, T = codes.shape + out = torch.full((K, T + max(delays)), fill_value, device=codes.device, dtype=torch.long) + for k, delay in enumerate(delays): + out[k, delay: delay + T] = codes[k] + return out + + +def _make_null(all_attributes: tp.Sequence[ConditionAttributes]) -> list[ConditionAttributes]: + # When using CFG, returns the null conditions. + return dropout_all_conditions(all_attributes) + + +@dataclass +class StateMachine: + """State machine that manipulates the `State` based on the model prediction. + In particular, every time the model predicts a `word` (see `TokenIds`) special token, + the state machine will pop the next word to synthesize and start feeding it. + The model is optionally equipped with a second input text stream providing a lookahead + into the future text. + + Args: + token_ids: special token values. + second_stream_ahead: if > 0, the model needs a second stream for lookahead. + max_padding: maximum number of padding tokens that can be sampled in a row. + initial_padding: number of padding tokens at the beginning, to prevent the first + word from being cut. + + """ + + token_ids: TokenIds + second_stream_ahead: int = 0 + max_padding: int = 6 + initial_padding: int = 2 + + def new_state(self, entries: tp.Sequence[Entry]) -> State: + state = State( + entries=deque(entries), + lookahead_queued=deque(), + remaining_padding=self.initial_padding, + forced_padding=self.initial_padding, + ) + return state + + def process(self, step: int, state: State, token: int) -> tuple[int, bool]: + """ + Process the output of the model. + + Args: + step: current step index. + state: state to act upon. + token: model prediction + + Returns: + - output_token: value to use as the text input for the model at the next step. + - consumed_new_word: True if a new word was consumed. + """ + consumed_new_word = False + if token not in [self.token_ids.new_word, self.token_ids.pad]: + token = self.token_ids.pad + + if state.queued: + # Some text tokens are yet to be fed, we must PAD. + token = self.token_ids.pad + elif state.forced_padding > 0: + # We are forced to pad, we must PAD. + token = self.token_ids.pad + elif state.remaining_padding <= 0: + # We are not allowed to pad, we must ask for a new WORD. + token = self.token_ids.new_word + + if token == self.token_ids.new_word: + if state.entries: + entry = state.entries.popleft() + state.consumption_times.append(step) + if entry.tokens: + consumed_new_word = True + state.transcript.append((entry.text, step)) + # We queue the tokens to be fed to the model. + state.queued.extend(entry.tokens) + if self.second_stream_ahead: + # We queue the tokens for the N+lookahead word into the second text stream. + state.lookahead_queued.extend(state.get_tokens_ahead(self.second_stream_ahead)) + # Entry contains a new word, we reset the max padding counter. + state.remaining_padding = self.max_padding + else: + # Entry is only here to insert a break, pretend the token was a PAD. + token = self.token_ids.pad + state.forced_padding = entry.padding + else: + token = self.token_ids.pad + if self.second_stream_ahead and state.end_step is None: + token = self.token_ids.new_word + # Trying to consume past the last word, we reached the end. + if state.end_step is None: + state.end_step = step + + output: int | None = None + if token == self.token_ids.pad: + # Decrement the counters for remaining and forced pads. + if state.remaining_padding > 0: + state.remaining_padding -= 1 + if state.forced_padding > 0: + state.forced_padding -= 1 + if state.queued: + # We have some text tokens to feed to the model. + output = state.queued.popleft() + else: + output = self.token_ids.pad + elif token == self.token_ids.new_word: + output = self.token_ids.new_word + elif token == self.token_ids.zero: + output = token + else: + raise RuntimeError(f"Invalid token {token}") + + if self.second_stream_ahead: + second = -1 + if output == self.token_ids.new_word: + # If sampled the `word` special token, we put it on the + # second text stream instead of the main one. + second = self.token_ids.new_word + if state.queued: + # This allows us to pass the current word tokens faster. + output = state.queued.popleft() + else: + output = self.token_ids.pad + elif state.lookahead_queued: + # Otherwise if we have some lookahead tokens we feed them. + second = state.lookahead_queued.popleft() + # Then we multiplex the two tokens. We add `+1` to `second` so that + # we can encode -1, which would translate to an all 0s embedding. + # This will get de-multiplexed in the embedding in lm.py. + output = (second + 1) * self.token_ids.card + output + + assert output is not None + return output, consumed_new_word + + +def script_to_entries(tokenizer: SentencePieceProcessor, token_ids: TokenIds, frame_rate: float, + script: tp.Sequence[str], multi_speaker: bool = True, padding_between: int = 0) -> list[Entry]: + """Process a given script into a list of `Entry` that will be consumed by the model. + + This function will perform some replacements such as removing some caracters such as ':', etc. + It also supports a single XML tag from SSML, namely ``. This allow the insertion + of a pause of roughly the requested duration. + + Args: + tokenizer: text tokenizer. + token_ids: See `TokenIds`. + frame_rate: frame rate of the audio codec. + script: list of turns, with each element indicating a change of turn. Starts with the main speaker. + Use an empty first turn to start with the other speaker. + multi_speaker: whether the model was trained to handle more than one speaker. + padding_between: amount of padding to force between words. Will make the model articulate + a bit better with values such as 1. + """ + speaker_tokens = [token_ids.main, token_ids.other] + last_speaker = None + entries = [] + + # break is indicated as e.g. + event_re = re.compile(r"(?:)|(?:\s+)") + + def _add_entry(idx: int, word: str): + nonlocal first_content, last_speaker + assert ' ' not in word + assert word + tokens = tokenizer.encode(word) # type: ignore + if first_content: + speaker = idx % len(speaker_tokens) + if multi_speaker and last_speaker != speaker: + last_speaker = speaker + tokens.insert(0, speaker_tokens[speaker]) + first_content = False + padding = 0 + if padding_between > 0: + padding = max(0, padding_between + len(tokens) - 1) + entries.append(Entry(tokens=tokens, text=word, padding=padding)) + + for idx, line in enumerate(script): + first_content = True + line = line.replace('’', "'") + line = line.replace(':', " ") + line = line.replace('(', "") + line = line.replace(')', "") + while line: + match = event_re.search(line) + if match is None: + break + word = line[:match.start()] + line = line[match.end():] + if word: + _add_entry(idx, word) + if match.group(1): + break_duration = float(match.group(1)) + padding = int(round(break_duration * frame_rate)) + entry = Entry(tokens=[], text='', padding=padding) + entries.append(entry) + if line: + _add_entry(idx, line) + return entries + + +@dataclass +class TTSResult: + """Represents the result of a run of the TTS model on a batch. + + Args: + frames: list of long tensors with shape `[B, 1 + Q, 1]` representing + the audio and text tokens for each step. Note that acoustic delay is already corrected + at that point. + logged_text_tokens: for debugging, list of tuples `(predicted_tokens, next_input_token)`. + end_steps: gives the last valid step in `frames` for each item in the batch, or `None` if the + full text could not be fully synthesized within the provided budget. + all_consumption_times: for each item in the batch, a list of steps at which individual entries + (see `Entry`) were consumed as input to the model. + all_transcripts: for each item in the batch, a list of pairs `(word, step)` indicating + at which step the given `word` in the transcript should appear. Divide by the frame rate + to obtain a time stamp. + """ + frames: list[torch.Tensor] + logged_text_tokens: list[list[tuple[int, int]]] + end_steps: list[int | None] + all_consumption_times: list[list[int]] + all_transcripts: list[list[tuple[str, int]]] + + +@dataclass +class TTSModel: + """Wrapper around a multi-stream language model, a mimi codec, and a text tokenizer that + provides the functionality of a TTS model. As an end-user, you should use `from_checkpoint_info` + rather than trying to build a TTSModel directly. + + Args: + lm: trained delayed streams model. + mimi: codec to use. + tokenizer: text tokenizer to use. + machine: TTS state machine to use, which depend on how the model was trained. + delay_steps: delay between the text and audio in steps. + max_speakers: maximum number of speakers in the cross attention for this model. + temp: temperature (for both text and audio). + cfg_coef: classifier free guidance coefficient. Note that some models were trained with + CFG distillation, e.g. CFG should not be used at inference time. + final_padding: how many steps to sample past the last word. + n_q: how many audio codebooks (e.g. RVQ levels) to generate. Trade off between quality and speed. + max_gen_length: will stop generating after that many steps even if the text has not been fully consumed. + padding_bonus: additive bonus for the padding logits, positive value will lead to slower speech. + kwargs: other arguments for `moshi.models.lm.LMGen`. + + """ + + # the following params will be automatically set by `from_checkpoint_info` + lm: LMModel + mimi: MimiModel + tokenizer: SentencePieceProcessor + + voice_suffix: str + voice_repo: str + + machine: StateMachine + delay_steps: int + max_speakers: int = 5 + + # The following params can be overriden to customize generation. + temp: float = 0.6 + cfg_coef: float = 1.0 + final_padding: int = 4 + n_q: int = 32 + max_gen_length: int = 30000 + padding_bonus: float = 0. + + @staticmethod + def from_checkpoint_info(checkpoint_info: loaders.CheckpointInfo, + initial_padding: int = 2, + max_padding: int = 8, + voice_repo: str = DEFAULT_DSM_TTS_VOICE_REPO, + device: torch.device | str = 'cpu', + dtype: torch.dtype = torch.bfloat16, **kwargs) -> 'TTSModel': + assert checkpoint_info.raw_config is not None + model_id = checkpoint_info.raw_config['model_id'] + voice_suffix = f".{model_id['sig']}@{model_id['epoch']}.safetensors" + + mimi = checkpoint_info.get_mimi(device=device) + tokenizer = checkpoint_info.get_text_tokenizer() + lm = checkpoint_info.get_moshi(device=device, dtype=dtype) + + token_ids = TokenIds(lm.text_card + 1) + delay_steps = int(checkpoint_info.tts_config['audio_delay'] * mimi.frame_rate) + second_stream_ahead = checkpoint_info.tts_config.get('second_stream_ahead', 0) + + machine = StateMachine( + token_ids=token_ids, second_stream_ahead=second_stream_ahead, + max_padding=max_padding, initial_padding=initial_padding) + tts_model = TTSModel( + lm=lm, mimi=mimi, tokenizer=tokenizer, + voice_suffix=voice_suffix, voice_repo=voice_repo, + machine=machine, delay_steps=delay_steps, + **kwargs) + mimi.set_num_codebooks(tts_model.n_q) + if not tts_model.multi_speaker: + tts_model.voice_suffix = '' + return tts_model + + @cached_property + def valid_cfg_conditionings(self) -> set[float]: + valid_cfg_conditionings = set() + if self.lm.condition_provider is not None and 'cfg' in self.lm.condition_provider.conditioners: + cfg_conditioner = self.lm.condition_provider.conditioners['cfg'] + assert isinstance(cfg_conditioner, LUTConditioner) + assert cfg_conditioner.tokenizer.possible_values is not None + valid_cfg_conditionings = set(float(x) for x in cfg_conditioner.tokenizer.possible_values) + return valid_cfg_conditionings + + @cached_property + def multi_speaker(self) -> bool: + if self.lm.condition_provider is None: + return False + return 'speaker_wavs' in self.lm.condition_provider.conditioners + + def prepare_script(self, script: tp.Sequence[str], padding_between: int = 0) -> list[Entry]: + """Wrapper around `script_to_entries`.""" + return script_to_entries( + self.tokenizer, self.machine.token_ids, self.mimi.frame_rate, script, + multi_speaker=self.multi_speaker, padding_between=padding_between) + + @torch.no_grad() + def generate(self, all_entries: tp.Sequence[tp.Sequence[Entry]], + attributes: tp.Sequence[ConditionAttributes], + prefixes: list[torch.Tensor] | None = None, + cfg_is_no_prefix: bool = True, + cfg_is_no_text: bool = True, + on_frame: tp.Optional[tp.Callable[[torch.Tensor], None]] = None, + **kwargs + ) -> TTSResult: + """Synthesize text to audio. Returns a `TTSResult`. + + Args: + all_entries: list with one item per batch item, consisting of a list of `Entry`, + obtained from `prepare_script`. + attributes: list of `ConditionAttributes` for speaker conditioning. + prefixes: this should be the list of the lengths up until when to mask for the CFG. + cfg_is_no_prefix: if true, the null logits are computed with a masked prefix. + cfg_is_no_text: if true, the null logits are computed without the text. + on_frame: a callback triggered when a frame of mimi codes is available, the frame + is a view on a pre-allocated tensor so has to be copied if you want to keep it. + **kwargs: passed to `moshi.models.lm.LMGen`. + """ + + if self.cfg_coef != 1.0: + if self.valid_cfg_conditionings: + raise ValueError( + "This model does not support direct CFG, but was trained with " + "CFG distillation. Pass instead `cfg_coef` to `make_condition_attributes`.") + nulled = _make_null(attributes) + attributes = list(attributes) + nulled + + assert self.lm.condition_provider is not None + prepared = self.lm.condition_provider.prepare(attributes) + condition_tensors = self.lm.condition_provider(prepared) + + states = [] + for entries in all_entries: + state = self.machine.new_state(entries) + states.append(state) + + cfg_is_masked_until = None + text_prefixes = None + audio_prefixes = None + device = self.lm.device + if prefixes is not None: + assert len(all_entries) == len(prefixes), f"Not enough prefixes, expected {len(all_entries)}." + if cfg_is_no_prefix: + cfg_is_masked_until = [] + text_prefixes = [] + audio_prefixes = [] + for prefix in prefixes: + if cfg_is_masked_until is not None: + cfg_is_masked_until.append(prefix.shape[-1] + self.delay_steps) + K, _ = prefix.shape + assert K == self.lm.num_codebooks + text_prefixes.append(deque(prefix[0].cpu().tolist())) + delays = [d + self.delay_steps for d in self.lm.delays[self.lm.audio_offset:]] + delayed = _delayed(prefix[self.lm.audio_offset:], delays, self.machine.token_ids.ungenerated) + delayed = delayed.to(device) + audio_prefixes.append(deque(delayed.t())) + + def _on_text_logits_hook(text_logits): + if self.padding_bonus: + text_logits[..., self.machine.token_ids.pad] += self.padding_bonus + return text_logits + + def _on_audio_hook(audio_tokens): + audio_offset = self.lm.audio_offset + delays = self.lm.delays + ungenerated = self.machine.token_ids.ungenerated + for q in range(audio_tokens.shape[1]): + delay = delays[q + audio_offset] + if offset < delay + self.delay_steps: + audio_tokens[:, q] = self.machine.token_ids.zero + if audio_prefixes is not None: + for b, audio_prefix in enumerate(audio_prefixes): + if audio_prefix: + audio_codes = audio_prefix.popleft() + mask = audio_codes != ungenerated + audio_tokens[b] = torch.where(mask, audio_codes, audio_tokens[b]) + + def _on_text_hook(text_tokens): + tokens = text_tokens.tolist() + out_tokens = [] + for b, (token, state, logged) in enumerate(zip(tokens, states, logged_text_tokens)): + if text_prefixes is not None and text_prefixes[b]: + out_token = text_prefixes[b].popleft() + else: + out_token, _ = self.machine.process(offset, state, token) + out_tokens.append(out_token) + logged.append((token, out_token)) + text_tokens[:] = torch.tensor(out_tokens, dtype=torch.long, device=text_tokens.device) + + self.lm.dep_q = self.n_q + lm_gen = LMGen( + self.lm, temp=self.temp, temp_text=self.temp, + cfg_coef=self.cfg_coef, condition_tensors=condition_tensors, + on_text_logits_hook=_on_text_logits_hook, on_text_hook=_on_text_hook, on_audio_hook=_on_audio_hook, + cfg_is_masked_until=cfg_is_masked_until, cfg_is_no_text=cfg_is_no_text, + **kwargs) + + logged_text_tokens = [[] for _ in states] + frames: list[torch.Tensor] = [] + + with lm_gen.streaming(len(states)): + for offset in range(self.max_gen_length): + if all(state.end_step is not None for state in states): + max_end_step = max(state.end_step for state in states) + if offset >= max_end_step + self.delay_steps + self.final_padding: + break + missing = self.lm.n_q - self.lm.dep_q + input_tokens = torch.full((len(states), missing, 1), self.machine.token_ids.zero, + dtype=torch.long, device=self.lm.device) + frame = lm_gen.step(input_tokens) + if frame is not None: + frames.append(frame.clone()) + if on_frame is not None: + on_frame(frame) + return TTSResult( + frames, logged_text_tokens, + [state.end_step for state in states], + [state.consumption_times for state in states], + [state.transcript for state in states]) + + def get_voice_path(self, voice_name: str) -> Path: + """Returns a local path given a voice name, potentially fetching the voice + from a HuggingFace repository. To retrieve a voice from another repo, you can also use + the `hf://REPO/PATH` syntax. + """ + file = loaders.hf_get(voice_name + self.voice_suffix, self.voice_repo, + check_local_file_exists=True) + return Path(file) + + def make_condition_attributes( + self, voices: list[Path], cfg_coef: float | None = None) -> ConditionAttributes: + """Given a list of pre computed voice embeddings, returns a ConditionAttributes. + + Args: + voices: list of file paths to pre computed voice embeddings, see `get_voice_path`. + cfg_coef: for model trained with CFG distillation, value of the CFG + to use as conditioning. Typically, values from 1. to 4. are supported + with 0.5 increments. + """ + if voices: + voice_tensor = None + mask = None + for idx in range(5): + if idx < len(voices): + emb = load_file(voices[idx], device='cpu')['speaker_wavs'] + assert emb.dim() == 3 + if voice_tensor is None: + voice_tensor = torch.zeros(1, self.max_speakers, emb.shape[2], emb.shape[1]) + if mask is None: + mask = torch.zeros(1, self.max_speakers, emb.shape[2], dtype=torch.bool) + voice_tensor[:, idx, :, :] = emb.transpose(1, 2) + mask[:, idx, :] = True + assert voice_tensor is not None + assert mask is not None + voice_tensor = voice_tensor.view(1, -1, voice_tensor.shape[-1]) + mask = mask.view(1, -1) + tensors = { + 'speaker_wavs': TensorCondition(voice_tensor, mask) + } + else: + tensors = {} + text: dict[str, str | None] = {'control': 'ok'} + if cfg_coef is None: + text['cfg'] = None + else: + if cfg_coef in self.valid_cfg_conditionings: + text['cfg'] = format(cfg_coef, '.1f') + else: + valids = ", ".join(str(x) for x in self.valid_cfg_conditionings) + raise ValueError(f"Unsupported value for cfg_coef, valid values are {valids}.") + return ConditionAttributes(text=text, tensor=tensors) + + def get_prefix(self, audio_path: Path) -> torch.Tensor: + wav, _ = sphn.read(audio_path, sample_rate=self.mimi.sample_rate) + with torch.no_grad(): + prefix = self.mimi.encode(torch.from_numpy(wav).to(device=self.lm.device)[None])[0, :, :-2] + null_text = torch.full_like(prefix[:1], self.machine.token_ids.zero) + prefix = torch.cat([null_text, prefix], dim=0) + return prefix diff --git a/moshi_src/moshi/modules/__init__.py b/moshi_src/moshi/modules/__init__.py new file mode 100644 index 0000000..c24cbf9 --- /dev/null +++ b/moshi_src/moshi/modules/__init__.py @@ -0,0 +1,23 @@ +# Copyright (c) Kyutai, all rights reserved. +# This source code is licensed under the license found in the +# LICENSE file in the root directory of this source tree. + +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the license found in the +# LICENSE file in the root directory of this source tree. +"""Modules used for building the models.""" + +# flake8: noqa +from .conv import ( + NormConv1d, + NormConvTranspose1d, + StreamingConv1d, + StreamingConvTranspose1d, + pad_for_conv1d, + pad1d, + unpad1d, +) +from .seanet import SEANetEncoder, SEANetDecoder +from .transformer import StreamingTransformer diff --git a/moshi_src/moshi/modules/conv.py b/moshi_src/moshi/modules/conv.py new file mode 100644 index 0000000..f36b5cd --- /dev/null +++ b/moshi_src/moshi/modules/conv.py @@ -0,0 +1,423 @@ +# Copyright (c) Kyutai, all rights reserved. +# This source code is licensed under the license found in the +# LICENSE file in the root directory of this source tree. + +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the license found in the +# LICENSE file in the root directory of this source tree. + +from dataclasses import dataclass +import itertools +import math +import typing as tp +import warnings + +import torch +from torch import nn +from torch.nn import functional as F +from torch.nn.utils import weight_norm + +from .streaming import StreamingModule, State + + +CONV_NORMALIZATIONS = frozenset(["none", "weight_norm"]) +M = tp.TypeVar('M', bound=nn.Module) + + +class TransposedLayerNorm(nn.Module): + """LayerNorm for [B, C, T] inputs.""" + + def __init__(self, **kwargs): + super().__init__() + self.layer_norm = nn.LayerNorm(**kwargs) + + def forward(self, x): + x = x.transpose(1, 2) + x = self.layer_norm(x) + return x.transpose(1, 2) + + +def apply_parametrization_norm(module: M, norm: str = "none") -> M: + assert norm in CONV_NORMALIZATIONS + if norm == "weight_norm": + return weight_norm(module) + else: + # We already check was in CONV_NORMALIZATION, so any other choice + # doesn't need reparametrization. + return module + + +def get_extra_padding_for_conv1d( + x: torch.Tensor, kernel_size: int, stride: int, padding_total: int = 0 +) -> int: + """See `pad_for_conv1d`.""" + length = x.shape[-1] + n_frames = (length - kernel_size + padding_total) / stride + 1 + ideal_length = (math.ceil(n_frames) - 1) * stride + (kernel_size - padding_total) + return ideal_length - length + + +def pad_for_conv1d( + x: torch.Tensor, kernel_size: int, stride: int, padding_total: int = 0 +): + """Pad for a convolution to make sure that the last window is full. + Extra padding is added at the end. This is required to ensure that we can rebuild + an output of the same length, as otherwise, even with padding, some time steps + might get removed. + For instance, with total padding = 4, kernel size = 4, stride = 2: + 0 0 1 2 3 4 5 0 0 # (0s are padding) + 1 2 3 # (output frames of a convolution, last 0 is never used) + 0 0 1 2 3 4 5 0 # (output of tr. conv., but pos. 5 is going to get removed as padding) + 1 2 3 4 # once you removed padding, we are missing one time step ! + """ + extra_padding = get_extra_padding_for_conv1d(x, kernel_size, stride, padding_total) + return F.pad(x, (0, extra_padding)) + + +def pad1d( + x: torch.Tensor, + paddings: tp.Tuple[int, int], + mode: str = "constant", + value: float = 0.0, +): + """Tiny wrapper around F.pad, just to allow for reflect padding on small input. + If this is the case, we insert extra 0 padding to the right before the reflection happen. + """ + length = x.shape[-1] + padding_left, padding_right = paddings + assert padding_left >= 0 and padding_right >= 0, (padding_left, padding_right) + if mode == "reflect": + max_pad = max(padding_left, padding_right) + extra_pad = 0 + if length <= max_pad: + extra_pad = max_pad - length + 1 + x = F.pad(x, (0, extra_pad)) + padded = F.pad(x, paddings, mode, value) + end = padded.shape[-1] - extra_pad + return padded[..., :end] + else: + return F.pad(x, paddings, mode, value) + + +def unpad1d(x: torch.Tensor, paddings: tp.Tuple[int, int]): + """Remove padding from x, handling properly zero padding. Only for 1d!""" + padding_left, padding_right = paddings + assert padding_left >= 0 and padding_right >= 0, (padding_left, padding_right) + assert (padding_left + padding_right) <= x.shape[-1] + end = x.shape[-1] - padding_right + return x[..., padding_left:end] + + +class NormConv1d(nn.Module): + """Wrapper around Conv1d and normalization applied to this conv + to provide a uniform interface across normalization approaches. + """ + + def __init__( + self, + *args, + causal: bool = False, + norm: str = "none", + norm_kwargs: tp.Dict[str, tp.Any] = {}, + **kwargs, + ): + super().__init__() + self.conv = apply_parametrization_norm( + nn.Conv1d(*args, **kwargs), norm + ) + self.norm_type = norm + + def forward(self, x): + x = self.conv(x) + return x + + +class NormConvTranspose1d(nn.Module): + """Wrapper around ConvTranspose1d and normalization applied to this conv + to provide a uniform interface across normalization approaches. + """ + + def __init__( + self, + *args, + causal: bool = False, + norm: str = "none", + norm_kwargs: tp.Dict[str, tp.Any] = {}, + **kwargs, + ): + super().__init__() + self.convtr = apply_parametrization_norm( + nn.ConvTranspose1d(*args, **kwargs), norm + ) + self.norm_type = norm + + def forward(self, x): + x = self.convtr(x) + return x + + +@dataclass +class _StreamingConv1dState(State): + previous: torch.Tensor + first: torch.Tensor + + def reset(self, reset_mask: torch.Tensor): + super().reset(reset_mask) + self.previous[:] = torch.where(reset_mask.view(-1, 1, 1), torch.zeros_like(self.previous), self.previous) + self.first[:] = torch.where(reset_mask, torch.ones_like(self.first), self.first) + + +class StreamingConv1d(StreamingModule[_StreamingConv1dState]): + """Conv1d with some builtin handling of asymmetric or causal padding + and normalization. + """ + + def __init__( + self, + in_channels: int, + out_channels: int, + kernel_size: int, + stride: int = 1, + dilation: int = 1, + groups: int = 1, + bias: bool = True, + causal: bool = False, + norm: str = "none", + norm_kwargs: tp.Dict[str, tp.Any] = {}, + pad_mode: str = "constant", + ): + super().__init__() + assert pad_mode in ['constant', 'replicate'], pad_mode + self.pad_mode = pad_mode + assert causal + # warn user on unusual setup between dilation and stride + if stride > 1 and dilation > 1: + warnings.warn( + "StreamingConv1d has been initialized with stride > 1 and dilation > 1" + f" (kernel_size={kernel_size} stride={stride}, dilation={dilation})." + ) + self.conv = NormConv1d( + in_channels, + out_channels, + kernel_size, + stride, + dilation=dilation, + groups=groups, + bias=bias, + causal=causal, + norm=norm, + norm_kwargs=norm_kwargs, + ) + + @property + def _stride(self) -> int: + return self.conv.conv.stride[0] + + @property + def _kernel_size(self) -> int: + return self.conv.conv.kernel_size[0] + + @property + def _effective_kernel_size(self) -> int: + dilation = self.conv.conv.dilation[0] + return ( + self._kernel_size - 1 + ) * dilation + 1 # effective kernel size with dilations + + @property + def _padding_total(self) -> int: + return self._effective_kernel_size - self._stride + + def _init_streaming_state(self, batch_size: int) -> _StreamingConv1dState: + stride = self._stride + # Effective kernel size accounting for dilation. + kernel = self._effective_kernel_size + param = next(iter(self.parameters())) + dtype = param.dtype + device = param.device + previous = torch.zeros(batch_size, self.conv.conv.in_channels, kernel - stride, + dtype=dtype, device=device) + first = torch.ones(batch_size, device=device, dtype=torch.bool) + return _StreamingConv1dState(batch_size, device, previous, first) + + def forward(self, x): + B, C, T = x.shape + S = self._stride + assert T > 0 and T % S == 0, "Steps must be multiple of stride" + state = self._streaming_state + if state is None: + state = self._init_streaming_state(B) + TP = state.previous.shape[-1] + if TP and self.pad_mode == 'replicate': + assert T >= TP, "Not enough content to pad streaming." + init = x[..., :1] + state.previous[:] = torch.where( + state.first.view(-1, 1, 1) & state.exec_mask.view(-1, 1, 1), + init, + state.previous) + if TP: + x = torch.cat([state.previous, x], dim=-1) + y = self.conv(x) + if TP: + state.previous[:] = torch.where( + state.exec_mask.view(-1, 1, 1), + x[..., -TP:], + state.previous) + if self.pad_mode == 'replicate': + state.first = torch.where( + state.exec_mask, + torch.zeros_like(state.first), + state.first, + ) + return y + + +@dataclass +class _StreamingConvTr1dState(State): + partial: torch.Tensor + + def reset(self, reset_mask: torch.Tensor): + super().reset(reset_mask) + self.partial[:] = torch.where( + reset_mask.view(-1, 1, 1), + torch.zeros_like(self.partial), + self.partial) + + +class StreamingConvTranspose1d(StreamingModule[_StreamingConvTr1dState]): + """ConvTranspose1d with some builtin handling of asymmetric or causal padding + and normalization. + """ + + def __init__( + self, + in_channels: int, + out_channels: int, + kernel_size: int, + stride: int = 1, + groups: int = 1, + bias: bool = True, + causal: bool = False, + norm: str = "none", + trim_right_ratio: float = 1.0, + norm_kwargs: tp.Dict[str, tp.Any] = {}, + ): + super().__init__() + assert trim_right_ratio == 1. + assert causal + self.convtr = NormConvTranspose1d( + in_channels, + out_channels, + kernel_size, + stride, + groups=groups, + bias=bias, + causal=causal, + norm=norm, + norm_kwargs=norm_kwargs, + ) + + @property + def _stride(self) -> int: + return self.convtr.convtr.stride[0] + + @property + def _kernel_size(self) -> int: + return self.convtr.convtr.kernel_size[0] + + def _init_streaming_state(self, batch_size: int) -> _StreamingConvTr1dState: + param = next(iter(self.parameters())) + dtype = param.dtype + device = param.device + K = self._kernel_size + S = self._stride + partial = torch.zeros(batch_size, self.convtr.convtr.out_channels, K - S, + device=device, dtype=dtype) + return _StreamingConvTr1dState(batch_size, device, partial) + + def forward(self, x): + B, C, T = x.shape + K = self._kernel_size + S = self._stride + state = self._streaming_state + + y = self.convtr(x) + if state is None: + y = unpad1d(y, (0, K - S)) + else: + PT = state.partial.shape[-1] + if PT > 0: + y[..., :PT] += state.partial + bias = self.convtr.convtr.bias + for_partial = y[..., -PT:] + if bias is not None: + for_partial -= bias[:, None] + state.partial[:] = torch.where( + state.exec_mask.view(-1, 1, 1), + for_partial, + state.partial) + y = y[..., :-PT] + return y + + +def test(): + torch.manual_seed(1234) + device = "cpu" + if torch.cuda.is_available(): + # Avoid the cuda optimizations that would take place on single precision + # floats for convolutions. + torch.backends.cudnn.enabled = True + torch.backends.cudnn.benchmark = False + torch.backends.cudnn.deterministic = True + torch.backends.cuda.matmul.allow_tf32 = False + torch.backends.cudnn.allow_tf32 = False + device = "cuda:0" + + kernel_sizes = [1, 3, 4, 8, 15, 16] + strides = [1, 2, 3, 4, 5, 6, 7, 8, 9] + chin = 6 + chout = 12 + + for kernel, stride in itertools.product(kernel_sizes, strides): + if stride > kernel: + continue + conv = StreamingConv1d(chin, chout, kernel, stride, causal=True).to(device) + convtr = StreamingConvTranspose1d(chout, chin, kernel, stride, causal=True).to(device) + + for frames in [1, 4, 8, 32, 54, 65, 128]: + print(f"ksize {kernel} strides {stride} frames {frames}") + batch_size = 3 + length = frames * stride + x = torch.randn(batch_size, chin, length).to(device) + y = conv(x) + z = convtr(y) + for chunk_frames in [1, 2, 8]: + if frames % chunk_frames != 0: + continue + ys = [] + zs = [] + chunk_length = chunk_frames * stride + with conv.streaming(batch_size), convtr.streaming(batch_size): + for offset in range(0, length, chunk_length): + chunk = x[..., offset : offset + chunk_length] + ys.append(conv(chunk)) + zs.append(convtr(ys[-1])) + y_stream = torch.cat(ys, dim=-1) + z_stream = torch.cat(zs, dim=-1) + y = y[..., : y_stream.shape[-1]] + z = z[..., : z_stream.shape[-1]] + assert y.shape == y_stream.shape, (y.shape, y_stream.shape) + delta = (y_stream - y).norm() / y.norm() + assert delta <= 1e-6, delta + assert frames == y_stream.shape[-1], (frames, y_stream.shape) + + assert z.shape == z_stream.shape, (z.shape, z_stream.shape) + delta = (z_stream - z).norm() / z.norm() + assert delta <= 1e-6, (delta, (z_stream - z).abs().mean(dim=(0, 1))) + + +if __name__ == "__main__": + with torch.no_grad(): + test() diff --git a/moshi_src/moshi/modules/conv_test.py b/moshi_src/moshi/modules/conv_test.py new file mode 100644 index 0000000..8bd3228 --- /dev/null +++ b/moshi_src/moshi/modules/conv_test.py @@ -0,0 +1,157 @@ +import functools +import torch +import torch.nn as nn +import pytest + +from .conv import StreamingConv1d, StreamingConvTranspose1d + + +torch.backends.cudnn.enabled = False # Disable cuDNN for deterministic behavior and for numerical stability + + +CONV1D_DATA = [ + # batch_size, in_channels, out_channels, seq_len, kernel_size + pytest.param( + 3, 4, 5, 10, 6, + id='small conv1d test 1', + ), + pytest.param( + 4, 5, 6, 10, 7, + id='small conv1d test 2', + ), + pytest.param( + 5, 6, 7, 10, 2, + id='small conv1d test 3', + ), + pytest.param( + 1, 512, 512, 256, 7, + id='large conv1d test 1', + ), +] + +CONV1D_TRANSPOSE_DATA = [ + # batch_size, in_channels, out_channels, seq_len, kernel_size, stride + pytest.param( + 3, 4, 5, 10, 6, 1, + id='small conv1d transpose test 1', + ), + pytest.param( + 4, 5, 6, 10, 7, 2, + id='small conv1d transpose test 2', + ), + pytest.param( + 5, 6, 7, 10, 4, 3, + id='small conv1d transpose test 3', + ), + pytest.param( + 1, 512, 512, 256, 7, 2, + id='large conv1d transpose test 1', + ), +] + + +def _init_weights(module, generator=None): + for name, param in module.named_parameters(): + if "weight" in name: + nn.init.xavier_uniform_(param, generator=generator) + elif "bias" in name: + nn.init.constant_(param, 0.0) + else: + nn.init.xavier_uniform_(param, generator=generator) + + +@pytest.mark.parametrize("batch_size, in_channels, out_channels, seq_len, kernel_size", CONV1D_DATA) +def test_conv1d(batch_size, in_channels, out_channels, seq_len, kernel_size): + """Test that StreamingConv1d() calls are causal. Having new inputs does not change the previous output.""" + assert seq_len > kernel_size + + layer = StreamingConv1d(in_channels, out_channels, kernel_size, causal=True, norm="none", pad_mode="constant") + + generator = torch.Generator() + generator = generator.manual_seed(41) + layer.apply(functools.partial(_init_weights, generator=generator)) + + shape = (batch_size, in_channels, seq_len,) + input_hidden_states = torch.rand(shape) + + expected_output = layer(input_hidden_states) + + for end_index in range(kernel_size, seq_len + 1): + actual_output = layer(input_hidden_states[..., :end_index]) + torch.testing.assert_close(actual_output, expected_output[..., :actual_output.shape[-1]], + msg=lambda original_msg: f"Failed at end_index={end_index}: \n{original_msg}") + + +@pytest.mark.parametrize("batch_size, in_channels, out_channels, seq_len, kernel_size", CONV1D_DATA) +def test_conv1d_streaming(batch_size, in_channels, out_channels, seq_len, kernel_size): + """Test that StreamingConv1d() streaming works as expected.""" + assert seq_len > kernel_size + + layer = StreamingConv1d(in_channels, out_channels, kernel_size, causal=True, norm="none", pad_mode="constant") + + generator = torch.Generator() + generator = generator.manual_seed(41) + layer.apply(functools.partial(_init_weights, generator=generator)) + + shape = (batch_size, in_channels, seq_len,) + input_hidden_states = torch.rand(shape) + expected_output = layer(input_hidden_states) + + start_index = 0 + actual_outputs = [] + with layer.streaming(batch_size=batch_size): + for end_index in range(kernel_size, seq_len + 1): + actual_output = layer(input_hidden_states[..., start_index:end_index]) + start_index = end_index + actual_outputs.append(actual_output) + actual_outputs = torch.cat(actual_outputs, dim=-1) + + torch.testing.assert_close(actual_outputs, expected_output) + + +@pytest.mark.parametrize("batch_size, in_channels, out_channels, seq_len, kernel_size, stride", CONV1D_TRANSPOSE_DATA) +def test_conv1d_transpose(batch_size, in_channels, out_channels, seq_len, kernel_size, stride): + """Test that StreamingConvTranspose1d() calls are causal. Having new inputs does not change the previous output.""" + assert seq_len > kernel_size + + layer = StreamingConvTranspose1d(in_channels, out_channels, kernel_size, stride, causal=True, norm="none") + + generator = torch.Generator() + generator = generator.manual_seed(41) + layer.apply(functools.partial(_init_weights, generator=generator)) + + shape = (batch_size, in_channels, seq_len,) + input_hidden_states = torch.rand(shape) + expected_output = layer(input_hidden_states) + + for end_index in range(kernel_size, seq_len + 1): + actual_output = layer(input_hidden_states[..., :end_index]) + torch.testing.assert_close(actual_output, expected_output[..., :actual_output.shape[-1]], + msg=lambda original_msg: f"Failed at end_index={end_index}: \n{original_msg}") + + +@pytest.mark.parametrize("batch_size, in_channels, out_channels, seq_len, kernel_size, stride", CONV1D_TRANSPOSE_DATA) +def test_conv1d_transpose_streaming(batch_size, in_channels, out_channels, seq_len, kernel_size, stride): + """Test that StreamingConvTranspose1d() streaming works as expected.""" + assert seq_len > kernel_size + + layer = StreamingConvTranspose1d(in_channels, out_channels, kernel_size, stride, causal=True, norm="none") + + generator = torch.Generator() + generator = generator.manual_seed(41) + layer.apply(functools.partial(_init_weights, generator=generator)) + + shape = (batch_size, in_channels, seq_len,) + input_hidden_states = torch.rand(shape) + expected_output = layer(input_hidden_states) + + start_index = 0 + actual_outputs = [] + with layer.streaming(batch_size=batch_size): + for end_index in range(kernel_size, seq_len + 1): + actual_output = layer(input_hidden_states[..., start_index:end_index]) + start_index = end_index + actual_outputs.append(actual_output) + actual_outputs = torch.cat(actual_outputs, dim=-1) + + torch.testing.assert_close(actual_outputs, expected_output) diff --git a/moshi_src/moshi/modules/gating.py b/moshi_src/moshi/modules/gating.py new file mode 100644 index 0000000..bd708ad --- /dev/null +++ b/moshi_src/moshi/modules/gating.py @@ -0,0 +1,115 @@ +# Copyright (c) Kyutai, all rights reserved. +# This source code is licensed under the license found in the +# LICENSE file in the root directory of this source tree. + +from contextlib import ExitStack +import torch +from torch import nn +from torch.nn import functional as F + +from ..utils.compile import torch_compile_lazy, no_compile + + + +def gating_forward_kernel( + weight_in: torch.Tensor, weight_out: torch.Tensor, activation, x: torch.Tensor +): + x = F.linear(x, weight_in) + B, T, _ = x.shape + x = x.view(B, T, 2, -1) + x = activation(x[..., 0, :]) * x[..., 1, :] + x = F.linear(x, weight_out) + return x + + +def gating_forward_generic( + linear_in: nn.Module, + linear_out: nn.Module, + activation, + x: torch.Tensor +): + x = linear_in(x) + B, T, _ = x.shape + x = x.view(B, T, 2, -1) + x = activation(x[..., 0, :]) * x[..., 1, :] + x = linear_out(x) + return x + + +class ActivationGating(nn.Module): + """ + Gating FFN layer, using the given activation. + Args: + dim (int): dimension of the input and output of the transformer. + activation (any callable Tensor to Tensor): activation function to use. + **factory_kwargs: other kwargs passed to the linear layer, in particular device and dtype. + """ + + _fsdp_final = True + + def __init__(self, dim: int, dim_feedforward: int, activation, quantized: bool = False, **factory_kwargs): + super().__init__() + # We should have 8 d^2 param, instead we will have + # 2 * h * d + h * d = 3 h * d = 8 d^2 + # so h = 8 d / 3 but following Hervé's advice we use 21 / 8 as an approx. + if dim_feedforward == 4 * dim: + hidden = (21 * dim) // 8 + else: + hidden = (2 * dim_feedforward) // 3 + + self.linear_in = nn.Linear(dim, 2 * hidden, bias=False, **factory_kwargs) + self.linear_out = nn.Linear(hidden, dim, bias=False, **factory_kwargs) + + # We try to follow the default PyTorch MHA convention, to easily compare results. + + self.activation = activation + + def forward(self, x: torch.Tensor): + if isinstance(self.linear_in, nn.Linear): + assert isinstance(self.linear_out, nn.Linear) + with ExitStack() as stack: + if self.training: + stack.enter_context(no_compile()) + return gating_forward_kernel( + self.linear_in.weight, self.linear_out.weight, self.activation, x + ) + else: + return gating_forward_generic( + self.linear_in, + self.linear_out, + self.activation, + x + ) + + +def _get_activation(name: str): + if name in ["sigmoid", "tanh", "relu"]: + return getattr(torch, name) + elif name in ["leaky_relu", "elu", "gelu", "silu", "mish", "softsign"]: + return getattr(torch.nn.functional, name) + elif name == "identity": + return torch.nn.Identity() + else: + raise ValueError(f"Unknown activation {name}") + + +def _make_gating( + name: str, dim: int, dim_feedforward: int, + **factory_kwargs +) -> nn.Module: + return ActivationGating( + dim, dim_feedforward, _get_activation(name), **factory_kwargs + ) + + +def make_gating( + name: str, dim: int, dim_feedforward: int, **factory_kwargs +) -> nn.Module: + gating = _make_gating(name, dim, dim_feedforward, **factory_kwargs) + if isinstance(gating.linear_in, nn.Linear): + max_params = 2 * dim * dim_feedforward + params = sum(p.numel() for p in gating.parameters()) + assert ( + params <= max_params + ), f"{name} gating has {params} params, max is {max_params}" + return gating diff --git a/moshi_src/moshi/modules/lora.py b/moshi_src/moshi/modules/lora.py new file mode 100644 index 0000000..b4731bc --- /dev/null +++ b/moshi_src/moshi/modules/lora.py @@ -0,0 +1,122 @@ +import torch +import torch.nn as nn + + +def replace_all_linear_with_lora(module, rank: int, scaling: float, device=None, dtype=None): + """ Recursively replace all Linear layers with LoRALinear layers.""" + for name, child in module.named_children(): + if isinstance(child, nn.Linear): + if device is None: + this_device = child.weight.device + else: + this_device = device + if dtype is None: + this_dtype = child.weight.dtype + else: + this_dtype = dtype + lora = LoRALinear(child.in_features, child.out_features, + rank, scaling, device=this_device, dtype=this_dtype) + lora.frozen_W = child + setattr(module, name, lora) + else: + replace_all_linear_with_lora(child, rank, scaling, device=device, dtype=dtype) + + +def replace_lora_with_linear(module): + """Recursively replace all LoRALinear layers with Linear layers.""" + for name, child in module.named_children(): + if isinstance(child, LoRALinear): + # Compute merged weights: W' = W + scaling * B @ A + merged_weight = child.frozen_W.weight.data + \ + child.scaling * (child.lora_B.weight @ child.lora_A.weight) + # Create a standard Linear layer with the same in/out features + new_linear = nn.Linear(child.frozen_W.in_features, + child.frozen_W.out_features, bias=False, + device=torch.device('meta'), + dtype=merged_weight.dtype) + new_linear.weight = nn.Parameter( + merged_weight, requires_grad=merged_weight.requires_grad) # Transfer merged weights + setattr(module, name, new_linear) # Replace the module + else: + replace_lora_with_linear(child) # Recursively process submodules + + +class LoRALinear(nn.Module): + """ + Implementation of: + - LoRA: https://arxiv.org/abs/2106.09685 + + Notes: + - Freezing is handled at the network level, not the layer level. + - Scaling factor controls relative importance of LoRA skip + connection versus original frozen weight. General guidance is + to keep it to 2.0 and sweep over learning rate when changing + the rank. + """ + + def __init__( + self, + in_features: int, + out_features: int, + rank: int, + scaling: float, + bias: bool = False, + device: torch.device | None = None, + dtype: torch.dtype = torch.bfloat16, + ): + super().__init__() + + self.in_features = in_features + self.out_features = out_features + assert not bias + self.bias = bias + self.rank = rank + self.scaling = scaling + + self.lora_A = nn.Linear( + self.in_features, + self.rank, + bias=self.bias, + device=device, + dtype=dtype, + ) + self.lora_B = nn.Linear( + self.rank, + self.out_features, + bias=self.bias, + device=device, + dtype=dtype, + ) + + self.frozen_W = nn.Linear(self.in_features, + self.out_features, + bias=self.bias, + device=device, + dtype=dtype) + + self._register_load_state_dict_pre_hook(LoRALinear._load_hook, with_module=True) + + def merge_weight(self): + with torch.no_grad(): + down_weight = self.lora_A.weight + up_weight = self.lora_B.weight + + weight = up_weight.mm(down_weight) * self.scaling + + weight += self.frozen_W.weight + return weight + + @staticmethod + def _load_hook(module, state_dict, prefix, *_): + key_name = prefix + "weight" + if key_name in state_dict: + w_ref = state_dict.pop(key_name) + state_dict[prefix + 'frozen_W.weight'] = w_ref + + def forward(self, x: torch.Tensor): + lora = self.lora_B(self.lora_A(x)) + return self.frozen_W(x) + lora * self.scaling + + def __repr__(self) -> str: + return "{}Linear(in_features={}, out_features={}, r={})".format( + "LoRA", self.in_features, self.out_features, self.rank) diff --git a/moshi_src/moshi/modules/resample.py b/moshi_src/moshi/modules/resample.py new file mode 100644 index 0000000..e260796 --- /dev/null +++ b/moshi_src/moshi/modules/resample.py @@ -0,0 +1,119 @@ +# Copyright (c) Kyutai, all rights reserved. +# This source code is licensed under the license found in the +# LICENSE file in the root directory of this source tree. + +import typing as tp + +from einops import rearrange +import torch +from torch import nn + +from .conv import StreamingConv1d, StreamingConvTranspose1d + + +class ConvDownsample1d(nn.Module): + """ + Downsampling by some integer amount `stride` using convolutions + with a kernel size of twice the stride. + If `causal` is True, the output uses a causal convolution. + """ + + def __init__( + self, + stride: int, + dimension: tp.Optional[int] = None, + causal: bool = False, + learnt: bool = False, + channel_wise: bool = False, + ): + super().__init__() + self.learnt = learnt + self.channel_wise = channel_wise + groups = 1 + if learnt: + assert dimension is not None, "Dimension required for learnt convolutions." + in_channels = dimension + out_channels = dimension + if channel_wise: + groups = dimension + else: + in_channels = 1 + out_channels = 1 + + self.conv = StreamingConv1d( + in_channels, + out_channels, + kernel_size=2 * stride, + stride=stride, + causal=causal, + groups=groups, + bias=False, + pad_mode="replicate", + ) + if not learnt: + actual_conv = self.conv.conv.conv + actual_conv.weight.requires_grad_(False) + actual_conv.weight.data.fill_(1.0 / (2 * stride)) + + def forward(self, x: torch.Tensor): + batch_size = len(x) + if not self.learnt: + x = rearrange(x, "b c t -> (b c) () t") + y = self.conv(x) + if not self.learnt: + y = rearrange(y, "(b c) () t -> b c t", b=batch_size) + return y + + +class ConvTrUpsample1d(nn.Module): + """ + Upsample by some integer amount `stride` using transposed convolutions. + """ + + def __init__( + self, + stride: int, + dimension: tp.Optional[int] = None, + causal: bool = False, + learnt: bool = False, + channel_wise: bool = False, + ): + super().__init__() + self.learnt = learnt + self.channel_wise = channel_wise + groups = 1 + if learnt: + assert dimension is not None, "Dimension required for learnt convolutions." + in_channels = dimension + out_channels = dimension + if channel_wise: + groups = dimension + else: + in_channels = 1 + out_channels = 1 + + self.convtr = StreamingConvTranspose1d( + in_channels, + out_channels, + kernel_size=2 * stride, + stride=stride, + causal=causal, + groups=groups, + bias=False, + ) + if not learnt: + actual_convtr = self.convtr.convtr.convtr + actual_convtr.weight.requires_grad_(False) + actual_convtr.weight.data.fill_(1.0) + + def forward(self, x: torch.Tensor): + batch_size = len(x) + if not self.learnt: + x = rearrange(x, "b c t -> (b c) () t") + y = self.convtr(x) + if not self.learnt: + x_for_normalization = torch.ones_like(x[:1]) + normalization = self.convtr(x_for_normalization) + y = y / normalization + y = rearrange(y, "(b c) () t -> b c t", b=batch_size) + return y diff --git a/moshi_src/moshi/modules/rope.py b/moshi_src/moshi/modules/rope.py new file mode 100644 index 0000000..0965501 --- /dev/null +++ b/moshi_src/moshi/modules/rope.py @@ -0,0 +1,90 @@ +# Copyright (c) Kyutai, all rights reserved. +# This source code is licensed under the license found in the +# LICENSE file in the root directory of this source tree. + +from torch import nn +import math +import torch +from ..utils.compile import torch_compile_lazy + + + +def apply_rope( + q: torch.Tensor, + k: torch.Tensor, + offset: torch.Tensor, + max_period: float = 10_000, + time_before_heads: bool = False, +): + """ + Args: + q (torch.Tensor): queries, shape `[B, T, H, D]`. + k (torch.Tensor): keys, shape `[B, T, H, D]`. + offset (int): current offset, e.g. when streaming. + max_period (float): maximum period for the cos and sin. + time_before_heads (bool): if True, expected [B, T, H, D], else [B, H, T ,D] + """ + + if time_before_heads: + B, T, H, D = q.shape + else: + B, H, T, D = q.shape + assert k.shape == q.shape + assert D > 0 + assert D % 2 == 0 + assert max_period > 0 + + ds = torch.arange(D // 2, device=q.device, dtype=torch.float32) + freqs = torch.exp(ds * (-math.log(max_period) * 2 / D)) + ts = offset.float().view(-1, 1) + torch.arange(T, device=q.device, dtype=torch.float32) + if time_before_heads: + ts = ts.view(B, -1, 1, 1) + else: + ts = ts.view(B, 1, -1, 1) + + dims = q.shape[:-1] + q = q.view(*dims, D // 2, 2) + k = k.view(*dims, D // 2, 2) + + # convention is `r` suffix is real part, `i` is imaginary. + qr = q[..., 0].float() + qi = q[..., 1].float() + + kr = k[..., 0].float() + ki = k[..., 1].float() + + rotr = torch.cos(freqs * ts) + roti = torch.sin(freqs * ts) + qor = qr * rotr - qi * roti + qoi = qr * roti + qi * rotr + + kor = kr * rotr - ki * roti + koi = kr * roti + ki * rotr + + dtype = q.dtype + qo = torch.stack([qor.to(dtype), qoi.to(dtype)], dim=-1) + ko = torch.stack([kor.to(dtype), koi.to(dtype)], dim=-1) + + return qo.view(*dims, D), ko.view(*dims, D) + + +class RotaryEmbedding(nn.Module): + """Rotary positional embedding (RoPE) from [Su et al 2022](https://arxiv.org/abs/2104.09864). + + Args: + max_period (float): Maximum period of the rotation frequencies. + """ + + def __init__(self, max_period: float = 10000.0): + super().__init__() + self.max_period = max_period + + def forward( + self, + q: torch.Tensor, + k: torch.Tensor, + offset: torch.Tensor, + time_before_heads: bool = False, + ): + """Apply rope rotation to query or key tensor.""" + return apply_rope(q, k, offset, self.max_period, time_before_heads) diff --git a/moshi_src/moshi/modules/seanet.py b/moshi_src/moshi/modules/seanet.py new file mode 100644 index 0000000..0129190 --- /dev/null +++ b/moshi_src/moshi/modules/seanet.py @@ -0,0 +1,392 @@ +# Copyright (c) Kyutai, all rights reserved. +# This source code is licensed under the license found in the +# LICENSE file in the root directory of this source tree. + +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the license found in the +# LICENSE file in the root directory of this source tree. + +import typing as tp + +import numpy as np +import torch.nn as nn + +from .conv import StreamingConv1d, StreamingConvTranspose1d +from .streaming import StreamingContainer + + +class SEANetResnetBlock(StreamingContainer): + """Residual block from SEANet model. + + Args: + dim (int): Dimension of the input/output. + kernel_sizes (list): List of kernel sizes for the convolutions. + dilations (list): List of dilations for the convolutions. + activation (str): Activation function. + activation_params (dict): Parameters to provide to the activation function. + norm (str): Normalization method. + norm_params (dict): Parameters to provide to the underlying normalization used along with the convolution. + causal (bool): Whether to use fully causal convolution. + pad_mode (str): Padding mode for the convolutions. + compress (int): Reduced dimensionality in residual branches (from Demucs v3). + true_skip (bool): Whether to use true skip connection or a simple + (streamable) convolution as the skip connection. + """ + + def __init__( + self, + dim: int, + kernel_sizes: tp.List[int] = [3, 1], + dilations: tp.List[int] = [1, 1], + activation: str = "ELU", + activation_params: dict = {"alpha": 1.0}, + norm: str = "none", + norm_params: tp.Dict[str, tp.Any] = {}, + causal: bool = False, + pad_mode: str = "reflect", + compress: int = 2, + true_skip: bool = True, + ): + super().__init__() + assert len(kernel_sizes) == len( + dilations + ), "Number of kernel sizes should match number of dilations" + act = getattr(nn, activation) + hidden = dim // compress + block = [] + for i, (kernel_size, dilation) in enumerate(zip(kernel_sizes, dilations)): + in_chs = dim if i == 0 else hidden + out_chs = dim if i == len(kernel_sizes) - 1 else hidden + block += [ + act(**activation_params), + StreamingConv1d( + in_chs, + out_chs, + kernel_size=kernel_size, + dilation=dilation, + norm=norm, + norm_kwargs=norm_params, + causal=causal, + pad_mode=pad_mode, + ), + ] + self.block = nn.Sequential(*block) + self.shortcut: nn.Module + if true_skip: + self.shortcut = nn.Identity() + else: + self.shortcut = StreamingConv1d( + dim, + dim, + kernel_size=1, + norm=norm, + norm_kwargs=norm_params, + causal=causal, + pad_mode=pad_mode, + ) + + def forward(self, x): + u, v = self.shortcut(x), self.block(x) + assert u.shape == v.shape, (u.shape, v.shape, x.shape) + return u + v + + +class SEANetEncoder(StreamingContainer): + """SEANet encoder. + + Args: + channels (int): Audio channels. + dimension (int): Intermediate representation dimension. + n_filters (int): Base width for the model. + n_residual_layers (int): nb of residual layers. + ratios (Sequence[int]): kernel size and stride ratios. The encoder uses downsampling ratios instead of + upsampling ratios, hence it will use the ratios in the reverse order to the ones specified here + that must match the decoder order. We use the decoder order as some models may only employ the decoder. + activation (str): Activation function. + activation_params (dict): Parameters to provide to the activation function. + norm (str): Normalization method. + norm_params (dict): Parameters to provide to the underlying normalization used along with the convolution. + kernel_size (int): Kernel size for the initial convolution. + last_kernel_size (int): Kernel size for the initial convolution. + residual_kernel_size (int): Kernel size for the residual layers. + dilation_base (int): How much to increase the dilation with each layer. + causal (bool): Whether to use fully causal convolution. + pad_mode (str): Padding mode for the convolutions. + true_skip (bool): Whether to use true skip connection or a simple + (streamable) convolution as the skip connection in the residual network blocks. + compress (int): Reduced dimensionality in residual branches (from Demucs v3). + disable_norm_outer_blocks (int): Number of blocks for which we don't apply norm. + For the encoder, it corresponds to the N first blocks. + mask_fn (nn.Module): Optional mask function to apply after convolution layers. + mask_position (int): Position of the mask function, with mask_position == 0 for the first convolution layer, + mask_position == 1 for the first conv block, etc. + """ + + def __init__( + self, + channels: int = 1, + dimension: int = 128, + n_filters: int = 32, + n_residual_layers: int = 3, + ratios: tp.List[int] = [8, 5, 4, 2], + activation: str = "ELU", + activation_params: dict = {"alpha": 1.0}, + norm: str = "none", + norm_params: tp.Dict[str, tp.Any] = {}, + kernel_size: int = 7, + last_kernel_size: int = 7, + residual_kernel_size: int = 3, + dilation_base: int = 2, + causal: bool = False, + pad_mode: str = "reflect", + true_skip: bool = True, + compress: int = 2, + disable_norm_outer_blocks: int = 0, + mask_fn: tp.Optional[nn.Module] = None, + mask_position: tp.Optional[int] = None, + ): + super().__init__() + self.channels = channels + self.dimension = dimension + self.n_filters = n_filters + self.ratios = list(reversed(ratios)) + del ratios + self.n_residual_layers = n_residual_layers + self.hop_length = int(np.prod(self.ratios)) + self.n_blocks = len(self.ratios) + 2 # first and last conv + residual blocks + self.disable_norm_outer_blocks = disable_norm_outer_blocks + assert ( + self.disable_norm_outer_blocks >= 0 and self.disable_norm_outer_blocks <= self.n_blocks + ), ( + "Number of blocks for which to disable norm is invalid." + "It should be lower or equal to the actual number of blocks in the network and greater or equal to 0." + ) + + act = getattr(nn, activation) + mult = 1 + model: tp.List[nn.Module] = [ + StreamingConv1d( + channels, + mult * n_filters, + kernel_size, + norm="none" if self.disable_norm_outer_blocks >= 1 else norm, + norm_kwargs=norm_params, + causal=causal, + pad_mode=pad_mode, + ) + ] + if mask_fn is not None and mask_position == 0: + model += [mask_fn] + # Downsample to raw audio scale + for i, ratio in enumerate(self.ratios): + block_norm = "none" if self.disable_norm_outer_blocks >= i + 2 else norm + # Add residual layers + for j in range(n_residual_layers): + model += [ + SEANetResnetBlock( + mult * n_filters, + kernel_sizes=[residual_kernel_size, 1], + dilations=[dilation_base**j, 1], + norm=block_norm, + norm_params=norm_params, + activation=activation, + activation_params=activation_params, + causal=causal, + pad_mode=pad_mode, + compress=compress, + true_skip=true_skip, + ) + ] + + # Add downsampling layers + model += [ + act(**activation_params), + StreamingConv1d( + mult * n_filters, + mult * n_filters * 2, + kernel_size=ratio * 2, + stride=ratio, + norm=block_norm, + norm_kwargs=norm_params, + causal=causal, + pad_mode=pad_mode, + ), + ] + mult *= 2 + if mask_fn is not None and mask_position == i + 1: + model += [mask_fn] + + model += [ + act(**activation_params), + StreamingConv1d( + mult * n_filters, + dimension, + last_kernel_size, + norm=( + "none" if self.disable_norm_outer_blocks == self.n_blocks else norm + ), + norm_kwargs=norm_params, + causal=causal, + pad_mode=pad_mode, + ), + ] + + self.model = nn.Sequential(*model) + + def forward(self, x): + return self.model(x) + + +class SEANetDecoder(StreamingContainer): + """SEANet decoder. + + Args: + channels (int): Audio channels. + dimension (int): Intermediate representation dimension. + n_filters (int): Base width for the model. + n_residual_layers (int): nb of residual layers. + ratios (Sequence[int]): kernel size and stride ratios. + activation (str): Activation function. + activation_params (dict): Parameters to provide to the activation function. + final_activation (str): Final activation function after all convolutions. + final_activation_params (dict): Parameters to provide to the activation function. + norm (str): Normalization method. + norm_params (dict): Parameters to provide to the underlying normalization used along with the convolution. + kernel_size (int): Kernel size for the initial convolution. + last_kernel_size (int): Kernel size for the initial convolution. + residual_kernel_size (int): Kernel size for the residual layers. + dilation_base (int): How much to increase the dilation with each layer. + causal (bool): Whether to use fully causal convolution. + pad_mode (str): Padding mode for the convolutions. + true_skip (bool): Whether to use true skip connection or a simple. + (streamable) convolution as the skip connection in the residual network blocks. + compress (int): Reduced dimensionality in residual branches (from Demucs v3). + disable_norm_outer_blocks (int): Number of blocks for which we don't apply norm. + For the decoder, it corresponds to the N last blocks. + trim_right_ratio (float): Ratio for trimming at the right of the transposed convolution under the causal setup. + If equal to 1.0, it means that all the trimming is done at the right. + """ + + def __init__( + self, + channels: int = 1, + dimension: int = 128, + n_filters: int = 32, + n_residual_layers: int = 3, + ratios: tp.List[int] = [8, 5, 4, 2], + activation: str = "ELU", + activation_params: dict = {"alpha": 1.0}, + final_activation: tp.Optional[str] = None, + final_activation_params: tp.Optional[dict] = None, + norm: str = "none", + norm_params: tp.Dict[str, tp.Any] = {}, + kernel_size: int = 7, + last_kernel_size: int = 7, + residual_kernel_size: int = 3, + dilation_base: int = 2, + causal: bool = False, + pad_mode: str = "reflect", + true_skip: bool = True, + compress: int = 2, + disable_norm_outer_blocks: int = 0, + trim_right_ratio: float = 1.0, + ): + super().__init__() + self.dimension = dimension + self.channels = channels + self.n_filters = n_filters + self.ratios = ratios + del ratios + self.n_residual_layers = n_residual_layers + self.hop_length = int(np.prod(self.ratios)) + self.n_blocks = len(self.ratios) + 2 # first and last conv + residual blocks + self.disable_norm_outer_blocks = disable_norm_outer_blocks + assert ( + self.disable_norm_outer_blocks >= 0 and self.disable_norm_outer_blocks <= self.n_blocks + ), ( + "Number of blocks for which to disable norm is invalid." + "It should be lower or equal to the actual number of blocks in the network and greater or equal to 0." + ) + + act = getattr(nn, activation) + mult = int(2 ** len(self.ratios)) + model: tp.List[nn.Module] = [ + StreamingConv1d( + dimension, + mult * n_filters, + kernel_size, + norm=( + "none" if self.disable_norm_outer_blocks == self.n_blocks else norm + ), + norm_kwargs=norm_params, + causal=causal, + pad_mode=pad_mode, + ) + ] + + # Upsample to raw audio scale + for i, ratio in enumerate(self.ratios): + block_norm = ( + "none" + if self.disable_norm_outer_blocks >= self.n_blocks - (i + 1) + else norm + ) + # Add upsampling layers + model += [ + act(**activation_params), + StreamingConvTranspose1d( + mult * n_filters, + mult * n_filters // 2, + kernel_size=ratio * 2, + stride=ratio, + norm=block_norm, + norm_kwargs=norm_params, + causal=causal, + trim_right_ratio=trim_right_ratio, + ), + ] + # Add residual layers + for j in range(n_residual_layers): + model += [ + SEANetResnetBlock( + mult * n_filters // 2, + kernel_sizes=[residual_kernel_size, 1], + dilations=[dilation_base**j, 1], + activation=activation, + activation_params=activation_params, + norm=block_norm, + norm_params=norm_params, + causal=causal, + pad_mode=pad_mode, + compress=compress, + true_skip=true_skip, + ) + ] + + mult //= 2 + + # Add final layers + model += [ + act(**activation_params), + StreamingConv1d( + n_filters, + channels, + last_kernel_size, + norm="none" if self.disable_norm_outer_blocks >= 1 else norm, + norm_kwargs=norm_params, + causal=causal, + pad_mode=pad_mode, + ), + ] + # Add optional final activation to decoder (eg. tanh) + if final_activation is not None: + final_act = getattr(nn, final_activation) + final_activation_params = final_activation_params or {} + model += [final_act(**final_activation_params)] + self.model = nn.Sequential(*model) + + def forward(self, z): + y = self.model(z) + return y diff --git a/moshi_src/moshi/modules/seanet_test.py b/moshi_src/moshi/modules/seanet_test.py new file mode 100644 index 0000000..ce8c321 --- /dev/null +++ b/moshi_src/moshi/modules/seanet_test.py @@ -0,0 +1,187 @@ +import functools +import torch +import torch.nn as nn +import pytest + +from .seanet import SEANetResnetBlock, SEANetDecoder + + +torch.backends.cudnn.enabled = False # Disable cuDNN for deterministic behavior and for numerical stability + + +SEANET_RESNET_DATA = [ + # batch_size, dim, res_layer_index, seq_len, kernel_size + pytest.param( + 3, 4, 1, 10, 6, + id='small resnet test 1', + ), + pytest.param( + 4, 5, 2, 10, 7, + id='small resnet test 2', + ), + pytest.param( + 5, 6, 4, 10, 2, + id='small resnet test 3', + ), + pytest.param( + 1, 512, 2, 256, 7, + id='large resnet test 1', + ), +] +NUM_TIMESTEPS_DATA = [ + pytest.param( + 1, + id='length 1', + ), + pytest.param( + 2, + id='length 2', + ), + pytest.param( + 10, + id='length 10', + ), + pytest.param( + 100, + id='length 100', + ), +] + +SEANET_KWARGS_DATA = [ + pytest.param( + { + "channels": 1, + "dimension": 8, + "causal": True, + "n_filters": 2, + "n_residual_layers": 1, + "activation": "ELU", + "compress": 2, + "dilation_base": 2, + "disable_norm_outer_blocks": 0, + "kernel_size": 7, + "residual_kernel_size": 3, + "last_kernel_size": 3, + # We train using weight_norm but then the weights are pre-processed for inference so + # that we can use a normal convolution. + "norm": "none", + "pad_mode": "constant", + "ratios": [5], + "true_skip": True, + }, + id='Tiny SEANet', + ), + + pytest.param( + { + "channels": 1, + "dimension": 512, + "causal": True, + "n_filters": 64, + "n_residual_layers": 1, + "activation": "ELU", + "compress": 2, + "dilation_base": 2, + "disable_norm_outer_blocks": 0, + "kernel_size": 7, + "residual_kernel_size": 3, + "last_kernel_size": 3, + # We train using weight_norm but then the weights are pre-processed for inference so + # that we can use a normal convolution. + "norm": "none", + "pad_mode": "constant", + "ratios": [8, 6, 5, 4], + "true_skip": True, + }, + id='Large SEANet', + ), +] + + +def _init_weights(module, generator=None): + for name, param in module.named_parameters(): + if "weight" in name: + nn.init.xavier_uniform_(param, generator=generator) + elif "bias" in name: + nn.init.constant_(param, 0.0) + else: + nn.init.xavier_uniform_(param, generator=generator) + + +@pytest.mark.parametrize("batch_size, dim, res_layer_index, seq_len, kernel_size", SEANET_RESNET_DATA) +def test_resnet(batch_size, dim, res_layer_index, seq_len, kernel_size): + """Test that SEANetResnetBlock() calls are causal. Having new inputs does not change the previous output.""" + assert seq_len > kernel_size + + dilation_base = 2 + layer = SEANetResnetBlock(dim=dim, dilations=[dilation_base**res_layer_index, 1], pad_mode="constant", causal=True) + + generator = torch.Generator() + generator = generator.manual_seed(41) + layer.apply(functools.partial(_init_weights, generator=generator)) + + shape = (batch_size, dim, seq_len,) + input_hidden_states = torch.rand(shape) + + expected_output = layer(input_hidden_states) + + for end_index in range(kernel_size, seq_len + 1): + actual_output = layer(input_hidden_states[..., :end_index]) + torch.testing.assert_close(actual_output, expected_output[..., :actual_output.shape[-1]], + msg=lambda original_msg: f"Failed at end_index={end_index}: \n{original_msg}") + + +@pytest.mark.parametrize("batch_size, dim, res_layer_index, seq_len, kernel_size", SEANET_RESNET_DATA) +def test_resnet_streaming(batch_size, dim, res_layer_index, seq_len, kernel_size): + """Test that SEANetResnetBlock() streaming works as expected.""" + assert seq_len > kernel_size + + dilation_base = 2 + layer = SEANetResnetBlock(dim=dim, dilations=[dilation_base**res_layer_index, 1], pad_mode="constant", causal=True) + + generator = torch.Generator() + generator = generator.manual_seed(41) + layer.apply(functools.partial(_init_weights, generator=generator)) + + shape = (batch_size, dim, seq_len,) + input_hidden_states = torch.rand(shape) + + expected_output = layer(input_hidden_states) + + start_index = 0 + actual_outputs = [] + with layer.streaming(batch_size=batch_size): + for end_index in range(kernel_size, seq_len + 1): + actual_output = layer(input_hidden_states[..., start_index:end_index]) + start_index = end_index + actual_outputs.append(actual_output) + actual_outputs = torch.cat(actual_outputs, dim=-1) + + torch.testing.assert_close(actual_outputs, expected_output) + + +@pytest.mark.parametrize("num_timesteps", NUM_TIMESTEPS_DATA) +@pytest.mark.parametrize("seanet_kwargs", SEANET_KWARGS_DATA) +def test_nonstreaming_causal_decode(num_timesteps, seanet_kwargs): + """Test that the SEANetDecoder does not depend on future inputs.""" + + device = 'cuda' if torch.cuda.is_available() else 'cpu' + decoder = SEANetDecoder(**seanet_kwargs).to(device=device) + + generator = torch.Generator(device=device) + generator = generator.manual_seed(41) + decoder.apply(functools.partial(_init_weights, generator=generator)) + + rand_generator = torch.Generator(device=device) + rand_generator.manual_seed(2147483647) + with torch.no_grad(): + # [B, K = 8, T] + codes = torch.randn(1, seanet_kwargs['dimension'], num_timesteps, generator=rand_generator, device=device) + expected_decoded = decoder(codes) + + num_timesteps = codes.shape[-1] + for t in range(num_timesteps): + current_codes = codes[..., :t + 1] + actual_decoded = decoder(current_codes) + torch.testing.assert_close(expected_decoded[..., :actual_decoded.shape[-1]], actual_decoded, + msg=lambda original_msg: f"Failed at t={t}: \n{original_msg}") diff --git a/moshi_src/moshi/modules/streaming.py b/moshi_src/moshi/modules/streaming.py new file mode 100644 index 0000000..4193dbe --- /dev/null +++ b/moshi_src/moshi/modules/streaming.py @@ -0,0 +1,217 @@ +# Copyright (c) Kyutai, all rights reserved. +# This source code is licensed under the license found in the +# LICENSE file in the root directory of this source tree. + +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the license found in the +# LICENSE file in the root directory of this source tree. + +""" +Streaming module API that should be implemented by all Streaming components, +""" + +import abc +from contextlib import ExitStack +from dataclasses import dataclass +import typing as tp +from torch import nn +import torch +from ..utils.compile import CUDAGraphed + + +@dataclass +class State(abc.ABC): + """Base State for streaming, requires to be resetable and also support the context + protocol. The state __enter__ and __exit__ will be called upon entering and exiting + streaming in the parent module, but not upon `reset()` calls. + """ + + batch_size: int + device: torch.device + + def __post_init__(self): + self.exec_mask = torch.ones(self.batch_size, dtype=torch.bool, device=self.device) + self._set_exec_mask_graphed: CUDAGraphed | None = None + + def set_exec_mask(self, exec_mask: torch.Tensor): + self.exec_mask[:] = exec_mask + + def reset(self, reset_mask: torch.Tensor) -> None: + self.exec_mask[:] = torch.where(reset_mask, torch.ones_like(self.exec_mask), self.exec_mask) + + def __enter__(self) -> None: + pass + + def __exit__(self, exc_type, exc_value, traceback) -> None: + pass + + +StateT = tp.TypeVar("StateT", bound=State) + + +class StreamingModule(abc.ABC, nn.Module, tp.Generic[StateT]): + """Common API for streaming components. + + Each streaming component has a streaming state, `self._streaming_state`, which is None by default. + + To set a streaming component in streaming state, use + + with module.streaming(batch_size): + ... + + This will automatically void the streaming state when exiting the context manager. + This also automatically propagates to all streaming children module. + When the streaming state is set, modules should store whatever state they need in there. + """ + def __init__(self) -> None: + super().__init__() + self._streaming_state: StateT | None = None + self._streaming_detached: bool = False + self._cached_children: list[tuple[str, StreamingModule]] | None = None + + @property + def is_streaming(self): + return self._streaming_state is not None + + def set_streaming_detached(self, streaming_detached: bool): + """If set to False, the default, this module and all submodules will switch to streaming mode + if a parent module is set to streaming mode. + If set to True, or in detach mode, only a direct call to this module `.streaming(...)` method + will set it into streaming mode, ignoring the changes from its parents. + + This is useful if streaming over two different dimensions, e.g. for the RQ-Transformer + with the inner Depth Transformer working on the dimension of the codebooks.""" + self._streaming_detached = streaming_detached + + def _apply_named_streaming(self, fn: tp.Any): + def _handle_module(prefix: str, module: nn.Module): + if isinstance(module, StreamingModule): + # If prefix is empty, we are the direct receiver of the streaming request, + # otherwise, we are inheriting from a parent and will stop if detached. + if module._streaming_detached and prefix != "": + return + assert self._cached_children is not None + self._cached_children.append((prefix, module)) + for name, child in module.named_children(): + if prefix: + new_prefix = prefix + "." + name + else: + new_prefix = name + _handle_module(new_prefix, child) + + if self._cached_children is None: + self._cached_children = [] + _handle_module("", self) + for name, child in self._cached_children: + fn(name, child) + + def _start_streaming(self, batch_size: int, exit_stack: ExitStack): + def _start_streaming(name: str, module: StreamingModule): + assert module._streaming_state is None, f"{name} is already streaming!" + state = module._init_streaming_state(batch_size) + exit_stack.enter_context(state) + module._streaming_state = state + + self._apply_named_streaming(_start_streaming) + + def _stop_streaming(self) -> None: + def _stop_streaming(name: str, module: StreamingModule): + module._streaming_state = None + + self._apply_named_streaming(_stop_streaming) + + @abc.abstractmethod + def _init_streaming_state(self, batch_size: int) -> StateT: ... + + def streaming_forever(self, batch_size: int): + self.streaming(batch_size).__enter__() + + def streaming(self, batch_size: int) -> ExitStack: + """Context manager to enter streaming mode. Reset streaming state on exit.""" + + exit_stack = ExitStack() + self._start_streaming(batch_size, exit_stack) + exit_stack.callback(self._stop_streaming) + return exit_stack + + def reset_streaming(self, reset_mask: torch.Tensor | None = None) -> None: + """Reset the streaming state.""" + + def _reset(name: str, module: StreamingModule): + state = module._streaming_state + if state is None: + raise ValueError( + f"Trying to reset streaming, but {name} wasn't streaming." + ) + state.reset(reset_mask) + + state = self._streaming_state + assert state is not None + if reset_mask is None: + reset_mask = torch.ones(state.batch_size, device=state.device, dtype=torch.bool) + else: + reset_mask = reset_mask.to(state.device) + self._apply_named_streaming(_reset) + + def get_streaming_state(self) -> dict[str, tp.Any]: + """Return the complete streaming state, including that of sub-modules.""" + state: dict[str, tp.Any] = {} + + def _add(name: str, module: StreamingModule): + state[name] = module._streaming_state + + self._apply_named_streaming(_add) + return state + + def set_streaming_state(self, state: dict[str, tp.Any]): + """Set the streaming state, including that of sub-modules.""" + state = dict(state) + + def _set(name: str, module: StreamingModule): + if name in state: + module._streaming_state = state[name] + state.pop(name) + else: + raise RuntimeError(f"Expected to find a streaming state for {name}.") + + self._apply_named_streaming(_set) + if state: + raise RuntimeError(f"Some states were not consumed: {list(state.keys())}") + + def set_exec_mask(self, exec_mask: torch.Tensor): + """Set the execution mask, a tensor of boolean of shape `(B,), indicating + for each batch item whether the internal state should be updated or not as if + real data had been received. + + This is useful for running desynchronized streams with batching, e.g. when + the mask is False for an entry, the internal state will be unchanged by the provided + data, e.g. will be on the next step as if the previous one had never happened. + + There is no magic here, each StreamingModule subclass is responsible for respecting + the exec_mask. + """ + + state = self._streaming_state + assert state is not None + exec_mask = exec_mask.to(state.device) + + def _set_exec_mask(exec_mask: torch.Tensor): + def _set_exec_mask_fn(name: str, module: StreamingModule): + state = module._streaming_state + assert state is not None + state.set_exec_mask(exec_mask) + self._apply_named_streaming(_set_exec_mask_fn) + + if state._set_exec_mask_graphed is None: + disable = state.device.type != 'cuda' + state._set_exec_mask_graphed = CUDAGraphed(_set_exec_mask, disable=disable) + + state._set_exec_mask_graphed(exec_mask) + + +class StreamingContainer(StreamingModule[State]): + def _init_streaming_state(self, batch_size: int) -> State: + device = next(iter(self.parameters())).device + return State(batch_size, device) diff --git a/moshi_src/moshi/modules/transformer.py b/moshi_src/moshi/modules/transformer.py new file mode 100644 index 0000000..5f70ccc --- /dev/null +++ b/moshi_src/moshi/modules/transformer.py @@ -0,0 +1,956 @@ +# Copyright (c) Kyutai, all rights reserved. +# This source code is licensed under the license found in the +# LICENSE file in the root directory of this source tree. + +""" +Transformer model, with streaming support, + CUDA Graphable. +Optimized for inference. + +See `StreamingTransformer` for more information. +""" + +from contextlib import ExitStack +from dataclasses import dataclass +import typing as tp +from einops import rearrange +import torch +import torch.nn as nn +from torch.nn import functional as F +from ..utils.compile import no_compile +from ..utils import quantize +from ..utils.quantize import replace_linear_with_qlinear +from .gating import make_gating +from .rope import RotaryEmbedding +from .streaming import StreamingModule, StreamingContainer, State +from .lora import LoRALinear +from torch.utils.checkpoint import checkpoint as torch_checkpoint + + +class LayerNormF32(nn.LayerNorm): + def forward(self, input: torch.Tensor) -> torch.Tensor: + x_f32 = input.float() + out_f32 = super().forward(x_f32) + return out_f32.to(input.dtype) + + +def _rms_norm( + x: torch.Tensor, + alpha: torch.Tensor, + dtype: tp.Optional[torch.dtype], + eps: float, +): + assert x.dim() == 3, f"RMSNorm expects 3D inputs but got {x.shape}" + x_dtype = x.dtype + if dtype is not None: + x = x.to(dtype) + var = eps + torch.mean(x**2, dim=2, keepdim=True) + y = (x * (alpha.to(var) * torch.rsqrt(var))).to(x_dtype) + return y + + +class RMSNorm(nn.Module): + def __init__( + self, + dim: int, + eps: float = 1e-5, + dtype: tp.Optional[torch.dtype] = None, + device=None, + ): + super().__init__() + self.eps = eps + self.dtype = dtype + self.alpha = nn.Parameter( + torch.full((1, 1, dim), 1.0, requires_grad=True, device=device, dtype=dtype) + ) + + def forward(self, x: torch.Tensor): + return _rms_norm(x, self.alpha, self.dtype, self.eps) + + +class LayerScale(nn.Module): + """Layer scale from [Touvron et al 2021] (https://arxiv.org/pdf/2103.17239.pdf). + This rescales diagonally the residual outputs close to 0, with a learnt scale. + + Args: + channels (int): Number of channels. + init (float): Initial scale. + channel_last (bool): If True, expect `[*, C]` shaped tensors, otherwise, `[*, C, T]`. + device (torch.device or str, optional): Device on which to initialize the module. + dtype (torch.dtype, optional): dtype to use to initialize the module. + """ + + def __init__( + self, + channels: int, + init: float = 1e-4, + channel_last: bool = True, + device=None, + dtype=None, + ): + super().__init__() + self.channel_last = channel_last + self.scale = nn.Parameter( + torch.full( + (channels,), init, requires_grad=True, device=device, dtype=dtype + ) + ) + + def forward(self, x: torch.Tensor): + if self.channel_last: + return self.scale * x + else: + return self.scale[:, None] * x + + +def create_norm_fn(norm_type: str, dim: int, **kwargs) -> nn.Module: + """Create normalization module for transformer encoder layer. + + Args: + norm_type (str): Normalization method. + dim (int): Dimension of the normalized layer. + **kwargs (dict): Additional parameters for normalization layer. + Returns: + nn.Module: Normalization module. + """ + if norm_type == "layer_norm": + return nn.LayerNorm(dim, eps=1e-5, **kwargs) + elif norm_type == "layer_norm_f32": + kwargs.pop("dtype", None) + return LayerNormF32(dim, eps=1e-8, **kwargs) + elif norm_type in {"rms_norm"}: + return RMSNorm(dim, eps=1e-5, **kwargs) + elif norm_type in {"rms_norm_f32"}: + kwargs.pop("dtype", None) + return RMSNorm(dim, eps=1e-8, dtype=torch.float, **kwargs) + else: + raise ValueError(f"Unknown norm type: {norm_type}") + + +def create_sin_embedding( + positions: torch.Tensor, + dim: int, + max_period: float = 10000, + dtype: torch.dtype = torch.float32, +) -> torch.Tensor: + """Create sinusoidal positional embedding, with shape `[B, T, C]`. + + Args: + positions (torch.Tensor): LongTensor of positions. + dim (int): Dimension of the embedding. + max_period (float): Maximum period of the cosine/sine functions. + dtype (torch.dtype or str): dtype to use to generate the embedding. + Returns: + torch.Tensor: Sinusoidal positional embedding. + """ + # We aim for BTC format + assert dim % 2 == 0 + half_dim = dim // 2 + positions = positions.to(dtype) + adim = torch.arange(half_dim, device=positions.device, dtype=dtype).view(1, 1, -1) + max_period_tensor = torch.full( + [], max_period, device=positions.device, dtype=dtype + ) # avoid sync point + phase = positions / (max_period_tensor ** (adim / (half_dim - 1))) + return torch.cat([torch.cos(phase), torch.sin(phase)], dim=-1) + + +def set_attention_context(model: nn.Module, context: tp.Optional[int] = None) -> None: + """Deactivates or changes the context span (in time steps) in a model. + Args: + model (nn.Module): model over which to look for attentions. + context (int or None): new temporary context value. + + ..Note:: this is not a context manager but a plain function changing the context forever. + Initially, it was a context manager, but that led to interesting bugs when using + activation checkpointing, with the context being inconsistent between the forward + and backward. + """ + for module in model.modules(): + if isinstance(module, StreamingMultiheadAttention): + module.context = context + + +class KVCacheResult(tp.NamedTuple): + keys: torch.Tensor + values: torch.Tensor + positions: torch.Tensor + + @staticmethod + def from_kv(keys: torch.Tensor, values: torch.Tensor) -> "KVCacheResult": + B, H, T, D = keys.shape + assert tuple(values.shape[:-1]) == (B, H, T) + positions = torch.arange(T, device=keys.device, dtype=torch.long) + return KVCacheResult(keys, values, positions.expand(B, -1)) + + +class RingKVCache: + """Efficient streaming KVCache to be compatible with Cuda Graph. + + Args: + batch_size (int): Batch size. + num_heads (int): Number of heads in the attention. + dim_per_head (int): Dimension per head. + device (torch.device): Device on which to initialize the cache. + dtype (torch.dtype): dtype to use for the cache. + """ + + def __init__( + self, + batch_size: int, + num_heads: int, + dim_per_head: int, + capacity: int, + respect_exec_mask: bool = True, + device: torch.device = torch.device("cuda"), + dtype: torch.dtype = torch.bfloat16, + ): + self.capacity = capacity + self.cache = torch.zeros( + (2, batch_size, num_heads, capacity, dim_per_head), + device=device, + dtype=dtype, + ) + self.respect_exec_mask = respect_exec_mask + if self.respect_exec_mask: + self.end_offset = torch.zeros(batch_size, device=device, dtype=torch.long) + else: + self.end_offset = torch.zeros(1, device=device, dtype=torch.long) + + def reset(self, reset_mask: torch.Tensor) -> None: + self.end_offset[:] = torch.where( + reset_mask, + torch.zeros_like(self.end_offset), + self.end_offset, + ) + + def complete(self, k: torch.Tensor, v: torch.Tensor, exec_mask: torch.Tensor) -> KVCacheResult: + assert k.shape[:-1] == v.shape[:-1], (k.shape, v.shape) + B, H, T, D = k.shape + assert T > 0 + indexes = torch.arange(T, device=self.end_offset.device, dtype=self.end_offset.dtype) + indexes = indexes + self.end_offset.view(-1, 1) + indexes = indexes % self.capacity + if self.respect_exec_mask: + # indexes is [B, T] + # k is [B, H, T, D] + # cache is [B, H, T', D] + this_indexes = indexes.view(B, 1, T, 1) + this_indexes = this_indexes.expand(-1, H, T, D) + self.cache[0].scatter_(2, this_indexes, k) + self.cache[1].scatter_(2, this_indexes, v) + else: + self.cache[0].index_copy_(2, indexes[0], k) + self.cache[1].index_copy_(2, indexes[0], v) + + keys = self.cache[0] + values = self.cache[1] + + indexes = torch.arange( + self.capacity, device=self.end_offset.device, dtype=torch.long + ) + + # end_index correspond to the actual index where the last value was written. + last_offset = self.end_offset.view(-1, 1) + T - 1 + end_index = last_offset % self.capacity + delta = indexes - end_index + + # We know that if `index == end_index`, then we should output `self.end_offset`. + # If `index = end_index - 1` we should output `self.end_offset - 1` + # If `index = end_index - n` we should output `self.end_offset - n` + # Now, for `index == end_index + 1` , we actually have the oldest entry in the cache, + # so we should output `end_index + 1 - self.capacity` + + positions = torch.where( + delta <= 0, + last_offset + delta, + last_offset + delta - self.capacity, + ) + if self.respect_exec_mask: + self.end_offset[:] = torch.where( + exec_mask, + self.end_offset + T, + self.end_offset) + else: + self.end_offset.add_(T) + invalid = indexes >= self.end_offset.view(-1, 1) + positions = torch.where(invalid, torch.full_like(positions, -1), positions) + + return KVCacheResult(keys, values, positions) + + +def apply_weights_per_step(modules: nn.ModuleList, schedule: list[int] | None, + x: torch.Tensor, offset: int | None) -> torch.Tensor: + """Utility to apply a multi linear layer to the given input. A multi linear layer + applies a different set of weight for each time step. + + Args: + modules (nn.ModuleList): apply weights per step. + schedule (list[int] or None): schedule for weight sharing. + x (torch.Tensor): Input tensor, with shape `[B, T, C]`. + offset (int): offset for the current time step, in particular for decoding, with + time steps provided one by one. + """ + + if len(modules) == 1: + return modules[0](x) + + assert offset is not None, "Out of sync execution with weights per step." + + ys: list[torch.Tensor] = [] + B, T, C = x.shape + for t in range(T): + module_index = t + offset + if schedule is not None: + module_index = schedule[module_index] + y = modules[module_index](x[:, t: t + 1]) + ys.append(y) + out = torch.cat(ys, 1) + return out + + +@dataclass +class _MHAState(State): + kv_cache: RingKVCache | None + offset: torch.Tensor + offset_cpu: int + k_cross: torch.Tensor | None = None + v_cross: torch.Tensor | None = None + + def reset(self, reset_mask: torch.Tensor): + super().reset(reset_mask) + self.offset[:] = torch.where(reset_mask, torch.zeros_like(self.offset), self.offset) + if self.kv_cache is not None: + self.kv_cache.reset(reset_mask) + self.offset_cpu = 0 + + +class StreamingMultiheadAttention(StreamingModule[_MHAState]): + """Similar to `nn.MultiheadAttention` but with support for streaming, causal evaluation. + + Args: + embed_dim (int): Dimension to project to. + num_heads (int): Number of heads. + causal (bool): Causal mask applied automatically. + context (int, optional): Number of time steps the attention can access to. + When causal, can access `context` time steps into the past, and when non causal, + can access `context // 2` steps in the past, and the same in the future. + rope (`RotaryEmbedding`, optional): Rope embedding to use. + weights_per_step (int): use different weights per time step. If non zero, should correspond to the + number of possible time steps. + weights_per_step_schedule (list[int] | None): if provided, some steps will share weights when + `weights_per_step` is True, e.g. step `I` will use weights `schedule[I]`. + cross_attention (bool): True if this is to be used as a cross attention. + device (torch.device, optional): Device on which to initialize. + dtype (torch.dtype, optional): dtype to use. + """ + + _fsdp_final = True + + def __init__( + self, + embed_dim: int, + num_heads: int, + causal: bool = False, + context: tp.Optional[int] = None, + rope: tp.Optional[RotaryEmbedding] = None, + weights_per_step: int = 0, + weights_per_step_schedule: list[int] | None = None, + cross_attention: bool = False, + cache_cross_attention: bool = True, + device=None, + dtype=None, + ): + super().__init__() + factory_kwargs = {"device": device, "dtype": dtype} + + self.embed_dim = embed_dim + self.causal = causal + self.context = context + self.rope = rope + self.num_heads = num_heads + self.weights_per_step = weights_per_step + self.weights_per_step_schedule = weights_per_step_schedule + self.cross_attention = cross_attention + self.cache_cross_attention = cache_cross_attention + if cross_attention: + assert not weights_per_step, "weights_per_step not supported for cross attention." + assert rope is None, "rope and cross_attention makes no sense." + assert not causal, "causal and cross attention makes no sense." + # We do not want to activate the streaming KVCache if we are a cross attention. + # self.set_streaming_detached(True) + + out_dim = 3 * embed_dim + mult = 1 + if weights_per_step: + if weights_per_step_schedule: + assert len(weights_per_step_schedule) == weights_per_step + mult = max(weights_per_step_schedule) + 1 + else: + mult = weights_per_step + self.mult = mult + + # Split in one linear per step + self.out_projs = nn.ModuleList( + [ + nn.Linear(embed_dim, embed_dim, bias=False, **factory_kwargs) + for _ in range(mult) + ] + ) + self.in_projs = nn.ModuleList( + [ + nn.Linear(embed_dim, out_dim, bias=False, **factory_kwargs) + for _ in range(mult) + ] + ) + + self._register_load_state_dict_pre_hook(StreamingMultiheadAttention._load_hook, with_module=True) + + @staticmethod + def _load_hook(module, state_dict, prefix, *_): + mappings = { + 'in_proj_weight': 'in_projs.{i}.weight', + 'in_proj.weight': 'in_projs.{i}.weight', + 'in_proj.lora_A.weight': 'in_projs.{i}.lora_A.weight', + 'in_proj.lora_B.weight': 'in_projs.{i}.lora_B.weight', + 'out_proj.weight': 'out_projs.{i}.weight', + 'out_proj.lora_A.weight': 'out_projs.{i}.lora_A.weight', + 'out_proj.lora_B.weight': 'out_projs.{i}.lora_B.weight', + } + + mult = module.mult + # _scb suffix is for quantized data. + for suffix in ['', '_scb']: + for source, target in mappings.items(): + this_source = prefix + source + suffix + if this_source in state_dict: + weight = state_dict[this_source] + _, *OD = weight.shape + weight = weight.view(mult, -1, *OD) + for i in range(mult): + this_target = prefix + target.format(i=i) + suffix + state_dict[this_target] = weight[i] + state_dict.pop(this_source) + + def _init_streaming_state(self, batch_size: int) -> _MHAState: + in_proj = self.in_projs[0] + if isinstance(in_proj, LoRALinear): + device = in_proj.lora_A.weight.device + dtype = in_proj.lora_A.weight.dtype + elif isinstance(in_proj, nn.Linear): + device = in_proj.weight.device + dtype = in_proj.weight.dtype + elif isinstance(in_proj, quantize.QLinear): + device = in_proj.weight.device + dtype = torch.float16 + else: + raise RuntimeError(f"Unknown type {type(in_proj)} for linear.") + + dim_per_head = self.embed_dim // self.num_heads + if self.cross_attention: + kv_cache = None + else: + if self.context is None: + if self.weights_per_step: + capacity = self.weights_per_step + else: + raise RuntimeError( + "Cannot create a streaming KVCache without a context to estimate capacity." + ) + else: + capacity = self.context + + kv_cache = RingKVCache( + batch_size, self.num_heads, dim_per_head, capacity, + respect_exec_mask=not self.weights_per_step, device=device, dtype=dtype + ) + return _MHAState( + batch_size, + device, + kv_cache, + offset=torch.zeros(batch_size, device=device, dtype=torch.long), + offset_cpu=0, + ) + + def _complete_kv(self, k, v) -> KVCacheResult: + state = self._streaming_state + if state is None or state.kv_cache is None: + return KVCacheResult.from_kv(k, v) + else: + return state.kv_cache.complete(k, v, state.exec_mask) + + def _compute_cross_attention( + self, key: torch.Tensor, value: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: + assert self.cross_attention + assert key is value + in_proj = self.in_projs[0] + assert in_proj.bias is None + assert isinstance(in_proj, nn.Linear) + dim = in_proj.weight.shape[0] // 3 + kv = nn.functional.linear(key, in_proj.weight[dim:]) + k, v = rearrange(kv, "b t (p h d) -> p b h t d", p=2, h=self.num_heads) + return k, v + + def update_streaming_cross_attention_src( + self, cross_attention_src: torch.Tensor) -> None: + state = self._streaming_state + assert state is not None + assert self.cross_attention + k, v = self._compute_cross_attention(cross_attention_src, cross_attention_src) + if state.k_cross is None: + state.k_cross = k + state.v_cross = v + else: + assert state.v_cross is not None + state.k_cross[:] = k + state.v_cross[:] = v + + def _get_cross_attention( + self, key: torch.Tensor, value: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: + state = self._streaming_state + if state is not None and state.k_cross is not None: + assert state.v_cross is not None + return state.k_cross, state.v_cross + k, v = self._compute_cross_attention(key, value) + if state is not None and self.cache_cross_attention: + state.k_cross = k + state.v_cross = v + return k, v + + def forward(self, query: torch.Tensor, key: torch.Tensor, value: torch.Tensor): + state = self._streaming_state + B, T = query.shape[:2] + + if state is None: + offset = torch.zeros(B, device=query.device, dtype=torch.long) + offset_cpu = 0 + else: + offset = state.offset + offset_cpu = state.offset_cpu + + if self.cross_attention: + assert len(self.in_projs) == 1 + in_proj = self.in_projs[0] + assert in_proj.bias is None + assert isinstance(in_proj, nn.Linear) + dim = in_proj.weight.shape[0] // 3 + q = nn.functional.linear(query, in_proj.weight[:dim]) + q = rearrange(q, "b t (h d) -> b h t d", h=self.num_heads) + k, v = self._get_cross_attention(key, value) + else: + projected = apply_weights_per_step( + self.in_projs, self.weights_per_step_schedule, query, offset_cpu) + + q, k, v = rearrange( + projected, "b t (p h d) -> p b h t d", p=3, h=self.num_heads + ) + if self.rope: + q, k = self.rope(q, k, offset, time_before_heads=False) + + k, v, pos_k = self._complete_kv(k, v) + pos_k = pos_k[:, None] + if self.causal: + pos_q = offset.view(-1, 1, 1) + torch.arange(T, device=q.device, dtype=torch.long).view( + -1, 1) + delta = pos_q - pos_k + attn_bias = (pos_k >= 0) & (delta >= 0) + if self.context is not None: + attn_bias = attn_bias & (delta < self.context) + attn_bias = attn_bias[:, None] + else: + attn_bias = None + x = F.scaled_dot_product_attention(q, k, v, attn_bias, dropout_p=0.0) + + x = rearrange(x, "b h t d -> b t (h d)") + x = apply_weights_per_step( + self.out_projs, self.weights_per_step_schedule, x, offset_cpu) + + if state is not None and not self.cross_attention: + state.offset[:] = torch.where( + state.exec_mask, + state.offset + T, + state.offset) + state.offset_cpu += T + return x + + +@dataclass +class _LayerState(State): + offset_cpu: int = 0 + + def reset(self, reset_mask: torch.Tensor): + super().reset(reset_mask) + self.offset_cpu = 0 + + +class StreamingTransformerLayer(StreamingModule[_LayerState]): + """TransformerLayer with Streaming / Causal support. + + Args: + d_model (int): Dimension of the data. + num_heads (int): Number of heads. + dim_feedforward (int): Intermediate dimension of FF module. + causal (bool): Causal mask applied automatically. + context (int, optional): Receptive field for the causal mask, infinite if None. + rope (`RotaryEmbedding`, optional): Rope embedding to use. + norm (str): Normalization to use. Currently, only 'layer_norm' is supported. + layer_scale (float, optional): If not None, LayerScale will be used with the given value as initial scale. + gating (str): if provided, replaces FFN with special gating, like GLU, GSiGLU etc. + weights_per_step (int): use different weights per time step. If non zero, should correspond to the + number of possible time steps. + weights_per_step_schedule (list[int] | None): if provided, some steps will share weights when + `weights_per_step` is True, e.g. step `I` will use weights `schedule[I]`. + skip_self_attn: If true, skips the self attention module and the norm + cross_attention (bool): If True, expect to get secondary input for cross-attention. + device (torch.device, optional): Device on which to initialize. + dtype (torch.dtype, optional): dtype to use. + """ + + _fsdp_final = True + + def __init__( + self, + d_model: int, + num_heads: int, + dim_feedforward: int | list[int] = 2048, + causal: bool = False, + context: tp.Optional[int] = None, + rope: tp.Optional[RotaryEmbedding] = None, + norm: str = "layer_norm", + layer_scale: tp.Optional[float] = None, + gating: str = "none", + weights_per_step: int = 0, + weights_per_step_schedule: list[int] | None = None, + activation=F.gelu, + skip_self_attn: bool = False, + cross_attention: bool = False, + device=None, + dtype=None, + ): + super().__init__() + factory_kwargs = {"device": device, "dtype": dtype} + # Redefine self_attn to our streaming multi-head attention + attn_kwargs: tp.Dict[str, tp.Any] = { + "embed_dim": d_model, + "num_heads": num_heads, + } + if not skip_self_attn: + self.self_attn: StreamingMultiheadAttention = StreamingMultiheadAttention( + causal=causal, + context=context, + rope=rope, + weights_per_step=weights_per_step, + weights_per_step_schedule=weights_per_step_schedule, + **attn_kwargs, # type: ignore + **factory_kwargs, # type: ignore + ) # type: ignore + self.norm1 = create_norm_fn(norm, d_model, **factory_kwargs) + self.norm2 = create_norm_fn(norm, d_model, **factory_kwargs) + # Redefine feedforward layers to expose bias parameter + self.weights_per_step = weights_per_step + self.weights_per_step_schedule = weights_per_step_schedule + self.gating: tp.Optional[nn.Module] = None + self.linear1: tp.Optional[nn.Module] = None + self.linear2: tp.Optional[nn.Module] = None + self.activation = activation + self.skip_self_attn = skip_self_attn + + num_weights = 1 + if weights_per_step is not None: + num_weights = weights_per_step + if weights_per_step_schedule is not None: + assert len(weights_per_step_schedule) == weights_per_step + num_weights = max(weights_per_step_schedule) + 1 + if isinstance(dim_feedforward, list): + assert dim_feedforward + assert len(dim_feedforward) == num_weights, ( + "Length of dim_feedforward must match weights_per_step," + f" got {len(dim_feedforward)} != {num_weights}" + ) + if gating == "none": + assert ( + not weights_per_step + ), "weights_per_step without gating not supported for now." + assert not isinstance( + dim_feedforward, list + ), "List dim_feedforward without gating not supported for now." + self.linear1 = nn.Linear( + d_model, dim_feedforward, bias=False, **factory_kwargs + ) + self.linear2 = nn.Linear( + dim_feedforward, d_model, bias=False, **factory_kwargs + ) + else: + self.linear1 = None + self.linear2 = None + if weights_per_step: + if isinstance(dim_feedforward, int): + dim_feedforward = [dim_feedforward] * num_weights + assert isinstance(dim_feedforward, list), dim_feedforward + self.gating = nn.ModuleList( + [ + make_gating(gating, d_model, dim, **factory_kwargs) + for dim in dim_feedforward + ] + ) + else: + assert isinstance(dim_feedforward, int) + self.gating = make_gating( + gating, d_model, dim_feedforward, **factory_kwargs + ) + + self.cross_attention: StreamingMultiheadAttention | None = None + if cross_attention: + self.cross_attention = StreamingMultiheadAttention( + cross_attention=True, **attn_kwargs, **factory_kwargs) # type: ignore + # Cross attention norm is always a layer norm, for no specific reason. + self.norm_cross = nn.LayerNorm(d_model, eps=1e-5, **factory_kwargs) # type: ignore + + self.layer_scale_1: nn.Module + self.layer_scale_2: nn.Module + if layer_scale is None: + self.layer_scale_1 = nn.Identity() + self.layer_scale_2 = nn.Identity() + if cross_attention: + self.layer_scale_cross = nn.Identity() + else: + self.layer_scale_1 = LayerScale(d_model, layer_scale, **factory_kwargs) # type: ignore + self.layer_scale_2 = LayerScale(d_model, layer_scale, **factory_kwargs) # type: ignore + if cross_attention: + self.layer_scale_cross = LayerScale(d_model, layer_scale, **factory_kwargs) # type: ignore + + def _init_streaming_state(self, batch_size: int) -> _LayerState: + device = next(iter(self.parameters())).device + return _LayerState(batch_size, device, offset_cpu=0) + + # feed forward block + def _ff_block(self, x: torch.Tensor) -> torch.Tensor: + state = self._streaming_state + offset = 0 + if state is not None: + offset = state.offset_cpu + x_orig = x + x = self.norm2(x) + if self.gating is None: + assert self.linear1 is not None + assert self.linear2 is not None + update = self.linear2(self.activation(self.linear1(x))) + else: + if self.weights_per_step: + assert isinstance(self.gating, nn.ModuleList) + update = apply_weights_per_step(self.gating, self.weights_per_step_schedule, x, offset) + else: + update = self.gating(x) + return x_orig.to(update) + self.layer_scale_2(update) + + def _sa_block(self, x: torch.Tensor): + if self.skip_self_attn: + return x + x_orig = x + x = self.norm1(x) + update = self.self_attn(x, x, x) + return x_orig.to(update) + self.layer_scale_1(update) + + def _cross_attention_block(self, x: torch.Tensor, + cross_attention_src: torch.Tensor) -> torch.Tensor: + assert self.cross_attention is not None + x_orig = x + x = self.norm_cross(x) + # queries are from src, keys and values from cross_attention_src. + update = self.cross_attention(x, cross_attention_src, cross_attention_src) + return x_orig + self.layer_scale_cross(update) + + def forward(self, x: torch.Tensor, cross_attention_src: torch.Tensor | None = None): + with ExitStack() as stack: + if x.device.type != 'cuda': + stack.enter_context(no_compile()) + x = self._sa_block(x) + if self.cross_attention is not None: + assert cross_attention_src is not None + x = self._cross_attention_block(x, cross_attention_src) + else: + assert cross_attention_src is None + x = self._ff_block(x) + state = self._streaming_state + if state: + state.offset_cpu += x.shape[1] + return x + + +@dataclass +class _TransformerState(State): + offsets: torch.Tensor + + def reset(self, reset_mask: torch.Tensor): + super().reset(reset_mask) + self.offsets[:] = torch.where(reset_mask, torch.zeros_like(self.offsets), self.offsets) + + +class StreamingTransformer(StreamingModule[_TransformerState]): + """Transformer with Streaming / Causal support. + + Args: + d_model (int): Dimension of the data. + num_heads (int): Number of heads. + dim_feedforward (int): Intermediate dimension of FF module. + causal (bool): Causal mask applied automatically. + context (int, optional): Receptive field for the causal mask, infinite if None. + layer_scale (float, optional): If not None, LayerScale will be used + with the given value as initial scale. + positional_embedding (str): Positional embedding strategy (sin, rope, sin_rope, or none). + max_period (float): Maximum period of the time embedding. + positional_scale (float): Scale of positional embedding, set to 0 to deactivate. + layer_class: (subclass of `StreamingTransformerLayer): class to use + to initialize the layers, allowing further customization outside of AudioCraft. + device (torch.device, optional): Device on which to initialize. + dtype (torch.dtype, optional): dtype to use. + **kwargs: See `StreamingTransformerLayer`. + """ + + def __init__( + self, + d_model: int, + num_heads: int, + num_layers: int, + dim_feedforward: int | list[int] = 2048, + causal: bool = False, + context: tp.Optional[int] = None, + positional_embedding: str = "sin", + max_period: float = 10_000, + positional_scale: float = 1.0, + betas: tp.Optional[tp.Tuple[float, float]] = None, + layer_class: tp.Type[StreamingTransformerLayer] = StreamingTransformerLayer, + quantize: bool = False, + checkpointing: bool = False, + device=None, + dtype=None, + **kwargs, + ): + super().__init__() + assert d_model % num_heads == 0 + + self.positional_embedding = positional_embedding + self.max_period = max_period + self.positional_scale = positional_scale + self.betas = betas + + assert positional_embedding in {"sin", "rope", "sin_rope", "none"} + self.rope: tp.Optional[RotaryEmbedding] = None + if self.positional_embedding in {"rope", "sin_rope"}: + self.rope = RotaryEmbedding(max_period=max_period) + + self.checkpointing = checkpointing + + self.layers = nn.ModuleList() + for _ in range(num_layers): + self.layers.append( + layer_class( + d_model=d_model, + num_heads=num_heads, + dim_feedforward=dim_feedforward, + causal=causal, + context=context, + rope=self.rope, + device=device, + dtype=dtype, + **kwargs, + ) + ) + if quantize: + # Quantizing layers one by one to avoid taking too much space during init. + self.layers[-1].to(device=device, dtype=dtype) + replace_linear_with_qlinear(self.layers[-1]) + + def _init_streaming_state(self, batch_size: int) -> _TransformerState: + device = next(self.parameters()).device + return _TransformerState(batch_size, device, offsets=torch.zeros(batch_size, device=device, dtype=torch.long)) + + def forward(self, x: torch.Tensor, *args, **kwargs): + B, T, C = x.shape + + dtype_input = x.dtype + state = self._streaming_state + if state is None: + offsets = torch.zeros(1, dtype=torch.long, device=x.device) + else: + offsets = state.offsets + + if self.positional_embedding in {"sin", "sin_rope"}: + positions = torch.arange(T, device=x.device).view(1, -1, 1) + positions = positions + offsets.view(-1, 1, 1) + pos_emb = create_sin_embedding( + positions, C, max_period=self.max_period, dtype=x.dtype + ) + x = x + self.positional_scale * pos_emb + + for layer in self.layers: + if self.checkpointing: + y = torch_checkpoint( + layer, x, *args, use_reentrant=False, + determinism_check='none', + preserve_rng_state=False, + **kwargs) + assert isinstance(y, torch.Tensor) + x = y + else: + x = layer(x, *args, **kwargs) + + if state is not None: + state.offsets[:] = torch.where( + state.exec_mask, + state.offsets + T, + state.offsets) + return x.to(dtype_input) + + +class ProjectedTransformer(StreamingContainer): + """Transformer with optional projections of the input and output to different dimensions when needed. + Supports multiple outputs. + + Args: + input_dimension (int): dimension of the input. + output_dimensions (tuple[int]): dimensions of the outputs. + d_model (int): inner dimension of the Transformer. + conv_layout (bool): If True, expects `[B, C, T]` shaped tensors, otherwise, `[B, T, C]`. + Similarly, the output will have the same layout. + """ + + def __init__( + self, + input_dimension: int, + output_dimensions: tp.Tuple[int, ...], + d_model: int, + *, + conv_layout: bool = False, + **kwargs, + ): + super().__init__() + self.transformer = StreamingTransformer(d_model=d_model, **kwargs) + self.input_dimension = input_dimension + self.output_dimensions = output_dimensions + self.conv_layout = conv_layout + self.input_proj = None + if d_model != input_dimension: + self.input_proj = nn.Linear(input_dimension, d_model, bias=False) + + self.output_projs = nn.ModuleList() + for output_dimension in output_dimensions: + if d_model == output_dimension: + self.output_projs.append(nn.Identity()) + else: + self.output_projs.append( + nn.Linear(d_model, output_dimension, bias=False) + ) + + def forward(self, x, *args, **kwargs): + if self.conv_layout: + x = x.transpose(1, 2) + if self.input_proj is not None: + x = self.input_proj(x) + z = self.transformer(x, *args, **kwargs) + ys = [] + for output_proj in self.output_projs: + y = output_proj(z) + if self.conv_layout: + y = y.transpose(1, 2) + ys.append(y) + return ys diff --git a/moshi_src/moshi/quantization/__init__.py b/moshi_src/moshi/quantization/__init__.py new file mode 100644 index 0000000..0e5e1b3 --- /dev/null +++ b/moshi_src/moshi/quantization/__init__.py @@ -0,0 +1,13 @@ +# Copyright (c) Kyutai, all rights reserved. +# This source code is licensed under the license found in the +# LICENSE file in the root directory of this source tree. + +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the license found in the +# LICENSE file in the root directory of this source tree. +"""RVQ.""" +# flake8: noqa +from .vq import ResidualVectorQuantizer, SplitResidualVectorQuantizer +from .base import BaseQuantizer, DummyQuantizer, QuantizedResult diff --git a/moshi_src/moshi/quantization/base.py b/moshi_src/moshi/quantization/base.py new file mode 100644 index 0000000..242d5f9 --- /dev/null +++ b/moshi_src/moshi/quantization/base.py @@ -0,0 +1,170 @@ +# Copyright (c) Kyutai, all rights reserved. +# This source code is licensed under the license found in the +# LICENSE file in the root directory of this source tree. + +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the license found in the +# LICENSE file in the root directory of this source tree. + +""" +Base class for all quantizers. +""" + +from dataclasses import dataclass, field +import typing as tp + +import torch +from torch import nn + + +@dataclass +class QuantizedResult: + x: torch.Tensor + codes: torch.Tensor + bandwidth: torch.Tensor # bandwidth in kb/s used, per batch item. + penalty: tp.Optional[torch.Tensor] = None + metrics: dict = field(default_factory=dict) + + +class BaseQuantizer(nn.Module): + """Base class for quantizers.""" + + def __init__(self): + super().__init__() + self._ema_frozen = False + + def forward(self, x: torch.Tensor, frame_rate: int) -> QuantizedResult: + """ + Given input tensor x, returns first the quantized (or approximately quantized) + representation along with quantized codes, bandwidth, and any penalty term for the loss. + Finally, this returns a dict of metrics to update logging etc. + Frame rate must be passed so that the bandwidth is properly computed. + """ + raise NotImplementedError() + + def encode(self, x: torch.Tensor) -> torch.Tensor: + """Encode a given input tensor with the specified sample rate at the given bandwidth.""" + raise NotImplementedError() + + def decode(self, codes: torch.Tensor) -> torch.Tensor: + """Decode the given codes to the quantized representation.""" + raise NotImplementedError() + + @property + def cardinality(self) -> int: + """Cardinality of each codebook.""" + raise NotImplementedError() + + @property + def total_codebooks(self) -> int: + """Total number of codebooks.""" + raise NotImplementedError() + + @property + def num_codebooks(self) -> int: + """Number of active codebooks.""" + raise NotImplementedError() + + @property + def semantic_quantizer(self) -> 'BaseQuantizer': + """This returns the quantizer that models the first level of the hierarchy (typically semantic). + + In this case, it's the quantizer itself. + """ + return self + + @property + def acoustic_quantizer(self) -> 'BaseQuantizer': + """This returns the quantizer that models the higher levels of the hierarchy (typically acoustic). + + In this case, it's the quantizer itself. + """ + return self + + def set_num_codebooks(self, n: int) -> None: + """Set the number of active codebooks.""" + raise NotImplementedError() + + @property + def ema_frozen(self) -> bool: + """Whether to apply ema to the codebooks.""" + return self._ema_frozen + + def ema_frozen_(self, ema_frozen: bool) -> None: + """Set whether ema should be applied to the codebooks.""" + self._ema_frozen = ema_frozen + + +class DummyQuantizer(BaseQuantizer): + """Fake quantizer that actually does not perform any quantization.""" + + def __init__( + self, + dimension: int, + input_dimension: tp.Optional[int] = None, + output_dimension: tp.Optional[int] = None, + ): + super().__init__() + self.dimension = dimension + self.input_dimension = input_dimension or dimension + self.output_dimension = output_dimension or dimension + self.input_proj: torch.nn.Module + self.output_proj: torch.nn.Module + if self.input_dimension == self.dimension: + self.input_proj = torch.nn.Identity() + else: + self.input_proj = torch.nn.Conv1d( + self.input_dimension, self.dimension, 1, bias=False + ) + if self.input_dimension == self.dimension: + self.output_proj = torch.nn.Identity() + else: + self.output_proj = torch.nn.Conv1d( + self.dimension, self.output_dimension, 1, bias=False + ) + + def forward(self, x: torch.Tensor, frame_rate: int): + q = x.unsqueeze(1) + x = self.output_proj(self.input_proj(x)) + return QuantizedResult( + x, q, torch.tensor(q.numel() * 32 * frame_rate / 1000 / len(x)).to(x) + ) + + def encode(self, x: torch.Tensor) -> torch.Tensor: + """Encode a given input tensor with the specified sample rate at the given bandwidth. + In the case of the DummyQuantizer, the codes are actually identical + to the input and resulting quantized representation as no quantization is done. + """ + x = self.input_proj(x) + return x.unsqueeze(1) + + def decode(self, codes: torch.Tensor) -> torch.Tensor: + """Decode the given codes to the quantized representation. + In the case of the DummyQuantizer, the codes are actually identical + to the input and resulting quantized representation as no quantization is done. + """ + y = codes.squeeze(1) + return self.output_proj(y) + + @property + def total_codebooks(self): + """Total number of codebooks.""" + return 1 + + @property + def num_codebooks(self): + """Total number of codebooks.""" + return self.total_codebooks + + def set_num_codebooks(self, n: int): + """Set the number of active codebooks.""" + raise AttributeError( + "Cannot override the number of codebooks for the dummy quantizer" + ) + + @property + def cardinality(self) -> int: + """Cardinality of each codebook.""" + return 1 diff --git a/moshi_src/moshi/quantization/core_vq.py b/moshi_src/moshi/quantization/core_vq.py new file mode 100644 index 0000000..09f9bbd --- /dev/null +++ b/moshi_src/moshi/quantization/core_vq.py @@ -0,0 +1,528 @@ +# Copyright (c) Kyutai, all rights reserved. +# This source code is licensed under the license found in the +# LICENSE file in the root directory of this source tree. + +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the license found in the +# LICENSE file in the root directory of this source tree. + +import math +import typing as tp + +from einops import rearrange, repeat +import torch +from torch import nn +from torch import distributed +import torch.nn.functional as F + + +class _CodebookForwardResult(tp.NamedTuple): + quantized: torch.Tensor + codes: torch.Tensor + metrics: tp.Dict[str, torch.Tensor] + + +class _VQForwardResult(tp.NamedTuple): + quantized: torch.Tensor + codes: torch.Tensor + loss: torch.Tensor + metrics: tp.Dict[str, torch.Tensor] + + +def _ema_inplace(moving_avg: torch.Tensor, new: torch.Tensor, decay: float) -> None: + moving_avg.data.mul_(decay).add_(new, alpha=(1 - decay)) + + +def _sample_vectors(samples: torch.Tensor, num: int) -> torch.Tensor: + num_samples, device = samples.shape[0], samples.device + + if num_samples >= num: + indices = torch.randperm(num_samples, device=device)[:num] + else: + indices = torch.randint(0, num_samples, (num,), device=device) + + return samples[indices] + + +def _compute_entropy(usage: torch.Tensor) -> torch.Tensor: + # Usage is some unnormalized distribution. + proba = usage / usage.sum() + p_log_p = torch.where( + proba == 0, zero_scalar(usage.device), proba * torch.log(proba) + ) + return -p_log_p.sum() + + +def _is_distributed() -> bool: + # Checks if we need to use distributed routines. + return distributed.is_initialized() and distributed.get_world_size() > 1 + + +def _average_tensors(tensors: tp.Sequence[torch.Tensor]) -> None: + if not _is_distributed(): + return + world_size = distributed.get_world_size() + handles = [] + for tensor in tensors: + handle = distributed.all_reduce( + tensor.data, op=distributed.ReduceOp.SUM, async_op=True) + handles.append(handle) + for tensor, handle in zip(tensors, handles): + handle.wait() + tensor.data /= world_size + + +def _run_kmeans(samples: torch.Tensor, num_clusters: int, num_iters: int = 50) -> tp.Tuple[torch.Tensor, torch.Tensor]: + # Kmeans algorithm used to initialize the codebooks. + dim = samples.shape[-1] + means = _sample_vectors(samples, num_clusters) + bins = None + + for _ in range(num_iters): + dists = torch.cdist(samples[None], means[None], p=2)[0] + buckets = dists.argmin(dim=-1) + bins = torch.bincount(buckets, minlength=num_clusters) + zero_mask = bins == 0 + bins.clamp_(min=1) + + new_means = torch.zeros_like(means) + new_means.scatter_add_(0, repeat(buckets, "n -> n d", d=dim), samples) + new_means /= bins[..., None] + resampled = _sample_vectors(samples, num_clusters) + means = torch.where(zero_mask[..., None], resampled, new_means) + + assert bins is not None + return means, bins + + +def zero_scalar(device) -> torch.Tensor: + """Returns a 0. value on the given device without introducing a synchronization point.""" + return torch.zeros([1], device=device)[0] + + +class EuclideanCodebook(nn.Module): + """Codebook with Euclidean distance. + + Args: + dim (int): Dimension. + codebook_size (int): Codebook size. + decay (float): Decay for exponential moving average over the codebooks. + epsilon (float): Epsilon value for numerical stability. + threshold_usage_ratio (float): Defines the threshold for the cluster usage under which a centroid + is replaced. This is expressed as a fraction of the usage a centroid would get under + a uniform distribution, so that it doesn't depend on the batch size etc. + replaced_usage_ratio (float): When replacing a centroid, use this as an initial centroid usage, + to avoid the centroid getting replaced too quickly. + check_unused_every (int): Check for unused centroids every `check_unused_every` iterations. + This is to avoid too many synchronization points. + + Buffers: + cluster_usage (torch.Tensor): EMA of the cluster usage per batch, e.g. this will + be dependent on the batch size etc. + embedding_sum (torch.Tensor): EMA of the sum of the assigned points to each cluster. + In particular, this can be normalized by `cluster_usage` to obtain the + actual cluster centroids. + """ + + def __init__( + self, + dim: int, + codebook_size: int, + decay: float = 0.99, + epsilon: float = 1e-5, + threshold_usage_ratio: float = 0.1, + replaced_usage_ratio: float = 1.0, + check_unused_every: int = 5, + ): + super().__init__() + self.decay = decay + + self.dim = dim + self.codebook_size = codebook_size + + self.epsilon = epsilon + self.threshold_usage_ratio = threshold_usage_ratio + self.replaced_usage_ratio = replaced_usage_ratio + self.check_unused_every = check_unused_every + self._next_unused_check = check_unused_every + self._cached_initialized = False + + self._initialized: torch.Tensor + self.cluster_usage: torch.Tensor + self.embedding_sum: torch.Tensor + self._embedding: torch.Tensor + self.register_buffer("_initialized", torch.tensor([False], dtype=torch.float)) + self.register_buffer("cluster_usage", torch.ones(codebook_size)) + embedding = torch.zeros(codebook_size, dim) + self.register_buffer("embedding_sum", embedding) + self.register_buffer("_embedding", None, persistent=False) + + def _load_from_state_dict(self, state_dict, prefix, *args, **kwargs) -> None: + # Mapping old names to new names + mappings = { + "inited": "_initialized", + "cluster_size": "cluster_usage", + "embed_avg": "embedding_sum", + "embed_sum": "embedding_sum", + } + for old_name, new_name in mappings.items(): + old_name = prefix + old_name + if old_name in state_dict: + value = state_dict.pop(old_name) + if new_name is not None: + state_dict[prefix + new_name] = value + super()._load_from_state_dict(state_dict, prefix, *args, **kwargs) + + @property + def embedding(self) -> torch.Tensor: + if self._embedding is None: + embedding = ( + self.embedding_sum / self.cluster_usage.clamp(min=self.epsilon)[:, None] + ) + self.register_buffer("_embedding", embedding, persistent=False) + return embedding + return self._embedding + + @property + def initialized(self) -> bool: + """Cached version of self._initialized, + This assumes that once the module is initialized, it will never go back to the uninitialized state.""" + if not self._cached_initialized: + self._cached_initialized = bool(self._initialized.item()) + return self._cached_initialized + + def _init_embedding(self, data: torch.Tensor) -> None: + # Initialize the codebook, e.g. using kmeans. + if self.initialized: + return + + rank = 0 + if _is_distributed(): + rank = distributed.get_rank() + # First gathering shapes in case not all GPUs have the same effective batch size. + # then gathering the actual content. + if rank == 0: + other_shapes: tp.List[torch.Size] = [None] * distributed.get_world_size() # type: ignore + distributed.gather_object(data.shape, other_shapes) + other_data: tp.List[torch.Tensor] = [ + torch.empty(shape, device=data.device, dtype=data.dtype) for shape in other_shapes] + distributed.gather(data, other_data) + data = torch.cat(other_data, dim=0) + else: + distributed.gather_object(data.shape) + distributed.gather(data) + if rank == 0: + embedding, cluster_usage = _run_kmeans(data, self.codebook_size) + self.embedding_sum.data.copy_(embedding * cluster_usage[:, None]) + self.cluster_usage.data.copy_(cluster_usage) + self._initialized.data.fill_(1) + # Make sure all buffers across workers are in sync after initialization + self._broadcast_buffers() + + def _broadcast_buffers(self) -> None: + if _is_distributed(): + for buffer in self.buffers(): + distributed.broadcast(buffer, 0) + + def _replace_expired_codes(self, samples: torch.Tensor, mask: torch.Tensor) -> None: + # Replaces expired centroids, as indicated by `mask` (a true value indicate the code needs to be replaced). + # The new codes are sampled from the batch `samples`. + new_vectors = _sample_vectors(samples, self.codebook_size) + replace_cluster_usage = ( + self.replaced_usage_ratio * self.cluster_usage.sum() / self.codebook_size + ) + self.embedding_sum[:] = torch.where( + mask[:, None], replace_cluster_usage * new_vectors, self.embedding_sum + ) + self.cluster_usage[:] = torch.where( + mask, replace_cluster_usage, self.cluster_usage + ) + + def _check_expired_codes(self, batch_samples: torch.Tensor) -> torch.Tensor: + # Checks whether some centroids are under utilized, and replace them if necessary. + if not self.initialized: + return zero_scalar(batch_samples.device) + + self._next_unused_check -= 1 + if self._next_unused_check > 0: + return zero_scalar(batch_samples.device) + # we don't check every iteration to avoid having too many sync points. + self._next_unused_check = self.check_unused_every + threshold_cluster_usage = self.threshold_usage_ratio * self.cluster_usage.sum() / self.codebook_size + expired_codes = self.cluster_usage < threshold_cluster_usage + + assert batch_samples.dim() == 2 + self._replace_expired_codes(batch_samples, mask=expired_codes) + self._broadcast_buffers() + + return expired_codes.float().mean() + + def _reshape_input(self, x: torch.Tensor) -> torch.Tensor: + # Flattens all the dimensions but the last one, e.g. return a vector of shape `[N, D]`. + x = rearrange(x, "... d -> (...) d") + return x + + def _reshape_codes(self, codes: torch.Tensor, shape: torch.Size) -> torch.Tensor: + return codes.view(*shape[:-1]) + + def _quantize(self, x: torch.Tensor) -> torch.Tensor: + # Projects each vector in `x` over the nearest centroid and return its index. + # `x` should be `[N, D]` with `N` the number of input vectors and `D` the dimension. + assert x.dim() == 2 + dists = torch.cdist(x[None], self.embedding[None], p=2)[0] + codes = dists.argmin(dim=-1) + return codes + + def encode(self, x: torch.Tensor) -> torch.Tensor: + """Given a tensor `x` of shape `[*, D]`, returns a tensor of integer codes of shape `[*]`. + The codes are defined as the indexes of the centroids nearest to each vector in `x`. + """ + assert x.dtype.is_floating_point, f"Input should be floats, got {x.dtype}" + shape = x.shape + x = self._reshape_input(x) + codes = self._quantize(x) + codes = self._reshape_codes(codes, shape) + return codes + + def decode(self, codes: torch.Tensor) -> torch.Tensor: + """Given a tensor of codes of shape `[*]`, returns a tensor of shape `[*, D]`, + corresponding to the centroids associated to each code index. + """ + assert ( + not codes.dtype.is_floating_point + ), f"Codes should be integers, got {codes.dtype}" + quantized = F.embedding(codes, self.embedding) + return quantized + + def forward( + self, x: torch.Tensor, initialize: bool = True + ) -> _CodebookForwardResult: + shape = x.shape + x = self._reshape_input(x) + + if self.training and initialize: + # If initialize is False, we are not allowed to initialize this layer + # and the rest of the code will operate on a 0 filled codebook. + # This is due to previous layers having used the batch to run kmeans init + # and thus, the residuals are mostly 0s. + self._init_embedding(x.detach()) + + flat_codes = self._quantize(x) + codes = self._reshape_codes(flat_codes, shape) + quantized = self.decode(codes) + metrics: tp.Dict[str, torch.Tensor] = {} + + if self.training: + # We do the expiry of the unused codes at this point as buffers are in sync + # and all the workers will take the same decision. + expired = self._check_expired_codes(x) + metrics['rvq_expired'] = expired + cluster_usage = torch.zeros_like(self.cluster_usage) + cluster_usage.scatter_add_( + 0, flat_codes, torch.ones_like(flat_codes, dtype=cluster_usage.dtype)) + _ema_inplace(self.cluster_usage, cluster_usage, self.decay) + + if self.initialized: + # We report the entropy normalized by that of the uniform distribution, + # This means the codebooks are optimally used when entropy=1. + metrics['rvq_entropy'] = _compute_entropy(self.cluster_usage) / math.log(self.codebook_size) + + embedding_sum = torch.zeros_like(self.embedding_sum) + embedding_sum.scatter_add_(0, repeat(flat_codes, "n -> n d", d=self.dim), x) + _ema_inplace(self.embedding_sum, embedding_sum, self.decay) + self.register_buffer('_embedding', None) + + return _CodebookForwardResult(quantized, codes, metrics) + + +class VectorQuantization(nn.Module): + """Vector quantization implementation. + Currently supports only euclidean distance. + + Args: + dim (int): Dimension + codebook_size (int): Codebook size + codebook_dim (int): Codebook dimension. If not defined, uses the specified dimension in dim. + decay (float): Decay for exponential moving average over the codebooks. + epsilon (float): Epsilon value for numerical stability. + threshold_usage_ratio (float): Defines the threshold for the cluster usage under which a centroid + is replaced. This is expressed as a fraction of the usage a centroid would get under + a uniform distribution, so that it doesn't depend on the batch size etc. + replaced_usage_ratio (float): When replacing a centroid, use this as an initial centroid usage, + to avoid the centroid getting replaced too quickly. + check_unused_every (int): Check for unused centroids every `check_unused_every` iterations. + This is to avoid too many synchronization points. + """ + + def __init__( + self, + dim: int, + codebook_size: int, + codebook_dim: tp.Optional[int] = None, + decay: float = 0.99, + epsilon: float = 1e-5, + threshold_usage_ratio: float = 0.1, + **kwargs, + ): + super().__init__() + if codebook_dim is None: + codebook_dim = dim + + requires_projection = codebook_dim != dim + self.project_in = ( + nn.Linear(dim, codebook_dim) if requires_projection else nn.Identity() + ) + self.project_out = ( + nn.Linear(codebook_dim, dim) if requires_projection else nn.Identity() + ) + self.epsilon = epsilon + self._codebook = EuclideanCodebook( + dim=codebook_dim, + codebook_size=codebook_size, + decay=decay, + epsilon=epsilon, + threshold_usage_ratio=threshold_usage_ratio, + **kwargs, + ) + self.codebook_size = codebook_size + + @property + def embedding(self): + return self._codebook.embedding + + @property + def initialized(self): + return self._codebook.initialized + + def _rearrange_input(self, x): + x = rearrange(x, "b d n -> b n d") + return x + + def _rearrange_output(self, quantized): + quantized = rearrange(quantized, "b n d -> b d n") + return quantized + + def encode(self, x: torch.Tensor) -> torch.Tensor: + """Encodes `x` into discrete integer codes.""" + x = self._rearrange_input(x) + x = self.project_in(x) + codes = self._codebook.encode(x) + return codes + + def decode(self, codes: torch.Tensor) -> torch.Tensor: + """Converts integer codes into quantized vectors.""" + quantized = self._codebook.decode(codes) + quantized = self.project_out(quantized) + quantized = self._rearrange_output(quantized) + return quantized + + def forward(self, x: torch.Tensor, initialize: bool = True) -> _VQForwardResult: + x = self._rearrange_input(x) + quantized, codes, metrics = self._codebook(x, initialize=initialize) + + if self.training: + quantized = x + (quantized - x).detach() + loss = F.mse_loss(x, quantized.detach()) + else: + loss = zero_scalar(x.device) + + quantized = self.project_out(quantized) + quantized = self._rearrange_output(quantized) + + return _VQForwardResult(quantized, codes, loss, metrics) + + +class ResidualVectorQuantization(nn.Module): + """Residual vector quantization implementation. + + Follows Algorithm 1. in https://arxiv.org/pdf/2107.03312.pdf + """ + + def __init__(self, *, num_quantizers: int, codebook_offset: int, **kwargs): + super().__init__() + self.layers = nn.ModuleList( + [VectorQuantization(**kwargs) for _ in range(num_quantizers)] + ) + self.codebook_offset = codebook_offset + + def forward( + self, x: torch.Tensor, n_q: tp.Optional[int] = None + ) -> _VQForwardResult: + """ + Args: + x (torch.Tensor): input tensor to quantize, of shape `[B, C, T]`. + n_q (int or None): if provided, number of codebook levels to use in RVQ. + """ + + quantized_out = zero_scalar(x.device) + residual = x + + all_losses = [] + all_codes = [] + all_metrics: tp.Dict[str, torch.Tensor] = {} + + n_q = n_q or len(self.layers) + previous_layer_is_initialized = True + + for i, layer in enumerate(self.layers[:n_q]): # type: ignore + if self.training: + this_layer_is_initialized = layer.initialized + # We only allow the kmeans initialization if the previous layer is already initialized from the previous + # iterations, this is to avoid learning the subsequent kmeans on the same batch, which would eventually + # lead to its exhaustion and running kmeans on 0 values. + quantized, codes, loss, metrics = layer( + residual, initialize=previous_layer_is_initialized + ) + if self.training: + previous_layer_is_initialized = this_layer_is_initialized # type: ignore + + quantized = quantized.detach() + residual = residual - quantized + quantized_out = quantized_out + quantized + + all_codes.append(codes) + all_losses.append(loss) + + for key, value in metrics.items(): + if key in all_metrics: + all_metrics[key] += value / n_q + else: + all_metrics[key] = value / n_q + all_metrics[key + f"_{i + self.codebook_offset}"] = value + + if self.training: + # Solving subtle bug with STE and RVQ: https://github.com/facebookresearch/encodec/issues/25 + quantized_out = x + (quantized_out - x).detach() + to_average = [] + for layer in self.layers: + assert isinstance(layer, VectorQuantization) + to_average += [layer._codebook.cluster_usage, layer._codebook.embedding_sum] + _average_tensors(to_average) + + out_losses, out_codes = map(torch.stack, (all_losses, all_codes)) + return _VQForwardResult(quantized_out, out_codes, out_losses, all_metrics) + + def encode(self, x: torch.Tensor, n_q: tp.Optional[int] = None) -> torch.Tensor: + """Encodes `x` into discrete integer codes. If `n_q` is provided, only uses the first `n_q` codebook levels.""" + residual = x + all_indices = [] + n_q = n_q or len(self.layers) + for layer in self.layers[:n_q]: # type: ignore + assert isinstance(layer, VectorQuantization) + indices = layer.encode(residual) + quantized = layer.decode(indices) + residual = residual - quantized + all_indices.append(indices) + out_indices = torch.stack(all_indices) + return out_indices + + def decode(self, codes: torch.Tensor) -> torch.Tensor: + """Converts the integer codes into quantized vectors.""" + quantized = zero_scalar(codes.device) + for idx, layer_codes in enumerate(codes): + layer = self.layers[idx] + assert isinstance(layer, VectorQuantization) + quantized = quantized + layer.decode(layer_codes) + return quantized diff --git a/moshi_src/moshi/quantization/vq.py b/moshi_src/moshi/quantization/vq.py new file mode 100644 index 0000000..bb39727 --- /dev/null +++ b/moshi_src/moshi/quantization/vq.py @@ -0,0 +1,318 @@ +# Copyright (c) Kyutai, all rights reserved. +# This source code is licensed under the license found in the +# LICENSE file in the root directory of this source tree. + +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the license found in the +# LICENSE file in the root directory of this source tree. + +import math +import random +import typing as tp + +import torch + +from .base import BaseQuantizer, QuantizedResult +from .core_vq import ResidualVectorQuantization + + +class ResidualVectorQuantizer(BaseQuantizer): + """Residual Vector Quantizer. + + Args: + dimension (int): Dimension of the codebooks. + input_dimension (None or int): dimension of the input, defaults to `dimension` if not provided. + output_dimension (None or int): dimension of the output, defaults to `dimension` if not provided. + n_q (int): Number of vector quantizers used. + q_dropout (bool): Random quantizer drop out at train time. + no_quantization_rate (float): Gives the probability of applying no quantization at all + at train time. The RVQ codebooks will still get the input value to learn the proper codebook. + bins (int): Codebook size. + decay (float): Decay for exponential moving average over the codebooks. + threshold_usage_ratio (float): Defines the threshold for the cluster usage under which a centroid + is replaced. This is expressed as a fraction of the usage a centroid would get under + a uniform distribution, so that it doesn't depend on the batch size etc. + replaced_usage_ratio (float): When replacing a centroid, use this as an initial centroid usage, + to avoid the centroid getting replaced too quickly. + codebook_offset (int): Offset to use for the codebook indices. This is useful when using multiple quantizers + such as in SplitResidualVectorQuantizer. + force_projection (bool): Whether to force input and output projections even when dimension is constant. + generator_seed (int or None): seed used to initialize the RNG used for no quantization. + """ + + def __init__( + self, + dimension: int = 128, + input_dimension: tp.Optional[int] = None, + output_dimension: tp.Optional[int] = None, + n_q: int = 8, + q_dropout: bool = False, + no_quantization_rate: float = 0.0, + bins: int = 1024, + decay: float = 0.99, + threshold_usage_ratio: float = 0.1, + replaced_usage_ratio: float = 1.0, + codebook_offset: int = 0, + force_projection: bool = False, + ): + super().__init__() + self.max_n_q = n_q + self.n_q = n_q + self.q_dropout = q_dropout + self.no_quantization_rate = no_quantization_rate + self.dimension = dimension + self.input_dimension = input_dimension or dimension + self.output_dimension = output_dimension or dimension + self.bins = bins + self.decay = decay + self.rng_dropout = random.Random(1234) + self.input_proj: torch.nn.Module + self.output_proj: torch.nn.Module + if self.input_dimension == self.dimension and not force_projection: + self.input_proj = torch.nn.Identity() + else: + self.input_proj = torch.nn.Conv1d( + self.input_dimension, self.dimension, 1, bias=False + ) + if self.output_dimension == self.dimension and not force_projection: + self.output_proj = torch.nn.Identity() + else: + self.output_proj = torch.nn.Conv1d( + self.dimension, self.output_dimension, 1, bias=False + ) + self.vq = ResidualVectorQuantization( + dim=self.dimension, + codebook_size=self.bins, + num_quantizers=self.n_q, + decay=self.decay, + threshold_usage_ratio=threshold_usage_ratio, + replaced_usage_ratio=replaced_usage_ratio, + codebook_offset=codebook_offset, + ) + + def forward(self, x: torch.Tensor, frame_rate: int): + """ + Args: + x (torch.Tensor): Input tensor of shape [B, C, T] with `C` number of channels. + frame_rate (int): frame rate of the input (e.g `T = frame_rate * duration`), used to compute + the bandwidth. + + Returns: + QuantizedResult: Quantized result with the following attributes: + - `x` (torch.Tensor): Quantized tensor of shape [B, C, T]. + - `codes` (torch.Tensor): Quantized codes of shape [B, K, T] with `K` number of codebooks. + - `bw` (torch.Tensor): Bandwidth of the quantized tensor in kbits per second. + - `penalty` (torch.Tensor): Commitment loss. + - `metrics` (dict): RVQ metrics, in particular rate of dead code replacement, and entropy. + """ + n_q = self.n_q + x = self.input_proj(x) + if self.training and self.q_dropout: + n_q = self.rng_dropout.randint(1, self.n_q) + bw_per_q = math.log2(self.bins) * frame_rate / 1000 + quantized, codes, commit_loss, metrics = self.vq(x, n_q=n_q) + B, _, _ = quantized.shape + if self.training and self.no_quantization_rate > 0: + mask = (torch.rand(B, 1, 1, device=x.device) <= self.no_quantization_rate).float() + quantized = x * mask + (1 - mask) * quantized + quantized = self.output_proj(quantized) + codes = codes.transpose(0, 1) + # codes is [B, K, T], with T frames, K nb of codebooks. + bw = torch.tensor(n_q * bw_per_q).to(x) + return QuantizedResult(quantized, codes, bw, penalty=torch.mean(commit_loss), metrics=metrics) + + def encode(self, x: torch.Tensor) -> torch.Tensor: + """Encode a given input tensor with the specified frame rate at the given bandwidth. + The RVQ encode method sets the appropriate number of quantizer to use + and returns indices for each quantizer. + """ + n_q = self.n_q + if x.shape[-1] == 0: + return torch.empty((x.shape[0], n_q, 0), device=x.device, dtype=torch.int64) + + x = self.input_proj(x) + codes = self.vq.encode(x, n_q=n_q) + codes = codes.transpose(0, 1) + # codes is [B, K, T], with T frames, K nb of codebooks. + return codes + + def decode(self, codes: torch.Tensor) -> torch.Tensor: + """Decode the given codes to the quantized representation.""" + # codes is [B, K, T], with T frames, K nb of codebooks, vq.decode expects [K, B, T]. + codes = codes.transpose(0, 1) + quantized = self.vq.decode(codes) + quantized = self.output_proj(quantized) + return quantized + + @property + def total_codebooks(self): + return self.max_n_q + + @property + def num_codebooks(self): + return self.n_q + + def set_num_codebooks(self, n: int): + assert n >= 0 and n <= self.max_n_q + self.n_q = n + + @property + def cardinality(self) -> int: + return self.bins + + +class SplitResidualVectorQuantizer(BaseQuantizer): + """Residual Vector Quantizer with separate projections for the first quantizer and the rest. + + Args: + n_q (int): Number of residual vector quantizers used. + n_semantic_q (int): Number of residual vector quantizers used for the semantic quantizer. + **kwargs: Arguments to the constructor of `ResidualVectorQuantizer` that are shared between both. + """ + + def __init__( + self, + *, + n_q: int = 8, + n_q_semantic: int = 1, + **kwargs, + ): + super().__init__() + assert n_q > n_q_semantic, ( + f"Number of quantizers {n_q} must be larger " + f"than the number of semantic quantizers {n_q_semantic}." + ) + self.max_n_q = n_q + self.n_q_semantic = n_q_semantic + self.n_q_acoustic = n_q - n_q_semantic + q_dropout = kwargs.pop("q_dropout", False) + self.rvq_first = ResidualVectorQuantizer( + n_q=n_q_semantic, force_projection=True, q_dropout=False, **kwargs + ) + self.rvq_rest = ResidualVectorQuantizer( + n_q=n_q - n_q_semantic, + codebook_offset=1, + force_projection=True, + q_dropout=q_dropout, + **kwargs, + ) + + def _renorm_and_add( + self, + first_val: torch.Tensor, + rest_val: torch.Tensor, + n_q_semantic: int, + n_q_acoustic: int, + ): + """Renormalizes values from `rvq_first` and `rvq_rest` and adds them. + + This allows correcting statistics that are normalized by the number of quantizers. To renormalize, we use the + number of quantizers that are actually used, e.g. taking into account quantizer dropout. + """ + n_q = n_q_semantic + n_q_acoustic + renorm_first_val = first_val * n_q_semantic / n_q + renorm_rest_val = rest_val * n_q_acoustic / n_q + return renorm_first_val + renorm_rest_val + + def forward(self, x: torch.Tensor, frame_rate: int): + """ + Args: + x (torch.Tensor): Input tensor of shape [B, C, T] with `C` number of channels. + frame_rate (int): frame rate of the input (e.g `T = frame_rate * duration`), used to compute + the bandwidth. + + Returns: + QuantizedResult: Quantized result with the following attributes: + - `x` (torch.Tensor): Quantized tensor of shape [B, C, T]. + - `codes` (torch.Tensor): Quantized codes of shape [B, K, T] with `K` number of codebooks. + - `bw` (torch.Tensor): Bandwidth of the quantized tensor in kbits per second. + - `penalty` (torch.Tensor): Commitment loss. + - `metrics` (dict): RVQ metrics, in particular rate of dead code replacement, and entropy. + """ + semantic_result = self.rvq_first(x, frame_rate) + if self.n_q == self.n_q_semantic: + return semantic_result + acoustic_result = self.rvq_rest(x, frame_rate) + full_quantized_emb = semantic_result.x + acoustic_result.x + full_quantized_codes = torch.cat( + [semantic_result.codes, acoustic_result.codes], dim=1 + ) + # This is the actual number of quantizers used, e.g. taking into account quantizer dropout. + n_q_semantic = semantic_result.codes.shape[1] + n_q_acoustic = acoustic_result.codes.shape[1] + full_quantized_bandwidth = semantic_result.bandwidth + acoustic_result.bandwidth + full_quantized_penalty = self._renorm_and_add( + semantic_result.penalty, acoustic_result.penalty, n_q_semantic, n_q_acoustic + ) + full_quantized_metrics = semantic_result.metrics + for key, value in acoustic_result.metrics.items(): + if key in full_quantized_metrics: + full_quantized_metrics[key] = self._renorm_and_add( + full_quantized_metrics[key], value, n_q_semantic, n_q_acoustic + ) + else: + full_quantized_metrics[key] = value + return QuantizedResult( + full_quantized_emb, + full_quantized_codes, + full_quantized_bandwidth, + penalty=full_quantized_penalty, + metrics=full_quantized_metrics, + ) + + def encode(self, x: torch.Tensor) -> torch.Tensor: + """Encode a given input tensor with the specified frame rate at the given bandwidth. + The RVQ encode method sets the appropriate number of quantizer to use + and returns indices for each quantizer. + """ + codes = self.rvq_first.encode(x) + if self.n_q > self.n_q_semantic: + acoustic_codes = self.rvq_rest.encode(x) + codes = torch.cat([codes, acoustic_codes], dim=1) + # codes is [B, K, T], with T frames, K nb of codebooks. + return codes + + def decode(self, codes: torch.Tensor) -> torch.Tensor: + """Decode the given codes to the quantized representation.""" + # codes is [B, K, T], with T frames, K nb of codebooks. + quantized = self.rvq_first.decode(codes[:, : self.n_q_semantic]) + if codes.shape[1] > self.n_q_semantic: + quantized += self.rvq_rest.decode(codes[:, self.n_q_semantic :]) + return quantized + + @property + def total_codebooks(self): + return self.rvq_first.max_n_q + self.rvq_rest.max_n_q + + @property + def num_codebooks(self): + return self.rvq_first.num_codebooks + self.rvq_rest.num_codebooks + + @property + def n_q(self): + return self.rvq_first.n_q + self.rvq_rest.n_q + + @property + def dimension(self): + return self.rvq_first.dimension + + @property + def semantic_quantizer(self) -> ResidualVectorQuantizer: + """This returns the quantizer that models the first level of the hierarchy (typically semantic).""" + return self.rvq_first + + @property + def acoustic_quantizer(self) -> ResidualVectorQuantizer: + """This returns the quantizer that models the higher levels of the hierarchy (typically acoustic).""" + return self.rvq_rest + + def set_num_codebooks(self, n: int): + assert n >= self.n_q_semantic and n <= self.total_codebooks + self.rvq_rest.set_num_codebooks(n - self.n_q_semantic) + + @property + def cardinality(self) -> int: + assert self.rvq_rest.cardinality == self.rvq_first.cardinality + return self.rvq_first.cardinality diff --git a/moshi_src/moshi/utils/__init__.py b/moshi_src/moshi/utils/__init__.py new file mode 100644 index 0000000..74bbb52 --- /dev/null +++ b/moshi_src/moshi/utils/__init__.py @@ -0,0 +1,10 @@ +# Copyright (c) Kyutai, all rights reserved. +# This source code is licensed under the license found in the +# LICENSE file in the root directory of this source tree. + +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the license found in the +# LICENSE file in the root directory of this source tree. +"""Utilities.""" diff --git a/moshi_src/moshi/utils/autocast.py b/moshi_src/moshi/utils/autocast.py new file mode 100644 index 0000000..e90efe0 --- /dev/null +++ b/moshi_src/moshi/utils/autocast.py @@ -0,0 +1,45 @@ +# Copyright (c) Kyutai, all rights reserved. +# This source code is licensed under the license found in the +# LICENSE file in the root directory of this source tree. + +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the license found in the +# LICENSE file in the root directory of this source tree. + +import torch + + +class TorchAutocast: + """TorchAutocast utility class. + Allows you to enable and disable autocast. This is specially useful + when dealing with different architectures and clusters with different + levels of support. + + Args: + enabled (bool): Whether to enable torch.autocast or not. + args: Additional args for torch.autocast. + kwargs: Additional kwargs for torch.autocast + """ + + def __init__(self, enabled: bool, *args, **kwargs): + self.autocast = torch.autocast(*args, **kwargs) if enabled else None + + def __enter__(self): + if self.autocast is None: + return + try: + self.autocast.__enter__() + except RuntimeError: + device = self.autocast.device + dtype = self.autocast.fast_dtype + raise RuntimeError( + f"There was an error autocasting with dtype={dtype} device={device}\n" + "If you are on the FAIR Cluster, you might need to use autocast_dtype=float16" + ) + + def __exit__(self, *args, **kwargs): + if self.autocast is None: + return + self.autocast.__exit__(*args, **kwargs) diff --git a/moshi_src/moshi/utils/compile.py b/moshi_src/moshi/utils/compile.py new file mode 100644 index 0000000..fe9c385 --- /dev/null +++ b/moshi_src/moshi/utils/compile.py @@ -0,0 +1,287 @@ +# Copyright (c) Kyutai, all rights reserved. +# This source code is licensed under the license found in the +# LICENSE file in the root directory of this source tree. + +""" +Provides some extra utilities around torch compile, in particular with a way +to fully deactivate it easily with a context manager. +Provides a simple activation checkpointing that is compatible with FSDP and torch compile. +Finally, provides some utilities for CUDA graphing functions. +""" +from contextlib import contextmanager +from functools import wraps +import inspect +import os +import typing as tp + +import torch +from torch import cuda + + +_compile_disabled: bool = False + + +@contextmanager +def no_compile(): + """Disable torch.compile locally. Now Pytorch 2.4 provides a function to do that.""" + global _compile_disabled + + prev_disabled = _compile_disabled + _compile_disabled = True + try: + yield + finally: + _compile_disabled = prev_disabled + + +def torch_compile_lazy(fun): + """torch.compile creates a huge pool of processes, even when not using the function at all, + e.g. with Dora. This can polute stderr when doing CTRL+C. So we do it in a lazy way. + """ + if os.environ.get("NO_TORCH_COMPILE"): + return fun + fun_compiled = None + + @wraps(fun) + def _wrapped(*args, **kwargs): + nonlocal fun_compiled + if _compile_disabled: + return fun(*args, **kwargs) + if fun_compiled is None: + fun_compiled = torch.compile(fun) + return fun_compiled(*args, **kwargs) + + return _wrapped + + +class Checkpoint(torch.autograd.Function): + @staticmethod + def forward(ctx, function, *args) -> tp.Any: + to_save = [] + ctx.others = [] + ctx.function = function + # Sources will indicate whether the arg in position N is + # a tensor stored in ctx.save_for_backward, or inside ctx.others. + ctx.sources = [] + new_args = [] + for arg in args: + if isinstance(arg, torch.Tensor): + to_save.append(arg) + ctx.sources.append("tensor") + new_args.append(arg.detach()) + else: + ctx.sources.append("other") + ctx.others.append(arg) + new_args.append(arg) + ctx.save_for_backward(*to_save) + # During the forward, we just make a pass with no gradient computed. + with torch.no_grad(): + res = function(*new_args) + return res + + @staticmethod + def backward(ctx, *grads) -> tp.Tuple[tp.Optional[torch.Tensor], ...]: + pseudo_tensors = [] + with torch.set_grad_enabled(True): + # We create leaf tensors to collect the output gradients. + # We call them pseudo_tensors because they are pretending to be the input + # to `function` but are not directly + for tensor in ctx.saved_tensors: + pseudo_tensor = tensor.detach() + pseudo_tensor.requires_grad_(True) + pseudo_tensors.append(pseudo_tensor) + pseudo_tensors_copy = list(pseudo_tensors) + args = [] + for source in ctx.sources: + if source == "other": + args.append(ctx.others.pop(0)) + else: + assert source == "tensor" + args.append(pseudo_tensors_copy.pop(0)) + res = ctx.function(*args) + # The second forward with grad computation allows us to connect the input leaf tensors + # inside pseudo_tensors, to the outputs of the function called. + if not isinstance(res, tuple): + res = (res,) + # Now we just ask Torch to compute the derivative of `res` given the gradient coming from above + # `grads`. The computed gradient will end up into the `pseudo_tensors` grad attributes. + torch.autograd.backward(res, grads) + out: tp.List[tp.Optional[torch.Tensor]] = [None] + for source in ctx.sources: + # We still need to output `None` values for non tensor parameters. + if source == "other": + out.append(None) + else: + assert source == "tensor" + out.append(pseudo_tensors.pop(0).grad) + return tuple(out) + + +def simple_checkpoint(module: torch.nn.Module, *args, **kwargs): + """Custom implementation of checkpointing in PyTorch as the builtin implementation is broken + when using torch compile. Only supports wrapping a `nn.Module` with a forward with no `*args` or `**kwargs`. + + https://github.com/pytorch/pytorch/issues/97436. + Should be resolved in nightlies, but it is quite fun and simple to code it ourselves. + """ + if hasattr(module, "_fsdp_wrapped_module"): + module_for_sig = module._fsdp_wrapped_module + else: + module_for_sig = module + assert isinstance(module_for_sig, torch.nn.Module) + sig = inspect.signature(module_for_sig.forward) + # We first flatten all arguments to use only *args, to make things easier and because + # torch.autograd.Function has weird support for kwargs. + bounded = sig.bind(*args, **kwargs) + new_args = [] + for name, param in sig.parameters.items(): + if param.kind in { + inspect.Parameter.VAR_POSITIONAL, + inspect.Parameter.VAR_KEYWORD, + }: + raise RuntimeError("simple_checkpoint doesn't support var args.") + if name not in bounded.arguments: + break + new_args.append(bounded.arguments[name]) + return Checkpoint.apply(module, *new_args) + + +_in_cuda_graph = False +_disable_cuda_graph = False + + +def in_cuda_graph() -> bool: + """Indicate whether we are in a function that is CUDA Graphed (or will be soon).""" + return _in_cuda_graph + + +@contextmanager +def _set_in_cuda_graph(): + global _in_cuda_graph + assert not _in_cuda_graph + _in_cuda_graph = True + try: + yield + finally: + _in_cuda_graph = False + + +def _is_cuda_graph_enabled() -> bool: + if _disable_cuda_graph: + return False + no_cuda_graph = os.environ.get("NO_CUDA_GRAPH", "") + if no_cuda_graph.lower() not in {"0", "no", "n", ""}: + return False + return True + + +@contextmanager +def no_cuda_graph(): + """Deactivate CUDA Graphing for all the calls in this context manager.""" + global _disable_cuda_graph + old_value = _disable_cuda_graph + _disable_cuda_graph = True + try: + yield + finally: + _disable_cuda_graph = old_value + + +class CUDAGraphed: + """Allow simple CUDA Graphing of a function. + + Args: + func: callable, taking any number of arguments. Its tensors arguments should + be top level args, not nested in structures (tuples, dicts, etc). Keyword + arguments are NOT supported for simplicity. + warmup_steps: how many call to make normally before CUDA Graphing. In particular, this + allows torch.compiled functions to get properly compiled. + disabled: if True, just call the func directly, useful to quickly deactivate on CPU. + """ + + def __init__(self, func: tp.Callable, warmup_steps: int = 1, disable: bool = False): + self.func = func + self.warmup_steps = warmup_steps + self.disable = disable + self._graph: cuda.CUDAGraph | None = None + self._output: tuple | None = None + self._args: tuple | None = None + + def reset(self, warmup_steps: int = 0) -> None: + """Reset the state, meaning the next call we get CUDA Graphed again. Useful if some + shapes have changed, or external state (e.g. KVCache) has changed.""" + self.warmup_steps = warmup_steps + self._graph = None + self._output = None + self._args = None + + def __call__(self, *args, **kwargs) -> tp.Any: + if kwargs: + raise RuntimeError("Named arguments not supported for now.") + if self.disable or not _is_cuda_graph_enabled() or in_cuda_graph(): + return self.func(*args, **kwargs) + + def _clone_tensors(args: tuple) -> tuple: + out: list = [] + for arg in args: + if isinstance(arg, torch.Tensor): + arg = arg.clone() + out.append(arg) + return tuple(out) + + def _match_values_copy_tensors(args: tuple, target_args: tuple) -> None: + if len(args) != len(target_args): + raise ValueError( + f"Expected {len(target_args)}, but got {args} for CUDA Graphed function." + ) + for idx, (source, target) in enumerate(zip(args, target_args)): + if isinstance(target, torch.Tensor): + if not isinstance(source, torch.Tensor): + raise ValueError( + f"Argument #{idx} was a tensor, and is no longer (now {source})." + ) + if source.shape != target.shape: + raise ValueError( + f"Argument #{idx} had shape {target.shape}, but got shape {source.shape}" + "When using CUDAGraph, every call must be done with exactly the same shapes. " + "Feel free to deactivate with the env variable NO_CUDA_GRAPH=1, or the decorator " + "`with no_cuda_graph():`" + ) + target.copy_(source) + else: + if isinstance(source, torch.Tensor): + raise ValueError( + f"Argument #{idx} was not a tensor {target}, but is now one." + ) + if source is not target and source != target: + raise ValueError( + f"Argument #{idx} changed value from {target} to {source}." + ) + + with _set_in_cuda_graph(): + # Prevent any one under us to try and CUDA Graph things. + if self._graph is None: + if self.warmup_steps <= 0: + self._graph = cuda.CUDAGraph() + # Making a copy just to ensure those are not used else where. + self._args = _clone_tensors(args) + with cuda.graph(self._graph): + self._output = self.func(*self._args) + # At this point nothing really happened, so we have to make it run for real. + self._graph.replay() + return self._output + else: + self.warmup_steps -= 1 + return self.func(*args) + else: + assert self._args is not None + _match_values_copy_tensors(args, self._args) + self._graph.replay() + return self._output + + +def cuda_graph(func: tp.Callable, warmup_steps: int = 1): + """Just calls `CUDAGraphed` on the given function.""" + if not _is_cuda_graph_enabled(): + return func + return CUDAGraphed(func, warmup_steps) diff --git a/moshi_src/moshi/utils/quantize.py b/moshi_src/moshi/utils/quantize.py new file mode 100644 index 0000000..70bb3b6 --- /dev/null +++ b/moshi_src/moshi/utils/quantize.py @@ -0,0 +1,57 @@ +# Copyright (c) Kyutai, all rights reserved. +# This source code is licensed under the license found in the +# LICENSE file in the root directory of this source tree. + +"""Quantization based on bitsandbytes, supporting only 8 bits for now. +We are taking from freedom from the intended use of bnb: +""" + +import torch +from torch import nn + + +class QLinear(nn.Module): + def __init__(self, linear: nn.Linear): + super().__init__() + from bitsandbytes import functional as bnbF # type: ignore + weight = linear.weight + assert weight.data.dtype.is_floating_point + assert linear.bias is None + CB, SCB, _ = bnbF.int8_vectorwise_quant(weight.data.to(torch.float16)) # type: ignore + self.weight = nn.Parameter(CB, requires_grad=False) + self.weight_scb = nn.Parameter(SCB, requires_grad=False) + + def forward(self, x): + import bitsandbytes as bnb # type: ignore + state = bnb.MatmulLtState() + state.CB = self.weight # type: ignore + assert isinstance(state.CB, torch.Tensor) + state.SCB = self.weight_scb # type: ignore + assert isinstance(state.SCB, torch.Tensor) + if state.SCB.dtype != torch.float: + raise RuntimeError( + "Expected `weight_scb` to have type float, but got bfloat16. " + "When using quantized models, care should be taken not to change the dtype of " + "the model once initialized.") + assert state.SCB.dtype == torch.float, state.SCB.dtype + state.has_fp16_weights = False + y = bnb.matmul(x.half(), state.CB, state=state) + assert isinstance(y, torch.Tensor) + return y + + +def replace_linear_with_qlinear(module): + """Recursively replace all Linear layers with QLinear layers.""" + for name, child in module.named_children(): + if isinstance(child, nn.Linear): + setattr(module, name, QLinear(child)) + elif isinstance(child, QLinear): + # Slight issue with the way we implement things: the scale param + # might get casted with the rest of the model to bfloat16, altough + # we most likely want to keep it as float. For the LM model we might call this function twice, + # first layer by layer to avoid to big of a memory usage, and second, at the end + # of the LM init, after all other modules are initialized and properly dtyped. + # In any case that should happen before loading the state dict to avoid a loss of precision. + child.float() + else: + replace_linear_with_qlinear(child) diff --git a/moshi_src/moshi/utils/sampling.py b/moshi_src/moshi/utils/sampling.py new file mode 100644 index 0000000..2a21d0a --- /dev/null +++ b/moshi_src/moshi/utils/sampling.py @@ -0,0 +1,127 @@ +# Copyright (c) Kyutai, all rights reserved. +# This source code is licensed under the license found in the +# LICENSE file in the root directory of this source tree. + +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the license found in the +# LICENSE file in the root directory of this source tree. + + +import torch + + +def multinomial( + input: torch.Tensor, num_samples: int, replacement=False, *, generator=None +): + """torch.multinomial with arbitrary number of dimensions, and number of candidates on the last dimension. + + Args: + input (torch.Tensor): The input tensor containing probabilities. + num_samples (int): Number of samples to draw. + replacement (bool): Whether to draw with replacement or not. + Keywords args: + generator (torch.Generator): A pseudorandom number generator for sampling. + Returns: + torch.Tensor: Last dimension contains num_samples indices + sampled from the multinomial probability distribution + located in the last dimension of tensor input. + """ + input_ = input.reshape(-1, input.shape[-1]) + # We should probably be able to remove this once the following PR has landed: + # https://github.com/pytorch/pytorch/pull/134818/files + # In the meantime, we specialize the case no-replacement, nsamples=1 so as not + # to have a synchronization point. + if replacement or num_samples != 1: + output_ = torch.multinomial( + input_, + num_samples=num_samples, + replacement=replacement, + generator=generator, + ) + else: + q = torch.empty_like(input_).exponential_(1, generator=generator) + q = input_ / q + output_ = q.argmax(dim=-1, keepdim=True) + output = output_.reshape(*list(input.shape[:-1]), -1) + return output + + +def sample_top_k(probs: torch.Tensor, k: int) -> torch.Tensor: + """Sample next token from top K values along the last dimension of the input probs tensor. + + Args: + probs (torch.Tensor): Input probabilities with token candidates on the last dimension. + k (int): The k in “top-k”. + Returns: + torch.Tensor: Sampled tokens. + """ + k = min(k, probs.shape[-1]) + probs, indices = torch.topk(probs, k, dim=-1) + next_token = multinomial(probs, num_samples=1) + next_token = indices.gather(-1, next_token) + return next_token + + +def sample_top_p(probs: torch.Tensor, p: float) -> torch.Tensor: + """Sample next token from top P probabilities along the last dimension of the input probs tensor. + + Args: + probs (torch.Tensor): Input probabilities with token candidates on the last dimension. + p (int): The p in “top-p”. + Returns: + torch.Tensor: Sampled tokens. + """ + probs_sort, probs_idx = torch.sort(probs, dim=-1, descending=True) + probs_sum = torch.cumsum(probs_sort, dim=-1) + mask = probs_sum - probs_sort > p + probs_sort *= (~mask).float() + probs_sort.div_(probs_sort.sum(dim=-1, keepdim=True)) + next_token = multinomial(probs_sort, num_samples=1) + next_token = torch.gather(probs_idx, -1, next_token) + return next_token + + +def sample_token( + logits: torch.Tensor, + use_sampling: bool = False, + temp: float = 1.0, + top_k: int = 0, + top_p: float = 0.0, +) -> torch.Tensor: + """Given logits of shape [*, Card], returns a LongTensor of shape [*].""" + # Apply softmax for sampling if temp > 0. Else, do greedy sampling to avoid zero division error. + if use_sampling and temp > 0.0: + probs = torch.softmax(logits / temp, dim=-1) + if top_p > 0.0: + next_token = sample_top_p(probs, p=top_p) + elif top_k > 0: + next_token = sample_top_k(probs, k=top_k) + else: + next_token = multinomial(probs, num_samples=1) + else: + next_token = torch.argmax(logits, dim=-1, keepdim=True) + assert next_token.shape[-1] == 1 + return next_token[..., 0] + + +if __name__ == "__main__": + torch.manual_seed(1234) + device = "cpu" + if torch.cuda.is_available(): + torch.backends.cuda.matmul.allow_tf32 = False + torch.backends.cudnn.allow_tf32 = False + device = "cuda:0" + + ps = torch.tensor([5.0, 2.0, 12.0, 6.0, 8.0, 1.0, 0.0, 4.0], device=device) + cnts = torch.zeros(ps.shape, dtype=torch.long, device=device) + total_samples = 1000 + for _ in range(total_samples): + vs = multinomial(ps, num_samples=1, replacement=False) + cnts[vs] += 1 + diff = cnts / cnts.sum() - ps / ps.sum() + max_diff = diff.abs().max().cpu().item() + print(ps / ps.sum()) + print(cnts / cnts.sum()) + assert max_diff < 1.5e-2 diff --git a/moshi_src/moshi/utils/utils.py b/moshi_src/moshi/utils/utils.py new file mode 100644 index 0000000..51eea20 --- /dev/null +++ b/moshi_src/moshi/utils/utils.py @@ -0,0 +1,52 @@ +import torch + +from .compile import torch_compile_lazy + + +@torch_compile_lazy +def cross_entropy( + logits: torch.Tensor, targets: torch.Tensor, mask: torch.Tensor, dtype=torch.float32, + logits_soft_clip: float | None = None) -> torch.Tensor: + """Compute cross entropy between multi-codebook targets and model's logits. + The cross entropy is computed per codebook to provide codebook-level cross entropy. + Valid timesteps for each of the codebook are pulled from the mask, where invalid + timesteps are set to 0. + + Args: + logits (torch.Tensor): Model's logits of shape [B, K, T, card]. + targets (torch.Tensor): Target codes, of shape [B, K, T]. + mask (torch.Tensor): Mask for valid target codes, of shape [B, K, T]. + dtype (type): Data type of the output cross entropy. + logits_soft_clip (float): Clipping value for the logits to avoid numerical instability. + Recommended value: 30.0. + Returns: + ce (torch.Tensor): Cross entropy [B, K, T] with type dtype. + """ + output_shape = targets.shape + assert logits.shape[:-1] == targets.shape + assert mask.shape == targets.shape + logits = logits.view(-1, logits.shape[-1]) + targets = targets.reshape(-1) + mask = mask.reshape(-1) + + safe_targets = torch.where( + mask, + targets, + torch.zeros(1, device=targets.device, dtype=targets.dtype), + ) + + # Chunking the conversion to float32 to avoid OOMs. + ce_chunks = [] + for logits_chunk, targets_chunk in zip(torch.chunk(logits, 4), torch.chunk(safe_targets, 4)): + logits_chunk = logits_chunk.to(dtype) + if logits_soft_clip is not None: + logits_chunk = logits_soft_clip * torch.tanh(logits_chunk / logits_soft_clip) + log_partition = torch.logsumexp(logits_chunk, dim=-1, keepdim=True) + + # For some reason, the PyTorch cross entropy is super slow with inputs with large cardinality (e.g. 32000) + # so we reimplement the cross entropy ourselves... + ce_chunks.append(log_partition - logits_chunk.gather(-1, targets_chunk[..., None])) + ce = torch.cat(ce_chunks, dim=0) + ce = ce[..., 0] + ce = torch.where(mask, ce, torch.zeros(1, device=ce.device, dtype=ce.dtype)) + return ce.view(output_shape) diff --git a/nodes.py b/nodes.py new file mode 100644 index 0000000..b6a6f74 --- /dev/null +++ b/nodes.py @@ -0,0 +1,184 @@ +import torch +import torch._dynamo +import sys +import os + + +os.environ['PYTORCH_CUDA_ALLOC_CONF'] = 'garbage_collection_threshold:0.1' +import folder_paths +from pathlib import Path +import json +import random +import numpy as np +import comfy.utils +from tqdm import tqdm + +# Add the correct moshi source directory to the Python path +sys.path.insert(0, os.path.join(os.path.dirname(__file__), "moshi_src")) + +from moshi.models.loaders import CheckpointInfo +from moshi.models.tts import TTSModel + +# Monkey-patch the problematic function in the moshi library +# This prevents a PyTorch compilation error on Windows by disabling +# the JIT compiler for this specific function. +try: + import moshi.modules.rope + torch._dynamo.disable(moshi.modules.rope.apply_rope) +except (ImportError, AttributeError) as e: + print(f"KyutaiTTS Node: Could not patch moshi.modules.rope.apply_rope. If you encounter an OverflowError, this may be the cause. Error: {e}") + +class KyutaiTTS: + @classmethod + def INPUT_TYPES(s): + return { + "required": { + "text": ("STRING", {"multiline": True, "default": "Hey there! How are you?"}), + "model_path": ("STRING", {"default": "", "multiline": False, "folder_input": True}), + "voice_model": (folder_paths.get_filename_list("loras"), ), + "device": (["cuda", "cpu"],), + "n_q": ("INT", {"default": 32}), + "temp": ("FLOAT", {"default": 0.6, "min": 0.0, "max": 1.0, "step": 0.1}), + "cfg_coef": ("FLOAT", {"default": 2.0, "min": 0.0, "max": 10.0, "step": 0.1}), + "padding_between": ("INT", {"default": 1}), + "seed": ("INT", {"default": 0, "min": 0, "max": 0xFFFFFFFF}), + } + } + + RETURN_TYPES = ("AUDIO", ) + FUNCTION = "generate" + CATEGORY = "Kyutai" + + def generate(self, text, model_path, voice_model, device, n_q, temp, cfg_coef, padding_between, seed): + + + def seed_all(seed): + torch.manual_seed(seed) + random.seed(seed) + np.random.seed(seed) + + seed_all(seed) + + device = torch.device(device) + + # Create CheckpointInfo from the local model path + full_model_path = model_path + if not full_model_path or not os.path.isdir(full_model_path): + raise FileNotFoundError(f"Model directory not found: {full_model_path}") + + # Define expected file names within the model directory + moshi_weights_path = os.path.join(full_model_path, "dsm_tts_1e68beda@240.safetensors") + if not os.path.exists(moshi_weights_path): + raise FileNotFoundError(f"Moshi weights (dsm_tts_1e68beda@240.safetensors) not found in {full_model_path}") + + mimi_weights_path = os.path.join(full_model_path, "tokenizer-e351c8d8-checkpoint125.safetensors") + if not os.path.exists(mimi_weights_path): + raise FileNotFoundError(f"Mimi weights (tokenizer-e351c8d8-checkpoint125.safetensors) not found in {full_model_path}") + + tokenizer_path = os.path.join(full_model_path, "tokenizer_spm_8k_en_fr_audio.model") + if not os.path.exists(tokenizer_path): + raise FileNotFoundError(f"Tokenizer (tokenizer_spm_8k_en_fr_audio.model) not found in {full_model_path}") + + config_path = os.path.join(full_model_path, "config.json") + if not os.path.exists(config_path): + raise FileNotFoundError(f"config.json not found in {full_model_path}") + + with open(config_path, 'r') as f: + raw_config = json.load(f) + + # Extract specific configs for CheckpointInfo and remove them from lm_config + tts_config = raw_config.get("tts_config", {}) + stt_config = raw_config.get("stt_config", {}) + lm_gen_config = raw_config.get("lm_gen_config", {}) + model_id = raw_config.get("model_id", {}) + model_type = raw_config.get("model_type", "moshi") # Extract model_type + + lm_config = dict(raw_config) # Create a copy for lm_config + + # Remove keys not meant for LMModel from lm_config + lm_config.pop("tts_config", None) + lm_config.pop("stt_config", None) + lm_config.pop("lm_gen_config", None) + lm_config.pop("model_id", None) + lm_config.pop("moshi_name", None) + lm_config.pop("mimi_name", None) + lm_config.pop("tokenizer_name", None) + lm_config.pop("lora_name", None) + lm_config.pop("model_type", None) # Remove model_type from lm_config + + checkpoint_info = CheckpointInfo( + moshi_weights=Path(moshi_weights_path), + mimi_weights=Path(mimi_weights_path), + tokenizer=Path(tokenizer_path), + lm_config=lm_config, + raw_config=raw_config, + tts_config=tts_config, + stt_config=stt_config, + lm_gen_config=lm_gen_config, + model_id=model_id, + model_type=model_type # Pass model_type to CheckpointInfo + ) + + tts_model = TTSModel.from_checkpoint_info( + checkpoint_info, n_q=n_q, temp=temp, device=device + ) + + entries = tts_model.prepare_script([text], padding_between=padding_between) + + voice_path = folder_paths.get_full_path("loras", voice_model) + if not voice_path or not os.path.exists(voice_path): + raise FileNotFoundError(f"Voice model not found: {voice_model}") + + condition_attributes = tts_model.make_condition_attributes( + [voice_path], cfg_coef=cfg_coef + ) + + # --- Step 1: Generate audio token frames --- + frames_list = [] + # A more accurate estimation including initial/final padding, model delays, and an empirical fudge factor. + initial_padding = tts_model.machine.initial_padding + final_padding = tts_model.final_padding + delay_steps = tts_model.delay_steps + # Add a small fudge factor for each word to account for un-predictable discretionary padding. + FUDGE_FACTOR_PER_WORD = 1 + word_steps = sum(len(entry.tokens) + entry.padding + FUDGE_FACTOR_PER_WORD for entry in entries) + total_steps = initial_padding + word_steps + delay_steps + final_padding + gen_pbar = comfy.utils.ProgressBar(total_steps) + with tqdm(total=total_steps, desc="Generating Tokens") as pbar_cmd_gen: + def _on_frame_collect(frame): + if (frame != -1).all(): + frames_list.append(frame.clone()) + # Update by 1 for each frame generated. + gen_pbar.update(1) + pbar_cmd_gen.update(1) + + all_entries = [entries] + all_condition_attributes = [condition_attributes] + with tts_model.mimi.streaming(len(all_entries)): + tts_model.generate(all_entries, all_condition_attributes, on_frame=_on_frame_collect) + + # --- Step 2: Decode frames to PCM audio --- + pcms = [] + if frames_list: + decode_pbar = comfy.utils.ProgressBar(len(frames_list)) + with tqdm(total=len(frames_list), desc="Decoding Audio") as pbar_cmd_decode: + for frame in frames_list: + pcm = tts_model.mimi.decode(frame[:, 1:, :]).cpu().numpy() + pcms.append(np.clip(pcm[0], -1, 1)[np.newaxis, :]) + decode_pbar.update(1) + pbar_cmd_decode.update(1) + + # --- Step 3: Concatenate audio chunks --- + audio = np.concatenate(pcms, axis=-1) if pcms else np.array([]) + + print(f"KyutaiTTS Node: Outputting audio with sample rate: {tts_model.mimi.sample_rate}") + # Return audio in the format expected by ComfyUI's AUDIO type + return ({"waveform": torch.from_numpy(audio), "sample_rate": tts_model.mimi.sample_rate},) + +NODE_CLASS_MAPPINGS = { + "KyutaiTTS": KyutaiTTS, +} + +NODE_DISPLAY_NAME_MAPPINGS = { + "KyutaiTTS": "KyutaiTTS", +} \ No newline at end of file diff --git a/requirements.txt b/requirements.txt new file mode 100644 index 0000000..7ce057e --- /dev/null +++ b/requirements.txt @@ -0,0 +1 @@ +sphn<0.2 \ No newline at end of file