\n 
\n ![](\"https://img.shields.io/badge/🕹️-Boson%20AI%20Playground-9C276A\"){kind=icon}
\n ![](\"https://img.shields.io/badge/🎮-HF%20Space%20Playground-8A2BE2\"){kind=icon}
\n ![](\"https://img.shields.io/badge/🤗-Checkpoints){kind=icon}
\n\n \n \n \n \n\n ![](\"figures/emergent-tts-emotions-win-rate.png\")
\n
![](\"figures/higgs_audio_v2_architecture_combined.png\")
\n\n\nHiggs Audio v2 adopts the \"generation variant\" depicted in the architecture figure above. Its strong performance is driven by three key technical innovations:\n- We developed an automated annotation pipeline that leverages multiple ASR models, sound event classification models, and our in-house audio understanding model. Using this pipeline, we cleaned and annotated 10 million hours audio data, which we refer to as **AudioVerse**. The in-house understanding model is finetuned on top of [Higgs Audio v1 Understanding](https://www.boson.ai/blog/higgs-audio), which adopts the \"understanding variant\" shown in the architecture figure.\n- We trained a unified audio tokenizer from scratch that captures both semantic and acoustic features. We also open-sourced our evaluation set on [HuggingFace](https://huggingface.co/datasets/bosonai/AudioTokenBench). Learn more in the [tokenizer blog](./tech_blogs/TOKENIZER_BLOG.md).\n- We proposed the DualFFN architecture, which enhances the LLM’s ability to model acoustics tokens with minimal computational overhead. See the [architecture blog](./tech_blogs/ARCHITECTURE_BLOG.md).\n\n## Evaluation\n\nHere's the performance of Higgs Audio v2 on four benchmarks, [Seed-TTS Eval](https://github.com/BytedanceSpeech/seed-tts-eval), [Emotional Speech Dataset (ESD)](https://paperswithcode.com/dataset/esd), [EmergentTTS-Eval](https://arxiv.org/abs/2505.23009), and Multi-speaker Eval:\n\n#### Seed-TTS Eval & ESD\n\nWe prompt Higgs Audio v2 with the reference text, reference audio, and target text for zero-shot TTS. We use the standard evaluation metrics from Seed-TTS Eval and ESD.\n\n| | SeedTTS-Eval| | ESD | |\n|------------------------------|--------|--------|---------|-------------------|\n| | WER ↓ | SIM ↑ | WER ↓ | SIM (emo2vec) ↑ |\n| Cosyvoice2 | 2.28 | 65.49 | 2.71 | 80.48 |\n| Qwen2.5-omni† | 2.33 | 64.10 | - | - |\n| ElevenLabs Multilingual V2 | **1.43** | 50.00 | 1.66 | 65.87 |\n| Higgs Audio v1 | 2.18 | 66.27 | **1.49** | 82.84 |\n| Higgs Audio v2 (base) | 2.44 | **67.70** | 1.78 | **86.13** |\n\n\n#### EmergentTTS-Eval (\"Emotions\" and \"Questions\")\n\nFollowing the [EmergentTTS-Eval Paper](https://arxiv.org/abs/2505.23009), we report the win-rate over \"gpt-4o-mini-tts\" with the \"alloy\" voice. The judge model is Gemini 2.5 Pro.\n\n| Model | Emotions (%) ↑ | Questions (%) ↑ |\n|------------------------------------|--------------|----------------|\n| Higgs Audio v2 (base) | **75.71%** | **55.71%** |\n| [gpt-4o-audio-preview†](https://platform.openai.com/docs/models/gpt-4o-audio-preview) | 61.64% | 47.85% |\n| [Hume.AI](https://www.hume.ai/research) | 61.60% | 43.21% |\n| **BASELINE:** [gpt-4o-mini-tts](https://platform.openai.com/docs/models/gpt-4o-mini-tts) | 50.00% | 50.00% |\n| [Qwen 2.5 Omni†](https://github.com/QwenLM/Qwen2.5-Omni) | 41.60% | 51.78% |\n| [minimax/speech-02-hd](https://replicate.com/minimax/speech-02-hd) | 40.86% | 47.32% |\n| [ElevenLabs Multilingual v2](https://elevenlabs.io/blog/eleven-multilingual-v2) | 30.35% | 39.46% |\n| [DeepGram Aura-2](https://deepgram.com/learn/introducing-aura-2-enterprise-text-to-speech) | 29.28% | 48.21% |\n| [Sesame csm-1B](https://github.com/SesameAILabs/csm) | 15.96% | 31.78% |\n\n'†' means using the strong-prompting method described in the paper.\n\n\n#### Multi-speaker Eval\n\nWe also designed a multi-speaker evaluation benchmark to evaluate the capability of Higgs Audio v2 for multi-speaker dialog generation. The benchmark contains three subsets\n\n- `two-speaker-conversation`: 1000 synthetic dialogues involving two speakers. We fix two reference audio clips to evaluate the model's ability in double voice cloning for utterances ranging from 4 to 10 dialogues between two randomly chosen persona.\n- `small talk (no ref)`: 250 synthetic dialogues curated in the same way as above, but are characterized by short utterances and a limited number of turns (4–6), we do not fix reference audios in this case and this set is designed to evaluate the model's ability to automatically assign appropriate voices to speakers.\n- `small talk (ref)`: 250 synthetic dialogues similar to above, but contains even shorter utterances as this set is meant to include reference clips in it's context, similar to `two-speaker-conversation`.\n\n\nWe report the word-error-rate (WER) and the geometric mean between intra-speaker similarity and inter-speaker dis-similarity on these three subsets. Other than Higgs Audio v2, we also evaluated [MoonCast](https://github.com/jzq2000/MoonCast) and [nari-labs/Dia-1.6B-0626](https://huggingface.co/nari-labs/Dia-1.6B-0626), two of the most popular open-source models capable of multi-speaker dialog generation. Results are summarized in the following table. We are not able to run [nari-labs/Dia-1.6B-0626](https://huggingface.co/nari-labs/Dia-1.6B-0626) on our \"two-speaker-conversation\" subset due to its strict limitation on the length of the utterances and output audio.\n\n| | two-speaker-conversation | |small talk | | small talk (no ref) | |\n| ---------------------------------------------- | -------------- | ------------------ | ---------- | -------------- | ------------------- | -------------- |\n| | WER ↓ | Mean Sim & Dis-sim ↑ | WER ↓ | Mean Sim & Dis-sim ↑ | WER ↓ | Mean Sim & Dis-sim ↑ |\n| [MoonCast](https://github.com/jzq2000/MoonCast) | 38.77 | 46.02 | **8.33** | 63.68 | 24.65 | 53.94 |\n| [nari-labs/Dia-1.6B-0626](https://huggingface.co/nari-labs/Dia-1.6B-0626) | \\- | \\- | 17.62 | 63.15 | 19.46 | **61.14** |\n| Higgs Audio v2 (base) | **18.88** | **51.95** | 11.89 | **67.92** | **14.65** | 55.28 |\n\n## Contribution and Support\n\nFor contribution and support guidelines, please see the support guidelines at [SUPPORT_GUIDELINES.md](SUPPORT_GUIDELINES.md).\n\n## Citation\n\nIf you feel the repository is helpful, please kindly cite as:\n\n```\n@misc{higgsaudio2025,\n author = {{Boson AI}},\n title = {{Higgs Audio V2: Redefining Expressiveness in Audio Generation}},\n year = {2025},\n howpublished = {\\url{https://github.com/boson-ai/higgs-audio}},\n note = {GitHub repository. Release blog available at \\url{https://www.boson.ai/blog/higgs-audio-v2}},\n}\n```\n\n## Third-Party Licenses\n\nThe `boson_multimodal/audio_processing/` directory contains code derived from third-party repositories, primarily from [xcodec](https://github.com/zhenye234/xcodec). Please see the [`LICENSE`](boson_multimodal/audio_processing/LICENSE) in that directory for complete attribution and licensing information.\n\n## We Are Hiring!\n\nIf you are passionate about multimodal AI, speech/audio models, or large-scale systems, \ncheck out our open positions at [Boson AI Careers](https://jobs.lever.co/bosonai).\n"
},
{
"path": "SUPPORT_GUIDELINES.md",
"content": "# Contribution & Support Guidelines\n\nThank you for your interest in this project! Before opening an issue, please take a moment to read the following guidelines:\n\n## Self-Check First\n- Write your question in **English** or include an English translation so the community can understand and assist you better.\n- Verify that you have **installed the correct version** of the package.\n- Check the GitHub [README](README.md), [Hugging Face Space](https://huggingface.co/spaces/smola/higgs_audio_v2), [Model Card](https://huggingface.co/bosonai/higgs-audio-v2-generation-3B-base) and existing issues — many questions already have answers.\n- Ensure your problem can be **reproduced** and is directly related to this project.\n\n## Asking Properly\n- Provide **clear reproduction steps / minimal code examples / error logs**.\n- Keep the issue title **concise and descriptive**, and include enough context in the body.\n- Avoid vague questions like *“It doesn’t work, what should I do?”* or *“Can you debug this for me?”*.\n\n## About Support\n- This is a **community-driven open source project**. Maintainers will respond when time allows.\n- There is **no obligation** to answer every request — please be patient and understanding.\n- For more reliable or timely support, consider:\n - Submitting a **Pull Request** to improve code or documentation.\n - Providing detailed context so that the community can help.\n\n## Code of Conduct\n- Be **respectful and polite**.\n- Do not spam or repeatedly demand responses.\n- Off-topic, vague, or inappropriate questions may be closed.\n"
},
{
"path": "boson_multimodal/__init__.py",
"content": ""
},
{
"path": "boson_multimodal/audio_processing/LICENSE",
"content": "Third-Party License Attribution for Audio Processing Module\n===========================================================\n\nThis directory contains code derived from multiple open-source projects. \nThe following sections detail the licenses and attributions for third-party code.\n\n## XCodec Repository\nThe code in this directory is derived from:\nhttps://github.com/zhenye234/xcodec\n\n## Individual File Attributions\n\n### Quantization Module (quantization/)\n- Several files contain code derived from Meta Platforms, Inc. and the vector-quantize-pytorch repository\n- Individual files contain their own license headers where applicable\n- The vector-quantize-pytorch portions are licensed under the MIT License\n\n## License Terms\n\n### MIT License (for applicable portions)\nPermission is hereby granted, free of charge, to any person obtaining a copy\nof this software and associated documentation files (the \"Software\"), to deal\nin the Software without restriction, including without limitation the rights\nto use, copy, modify, merge, publish, distribute, sublicense, and/or sell\ncopies of the Software, and to permit persons to whom the Software is\nfurnished to do so, subject to the following conditions:\n\nThe above copyright notice and this permission notice shall be included in all\ncopies or substantial portions of the Software.\n\nTHE SOFTWARE IS PROVIDED \"AS IS\", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR\nIMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,\nFITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE\nAUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER\nLIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,\nOUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE\nSOFTWARE.\n\n## Attribution Requirements\nWhen using this code, please ensure proper attribution to:\n1. The original xcodec repository: https://github.com/zhenye234/xcodec\n2. Any other repositories mentioned in individual file headers\n3. This derivative work and its modifications\n\n## Disclaimer\nThis directory contains modified versions of the original code. Please refer to\nthe original repositories for the canonical implementations and their specific\nlicense terms.\n\nFor any questions about licensing or attribution, please check the individual\nfile headers and the original source repositories. "
},
{
"path": "boson_multimodal/audio_processing/descriptaudiocodec/__init__.py",
"content": ""
},
{
"path": "boson_multimodal/audio_processing/descriptaudiocodec/dac/model/base.py",
"content": "import math\nfrom dataclasses import dataclass\nfrom pathlib import Path\nfrom typing import Union\n\nimport numpy as np\nimport torch\nimport tqdm\nfrom audiotools import AudioSignal\nfrom torch import nn\n\nSUPPORTED_VERSIONS = [\"1.0.0\"]\n\n\n@dataclass\nclass DACFile:\n codes: torch.Tensor\n\n # Metadata\n chunk_length: int\n original_length: int\n input_db: float\n channels: int\n sample_rate: int\n padding: bool\n dac_version: str\n\n def save(self, path):\n artifacts = {\n \"codes\": self.codes.numpy().astype(np.uint16),\n \"metadata\": {\n \"input_db\": self.input_db.numpy().astype(np.float32),\n \"original_length\": self.original_length,\n \"sample_rate\": self.sample_rate,\n \"chunk_length\": self.chunk_length,\n \"channels\": self.channels,\n \"padding\": self.padding,\n \"dac_version\": SUPPORTED_VERSIONS[-1],\n },\n }\n path = Path(path).with_suffix(\".dac\")\n with open(path, \"wb\") as f:\n np.save(f, artifacts)\n return path\n\n @classmethod\n def load(cls, path):\n artifacts = np.load(path, allow_pickle=True)[()]\n codes = torch.from_numpy(artifacts[\"codes\"].astype(int))\n if artifacts[\"metadata\"].get(\"dac_version\", None) not in SUPPORTED_VERSIONS:\n raise RuntimeError(f\"Given file {path} can't be loaded with this version of descript-audio-codec.\")\n return cls(codes=codes, **artifacts[\"metadata\"])\n\n\nclass CodecMixin:\n @property\n def padding(self):\n if not hasattr(self, \"_padding\"):\n self._padding = True\n return self._padding\n\n @padding.setter\n def padding(self, value):\n assert isinstance(value, bool)\n\n layers = [l for l in self.modules() if isinstance(l, (nn.Conv1d, nn.ConvTranspose1d))]\n\n for layer in layers:\n if value:\n if hasattr(layer, \"original_padding\"):\n layer.padding = layer.original_padding\n else:\n layer.original_padding = layer.padding\n layer.padding = tuple(0 for _ in range(len(layer.padding)))\n\n self._padding = value\n\n def get_delay(self):\n # Any number works here, delay is invariant to input length\n l_out = self.get_output_length(0)\n L = l_out\n\n layers = []\n for layer in self.modules():\n if isinstance(layer, (nn.Conv1d, nn.ConvTranspose1d)):\n layers.append(layer)\n\n for layer in reversed(layers):\n d = layer.dilation[0]\n k = layer.kernel_size[0]\n s = layer.stride[0]\n\n if isinstance(layer, nn.ConvTranspose1d):\n L = ((L - d * (k - 1) - 1) / s) + 1\n elif isinstance(layer, nn.Conv1d):\n L = (L - 1) * s + d * (k - 1) + 1\n\n L = math.ceil(L)\n\n l_in = L\n\n return (l_in - l_out) // 2\n\n def get_output_length(self, input_length):\n L = input_length\n # Calculate output length\n for layer in self.modules():\n if isinstance(layer, (nn.Conv1d, nn.ConvTranspose1d)):\n d = layer.dilation[0]\n k = layer.kernel_size[0]\n s = layer.stride[0]\n\n if isinstance(layer, nn.Conv1d):\n L = ((L - d * (k - 1) - 1) / s) + 1\n elif isinstance(layer, nn.ConvTranspose1d):\n L = (L - 1) * s + d * (k - 1) + 1\n\n L = math.floor(L)\n return L\n\n @torch.no_grad()\n def compress(\n self,\n audio_path_or_signal: Union[str, Path, AudioSignal],\n win_duration: float = 1.0,\n verbose: bool = False,\n normalize_db: float = -16,\n n_quantizers: int = None,\n ) -> DACFile:\n \"\"\"Processes an audio signal from a file or AudioSignal object into\n discrete codes. This function processes the signal in short windows,\n using constant GPU memory.\n\n Parameters\n ----------\n audio_path_or_signal : Union[str, Path, AudioSignal]\n audio signal to reconstruct\n win_duration : float, optional\n window duration in seconds, by default 5.0\n verbose : bool, optional\n by default False\n normalize_db : float, optional\n normalize db, by default -16\n\n Returns\n -------\n DACFile\n Object containing compressed codes and metadata\n required for decompression\n \"\"\"\n audio_signal = audio_path_or_signal\n if isinstance(audio_signal, (str, Path)):\n audio_signal = AudioSignal.load_from_file_with_ffmpeg(str(audio_signal))\n\n self.eval()\n original_padding = self.padding\n original_device = audio_signal.device\n\n audio_signal = audio_signal.clone()\n original_sr = audio_signal.sample_rate\n\n resample_fn = audio_signal.resample\n loudness_fn = audio_signal.loudness\n\n # If audio is > 10 minutes long, use the ffmpeg versions\n if audio_signal.signal_duration >= 10 * 60 * 60:\n resample_fn = audio_signal.ffmpeg_resample\n loudness_fn = audio_signal.ffmpeg_loudness\n\n original_length = audio_signal.signal_length\n resample_fn(self.sample_rate)\n input_db = loudness_fn()\n\n if normalize_db is not None:\n audio_signal.normalize(normalize_db)\n audio_signal.ensure_max_of_audio()\n\n nb, nac, nt = audio_signal.audio_data.shape\n audio_signal.audio_data = audio_signal.audio_data.reshape(nb * nac, 1, nt)\n win_duration = audio_signal.signal_duration if win_duration is None else win_duration\n\n if audio_signal.signal_duration <= win_duration:\n # Unchunked compression (used if signal length < win duration)\n self.padding = True\n n_samples = nt\n hop = nt\n else:\n # Chunked inference\n self.padding = False\n # Zero-pad signal on either side by the delay\n audio_signal.zero_pad(self.delay, self.delay)\n n_samples = int(win_duration * self.sample_rate)\n # Round n_samples to nearest hop length multiple\n n_samples = int(math.ceil(n_samples / self.hop_length) * self.hop_length)\n hop = self.get_output_length(n_samples)\n\n codes = []\n range_fn = range if not verbose else tqdm.trange\n\n for i in range_fn(0, nt, hop):\n x = audio_signal[..., i : i + n_samples]\n x = x.zero_pad(0, max(0, n_samples - x.shape[-1]))\n\n audio_data = x.audio_data.to(self.device)\n audio_data = self.preprocess(audio_data, self.sample_rate)\n _, c, _, _, _ = self.encode(audio_data, n_quantizers)\n codes.append(c.to(original_device))\n chunk_length = c.shape[-1]\n\n codes = torch.cat(codes, dim=-1)\n\n dac_file = DACFile(\n codes=codes,\n chunk_length=chunk_length,\n original_length=original_length,\n input_db=input_db,\n channels=nac,\n sample_rate=original_sr,\n padding=self.padding,\n dac_version=SUPPORTED_VERSIONS[-1],\n )\n\n if n_quantizers is not None:\n codes = codes[:, :n_quantizers, :]\n\n self.padding = original_padding\n return dac_file\n\n @torch.no_grad()\n def decompress(\n self,\n obj: Union[str, Path, DACFile],\n verbose: bool = False,\n ) -> AudioSignal:\n \"\"\"Reconstruct audio from a given .dac file\n\n Parameters\n ----------\n obj : Union[str, Path, DACFile]\n .dac file location or corresponding DACFile object.\n verbose : bool, optional\n Prints progress if True, by default False\n\n Returns\n -------\n AudioSignal\n Object with the reconstructed audio\n \"\"\"\n self.eval()\n if isinstance(obj, (str, Path)):\n obj = DACFile.load(obj)\n\n original_padding = self.padding\n self.padding = obj.padding\n\n range_fn = range if not verbose else tqdm.trange\n codes = obj.codes\n original_device = codes.device\n chunk_length = obj.chunk_length\n recons = []\n\n for i in range_fn(0, codes.shape[-1], chunk_length):\n c = codes[..., i : i + chunk_length].to(self.device)\n z = self.quantizer.from_codes(c)[0]\n r = self.decode(z)\n recons.append(r.to(original_device))\n\n recons = torch.cat(recons, dim=-1)\n recons = AudioSignal(recons, self.sample_rate)\n\n resample_fn = recons.resample\n loudness_fn = recons.loudness\n\n # If audio is > 10 minutes long, use the ffmpeg versions\n if recons.signal_duration >= 10 * 60 * 60:\n resample_fn = recons.ffmpeg_resample\n loudness_fn = recons.ffmpeg_loudness\n\n recons.normalize(obj.input_db)\n resample_fn(obj.sample_rate)\n recons = recons[..., : obj.original_length]\n loudness_fn()\n recons.audio_data = recons.audio_data.reshape(-1, obj.channels, obj.original_length)\n\n self.padding = original_padding\n return recons\n"
},
{
"path": "boson_multimodal/audio_processing/descriptaudiocodec/dac/model/dac.py",
"content": "import math\nfrom typing import List\nfrom typing import Union\n\nimport numpy as np\nimport torch\nfrom audiotools import AudioSignal\nfrom audiotools.ml import BaseModel\nfrom torch import nn\n\nfrom .base import CodecMixin\nfrom dac.nn.layers import Snake1d\nfrom dac.nn.layers import WNConv1d\nfrom dac.nn.layers import WNConvTranspose1d\nfrom dac.nn.quantize import ResidualVectorQuantize\n\n\ndef init_weights(m):\n if isinstance(m, nn.Conv1d):\n nn.init.trunc_normal_(m.weight, std=0.02)\n nn.init.constant_(m.bias, 0)\n\n\nclass ResidualUnit(nn.Module):\n def __init__(self, dim: int = 16, dilation: int = 1):\n super().__init__()\n pad = ((7 - 1) * dilation) // 2\n self.block = nn.Sequential(\n Snake1d(dim),\n WNConv1d(dim, dim, kernel_size=7, dilation=dilation, padding=pad),\n Snake1d(dim),\n WNConv1d(dim, dim, kernel_size=1),\n )\n\n def forward(self, x):\n y = self.block(x)\n pad = (x.shape[-1] - y.shape[-1]) // 2\n if pad > 0:\n x = x[..., pad:-pad]\n return x + y\n\n\nclass EncoderBlock(nn.Module):\n def __init__(self, dim: int = 16, stride: int = 1):\n super().__init__()\n self.block = nn.Sequential(\n ResidualUnit(dim // 2, dilation=1),\n ResidualUnit(dim // 2, dilation=3),\n ResidualUnit(dim // 2, dilation=9),\n Snake1d(dim // 2),\n WNConv1d(\n dim // 2,\n dim,\n kernel_size=2 * stride,\n stride=stride,\n padding=math.ceil(stride / 2),\n ),\n )\n\n def forward(self, x):\n return self.block(x)\n\n\nclass Encoder(nn.Module):\n def __init__(\n self,\n d_model: int = 64,\n strides: list = [2, 4, 8, 8],\n d_latent: int = 256,\n ):\n super().__init__()\n # Create first convolution\n self.block = [WNConv1d(1, d_model, kernel_size=7, padding=3)]\n\n # Create EncoderBlocks that double channels as they downsample by `stride`\n for stride in strides:\n d_model *= 2\n self.block += [EncoderBlock(d_model, stride=stride)]\n\n # Create last convolution\n self.block += [\n Snake1d(d_model),\n WNConv1d(d_model, d_latent, kernel_size=3, padding=1),\n ]\n\n # Wrap black into nn.Sequential\n self.block = nn.Sequential(*self.block)\n self.enc_dim = d_model\n\n def forward(self, x):\n return self.block(x)\n\n\nclass DecoderBlock(nn.Module):\n def __init__(self, input_dim: int = 16, output_dim: int = 8, stride: int = 1, out_pad=0):\n super().__init__()\n self.block = nn.Sequential(\n Snake1d(input_dim),\n WNConvTranspose1d(\n input_dim,\n output_dim,\n kernel_size=2 * stride,\n stride=stride,\n padding=math.ceil(stride / 2),\n output_padding=stride % 2, # out_pad,\n ),\n ResidualUnit(output_dim, dilation=1),\n ResidualUnit(output_dim, dilation=3),\n ResidualUnit(output_dim, dilation=9),\n )\n\n def forward(self, x):\n return self.block(x)\n\n\nclass Decoder(nn.Module):\n def __init__(\n self,\n input_channel,\n channels,\n rates,\n d_out: int = 1,\n ):\n super().__init__()\n\n # Add first conv layer\n layers = [WNConv1d(input_channel, channels, kernel_size=7, padding=3)]\n\n # Add upsampling + MRF blocks\n for i, stride in enumerate(rates):\n input_dim = channels // 2**i\n output_dim = channels // 2 ** (i + 1)\n if i == 1:\n out_pad = 1\n else:\n out_pad = 0\n layers += [DecoderBlock(input_dim, output_dim, stride, out_pad)]\n\n # Add final conv layer\n layers += [\n Snake1d(output_dim),\n WNConv1d(output_dim, d_out, kernel_size=7, padding=3),\n # nn.Tanh(),\n ]\n\n self.model = nn.Sequential(*layers)\n\n def forward(self, x):\n return self.model(x)\n\n\nclass DAC(BaseModel, CodecMixin):\n def __init__(\n self,\n encoder_dim: int = 64,\n encoder_rates: List[int] = [2, 4, 8, 8],\n latent_dim: int = None,\n decoder_dim: int = 1536,\n decoder_rates: List[int] = [8, 8, 4, 2],\n n_codebooks: int = 9,\n codebook_size: int = 1024,\n codebook_dim: Union[int, list] = 8,\n quantizer_dropout: bool = False,\n sample_rate: int = 44100,\n ):\n super().__init__()\n\n self.encoder_dim = encoder_dim\n self.encoder_rates = encoder_rates\n self.decoder_dim = decoder_dim\n self.decoder_rates = decoder_rates\n self.sample_rate = sample_rate\n\n if latent_dim is None:\n latent_dim = encoder_dim * (2 ** len(encoder_rates))\n\n self.latent_dim = latent_dim\n\n self.hop_length = np.prod(encoder_rates)\n self.encoder = Encoder(encoder_dim, encoder_rates, latent_dim)\n\n self.n_codebooks = n_codebooks\n self.codebook_size = codebook_size\n self.codebook_dim = codebook_dim\n self.quantizer = ResidualVectorQuantize(\n input_dim=latent_dim,\n n_codebooks=n_codebooks,\n codebook_size=codebook_size,\n codebook_dim=codebook_dim,\n quantizer_dropout=quantizer_dropout,\n )\n\n self.decoder = Decoder(\n latent_dim,\n decoder_dim,\n decoder_rates,\n )\n self.sample_rate = sample_rate\n self.apply(init_weights)\n\n self.delay = self.get_delay()\n\n def preprocess(self, audio_data, sample_rate):\n if sample_rate is None:\n sample_rate = self.sample_rate\n assert sample_rate == self.sample_rate\n\n length = audio_data.shape[-1]\n right_pad = math.ceil(length / self.hop_length) * self.hop_length - length\n audio_data = nn.functional.pad(audio_data, (0, right_pad))\n\n return audio_data\n\n def encode(\n self,\n audio_data: torch.Tensor,\n n_quantizers: int = None,\n ):\n \"\"\"Encode given audio data and return quantized latent codes\n\n Parameters\n ----------\n audio_data : Tensor[B x 1 x T]\n Audio data to encode\n n_quantizers : int, optional\n Number of quantizers to use, by default None\n If None, all quantizers are used.\n\n Returns\n -------\n dict\n A dictionary with the following keys:\n \"z\" : Tensor[B x D x T]\n Quantized continuous representation of input\n \"codes\" : Tensor[B x N x T]\n Codebook indices for each codebook\n (quantized discrete representation of input)\n \"latents\" : Tensor[B x N*D x T]\n Projected latents (continuous representation of input before quantization)\n \"vq/commitment_loss\" : Tensor[1]\n Commitment loss to train encoder to predict vectors closer to codebook\n entries\n \"vq/codebook_loss\" : Tensor[1]\n Codebook loss to update the codebook\n \"length\" : int\n Number of samples in input audio\n \"\"\"\n z = self.encoder(audio_data)\n z, codes, latents, commitment_loss, codebook_loss = self.quantizer(z, n_quantizers)\n return z, codes, latents, commitment_loss, codebook_loss\n\n def decode(self, z: torch.Tensor):\n \"\"\"Decode given latent codes and return audio data\n\n Parameters\n ----------\n z : Tensor[B x D x T]\n Quantized continuous representation of input\n length : int, optional\n Number of samples in output audio, by default None\n\n Returns\n -------\n dict\n A dictionary with the following keys:\n \"audio\" : Tensor[B x 1 x length]\n Decoded audio data.\n \"\"\"\n return self.decoder(z)\n\n def forward(\n self,\n audio_data: torch.Tensor,\n sample_rate: int = None,\n n_quantizers: int = None,\n ):\n \"\"\"Model forward pass\n\n Parameters\n ----------\n audio_data : Tensor[B x 1 x T]\n Audio data to encode\n sample_rate : int, optional\n Sample rate of audio data in Hz, by default None\n If None, defaults to `self.sample_rate`\n n_quantizers : int, optional\n Number of quantizers to use, by default None.\n If None, all quantizers are used.\n\n Returns\n -------\n dict\n A dictionary with the following keys:\n \"z\" : Tensor[B x D x T]\n Quantized continuous representation of input\n \"codes\" : Tensor[B x N x T]\n Codebook indices for each codebook\n (quantized discrete representation of input)\n \"latents\" : Tensor[B x N*D x T]\n Projected latents (continuous representation of input before quantization)\n \"vq/commitment_loss\" : Tensor[1]\n Commitment loss to train encoder to predict vectors closer to codebook\n entries\n \"vq/codebook_loss\" : Tensor[1]\n Codebook loss to update the codebook\n \"length\" : int\n Number of samples in input audio\n \"audio\" : Tensor[B x 1 x length]\n Decoded audio data.\n \"\"\"\n length = audio_data.shape[-1]\n audio_data = self.preprocess(audio_data, sample_rate)\n z, codes, latents, commitment_loss, codebook_loss = self.encode(audio_data, n_quantizers)\n\n x = self.decode(z)\n return {\n \"audio\": x[..., :length],\n \"z\": z,\n \"codes\": codes,\n \"latents\": latents,\n \"vq/commitment_loss\": commitment_loss,\n \"vq/codebook_loss\": codebook_loss,\n }\n\n\nif __name__ == \"__main__\":\n import numpy as np\n from functools import partial\n\n model = DAC().to(\"cpu\")\n\n for n, m in model.named_modules():\n o = m.extra_repr()\n p = sum([np.prod(p.size()) for p in m.parameters()])\n fn = lambda o, p: o + f\" {p / 1e6:<.3f}M params.\"\n setattr(m, \"extra_repr\", partial(fn, o=o, p=p))\n print(model)\n print(\"Total # of params: \", sum([np.prod(p.size()) for p in model.parameters()]))\n\n length = 88200 * 2\n x = torch.randn(1, 1, length).to(model.device)\n x.requires_grad_(True)\n x.retain_grad()\n\n # Make a forward pass\n out = model(x)[\"audio\"]\n print(\"Input shape:\", x.shape)\n print(\"Output shape:\", out.shape)\n\n # Create gradient variable\n grad = torch.zeros_like(out)\n grad[:, :, grad.shape[-1] // 2] = 1\n\n # Make a backward pass\n out.backward(grad)\n\n # Check non-zero values\n gradmap = x.grad.squeeze(0)\n gradmap = (gradmap != 0).sum(0) # sum across features\n rf = (gradmap != 0).sum()\n\n print(f\"Receptive field: {rf.item()}\")\n\n x = AudioSignal(torch.randn(1, 1, 44100 * 60), 44100)\n model.decompress(model.compress(x, verbose=True), verbose=True)\n"
},
{
"path": "boson_multimodal/audio_processing/descriptaudiocodec/dac/nn/layers.py",
"content": "import numpy as np\nimport torch\nimport torch.nn as nn\nimport torch.nn.functional as F\nfrom einops import rearrange\nfrom torch.nn.utils import weight_norm\n\n\ndef WNConv1d(*args, **kwargs):\n return weight_norm(nn.Conv1d(*args, **kwargs))\n\n\ndef WNConvTranspose1d(*args, **kwargs):\n return weight_norm(nn.ConvTranspose1d(*args, **kwargs))\n\n\n# Scripting this brings model speed up 1.4x\n@torch.jit.script\ndef snake(x, alpha):\n shape = x.shape\n x = x.reshape(shape[0], shape[1], -1)\n x = x + (alpha + 1e-9).reciprocal() * torch.sin(alpha * x).pow(2)\n x = x.reshape(shape)\n return x\n\n\nclass Snake1d(nn.Module):\n def __init__(self, channels):\n super().__init__()\n self.alpha = nn.Parameter(torch.ones(1, channels, 1))\n\n def forward(self, x):\n return snake(x, self.alpha)\n"
},
{
"path": "boson_multimodal/audio_processing/descriptaudiocodec/dac/nn/quantize.py",
"content": "from typing import Union\n\nimport numpy as np\nimport torch\nimport torch.nn as nn\nimport torch.nn.functional as F\nfrom einops import rearrange\nfrom torch.nn.utils import weight_norm\n\nfrom dac.nn.layers import WNConv1d\n\n\nclass VectorQuantize(nn.Module):\n \"\"\"\n Implementation of VQ similar to Karpathy's repo:\n https://github.com/karpathy/deep-vector-quantization\n Additionally uses following tricks from Improved VQGAN\n (https://arxiv.org/pdf/2110.04627.pdf):\n 1. Factorized codes: Perform nearest neighbor lookup in low-dimensional space\n for improved codebook usage\n 2. l2-normalized codes: Converts euclidean distance to cosine similarity which\n improves training stability\n \"\"\"\n\n def __init__(self, input_dim: int, codebook_size: int, codebook_dim: int):\n super().__init__()\n self.codebook_size = codebook_size\n self.codebook_dim = codebook_dim\n\n self.in_proj = WNConv1d(input_dim, codebook_dim, kernel_size=1)\n self.out_proj = WNConv1d(codebook_dim, input_dim, kernel_size=1)\n self.codebook = nn.Embedding(codebook_size, codebook_dim)\n\n def forward(self, z):\n \"\"\"Quantized the input tensor using a fixed codebook and returns\n the corresponding codebook vectors\n\n Parameters\n ----------\n z : Tensor[B x D x T]\n\n Returns\n -------\n Tensor[B x D x T]\n Quantized continuous representation of input\n Tensor[1]\n Commitment loss to train encoder to predict vectors closer to codebook\n entries\n Tensor[1]\n Codebook loss to update the codebook\n Tensor[B x T]\n Codebook indices (quantized discrete representation of input)\n Tensor[B x D x T]\n Projected latents (continuous representation of input before quantization)\n \"\"\"\n\n # Factorized codes (ViT-VQGAN) Project input into low-dimensional space\n z_e = self.in_proj(z) # z_e : (B x D x T)\n z_q, indices = self.decode_latents(z_e)\n\n commitment_loss = F.mse_loss(z_e, z_q.detach(), reduction=\"none\").mean([1, 2])\n codebook_loss = F.mse_loss(z_q, z_e.detach(), reduction=\"none\").mean([1, 2])\n\n z_q = z_e + (z_q - z_e).detach() # noop in forward pass, straight-through gradient estimator in backward pass\n\n z_q = self.out_proj(z_q)\n\n return z_q, commitment_loss, codebook_loss, indices, z_e\n\n def embed_code(self, embed_id):\n return F.embedding(embed_id, self.codebook.weight)\n\n def decode_code(self, embed_id):\n return self.embed_code(embed_id).transpose(1, 2)\n\n def decode_latents(self, latents):\n encodings = rearrange(latents, \"b d t -> (b t) d\")\n codebook = self.codebook.weight # codebook: (N x D)\n\n # L2 normalize encodings and codebook (ViT-VQGAN)\n encodings = F.normalize(encodings)\n codebook = F.normalize(codebook)\n\n # Compute euclidean distance with codebook\n dist = (\n encodings.pow(2).sum(1, keepdim=True)\n - 2 * encodings @ codebook.t()\n + codebook.pow(2).sum(1, keepdim=True).t()\n )\n indices = rearrange((-dist).max(1)[1], \"(b t) -> b t\", b=latents.size(0))\n z_q = self.decode_code(indices)\n return z_q, indices\n\n\nclass ResidualVectorQuantize(nn.Module):\n \"\"\"\n Introduced in SoundStream: An end2end neural audio codec\n https://arxiv.org/abs/2107.03312\n \"\"\"\n\n def __init__(\n self,\n input_dim: int = 512,\n n_codebooks: int = 9,\n codebook_size: int = 1024,\n codebook_dim: Union[int, list] = 8,\n quantizer_dropout: float = 0.0,\n ):\n super().__init__()\n if isinstance(codebook_dim, int):\n codebook_dim = [codebook_dim for _ in range(n_codebooks)]\n\n self.n_codebooks = n_codebooks\n self.codebook_dim = codebook_dim\n self.codebook_size = codebook_size\n\n self.quantizers = nn.ModuleList(\n [VectorQuantize(input_dim, codebook_size, codebook_dim[i]) for i in range(n_codebooks)]\n )\n self.quantizer_dropout = quantizer_dropout\n\n def forward(self, z, n_quantizers: int = None):\n \"\"\"Quantized the input tensor using a fixed set of `n` codebooks and returns\n the corresponding codebook vectors\n Parameters\n ----------\n z : Tensor[B x D x T]\n n_quantizers : int, optional\n No. of quantizers to use\n (n_quantizers < self.n_codebooks ex: for quantizer dropout)\n Note: if `self.quantizer_dropout` is True, this argument is ignored\n when in training mode, and a random number of quantizers is used.\n Returns\n -------\n dict\n A dictionary with the following keys:\n\n \"z\" : Tensor[B x D x T]\n Quantized continuous representation of input\n \"codes\" : Tensor[B x N x T]\n Codebook indices for each codebook\n (quantized discrete representation of input)\n \"latents\" : Tensor[B x N*D x T]\n Projected latents (continuous representation of input before quantization)\n \"vq/commitment_loss\" : Tensor[1]\n Commitment loss to train encoder to predict vectors closer to codebook\n entries\n \"vq/codebook_loss\" : Tensor[1]\n Codebook loss to update the codebook\n \"\"\"\n z_q = 0\n residual = z\n commitment_loss = 0\n codebook_loss = 0\n\n codebook_indices = []\n latents = []\n\n if n_quantizers is None:\n n_quantizers = self.n_codebooks\n if self.training:\n n_quantizers = torch.ones((z.shape[0],)) * self.n_codebooks + 1\n dropout = torch.randint(1, self.n_codebooks + 1, (z.shape[0],))\n n_dropout = int(z.shape[0] * self.quantizer_dropout)\n n_quantizers[:n_dropout] = dropout[:n_dropout]\n n_quantizers = n_quantizers.to(z.device)\n\n for i, quantizer in enumerate(self.quantizers):\n if self.training is False and i >= n_quantizers:\n break\n\n z_q_i, commitment_loss_i, codebook_loss_i, indices_i, z_e_i = quantizer(residual)\n\n # Create mask to apply quantizer dropout\n mask = torch.full((z.shape[0],), fill_value=i, device=z.device) < n_quantizers\n z_q = z_q + z_q_i * mask[:, None, None]\n residual = residual - z_q_i\n\n # Sum losses\n commitment_loss += (commitment_loss_i * mask).mean()\n codebook_loss += (codebook_loss_i * mask).mean()\n\n codebook_indices.append(indices_i)\n latents.append(z_e_i)\n\n codes = torch.stack(codebook_indices, dim=1)\n latents = torch.cat(latents, dim=1)\n\n return z_q, codes, latents, commitment_loss, codebook_loss\n\n def from_codes(self, codes: torch.Tensor):\n \"\"\"Given the quantized codes, reconstruct the continuous representation\n Parameters\n ----------\n codes : Tensor[B x N x T]\n Quantized discrete representation of input\n Returns\n -------\n Tensor[B x D x T]\n Quantized continuous representation of input\n \"\"\"\n z_q = 0.0\n z_p = []\n n_codebooks = codes.shape[1]\n for i in range(n_codebooks):\n z_p_i = self.quantizers[i].decode_code(codes[:, i, :])\n z_p.append(z_p_i)\n\n z_q_i = self.quantizers[i].out_proj(z_p_i)\n z_q = z_q + z_q_i\n return z_q, torch.cat(z_p, dim=1), codes\n\n def from_latents(self, latents: torch.Tensor):\n \"\"\"Given the unquantized latents, reconstruct the\n continuous representation after quantization.\n\n Parameters\n ----------\n latents : Tensor[B x N x T]\n Continuous representation of input after projection\n\n Returns\n -------\n Tensor[B x D x T]\n Quantized representation of full-projected space\n Tensor[B x D x T]\n Quantized representation of latent space\n \"\"\"\n z_q = 0\n z_p = []\n codes = []\n dims = np.cumsum([0] + [q.codebook_dim for q in self.quantizers])\n\n n_codebooks = np.where(dims <= latents.shape[1])[0].max(axis=0, keepdims=True)[0]\n for i in range(n_codebooks):\n j, k = dims[i], dims[i + 1]\n z_p_i, codes_i = self.quantizers[i].decode_latents(latents[:, j:k, :])\n z_p.append(z_p_i)\n codes.append(codes_i)\n\n z_q_i = self.quantizers[i].out_proj(z_p_i)\n z_q = z_q + z_q_i\n\n return z_q, torch.cat(z_p, dim=1), torch.stack(codes, dim=1)\n\n\nif __name__ == \"__main__\":\n rvq = ResidualVectorQuantize(quantizer_dropout=True)\n x = torch.randn(16, 512, 80)\n y = rvq(x)\n print(y[\"latents\"].shape)\n"
},
{
"path": "boson_multimodal/audio_processing/higgs_audio_tokenizer.py",
"content": "# Based on code from: https://github.com/zhenye234/xcodec\n# Licensed under MIT License\n# Modifications by BosonAI\n\nimport math\nimport os\nimport torch\nimport torch.nn as nn\nimport torch.nn.functional as F\nfrom typing import Optional, Union, Sequence\nimport numpy as np\nfrom transformers import AutoModel\nimport torchaudio\nimport json\nimport librosa\nfrom huggingface_hub import snapshot_download\n\nfrom vector_quantize_pytorch import ResidualFSQ\nfrom .descriptaudiocodec.dac.model import dac as dac2\nfrom .quantization.vq import ResidualVectorQuantizer\nfrom .semantic_module import Encoder, Decoder\n\n\nclass EncodedResult:\n def __init__(self, audio_codes):\n self.audio_codes = audio_codes\n\n\nclass HiggsAudioFeatureExtractor(nn.Module):\n def __init__(self, sampling_rate=16000):\n super().__init__()\n self.sampling_rate = sampling_rate\n\n def forward(self, raw_audio, sampling_rate=16000, return_tensors=\"pt\"):\n # Convert from librosa to torch\n audio_signal = torch.tensor(raw_audio)\n audio_signal = audio_signal.unsqueeze(0)\n if len(audio_signal.shape) < 3:\n audio_signal = audio_signal.unsqueeze(0)\n return {\"input_values\": audio_signal}\n\n\nclass HiggsAudioTokenizer(nn.Module):\n def __init__(\n self,\n n_filters: int = 32,\n D: int = 128,\n target_bandwidths: Sequence[Union[int, float]] = [1, 1.5, 2, 4, 6],\n ratios: Sequence[int] = [8, 5, 4, 2], # downsampling by 320\n sample_rate: int = 16000,\n bins: int = 1024,\n n_q: int = 8,\n codebook_dim: int = None,\n normalize: bool = False,\n causal: bool = False,\n semantic_techer: str = \"hubert_base_general\",\n last_layer_semantic: bool = True,\n merge_mode: str = \"concat\",\n downsample_mode: str = \"step_down\",\n semantic_mode: str = \"classic\",\n vq_scale: int = 1,\n semantic_sample_rate: int = None,\n device: str = \"cuda\",\n ):\n super().__init__()\n self.hop_length = np.prod(ratios)\n self.semantic_techer = semantic_techer\n\n self.frame_rate = math.ceil(sample_rate / np.prod(ratios)) # 50 Hz\n\n self.target_bandwidths = target_bandwidths\n self.n_q = n_q\n self.sample_rate = sample_rate\n self.encoder = dac2.Encoder(64, ratios, D)\n\n self.decoder_2 = dac2.Decoder(D, 1024, ratios)\n self.last_layer_semantic = last_layer_semantic\n self.device = device\n if semantic_techer == \"hubert_base\":\n self.semantic_model = AutoModel.from_pretrained(\"facebook/hubert-base-ls960\")\n self.semantic_sample_rate = 16000\n self.semantic_dim = 768\n self.encoder_semantic_dim = 768\n\n elif semantic_techer == \"wavlm_base_plus\":\n self.semantic_model = AutoModel.from_pretrained(\"microsoft/wavlm-base-plus\")\n self.semantic_sample_rate = 16000\n self.semantic_dim = 768\n self.encoder_semantic_dim = 768\n\n elif semantic_techer == \"hubert_base_general\":\n self.semantic_model = AutoModel.from_pretrained(\"bosonai/hubert_base\", trust_remote_code=True)\n self.semantic_sample_rate = 16000\n self.semantic_dim = 768\n self.encoder_semantic_dim = 768\n\n # Overwrite semantic model sr to ensure semantic_downsample_factor is an integer\n if semantic_sample_rate is not None:\n self.semantic_sample_rate = semantic_sample_rate\n\n self.semantic_model.eval()\n\n # make the semantic model parameters do not need gradient\n for param in self.semantic_model.parameters():\n param.requires_grad = False\n\n self.semantic_downsample_factor = int(self.hop_length / (self.sample_rate / self.semantic_sample_rate) / 320)\n\n self.quantizer_dim = int((D + self.encoder_semantic_dim) // vq_scale)\n self.encoder_semantic = Encoder(input_channels=self.semantic_dim, encode_channels=self.encoder_semantic_dim)\n self.decoder_semantic = Decoder(\n code_dim=self.encoder_semantic_dim, output_channels=self.semantic_dim, decode_channels=self.semantic_dim\n )\n\n # out_D=D+768\n if isinstance(bins, int): # RVQ\n self.quantizer = ResidualVectorQuantizer(\n dimension=self.quantizer_dim, codebook_dim=codebook_dim, n_q=n_q, bins=bins\n )\n self.quantizer_type = \"RVQ\"\n else: # RFSQ\n self.quantizer = ResidualFSQ(dim=self.quantizer_dim, levels=bins, num_quantizers=n_q)\n self.quantizer_type = \"RFSQ\"\n\n self.fc_prior = nn.Linear(D + self.encoder_semantic_dim, self.quantizer_dim)\n self.fc_post1 = nn.Linear(self.quantizer_dim, self.encoder_semantic_dim)\n self.fc_post2 = nn.Linear(self.quantizer_dim, D)\n\n self.downsample_mode = downsample_mode\n if downsample_mode == \"avg\":\n self.semantic_pooling = nn.AvgPool1d(\n kernel_size=self.semantic_downsample_factor, stride=self.semantic_downsample_factor\n )\n\n self.audio_tokenizer_feature_extractor = HiggsAudioFeatureExtractor(sampling_rate=self.sample_rate)\n\n @property\n def tps(self):\n return self.frame_rate\n\n @property\n def sampling_rate(self):\n return self.sample_rate\n\n @property\n def num_codebooks(self):\n return self.n_q\n\n @property\n def codebook_size(self):\n return self.quantizer_dim\n\n def get_last_layer(self):\n return self.decoder.layers[-1].weight\n\n def calculate_rec_loss(self, rec, target):\n target = target / target.norm(dim=-1, keepdim=True)\n rec = rec / rec.norm(dim=-1, keepdim=True)\n rec_loss = (1 - (target * rec).sum(-1)).mean()\n\n return rec_loss\n\n @torch.no_grad()\n def get_regress_target(self, x):\n x = torchaudio.functional.resample(x, self.sample_rate, self.semantic_sample_rate)\n\n if (\n self.semantic_techer == \"hubert_base\"\n or self.semantic_techer == \"hubert_base_general\"\n or self.semantic_techer == \"wavlm_base_plus\"\n ):\n x = x[:, 0, :]\n x = F.pad(x, (160, 160))\n target = self.semantic_model(x, output_hidden_states=True).hidden_states\n target = torch.stack(target, dim=1) # .transpose(-1, -2)#.flatten(start_dim=1, end_dim=2)\n\n # average for all layers\n target = target.mean(1)\n # target = target[9]\n # if self.hop_length > 320:\n # target = self.semantic_pooling(target.transpose(1, 2)).transpose(1, 2)\n\n elif self.semantic_techer == \"w2v_bert2\":\n target = self.semantic_model(x)\n\n elif self.semantic_techer.startswith(\"whisper\"):\n if self.last_layer_semantic:\n target = self.semantic_model(x, avg_layers=False)\n else:\n target = self.semantic_model(x, avg_layers=True)\n\n elif self.semantic_techer.startswith(\"mert_music\"):\n if self.last_layer_semantic:\n target = self.semantic_model(x, avg_layers=False)\n else:\n target = self.semantic_model(x, avg_layers=True)\n\n elif self.semantic_techer.startswith(\"qwen_audio_omni\"):\n target = self.semantic_model(x)\n\n if self.downsample_mode == \"step_down\":\n if self.semantic_downsample_factor > 1:\n target = target[:, :: self.semantic_downsample_factor, :]\n\n elif self.downsample_mode == \"avg\":\n target = self.semantic_pooling(target.transpose(1, 2)).transpose(1, 2)\n return target\n\n def forward(self, x: torch.Tensor, bw: int):\n e_semantic_input = self.get_regress_target(x).detach()\n\n e_semantic = self.encoder_semantic(e_semantic_input.transpose(1, 2))\n e_acoustic = self.encoder(x)\n\n e = torch.cat([e_acoustic, e_semantic], dim=1)\n\n e = self.fc_prior(e.transpose(1, 2))\n\n if self.quantizer_type == \"RVQ\":\n e = e.transpose(1, 2)\n quantized, codes, bandwidth, commit_loss = self.quantizer(e, self.frame_rate, bw)\n quantized = quantized.transpose(1, 2)\n else:\n quantized, codes = self.quantizer(e)\n commit_loss = torch.tensor(0.0)\n\n quantized_semantic = self.fc_post1(quantized).transpose(1, 2)\n quantized_acoustic = self.fc_post2(quantized).transpose(1, 2)\n\n o = self.decoder_2(quantized_acoustic)\n\n o_semantic = self.decoder_semantic(quantized_semantic)\n semantic_recon_loss = F.mse_loss(e_semantic_input.transpose(1, 2).detach(), o_semantic)\n\n return o, commit_loss, semantic_recon_loss, None\n\n def encode(self, audio_path_or_wv, sr=None, loudness_normalize=False, loudness_threshold=-23.0):\n if isinstance(audio_path_or_wv, str):\n wv, sr = librosa.load(audio_path_or_wv, mono=True, sr=None)\n else:\n wv = audio_path_or_wv\n assert sr is not None\n if loudness_normalize:\n import pyloudnorm as pyln\n\n meter = pyln.Meter(sr)\n l = meter.integrated_loudness(wv)\n wv = pyln.normalize.loudness(wv, l, loudness_threshold)\n if sr != self.sampling_rate:\n wv = librosa.resample(wv, orig_sr=sr, target_sr=self.sampling_rate)\n if self.audio_tokenizer_feature_extractor is not None:\n inputs = self.audio_tokenizer_feature_extractor(\n raw_audio=wv, sampling_rate=self.audio_tokenizer_feature_extractor.sampling_rate, return_tensors=\"pt\"\n )\n input_values = inputs[\"input_values\"].to(self.device)\n else:\n input_values = torch.from_numpy(wv).float().unsqueeze(0)\n with torch.no_grad():\n encoder_outputs = self._xcodec_encode(input_values)\n vq_code = encoder_outputs.audio_codes[0]\n return vq_code\n\n def _xcodec_encode(self, x: torch.Tensor, target_bw: Optional[int] = None) -> torch.Tensor:\n bw = target_bw\n\n e_semantic_input = self.get_regress_target(x).detach()\n\n e_semantic = self.encoder_semantic(e_semantic_input.transpose(1, 2))\n e_acoustic = self.encoder(x)\n\n if e_acoustic.shape[2] != e_semantic.shape[2]:\n pad_size = 160 * self.semantic_downsample_factor\n e_acoustic = self.encoder(F.pad(x[:, 0, :], (pad_size, pad_size)).unsqueeze(0))\n\n if e_acoustic.shape[2] != e_semantic.shape[2]:\n if e_acoustic.shape[2] > e_semantic.shape[2]:\n e_acoustic = e_acoustic[:, :, : e_semantic.shape[2]]\n else:\n e_semantic = e_semantic[:, :, : e_acoustic.shape[2]]\n\n e = torch.cat([e_acoustic, e_semantic], dim=1)\n\n e = self.fc_prior(e.transpose(1, 2))\n\n if self.quantizer_type == \"RVQ\":\n e = e.transpose(1, 2)\n quantized, codes, bandwidth, commit_loss = self.quantizer(e, self.frame_rate, bw)\n codes = codes.permute(1, 0, 2)\n else:\n quantized, codes = self.quantizer(e)\n codes = codes.permute(0, 2, 1)\n\n # return codes\n return EncodedResult(codes)\n\n def decode(self, vq_code: torch.Tensor) -> torch.Tensor:\n vq_code = vq_code.to(self.device)\n\n if self.quantizer_type == \"RVQ\":\n vq_code = vq_code.permute(1, 0, 2)\n quantized = self.quantizer.decode(vq_code)\n quantized = quantized.transpose(1, 2)\n else:\n vq_code = vq_code.permute(0, 2, 1)\n quantized = self.quantizer.get_output_from_indices(vq_code)\n quantized_acoustic = self.fc_post2(quantized).transpose(1, 2)\n\n o = self.decoder_2(quantized_acoustic)\n return o.detach().cpu().numpy()\n\n\ndef load_higgs_audio_tokenizer(tokenizer_name_or_path, device=\"cuda\"):\n is_local = os.path.exists(tokenizer_name_or_path)\n if not is_local:\n tokenizer_path = snapshot_download(tokenizer_name_or_path)\n else:\n tokenizer_path = tokenizer_name_or_path\n config_path = os.path.join(tokenizer_path, \"config.json\")\n model_path = os.path.join(tokenizer_path, \"model.pth\")\n config = json.load(open(config_path))\n model = HiggsAudioTokenizer(\n **config,\n device=device,\n )\n parameter_dict = torch.load(model_path, map_location=device)\n model.load_state_dict(parameter_dict, strict=False)\n model.to(device)\n model.eval()\n return model\n"
},
{
"path": "boson_multimodal/audio_processing/quantization/__init__.py",
"content": "# Copyright (c) Meta Platforms, Inc. and affiliates.\n# All rights reserved.\n#\n# This source code is licensed under the license found in the\n# LICENSE file in the root directory of this source tree.\n\n# flake8: noqa\nfrom .vq import QuantizedResult, ResidualVectorQuantizer\n"
},
{
"path": "boson_multimodal/audio_processing/quantization/ac.py",
"content": "# Copyright (c) Meta Platforms, Inc. and affiliates.\n# All rights reserved.\n#\n# This source code is licensed under the license found in the\n# LICENSE file in the root directory of this source tree.\n\n\"\"\"Arithmetic coder.\"\"\"\n\nimport io\nimport math\nimport random\nimport typing as tp\nimport torch\n\nfrom ..binary import BitPacker, BitUnpacker\n\n\ndef build_stable_quantized_cdf(\n pdf: torch.Tensor, total_range_bits: int, roundoff: float = 1e-8, min_range: int = 2, check: bool = True\n) -> torch.Tensor:\n \"\"\"Turn the given PDF into a quantized CDF that splits\n [0, 2 ** self.total_range_bits - 1] into chunks of size roughly proportional\n to the PDF.\n\n Args:\n pdf (torch.Tensor): probability distribution, shape should be `[N]`.\n total_range_bits (int): see `ArithmeticCoder`, the typical range we expect\n during the coding process is `[0, 2 ** total_range_bits - 1]`.\n roundoff (float): will round the pdf up to that level to remove difference coming\n from e.g. evaluating the Language Model on different architectures.\n min_range (int): minimum range width. Should always be at least 2 for numerical\n stability. Use this to avoid pathological behavior is a value\n that is expected to be rare actually happens in real life.\n check (bool): if True, checks that nothing bad happened, can be deactivated for speed.\n \"\"\"\n pdf = pdf.detach()\n if roundoff:\n pdf = (pdf / roundoff).floor() * roundoff\n # interpolate with uniform distribution to achieve desired minimum probability.\n total_range = 2**total_range_bits\n cardinality = len(pdf)\n alpha = min_range * cardinality / total_range\n assert alpha <= 1, \"you must reduce min_range\"\n ranges = (((1 - alpha) * total_range) * pdf).floor().long()\n ranges += min_range\n quantized_cdf = torch.cumsum(ranges, dim=-1)\n if min_range < 2:\n raise ValueError(\"min_range must be at least 2.\")\n if check:\n assert quantized_cdf[-1] <= 2**total_range_bits, quantized_cdf[-1]\n if ((quantized_cdf[1:] - quantized_cdf[:-1]) < min_range).any() or quantized_cdf[0] < min_range:\n raise ValueError(\"You must increase your total_range_bits.\")\n return quantized_cdf\n\n\nclass ArithmeticCoder:\n \"\"\"ArithmeticCoder,\n Let us take a distribution `p` over `N` symbols, and assume we have a stream\n of random variables `s_t` sampled from `p`. Let us assume that we have a budget\n of `B` bits that we can afford to write on device. There are `2**B` possible numbers,\n corresponding to the range `[0, 2 ** B - 1]`. We can map each of those number to a single\n sequence `(s_t)` by doing the following:\n\n 1) Initialize the current range to` [0 ** 2 B - 1]`.\n 2) For each time step t, split the current range into contiguous chunks,\n one for each possible outcome, with size roughly proportional to `p`.\n For instance, if `p = [0.75, 0.25]`, and the range is `[0, 3]`, the chunks\n would be `{[0, 2], [3, 3]}`.\n 3) Select the chunk corresponding to `s_t`, and replace the current range with this.\n 4) When done encoding all the values, just select any value remaining in the range.\n\n You will notice that this procedure can fail: for instance if at any point in time\n the range is smaller than `N`, then we can no longer assign a non-empty chunk to each\n possible outcome. Intuitively, the more likely a value is, the less the range width\n will reduce, and the longer we can go on encoding values. This makes sense: for any efficient\n coding scheme, likely outcomes would take less bits, and more of them can be coded\n with a fixed budget.\n\n In practice, we do not know `B` ahead of time, but we have a way to inject new bits\n when the current range decreases below a given limit (given by `total_range_bits`), without\n having to redo all the computations. If we encode mostly likely values, we will seldom\n need to inject new bits, but a single rare value can deplete our stock of entropy!\n\n In this explanation, we assumed that the distribution `p` was constant. In fact, the present\n code works for any sequence `(p_t)` possibly different for each timestep.\n We also assume that `s_t ~ p_t`, but that doesn't need to be true, although the smaller\n the KL between the true distribution and `p_t`, the most efficient the coding will be.\n\n Args:\n fo (IO[bytes]): file-like object to which the bytes will be written to.\n total_range_bits (int): the range `M` described above is `2 ** total_range_bits.\n Any time the current range width fall under this limit, new bits will\n be injected to rescale the initial range.\n \"\"\"\n\n def __init__(self, fo: tp.IO[bytes], total_range_bits: int = 24):\n assert total_range_bits <= 30\n self.total_range_bits = total_range_bits\n self.packer = BitPacker(bits=1, fo=fo) # we push single bits at a time.\n self.low: int = 0\n self.high: int = 0\n self.max_bit: int = -1\n self._dbg: tp.List[tp.Any] = []\n self._dbg2: tp.List[tp.Any] = []\n\n @property\n def delta(self) -> int:\n \"\"\"Return the current range width.\"\"\"\n return self.high - self.low + 1\n\n def _flush_common_prefix(self):\n # If self.low and self.high start with the sames bits,\n # those won't change anymore as we always just increase the range\n # by powers of 2, and we can flush them out to the bit stream.\n assert self.high >= self.low, (self.low, self.high)\n assert self.high < 2 ** (self.max_bit + 1)\n while self.max_bit >= 0:\n b1 = self.low >> self.max_bit\n b2 = self.high >> self.max_bit\n if b1 == b2:\n self.low -= b1 << self.max_bit\n self.high -= b1 << self.max_bit\n assert self.high >= self.low, (self.high, self.low, self.max_bit)\n assert self.low >= 0\n self.max_bit -= 1\n self.packer.push(b1)\n else:\n break\n\n def push(self, symbol: int, quantized_cdf: torch.Tensor):\n \"\"\"Push the given symbol on the stream, flushing out bits\n if possible.\n\n Args:\n symbol (int): symbol to encode with the AC.\n quantized_cdf (torch.Tensor): use `build_stable_quantized_cdf`\n to build this from your pdf estimate.\n \"\"\"\n while self.delta < 2**self.total_range_bits:\n self.low *= 2\n self.high = self.high * 2 + 1\n self.max_bit += 1\n\n range_low = 0 if symbol == 0 else quantized_cdf[symbol - 1].item()\n range_high = quantized_cdf[symbol].item() - 1\n effective_low = int(math.ceil(range_low * (self.delta / (2**self.total_range_bits))))\n effective_high = int(math.floor(range_high * (self.delta / (2**self.total_range_bits))))\n assert self.low <= self.high\n self.high = self.low + effective_high\n self.low = self.low + effective_low\n assert self.low <= self.high, (effective_low, effective_high, range_low, range_high)\n self._dbg.append((self.low, self.high))\n self._dbg2.append((self.low, self.high))\n outs = self._flush_common_prefix()\n assert self.low <= self.high\n assert self.max_bit >= -1\n assert self.max_bit <= 61, self.max_bit\n return outs\n\n def flush(self):\n \"\"\"Flush the remaining information to the stream.\"\"\"\n while self.max_bit >= 0:\n b1 = (self.low >> self.max_bit) & 1\n self.packer.push(b1)\n self.max_bit -= 1\n self.packer.flush()\n\n\nclass ArithmeticDecoder:\n \"\"\"ArithmeticDecoder, see `ArithmeticCoder` for a detailed explanation.\n\n Note that this must be called with **exactly** the same parameters and sequence\n of quantized cdf as the arithmetic encoder or the wrong values will be decoded.\n\n If the AC encoder current range is [L, H], with `L` and `H` having the some common\n prefix (i.e. the same most significant bits), then this prefix will be flushed to the stream.\n For instances, having read 3 bits `b1 b2 b3`, we know that `[L, H]` is contained inside\n `[b1 b2 b3 0 ... 0 b1 b3 b3 1 ... 1]`. Now this specific sub-range can only be obtained\n for a specific sequence of symbols and a binary-search allows us to decode those symbols.\n At some point, the prefix `b1 b2 b3` will no longer be sufficient to decode new symbols,\n and we will need to read new bits from the stream and repeat the process.\n\n \"\"\"\n\n def __init__(self, fo: tp.IO[bytes], total_range_bits: int = 24):\n self.total_range_bits = total_range_bits\n self.low: int = 0\n self.high: int = 0\n self.current: int = 0\n self.max_bit: int = -1\n self.unpacker = BitUnpacker(bits=1, fo=fo) # we pull single bits at a time.\n # Following is for debugging\n self._dbg: tp.List[tp.Any] = []\n self._dbg2: tp.List[tp.Any] = []\n self._last: tp.Any = None\n\n @property\n def delta(self) -> int:\n return self.high - self.low + 1\n\n def _flush_common_prefix(self):\n # Given the current range [L, H], if both have a common prefix,\n # we know we can remove it from our representation to avoid handling large numbers.\n while self.max_bit >= 0:\n b1 = self.low >> self.max_bit\n b2 = self.high >> self.max_bit\n if b1 == b2:\n self.low -= b1 << self.max_bit\n self.high -= b1 << self.max_bit\n self.current -= b1 << self.max_bit\n assert self.high >= self.low\n assert self.low >= 0\n self.max_bit -= 1\n else:\n break\n\n def pull(self, quantized_cdf: torch.Tensor) -> tp.Optional[int]:\n \"\"\"Pull a symbol, reading as many bits from the stream as required.\n This returns `None` when the stream has been exhausted.\n\n Args:\n quantized_cdf (torch.Tensor): use `build_stable_quantized_cdf`\n to build this from your pdf estimate. This must be **exatly**\n the same cdf as the one used at encoding time.\n \"\"\"\n while self.delta < 2**self.total_range_bits:\n bit = self.unpacker.pull()\n if bit is None:\n return None\n self.low *= 2\n self.high = self.high * 2 + 1\n self.current = self.current * 2 + bit\n self.max_bit += 1\n\n def bin_search(low_idx: int, high_idx: int):\n # Binary search is not just for coding interviews :)\n if high_idx < low_idx:\n raise RuntimeError(\"Binary search failed\")\n mid = (low_idx + high_idx) // 2\n range_low = quantized_cdf[mid - 1].item() if mid > 0 else 0\n range_high = quantized_cdf[mid].item() - 1\n effective_low = int(math.ceil(range_low * (self.delta / (2**self.total_range_bits))))\n effective_high = int(math.floor(range_high * (self.delta / (2**self.total_range_bits))))\n low = effective_low + self.low\n high = effective_high + self.low\n if self.current >= low:\n if self.current <= high:\n return (mid, low, high, self.current)\n else:\n return bin_search(mid + 1, high_idx)\n else:\n return bin_search(low_idx, mid - 1)\n\n self._last = (self.low, self.high, self.current, self.max_bit)\n sym, self.low, self.high, self.current = bin_search(0, len(quantized_cdf) - 1)\n self._dbg.append((self.low, self.high, self.current))\n self._flush_common_prefix()\n self._dbg2.append((self.low, self.high, self.current))\n\n return sym\n\n\ndef test():\n torch.manual_seed(1234)\n random.seed(1234)\n for _ in range(4):\n pdfs = []\n cardinality = random.randrange(4000)\n steps = random.randrange(100, 500)\n fo = io.BytesIO()\n encoder = ArithmeticCoder(fo)\n symbols = []\n for step in range(steps):\n pdf = torch.softmax(torch.randn(cardinality), dim=0)\n pdfs.append(pdf)\n q_cdf = build_stable_quantized_cdf(pdf, encoder.total_range_bits)\n symbol = torch.multinomial(pdf, 1).item()\n symbols.append(symbol)\n encoder.push(symbol, q_cdf)\n encoder.flush()\n\n fo.seek(0)\n decoder = ArithmeticDecoder(fo)\n for idx, (pdf, symbol) in enumerate(zip(pdfs, symbols)):\n q_cdf = build_stable_quantized_cdf(pdf, encoder.total_range_bits)\n decoded_symbol = decoder.pull(q_cdf)\n assert decoded_symbol == symbol, idx\n assert decoder.pull(torch.zeros(1)) is None\n\n\nif __name__ == \"__main__\":\n test()\n"
},
{
"path": "boson_multimodal/audio_processing/quantization/core_vq.py",
"content": "# Copyright (c) Meta Platforms, Inc. and affiliates.\n# All rights reserved.\n#\n# This source code is licensed under the license found in the\n# LICENSE file in the root directory of this source tree.\n#\n# This implementation is inspired from\n# https://github.com/lucidrains/vector-quantize-pytorch\n# which is released under MIT License. Hereafter, the original license:\n# MIT License\n#\n# Copyright (c) 2020 Phil Wang\n#\n# Permission is hereby granted, free of charge, to any person obtaining a copy\n# of this software and associated documentation files (the \"Software\"), to deal\n# in the Software without restriction, including without limitation the rights\n# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell\n# copies of the Software, and to permit persons to whom the Software is\n# furnished to do so, subject to the following conditions:\n#\n# The above copyright notice and this permission notice shall be included in all\n# copies or substantial portions of the Software.\n#\n# THE SOFTWARE IS PROVIDED \"AS IS\", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR\n# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,\n# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE\n# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER\n# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,\n# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE\n# SOFTWARE.\n\n\"\"\"Core vector quantization implementation.\"\"\"\n\nimport typing as tp\n\nfrom einops import rearrange, repeat\nimport torch\nfrom torch import nn\nimport torch.nn.functional as F\n\nfrom xcodec.quantization.distrib import broadcast_tensors, rank\n\n\ndef default(val: tp.Any, d: tp.Any) -> tp.Any:\n return val if val is not None else d\n\n\ndef ema_inplace(moving_avg, new, decay: float):\n moving_avg.data.mul_(decay).add_(new, alpha=(1 - decay))\n\n\ndef laplace_smoothing(x, n_categories: int, epsilon: float = 1e-5):\n return (x + epsilon) / (x.sum() + n_categories * epsilon)\n\n\ndef uniform_init(*shape: int):\n t = torch.empty(shape)\n nn.init.kaiming_uniform_(t)\n return t\n\n\ndef sample_vectors(samples, num: int):\n num_samples, device = samples.shape[0], samples.device\n\n if num_samples >= num:\n indices = torch.randperm(num_samples, device=device)[:num]\n else:\n indices = torch.randint(0, num_samples, (num,), device=device)\n\n return samples[indices]\n\n\ndef kmeans(samples, num_clusters: int, num_iters: int = 10):\n dim, dtype = samples.shape[-1], samples.dtype\n\n means = sample_vectors(samples, num_clusters)\n\n for _ in range(num_iters):\n diffs = rearrange(samples, \"n d -> n () d\") - rearrange(means, \"c d -> () c d\")\n dists = -(diffs**2).sum(dim=-1)\n\n buckets = dists.max(dim=-1).indices\n bins = torch.bincount(buckets, minlength=num_clusters)\n zero_mask = bins == 0\n bins_min_clamped = bins.masked_fill(zero_mask, 1)\n\n new_means = buckets.new_zeros(num_clusters, dim, dtype=dtype)\n new_means.scatter_add_(0, repeat(buckets, \"n -> n d\", d=dim), samples)\n new_means = new_means / bins_min_clamped[..., None]\n\n means = torch.where(zero_mask[..., None], means, new_means)\n\n return means, bins\n\n\nclass EuclideanCodebook(nn.Module):\n \"\"\"Codebook with Euclidean distance.\n Args:\n dim (int): Dimension.\n codebook_size (int): Codebook size.\n kmeans_init (bool): Whether to use k-means to initialize the codebooks.\n If set to true, run the k-means algorithm on the first training batch and use\n the learned centroids as initialization.\n kmeans_iters (int): Number of iterations used for k-means algorithm at initialization.\n decay (float): Decay for exponential moving average over the codebooks.\n epsilon (float): Epsilon value for numerical stability.\n threshold_ema_dead_code (int): Threshold for dead code expiration. Replace any codes\n that have an exponential moving average cluster size less than the specified threshold with\n randomly selected vector from the current batch.\n \"\"\"\n\n def __init__(\n self,\n dim: int,\n codebook_size: int,\n kmeans_init: int = False,\n kmeans_iters: int = 10,\n decay: float = 0.99,\n epsilon: float = 1e-5,\n threshold_ema_dead_code: int = 2,\n ):\n super().__init__()\n self.decay = decay\n init_fn: tp.Union[tp.Callable[..., torch.Tensor], tp.Any] = uniform_init if not kmeans_init else torch.zeros\n embed = init_fn(codebook_size, dim)\n\n self.codebook_size = codebook_size\n\n self.kmeans_iters = kmeans_iters\n self.epsilon = epsilon\n self.threshold_ema_dead_code = threshold_ema_dead_code\n\n self.register_buffer(\"inited\", torch.Tensor([not kmeans_init]))\n self.register_buffer(\"cluster_size\", torch.zeros(codebook_size))\n self.register_buffer(\"embed\", embed)\n self.register_buffer(\"embed_avg\", embed.clone())\n\n @torch.jit.ignore\n def init_embed_(self, data):\n if self.inited:\n return\n\n embed, cluster_size = kmeans(data, self.codebook_size, self.kmeans_iters)\n self.embed.data.copy_(embed)\n self.embed_avg.data.copy_(embed.clone())\n self.cluster_size.data.copy_(cluster_size)\n self.inited.data.copy_(torch.Tensor([True]))\n # Make sure all buffers across workers are in sync after initialization\n broadcast_tensors(self.buffers())\n\n def replace_(self, samples, mask):\n modified_codebook = torch.where(mask[..., None], sample_vectors(samples, self.codebook_size), self.embed)\n self.embed.data.copy_(modified_codebook)\n\n def expire_codes_(self, batch_samples):\n if self.threshold_ema_dead_code == 0:\n return\n\n expired_codes = self.cluster_size < self.threshold_ema_dead_code\n if not torch.any(expired_codes):\n return\n\n batch_samples = rearrange(batch_samples, \"... d -> (...) d\")\n self.replace_(batch_samples, mask=expired_codes)\n broadcast_tensors(self.buffers())\n\n def preprocess(self, x):\n x = rearrange(x, \"... d -> (...) d\")\n return x\n\n def quantize(self, x):\n embed = self.embed.t()\n dist = -(x.pow(2).sum(1, keepdim=True) - 2 * x @ embed + embed.pow(2).sum(0, keepdim=True))\n embed_ind = dist.max(dim=-1).indices\n return embed_ind\n\n def postprocess_emb(self, embed_ind, shape):\n return embed_ind.view(*shape[:-1])\n\n def dequantize(self, embed_ind):\n quantize = F.embedding(embed_ind, self.embed) # get embedding based on index\n return quantize\n\n def encode(self, x):\n shape = x.shape\n # pre-process\n x = self.preprocess(x)\n # quantize\n embed_ind = self.quantize(x) # get index based on Euclidean distance\n # post-process\n embed_ind = self.postprocess_emb(embed_ind, shape)\n return embed_ind\n\n def decode(self, embed_ind):\n quantize = self.dequantize(embed_ind)\n return quantize\n\n def forward(self, x):\n shape, dtype = x.shape, x.dtype\n x = self.preprocess(x)\n\n self.init_embed_(x)\n\n embed_ind = self.quantize(x)\n embed_onehot = F.one_hot(embed_ind, self.codebook_size).type(dtype)\n embed_ind = self.postprocess_emb(embed_ind, shape)\n quantize = self.dequantize(embed_ind)\n\n if self.training:\n # We do the expiry of code at that point as buffers are in sync\n # and all the workers will take the same decision.\n self.expire_codes_(x)\n ema_inplace(self.cluster_size, embed_onehot.sum(0), self.decay)\n embed_sum = x.t() @ embed_onehot\n ema_inplace(self.embed_avg, embed_sum.t(), self.decay)\n cluster_size = (\n laplace_smoothing(self.cluster_size, self.codebook_size, self.epsilon) * self.cluster_size.sum()\n )\n embed_normalized = self.embed_avg / cluster_size.unsqueeze(1)\n self.embed.data.copy_(embed_normalized)\n\n return quantize, embed_ind\n\n\nclass VectorQuantization(nn.Module):\n \"\"\"Vector quantization implementation.\n Currently supports only euclidean distance.\n Args:\n dim (int): Dimension\n codebook_size (int): Codebook size\n codebook_dim (int): Codebook dimension. If not defined, uses the specified dimension in dim.\n decay (float): Decay for exponential moving average over the codebooks.\n epsilon (float): Epsilon value for numerical stability.\n kmeans_init (bool): Whether to use kmeans to initialize the codebooks.\n kmeans_iters (int): Number of iterations used for kmeans initialization.\n threshold_ema_dead_code (int): Threshold for dead code expiration. Replace any codes\n that have an exponential moving average cluster size less than the specified threshold with\n randomly selected vector from the current batch.\n commitment_weight (float): Weight for commitment loss.\n \"\"\"\n\n def __init__(\n self,\n dim: int,\n codebook_size: int,\n codebook_dim: tp.Optional[int] = None,\n decay: float = 0.99,\n epsilon: float = 1e-5,\n kmeans_init: bool = True,\n kmeans_iters: int = 50,\n threshold_ema_dead_code: int = 2,\n commitment_weight: float = 1.0,\n ):\n super().__init__()\n _codebook_dim: int = default(codebook_dim, dim)\n\n requires_projection = _codebook_dim != dim\n self.project_in = nn.Linear(dim, _codebook_dim) if requires_projection else nn.Identity()\n self.project_out = nn.Linear(_codebook_dim, dim) if requires_projection else nn.Identity()\n\n self.epsilon = epsilon\n self.commitment_weight = commitment_weight\n\n self._codebook = EuclideanCodebook(\n dim=_codebook_dim,\n codebook_size=codebook_size,\n kmeans_init=kmeans_init,\n kmeans_iters=kmeans_iters,\n decay=decay,\n epsilon=epsilon,\n threshold_ema_dead_code=threshold_ema_dead_code,\n )\n self.codebook_size = codebook_size\n\n @property\n def codebook(self):\n return self._codebook.embed\n\n def encode(self, x):\n x = rearrange(x, \"b d n -> b n d\")\n x = self.project_in(x)\n embed_in = self._codebook.encode(x)\n return embed_in\n\n def decode(self, embed_ind):\n quantize = self._codebook.decode(embed_ind)\n quantize = self.project_out(quantize)\n quantize = rearrange(quantize, \"b n d -> b d n\")\n return quantize\n\n def forward(self, x):\n device = x.device\n x = rearrange(x, \"b d n -> b n d\")\n x = self.project_in(x)\n\n quantize, embed_ind = self._codebook(x)\n\n if self.training:\n quantize = x + (quantize - x).detach()\n\n loss = torch.tensor([0.0], device=device, requires_grad=self.training)\n\n if self.training:\n if self.commitment_weight > 0:\n commit_loss = F.mse_loss(quantize.detach(), x)\n loss = loss + commit_loss * self.commitment_weight\n\n quantize = self.project_out(quantize)\n quantize = rearrange(quantize, \"b n d -> b d n\")\n return quantize, embed_ind, loss\n\n\nclass ResidualVectorQuantization(nn.Module):\n \"\"\"Residual vector quantization implementation.\n Follows Algorithm 1. in https://arxiv.org/pdf/2107.03312.pdf\n \"\"\"\n\n def __init__(self, *, num_quantizers, **kwargs):\n super().__init__()\n self.layers = nn.ModuleList([VectorQuantization(**kwargs) for _ in range(num_quantizers)])\n\n def forward(self, x, n_q: tp.Optional[int] = None):\n quantized_out = 0.0\n residual = x\n\n all_losses = []\n all_indices = []\n\n n_q = n_q or len(self.layers)\n\n for layer in self.layers[:n_q]:\n quantized, indices, loss = layer(residual)\n residual = residual - quantized\n quantized_out = quantized_out + quantized\n\n all_indices.append(indices)\n all_losses.append(loss)\n\n out_losses, out_indices = map(torch.stack, (all_losses, all_indices))\n return quantized_out, out_indices, out_losses\n\n def encode(self, x: torch.Tensor, n_q: tp.Optional[int] = None) -> torch.Tensor:\n residual = x\n all_indices = []\n n_q = n_q or len(self.layers)\n for layer in self.layers[:n_q]:\n indices = layer.encode(residual)\n quantized = layer.decode(indices)\n residual = residual - quantized\n all_indices.append(indices)\n out_indices = torch.stack(all_indices)\n return out_indices\n\n def decode(self, q_indices: torch.Tensor) -> torch.Tensor:\n quantized_out = torch.tensor(0.0, device=q_indices.device)\n for i, indices in enumerate(q_indices):\n layer = self.layers[i]\n quantized = layer.decode(indices)\n quantized_out = quantized_out + quantized\n return quantized_out\n"
},
{
"path": "boson_multimodal/audio_processing/quantization/core_vq_lsx_version.py",
"content": "# Copyright (c)\n#\n# This source code is licensed under the license found in the\n# LICENSE file in the root directory of this source tree.\n# This implementation is inspired from\n# https://github.com/rosinality/vq-vae-2-pytorch/blob/master/vqvae.py and\n# https://github.com/clementchadebec/benchmark_VAE/blob/dfa0dcf6c79172df5d27769c09c860c42008baaa/src/pythae/models/vq_vae/vq_vae_utils.py#L81\n#\n# Copyright (c) Meta Platforms, Inc. and affiliates.\n# All rights reserved.\n#\n# This source code is licensed under the license found in the\n# LICENSE file in the root directory of this source tree.\n#\n# This implementation is inspired from\n# https://github.com/lucidrains/vector-quantize-pytorch\n# which is released under MIT License. Hereafter, the original license:\n# MIT License\n#\n# Copyright (c) 2020 Phil Wang\n#\n# Permission is hereby granted, free of charge, to any person obtaining a copy\n# of this software and associated documentation files (the \"Software\"), to deal\n# in the Software without restriction, including without limitation the rights\n# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell\n# copies of the Software, and to permit persons to whom the Software is\n# furnished to do so, subject to the following conditions:\n#\n# The above copyright notice and this permission notice shall be included in all\n# copies or substantial portions of the Software.\n#\n# THE SOFTWARE IS PROVIDED \"AS IS\", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR\n# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,\n# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE\n# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER\n# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,\n# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE\n# SOFTWARE.\n\n\"\"\"Core vector quantization implementation.\"\"\"\n\nimport typing as tp\n\nfrom einops import rearrange\nimport torch\nfrom torch import nn\nimport torch.nn.functional as F\nimport torch.distributed as dist\n\nfrom .distrib import broadcast_tensors, is_distributed\nfrom .ddp_utils import SyncFunction\n\n\ndef default(val: tp.Any, d: tp.Any) -> tp.Any:\n return val if val is not None else d\n\n\ndef ema_inplace(moving_avg, new, decay: float):\n moving_avg.data.mul_(decay).add_(new, alpha=(1 - decay))\n\n\ndef laplace_smoothing(x, n_categories: int, epsilon: float = 1e-5):\n return (x + epsilon) / (x.sum() + n_categories * epsilon)\n\n\ndef uniform_init(*shape: int):\n t = torch.empty(shape)\n nn.init.kaiming_uniform_(t)\n return t\n\n\ndef sample_vectors(samples, num: int):\n num_samples, device = samples.shape[0], samples.device\n\n if num_samples >= num:\n indices = torch.randperm(num_samples, device=device)[:num]\n else:\n indices = torch.randint(0, num_samples, (num,), device=device)\n\n return samples[indices]\n\n\ndef kmeans(samples, num_clusters: int, num_iters: int = 10, frames_to_use: int = 10_000, batch_size: int = 64):\n \"\"\"\n Memory-efficient K-means clustering.\n Args:\n samples (tensor): shape [N, D]\n num_clusters (int): number of centroids.\n num_iters (int): number of iterations.\n frames_to_use (int): subsample size from total samples.\n batch_size (int): batch size used in distance computation.\n Returns:\n means: [num_clusters, D]\n bins: [num_clusters] (number of points per cluster)\n \"\"\"\n N, D = samples.shape\n dtype, device = samples.dtype, samples.device\n\n if frames_to_use < N:\n indices = torch.randperm(N, device=device)[:frames_to_use]\n samples = samples[indices]\n\n means = sample_vectors(samples, num_clusters)\n\n for _ in range(num_iters):\n # Store cluster assignments\n all_assignments = []\n\n for i in range(0, samples.shape[0], batch_size):\n batch = samples[i : i + batch_size] # [B, D]\n dists = torch.cdist(batch, means, p=2) # [B, C]\n assignments = dists.argmin(dim=1) # [B]\n all_assignments.append(assignments)\n\n buckets = torch.cat(all_assignments, dim=0) # [N]\n bins = torch.bincount(buckets, minlength=num_clusters)\n zero_mask = bins == 0\n bins_min_clamped = bins.masked_fill(zero_mask, 1)\n\n # Compute new means\n new_means = torch.zeros_like(means)\n for i in range(num_clusters):\n mask = buckets == i\n if mask.any():\n new_means[i] = samples[mask].mean(dim=0)\n\n means = torch.where(zero_mask[:, None], means, new_means)\n\n return means, bins\n\n\nclass EuclideanCodebook(nn.Module):\n \"\"\"Codebook with Euclidean distance.\n Args:\n dim (int): Dimension.\n codebook_size (int): Codebook size.\n kmeans_init (bool): Whether to use k-means to initialize the codebooks.\n If set to true, run the k-means algorithm on the first training batch and use\n the learned centroids as initialization.\n kmeans_iters (int): Number of iterations used for k-means algorithm at initialization.\n decay (float): Decay for exponential moving average over the codebooks.\n epsilon (float): Epsilon value for numerical stability.\n threshold_ema_dead_code (int): Threshold for dead code expiration. Replace any codes\n that have an exponential moving average cluster size less than the specified threshold with\n randomly selected vector from the current batch.\n \"\"\"\n\n def __init__(\n self,\n dim: int,\n codebook_size: int,\n kmeans_init: int = False,\n kmeans_iters: int = 10,\n decay: float = 0.99,\n epsilon: float = 1e-5,\n threshold_ema_dead_code: int = 2,\n ):\n super().__init__()\n self.decay = decay\n init_fn: tp.Union[tp.Callable[..., torch.Tensor], tp.Any] = uniform_init if not kmeans_init else torch.zeros\n embed = init_fn(codebook_size, dim)\n\n self.codebook_size = codebook_size\n\n self.kmeans_iters = kmeans_iters\n self.epsilon = epsilon\n self.threshold_ema_dead_code = threshold_ema_dead_code\n\n # Flag variable to indicate whether the codebook is initialized\n self.register_buffer(\"inited\", torch.Tensor([not kmeans_init]))\n # Runing EMA cluster size/count: N_i^t in eq. (6) in vqvae paper\n self.register_buffer(\"cluster_size\", torch.zeros(codebook_size))\n # Codebook\n self.register_buffer(\"embed\", embed)\n # EMA codebook: eq. (7) in vqvae paper\n self.register_buffer(\"embed_avg\", embed.clone())\n\n @torch.jit.ignore\n def init_embed_(self, data):\n \"\"\"Initialize codebook.\n Args:\n data (tensor): [B * T, D].\n \"\"\"\n if self.inited:\n return\n\n ## NOTE (snippet added by Songxiang Liu): gather data from all gpus\n if dist.is_available() and dist.is_initialized():\n # [B * T * world_size, D]\n data = SyncFunction.apply(data)\n\n embed, cluster_size = kmeans(data, self.codebook_size, self.kmeans_iters)\n self.embed.data.copy_(embed)\n self.embed_avg.data.copy_(embed.clone())\n self.cluster_size.data.copy_(cluster_size)\n self.inited.data.copy_(torch.Tensor([True]))\n # Make sure all buffers across workers are in sync after initialization\n broadcast_tensors(self.buffers())\n\n def replace_(self, samples, mask):\n modified_codebook = torch.where(mask[..., None], sample_vectors(samples, self.codebook_size), self.embed)\n self.embed.data.copy_(modified_codebook)\n\n def expire_codes_(self, batch_samples):\n if self.threshold_ema_dead_code == 0:\n return\n\n expired_codes = self.cluster_size < self.threshold_ema_dead_code\n if not torch.any(expired_codes):\n return\n\n ## NOTE (snippet added by Songxiang Liu): gather data from all gpus\n if is_distributed():\n # [B * T * world_size, D]\n batch_samples = SyncFunction.apply(batch_samples)\n\n batch_samples = rearrange(batch_samples, \"... d -> (...) d\")\n self.replace_(batch_samples, mask=expired_codes)\n broadcast_tensors(self.buffers())\n\n def preprocess(self, x):\n x = rearrange(x, \"... d -> (...) d\")\n return x\n\n def quantize(self, x):\n embed = self.embed.t()\n dist = -(x.pow(2).sum(1, keepdim=True) - 2 * x @ embed + embed.pow(2).sum(0, keepdim=True))\n embed_ind = dist.max(dim=-1).indices\n return embed_ind\n\n def postprocess_emb(self, embed_ind, shape):\n return embed_ind.view(*shape[:-1])\n\n def dequantize(self, embed_ind):\n quantize = F.embedding(embed_ind, self.embed)\n return quantize\n\n def encode(self, x):\n shape = x.shape\n # pre-process\n x = self.preprocess(x) # [B, T, D] -> [B*T, D]\n # quantize\n embed_ind = self.quantize(x)\n # post-process\n embed_ind = self.postprocess_emb(embed_ind, shape)\n return embed_ind\n\n def decode(self, embed_ind):\n quantize = self.dequantize(embed_ind)\n return quantize\n\n def forward(self, x):\n # shape: [B, T, D]\n shape, dtype = x.shape, x.dtype\n x = self.preprocess(x) # [B, T, D] -> [B*T, D]\n\n # Initialize codebook\n self.init_embed_(x)\n\n embed_ind = self.quantize(x) # [B*T,]\n embed_onehot = F.one_hot(embed_ind, self.codebook_size).type(dtype) # [B*T, cb-size]\n embed_ind = self.postprocess_emb(embed_ind, shape) # [B, T]\n quantize = self.dequantize(embed_ind) # [B, T, D]\n\n if self.training:\n ### Update codebook by EMA\n embed_onehot_sum = embed_onehot.sum(0) # [cb-size,]\n embed_sum = x.t() @ embed_onehot # [D, cb-size]\n if is_distributed():\n dist.all_reduce(embed_onehot_sum)\n dist.all_reduce(embed_sum)\n # Update ema cluster count N_i^t, eq. (6) in vqvae paper\n self.cluster_size.data.mul_(self.decay).add_(embed_onehot_sum, alpha=1 - self.decay)\n # Update ema embed: eq. (7) in vqvae paper\n self.embed_avg.data.mul_(self.decay).add_(embed_sum.t(), alpha=1 - self.decay)\n # apply laplace smoothing\n n = self.cluster_size.sum()\n cluster_size = (self.cluster_size + self.epsilon) / (n + self.codebook_size * self.epsilon) * n\n # Update ema embed: eq. (8) in vqvae paper\n embed_normalized = self.embed_avg / cluster_size.unsqueeze(1)\n self.embed.data.copy_(embed_normalized)\n\n # We do the expiry of code at that point as buffers are in sync\n # and all the workers will take the same decision.\n self.expire_codes_(x)\n\n return quantize, embed_ind\n\n\nclass VectorQuantization(nn.Module):\n \"\"\"Vector quantization implementation.\n Currently supports only euclidean distance.\n Args:\n dim (int): Dimension\n codebook_size (int): Codebook size\n codebook_dim (int): Codebook dimension. If not defined, uses the specified dimension in dim.\n decay (float): Decay for exponential moving average over the codebooks.\n epsilon (float): Epsilon value for numerical stability.\n kmeans_init (bool): Whether to use kmeans to initialize the codebooks.\n kmeans_iters (int): Number of iterations used for kmeans initialization.\n threshold_ema_dead_code (int): Threshold for dead code expiration. Replace any codes\n that have an exponential moving average cluster size less than the specified threshold with\n randomly selected vector from the current batch.\n commitment_weight (float): Weight for commitment loss.\n \"\"\"\n\n def __init__(\n self,\n dim: int,\n codebook_size: int,\n codebook_dim: tp.Optional[int] = None,\n decay: float = 0.99,\n epsilon: float = 1e-5,\n kmeans_init: bool = True,\n kmeans_iters: int = 50,\n threshold_ema_dead_code: int = 2,\n commitment_weight: float = 1.0,\n ):\n super().__init__()\n _codebook_dim: int = default(codebook_dim, dim)\n\n requires_projection = _codebook_dim != dim\n self.project_in = nn.Linear(dim, _codebook_dim) if requires_projection else nn.Identity()\n self.project_out = nn.Linear(_codebook_dim, dim) if requires_projection else nn.Identity()\n\n self.epsilon = epsilon\n self.commitment_weight = commitment_weight\n\n self._codebook = EuclideanCodebook(\n dim=_codebook_dim,\n codebook_size=codebook_size,\n kmeans_init=kmeans_init,\n kmeans_iters=kmeans_iters,\n decay=decay,\n epsilon=epsilon,\n threshold_ema_dead_code=threshold_ema_dead_code,\n )\n self.codebook_size = codebook_size\n\n @property\n def codebook(self):\n return self._codebook.embed\n\n def encode(self, x):\n x = rearrange(x, \"b d n -> b n d\")\n x = self.project_in(x)\n embed_in = self._codebook.encode(x)\n return embed_in\n\n def decode(self, embed_ind):\n quantize = self._codebook.decode(embed_ind)\n quantize = self.project_out(quantize)\n quantize = rearrange(quantize, \"b n d -> b d n\")\n return quantize\n\n def forward(self, x):\n device = x.device\n x = x.transpose(1, 2).contiguous() # [b d n] -> [b n d]\n x = self.project_in(x)\n\n quantize, embed_ind = self._codebook(x)\n\n if self.training:\n quantize = x + (quantize - x).detach()\n\n loss = torch.tensor([0.0], device=device, requires_grad=self.training)\n\n if self.training:\n if self.commitment_weight > 0:\n commit_loss = F.mse_loss(quantize.detach(), x)\n loss = loss + commit_loss * self.commitment_weight\n\n quantize = self.project_out(quantize)\n quantize = quantize.transpose(1, 2).contiguous() # [b n d] -> [b d n]\n return quantize, embed_ind, loss\n\n\nclass ResidualVectorQuantization(nn.Module):\n \"\"\"Residual vector quantization implementation.\n Follows Algorithm 1. in https://arxiv.org/pdf/2107.03312.pdf\n \"\"\"\n\n def __init__(self, *, num_quantizers, **kwargs):\n super().__init__()\n self.layers = nn.ModuleList([VectorQuantization(**kwargs) for _ in range(num_quantizers)])\n\n def forward(self, x, n_q: tp.Optional[int] = None):\n quantized_out = 0.0\n residual = x\n\n all_losses = []\n all_indices = []\n\n n_q = n_q or len(self.layers)\n\n for layer in self.layers[:n_q]:\n quantized, indices, loss = layer(residual)\n residual = residual - quantized\n quantized_out = quantized_out + quantized\n\n all_indices.append(indices)\n all_losses.append(loss)\n\n out_losses, out_indices = map(torch.stack, (all_losses, all_indices))\n return quantized_out, out_indices, out_losses\n\n def encode(self, x: torch.Tensor, n_q: tp.Optional[int] = None) -> torch.Tensor:\n residual = x\n all_indices = []\n n_q = n_q or len(self.layers)\n for layer in self.layers[:n_q]:\n indices = layer.encode(residual)\n quantized = layer.decode(indices)\n residual = residual - quantized\n all_indices.append(indices)\n out_indices = torch.stack(all_indices)\n return out_indices\n\n def decode(self, q_indices: torch.Tensor) -> torch.Tensor:\n quantized_out = torch.tensor(0.0, device=q_indices.device)\n for i, indices in enumerate(q_indices):\n layer = self.layers[i]\n quantized = layer.decode(indices)\n quantized_out = quantized_out + quantized\n return quantized_out\n"
},
{
"path": "boson_multimodal/audio_processing/quantization/ddp_utils.py",
"content": "import logging\nimport random\nimport subprocess\nfrom datetime import datetime\n\nimport numpy as np\nimport torch\nimport torch.distributed as dist\nfrom torch.nn.parallel import DistributedDataParallel\nfrom torch.nn.parallel.distributed import _find_tensors\nimport torch.optim\nimport torch.utils.data\nfrom packaging import version\nfrom omegaconf import OmegaConf\n\n\ndef set_random_seed(seed):\n random.seed(seed)\n np.random.seed(seed)\n torch.manual_seed(seed)\n torch.cuda.manual_seed_all(seed)\n\n\ndef is_logging_process():\n return not dist.is_initialized() or dist.get_rank() == 0\n\n\ndef get_logger(cfg, name=None):\n # log_file_path is used when unit testing\n if is_logging_process():\n logging.config.dictConfig(OmegaConf.to_container(cfg.job_logging_config, resolve=True))\n return logging.getLogger(name)\n\n\n# from https://github.com/Lightning-AI/lightning-bolts/blob/5d61197cd2f491f69e238137a5edabe80ae14ad9/pl_bolts/models/self_supervised/simclr/simclr_module.py#L20\nclass SyncFunction(torch.autograd.Function):\n @staticmethod\n # @torch.no_grad()\n def forward(ctx, tensor):\n ctx.batch_size = tensor.shape[0]\n\n gathered_tensor = [torch.zeros_like(tensor) for _ in range(torch.distributed.get_world_size())]\n\n torch.distributed.all_gather(gathered_tensor, tensor)\n gathered_tensor = torch.cat(gathered_tensor, 0)\n\n return gathered_tensor\n\n @staticmethod\n def backward(ctx, grad_output):\n grad_input = grad_output.clone()\n torch.distributed.all_reduce(grad_input, op=torch.distributed.ReduceOp.SUM, async_op=False)\n\n idx_from = torch.distributed.get_rank() * ctx.batch_size\n idx_to = (torch.distributed.get_rank() + 1) * ctx.batch_size\n return grad_input[idx_from:idx_to]\n\n\ndef get_timestamp():\n return datetime.now().strftime(\"%y%m%d-%H%M%S\")\n\n\ndef get_commit_hash():\n message = subprocess.check_output([\"git\", \"rev-parse\", \"--short\", \"HEAD\"])\n return message.strip().decode(\"utf-8\")\n\n\nclass DDP(DistributedDataParallel):\n \"\"\"\n Override the forward call in lightning so it goes to training and validation step respectively\n \"\"\"\n\n def forward(self, *inputs, **kwargs): # pragma: no cover\n if version.parse(torch.__version__[:6]) < version.parse(\"1.11\"):\n self._sync_params()\n inputs, kwargs = self.scatter(inputs, kwargs, self.device_ids)\n assert len(self.device_ids) == 1\n if self.module.training:\n output = self.module.training_step(*inputs[0], **kwargs[0])\n elif self.module.testing:\n output = self.module.test_step(*inputs[0], **kwargs[0])\n else:\n output = self.module.validation_step(*inputs[0], **kwargs[0])\n if torch.is_grad_enabled():\n # We'll return the output object verbatim since it is a freeform\n # object. We need to find any tensors in this object, though,\n # because we need to figure out which parameters were used during\n # this forward pass, to ensure we short circuit reduction for any\n # unused parameters. Only if `find_unused_parameters` is set.\n if self.find_unused_parameters:\n self.reducer.prepare_for_backward(list(_find_tensors(output)))\n else:\n self.reducer.prepare_for_backward([])\n else:\n from torch.nn.parallel.distributed import (\n logging,\n Join,\n _DDPSink,\n _tree_flatten_with_rref,\n _tree_unflatten_with_rref,\n )\n\n with torch.autograd.profiler.record_function(\"DistributedDataParallel.forward\"):\n if torch.is_grad_enabled() and self.require_backward_grad_sync:\n self.logger.set_runtime_stats_and_log()\n self.num_iterations += 1\n self.reducer.prepare_for_forward()\n\n # Notify the join context that this process has not joined, if\n # needed\n work = Join.notify_join_context(self)\n if work:\n self.reducer._set_forward_pass_work_handle(work, self._divide_by_initial_world_size)\n\n # Calling _rebuild_buckets before forward compuation,\n # It may allocate new buckets before deallocating old buckets\n # inside _rebuild_buckets. To save peak memory usage,\n # call _rebuild_buckets before the peak memory usage increases\n # during forward computation.\n # This should be called only once during whole training period.\n if torch.is_grad_enabled() and self.reducer._rebuild_buckets():\n logging.info(\"Reducer buckets have been rebuilt in this iteration.\")\n self._has_rebuilt_buckets = True\n\n # sync params according to location (before/after forward) user\n # specified as part of hook, if hook was specified.\n buffer_hook_registered = hasattr(self, \"buffer_hook\")\n if self._check_sync_bufs_pre_fwd():\n self._sync_buffers()\n\n if self._join_config.enable:\n # Notify joined ranks whether they should sync in backwards pass or not.\n self._check_global_requires_backward_grad_sync(is_joined_rank=False)\n\n inputs, kwargs = self.scatter(inputs, kwargs, self.device_ids)\n if self.module.training:\n output = self.module.training_step(*inputs[0], **kwargs[0])\n elif self.module.testing:\n output = self.module.test_step(*inputs[0], **kwargs[0])\n else:\n output = self.module.validation_step(*inputs[0], **kwargs[0])\n\n # sync params according to location (before/after forward) user\n # specified as part of hook, if hook was specified.\n if self._check_sync_bufs_post_fwd():\n self._sync_buffers()\n\n if torch.is_grad_enabled() and self.require_backward_grad_sync:\n self.require_forward_param_sync = True\n # We'll return the output object verbatim since it is a freeform\n # object. We need to find any tensors in this object, though,\n # because we need to figure out which parameters were used during\n # this forward pass, to ensure we short circuit reduction for any\n # unused parameters. Only if `find_unused_parameters` is set.\n if self.find_unused_parameters and not self.static_graph:\n # Do not need to populate this for static graph.\n self.reducer.prepare_for_backward(list(_find_tensors(output)))\n else:\n self.reducer.prepare_for_backward([])\n else:\n self.require_forward_param_sync = False\n\n # TODO: DDPSink is currently enabled for unused parameter detection and\n # static graph training for first iteration.\n if (self.find_unused_parameters and not self.static_graph) or (\n self.static_graph and self.num_iterations == 1\n ):\n state_dict = {\n \"static_graph\": self.static_graph,\n \"num_iterations\": self.num_iterations,\n }\n\n output_tensor_list, treespec, output_is_rref = _tree_flatten_with_rref(output)\n output_placeholders = [None for _ in range(len(output_tensor_list))]\n # Do not touch tensors that have no grad_fn, which can cause issues\n # such as https://github.com/pytorch/pytorch/issues/60733\n for i, output in enumerate(output_tensor_list):\n if torch.is_tensor(output) and output.grad_fn is None:\n output_placeholders[i] = output\n\n # When find_unused_parameters=True, makes tensors which require grad\n # run through the DDPSink backward pass. When not all outputs are\n # used in loss, this makes those corresponding tensors receive\n # undefined gradient which the reducer then handles to ensure\n # param.grad field is not touched and we don't error out.\n passthrough_tensor_list = _DDPSink.apply(\n self.reducer,\n state_dict,\n *output_tensor_list,\n )\n for i in range(len(output_placeholders)):\n if output_placeholders[i] is None:\n output_placeholders[i] = passthrough_tensor_list[i]\n\n # Reconstruct output data structure.\n output = _tree_unflatten_with_rref(output_placeholders, treespec, output_is_rref)\n return output\n"
},
{
"path": "boson_multimodal/audio_processing/quantization/distrib.py",
"content": "# Copyright (c) Meta Platforms, Inc. and affiliates.\n# All rights reserved.\n#\n# This source code is licensed under the license found in the\n# LICENSE file in the root directory of this source tree.\n\n\"\"\"Torch distributed utilities.\"\"\"\n\nimport typing as tp\n\nimport torch\n\n\ndef rank():\n if torch.distributed.is_initialized():\n return torch.distributed.get_rank()\n else:\n return 0\n\n\ndef world_size():\n if torch.distributed.is_initialized():\n return torch.distributed.get_world_size()\n else:\n return 1\n\n\ndef is_distributed():\n return world_size() > 1\n\n\ndef all_reduce(tensor: torch.Tensor, op=torch.distributed.ReduceOp.SUM):\n if is_distributed():\n return torch.distributed.all_reduce(tensor, op)\n\n\ndef _is_complex_or_float(tensor):\n return torch.is_floating_point(tensor) or torch.is_complex(tensor)\n\n\ndef _check_number_of_params(params: tp.List[torch.Tensor]):\n # utility function to check that the number of params in all workers is the same,\n # and thus avoid a deadlock with distributed all reduce.\n if not is_distributed() or not params:\n return\n # print('params[0].device ', params[0].device)\n tensor = torch.tensor([len(params)], device=params[0].device, dtype=torch.long)\n all_reduce(tensor)\n if tensor.item() != len(params) * world_size():\n # If not all the workers have the same number, for at least one of them,\n # this inequality will be verified.\n raise RuntimeError(\n f\"Mismatch in number of params: ours is {len(params)}, at least one worker has a different one.\"\n )\n\n\ndef broadcast_tensors(tensors: tp.Iterable[torch.Tensor], src: int = 0):\n \"\"\"Broadcast the tensors from the given parameters to all workers.\n This can be used to ensure that all workers have the same model to start with.\n \"\"\"\n if not is_distributed():\n return\n tensors = [tensor for tensor in tensors if _is_complex_or_float(tensor)]\n _check_number_of_params(tensors)\n handles = []\n for tensor in tensors:\n handle = torch.distributed.broadcast(tensor.data, src=src, async_op=True)\n handles.append(handle)\n for handle in handles:\n handle.wait()\n\n\ndef sync_buffer(buffers, average=True):\n \"\"\"\n Sync grad for buffers. If average is False, broadcast instead of averaging.\n \"\"\"\n if not is_distributed():\n return\n handles = []\n for buffer in buffers:\n if torch.is_floating_point(buffer.data):\n if average:\n handle = torch.distributed.all_reduce(buffer.data, op=torch.distributed.ReduceOp.SUM, async_op=True)\n else:\n handle = torch.distributed.broadcast(buffer.data, src=0, async_op=True)\n handles.append((buffer, handle))\n for buffer, handle in handles:\n handle.wait()\n if average:\n buffer.data /= world_size\n\n\ndef sync_grad(params):\n \"\"\"\n Simpler alternative to DistributedDataParallel, that doesn't rely\n on any black magic. For simple models it can also be as fast.\n Just call this on your model parameters after the call to backward!\n \"\"\"\n if not is_distributed():\n return\n handles = []\n for p in params:\n if p.grad is not None:\n handle = torch.distributed.all_reduce(p.grad.data, op=torch.distributed.ReduceOp.SUM, async_op=True)\n handles.append((p, handle))\n for p, handle in handles:\n handle.wait()\n p.grad.data /= world_size()\n\n\ndef average_metrics(metrics: tp.Dict[str, float], count=1.0):\n \"\"\"Average a dictionary of metrics across all workers, using the optional\n `count` as unormalized weight.\n \"\"\"\n if not is_distributed():\n return metrics\n keys, values = zip(*metrics.items())\n device = \"cuda\" if torch.cuda.is_available() else \"cpu\"\n tensor = torch.tensor(list(values) + [1], device=device, dtype=torch.float32)\n tensor *= count\n all_reduce(tensor)\n averaged = (tensor[:-1] / tensor[-1]).cpu().tolist()\n return dict(zip(keys, averaged))\n"
},
{
"path": "boson_multimodal/audio_processing/quantization/vq.py",
"content": "# Copyright (c) Meta Platforms, Inc. and affiliates.\n# All rights reserved.\n#\n# This source code is licensed under the license found in the\n# LICENSE file in the root directory of this source tree.\n\n\"\"\"Residual vector quantizer implementation.\"\"\"\n\nfrom dataclasses import dataclass, field\nimport math\nimport typing as tp\n\nimport torch\nfrom torch import nn\n\n# from .core_vq import ResidualVectorQuantization\nfrom .core_vq_lsx_version import ResidualVectorQuantization\n\n\n@dataclass\nclass QuantizedResult:\n quantized: torch.Tensor\n codes: torch.Tensor\n bandwidth: torch.Tensor # bandwidth in kb/s used, per batch item.\n penalty: tp.Optional[torch.Tensor] = None\n metrics: dict = field(default_factory=dict)\n\n\nclass ResidualVectorQuantizer(nn.Module):\n \"\"\"Residual Vector Quantizer.\n Args:\n dimension (int): Dimension of the codebooks.\n n_q (int): Number of residual vector quantizers used.\n bins (int): Codebook size.\n decay (float): Decay for exponential moving average over the codebooks.\n kmeans_init (bool): Whether to use kmeans to initialize the codebooks.\n kmeans_iters (int): Number of iterations used for kmeans initialization.\n threshold_ema_dead_code (int): Threshold for dead code expiration. Replace any codes\n that have an exponential moving average cluster size less than the specified threshold with\n randomly selected vector from the current batch.\n \"\"\"\n\n def __init__(\n self,\n dimension: int = 256,\n codebook_dim: int = None,\n n_q: int = 8,\n bins: int = 1024,\n decay: float = 0.99,\n kmeans_init: bool = True,\n kmeans_iters: int = 50,\n threshold_ema_dead_code: int = 2,\n ):\n super().__init__()\n self.n_q = n_q\n self.dimension = dimension\n self.codebook_dim = codebook_dim\n self.bins = bins\n self.decay = decay\n self.kmeans_init = kmeans_init\n self.kmeans_iters = kmeans_iters\n self.threshold_ema_dead_code = threshold_ema_dead_code\n self.vq = ResidualVectorQuantization(\n dim=self.dimension,\n codebook_dim=self.codebook_dim,\n codebook_size=self.bins,\n num_quantizers=self.n_q,\n decay=self.decay,\n kmeans_init=self.kmeans_init,\n kmeans_iters=self.kmeans_iters,\n threshold_ema_dead_code=self.threshold_ema_dead_code,\n )\n\n def forward(self, x: torch.Tensor, sample_rate: int, bandwidth: tp.Optional[float] = None): # -> QuantizedResult:\n \"\"\"Residual vector quantization on the given input tensor.\n Args:\n x (torch.Tensor): Input tensor.\n sample_rate (int): Sample rate of the input tensor.\n bandwidth (float): Target bandwidth.\n Returns:\n QuantizedResult:\n The quantized (or approximately quantized) representation with\n the associated bandwidth and any penalty term for the loss.\n \"\"\"\n bw_per_q = self.get_bandwidth_per_quantizer(sample_rate)\n n_q = self.get_num_quantizers_for_bandwidth(sample_rate, bandwidth)\n quantized, codes, commit_loss = self.vq(x, n_q=n_q)\n bw = torch.tensor(n_q * bw_per_q).to(x)\n return quantized, codes, bw, torch.mean(commit_loss)\n # return QuantizedResult(quantized, codes, bw, penalty=torch.mean(commit_loss))\n\n def get_num_quantizers_for_bandwidth(self, sample_rate: int, bandwidth: tp.Optional[float] = None) -> int:\n \"\"\"Return n_q based on specified target bandwidth.\"\"\"\n bw_per_q = self.get_bandwidth_per_quantizer(sample_rate)\n n_q = self.n_q\n if bandwidth and bandwidth > 0.0:\n n_q = int(max(1, math.floor(bandwidth / bw_per_q)))\n return n_q\n\n def get_bandwidth_per_quantizer(self, sample_rate: int):\n \"\"\"Return bandwidth per quantizer for a given input sample rate.\"\"\"\n return math.log2(self.bins) * sample_rate / 1000\n\n def encode(self, x: torch.Tensor, sample_rate: int, bandwidth: tp.Optional[float] = None) -> torch.Tensor:\n \"\"\"Encode a given input tensor with the specified sample rate at the given bandwidth.\n The RVQ encode method sets the appropriate number of quantizer to use\n and returns indices for each quantizer.\n \"\"\"\n n_q = self.get_num_quantizers_for_bandwidth(sample_rate, bandwidth)\n codes = self.vq.encode(x, n_q=n_q)\n return codes\n\n def decode(self, codes: torch.Tensor) -> torch.Tensor:\n \"\"\"Decode the given codes to the quantized representation.\"\"\"\n quantized = self.vq.decode(codes)\n return quantized\n"
},
{
"path": "boson_multimodal/audio_processing/semantic_module.py",
"content": "# Based on code from: https://github.com/zhenye234/xcodec\n# Licensed under MIT License\n# Modifications by BosonAI\n\nimport torch\nimport torch.nn as nn\n\n\nclass Conv1d1x1(nn.Conv1d):\n \"\"\"1x1 Conv1d.\"\"\"\n\n def __init__(self, in_channels, out_channels, bias=True):\n super(Conv1d1x1, self).__init__(in_channels, out_channels, kernel_size=1, bias=bias)\n\n\nclass Conv1d(nn.Module):\n def __init__(\n self,\n in_channels: int,\n out_channels: int,\n kernel_size: int,\n stride: int = 1,\n padding: int = -1,\n dilation: int = 1,\n groups: int = 1,\n bias: bool = True,\n ):\n super().__init__()\n self.in_channels = in_channels\n self.out_channels = out_channels\n self.kernel_size = kernel_size\n if padding < 0:\n padding = (kernel_size - 1) // 2 * dilation\n self.dilation = dilation\n self.conv = nn.Conv1d(\n in_channels=in_channels,\n out_channels=out_channels,\n kernel_size=kernel_size,\n stride=stride,\n padding=padding,\n dilation=dilation,\n groups=groups,\n bias=bias,\n )\n\n def forward(self, x):\n \"\"\"\n Args:\n x (Tensor): Float tensor variable with the shape (B, C, T).\n Returns:\n Tensor: Float tensor variable with the shape (B, C, T).\n \"\"\"\n x = self.conv(x)\n return x\n\n\nclass ResidualUnit(nn.Module):\n def __init__(\n self,\n in_channels: int,\n out_channels: int,\n kernel_size=3,\n dilation=1,\n bias=False,\n nonlinear_activation=\"ELU\",\n nonlinear_activation_params={},\n ):\n super().__init__()\n self.activation = getattr(nn, nonlinear_activation)(**nonlinear_activation_params)\n self.conv1 = Conv1d(\n in_channels=in_channels,\n out_channels=out_channels,\n kernel_size=kernel_size,\n stride=1,\n dilation=dilation,\n bias=bias,\n )\n self.conv2 = Conv1d1x1(out_channels, out_channels, bias)\n\n def forward(self, x):\n y = self.conv1(self.activation(x))\n y = self.conv2(self.activation(y))\n return x + y\n\n\nclass ConvTranspose1d(nn.Module):\n def __init__(\n self,\n in_channels: int,\n out_channels: int,\n kernel_size: int,\n stride: int,\n padding=-1,\n output_padding=-1,\n groups=1,\n bias=True,\n ):\n super().__init__()\n if padding < 0:\n padding = (stride + 1) // 2\n if output_padding < 0:\n output_padding = 1 if stride % 2 else 0\n self.deconv = nn.ConvTranspose1d(\n in_channels=in_channels,\n out_channels=out_channels,\n kernel_size=kernel_size,\n stride=stride,\n padding=padding,\n output_padding=output_padding,\n groups=groups,\n bias=bias,\n )\n\n def forward(self, x):\n \"\"\"\n Args:\n x (Tensor): Float tensor variable with the shape (B, C, T).\n Returns:\n Tensor: Float tensor variable with the shape (B, C', T').\n \"\"\"\n x = self.deconv(x)\n return x\n\n\nclass EncoderBlock(nn.Module):\n def __init__(\n self, in_channels: int, out_channels: int, stride: int, dilations=(1, 1), unit_kernel_size=3, bias=True\n ):\n super().__init__()\n self.res_units = torch.nn.ModuleList()\n for dilation in dilations:\n self.res_units += [ResidualUnit(in_channels, in_channels, kernel_size=unit_kernel_size, dilation=dilation)]\n self.num_res = len(self.res_units)\n\n self.conv = Conv1d(\n in_channels=in_channels,\n out_channels=out_channels,\n kernel_size=3 if stride == 1 else (2 * stride), # special case: stride=1, do not use kernel=2\n stride=stride,\n bias=bias,\n )\n\n def forward(self, x):\n for idx in range(self.num_res):\n x = self.res_units[idx](x)\n x = self.conv(x)\n return x\n\n\nclass Encoder(nn.Module):\n def __init__(\n self,\n input_channels: int,\n encode_channels: int,\n channel_ratios=(1, 1),\n strides=(1, 1),\n kernel_size=3,\n bias=True,\n block_dilations=(1, 1),\n unit_kernel_size=3,\n ):\n super().__init__()\n assert len(channel_ratios) == len(strides)\n\n self.conv = Conv1d(\n in_channels=input_channels, out_channels=encode_channels, kernel_size=kernel_size, stride=1, bias=False\n )\n self.conv_blocks = torch.nn.ModuleList()\n in_channels = encode_channels\n for idx, stride in enumerate(strides):\n out_channels = int(encode_channels * channel_ratios[idx]) # could be float\n self.conv_blocks += [\n EncoderBlock(\n in_channels,\n out_channels,\n stride,\n dilations=block_dilations,\n unit_kernel_size=unit_kernel_size,\n bias=bias,\n )\n ]\n in_channels = out_channels\n self.num_blocks = len(self.conv_blocks)\n self.out_channels = out_channels\n\n def forward(self, x):\n x = self.conv(x)\n for i in range(self.num_blocks):\n x = self.conv_blocks[i](x)\n return x\n\n\nclass DecoderBlock(nn.Module):\n \"\"\"Decoder block (no up-sampling)\"\"\"\n\n def __init__(\n self, in_channels: int, out_channels: int, stride: int, dilations=(1, 1), unit_kernel_size=3, bias=True\n ):\n super().__init__()\n\n if stride == 1:\n self.conv = Conv1d(\n in_channels=in_channels,\n out_channels=out_channels,\n kernel_size=3, # fix kernel=3 when stride=1 for unchanged shape\n stride=stride,\n bias=bias,\n )\n else:\n self.conv = ConvTranspose1d(\n in_channels=in_channels,\n out_channels=out_channels,\n kernel_size=(2 * stride),\n stride=stride,\n bias=bias,\n )\n\n self.res_units = torch.nn.ModuleList()\n for idx, dilation in enumerate(dilations):\n self.res_units += [\n ResidualUnit(out_channels, out_channels, kernel_size=unit_kernel_size, dilation=dilation)\n ]\n self.num_res = len(self.res_units)\n\n def forward(self, x):\n x = self.conv(x)\n for idx in range(self.num_res):\n x = self.res_units[idx](x)\n return x\n\n\nclass Decoder(nn.Module):\n def __init__(\n self,\n code_dim: int,\n output_channels: int,\n decode_channels: int,\n channel_ratios=(1, 1),\n strides=(1, 1),\n kernel_size=3,\n bias=True,\n block_dilations=(1, 1),\n unit_kernel_size=3,\n ):\n super().__init__()\n assert len(channel_ratios) == len(strides)\n\n self.conv1 = Conv1d(\n in_channels=code_dim,\n out_channels=int(decode_channels * channel_ratios[0]),\n kernel_size=kernel_size,\n stride=1,\n bias=False,\n )\n\n self.conv_blocks = torch.nn.ModuleList()\n for idx, stride in enumerate(strides):\n in_channels = int(decode_channels * channel_ratios[idx])\n if idx < (len(channel_ratios) - 1):\n out_channels = int(decode_channels * channel_ratios[idx + 1])\n else:\n out_channels = decode_channels\n self.conv_blocks += [\n DecoderBlock(\n in_channels,\n out_channels,\n stride,\n dilations=block_dilations,\n unit_kernel_size=unit_kernel_size,\n bias=bias,\n )\n ]\n self.num_blocks = len(self.conv_blocks)\n\n self.conv2 = Conv1d(out_channels, output_channels, kernel_size, 1, bias=False)\n\n def forward(self, z):\n x = self.conv1(z)\n for i in range(self.num_blocks):\n x = self.conv_blocks[i](x)\n x = self.conv2(x)\n return x\n"
},
{
"path": "boson_multimodal/constants.py",
"content": "AUDIO_IN_TOKEN = \"<|AUDIO|>\"\nAUDIO_OUT_TOKEN = \"<|AUDIO_OUT|>\"\nEOS_TOKEN = \"<|end_of_text|>\"\n"
},
{
"path": "boson_multimodal/data_collator/__init__.py",
"content": ""
},
{
"path": "boson_multimodal/data_collator/higgs_audio_collator.py",
"content": "import librosa\nimport torch\nimport torch.nn.functional as F\nimport math\nfrom typing import List, Tuple\n\nfrom dataclasses import dataclass\nfrom typing import List, Optional\nfrom transformers.models.whisper.processing_whisper import WhisperProcessor\n\nfrom ..dataset.chatml_dataset import ChatMLDatasetSample\nfrom ..model.higgs_audio.utils import build_delay_pattern_mask\n\n\ndef _ceil_to_nearest(n, round_to):\n return (n + round_to - 1) // round_to * round_to\n\n\ndef _ceil_to_next_power_of_two(self, x):\n return 1 if x == 0 else 2 ** (x - 1).bit_length()\n\n\n@dataclass\nclass HiggsAudioBatchInput:\n input_ids: torch.LongTensor # shape (bsz, seq_len).\n attention_mask: torch.Tensor # shape (bsz, seq_len).\n audio_features: Optional[torch.Tensor] # shape (num_audio_in, feature_dim, max_mel_seq_len).\n audio_feature_attention_mask: Optional[torch.Tensor] # shape (num_audio_in, max_mel_seq_len).\n audio_out_ids: Optional[torch.LongTensor] # shape (num_codebooks, audio_out_total_length)\n audio_out_ids_start: Optional[torch.LongTensor] # shape (num_audio_out,)\n # The audio_out_ids_start_group_loc has the same length as audio_out_ids_start. It is used to recover group location in a batch for an audio segment\n # Currently, we concatenante audio segments along dim 0 to handle variadic audio segment length. However, in the alignment stage, we need the location information\n # For example,\n # audio_out_ids_start = [0, 2, 4, 8]; and the first two audio segments come from the same sample in a batch, and other two come from different samples.\n # This is a batch of 3 samples, then we will have the group location as:\n # audio_out_ids_start_group_loc = [0, 0, 1, 2]\n audio_out_ids_start_group_loc: Optional[\n torch.LongTensor\n ] # shape (num_audio_out,), specify which a sample's group location in the batch\n audio_in_ids: Optional[torch.LongTensor] # shape (num_codebooks, audio_in_total_length)\n audio_in_ids_start: Optional[torch.LongTensor] # shape (num_audio_in,)\n label_ids: Optional[torch.LongTensor] # shape (bsz, seq_len)\n label_audio_ids: Optional[torch.LongTensor] # shape (num_codebooks, audio_out_total_length)\n reward: Optional[float] = None\n\n\nclass HiggsAudioSampleCollator:\n \"\"\"Sample collator for Higgs-Audio model.\n\n Args:\n whisper_processor (WhisperProcessor): The whisper processor.\n audio_in_token_id (int): The token id for audio-in.\n audio_out_token_id (int): The token id for audio-out.\n pad_token_id (int): The token id for padding.\n audio_stream_bos_id (int): The token id for audio-stream beginning of sentence.\n audio_stream_eos_id (int): The token id for audio-stream end of sentence.\n round_to (int): The round-to value.\n pad_left (bool): Whether to pad left.\n return_audio_in_tokens (bool): Whether to return audio-in tokens.\n use_delay_pattern (bool): Whether to use delay pattern.\n disable_audio_codes_transform (bool): Whether to add bos and eos tokens to audio codes.\n chunk_size_seconds (int): The chunk size in seconds.\n add_new_bos_eos_for_long_chunk (bool): Whether to add new bos and eos tokens for long chunks.\n mask_audio_out_token_label (bool): Whether to always mask the label associated with <|AUDIO_OUT|> token. Since we will always have `<|AUDIO_OUT|>` after `<|audio_bos|>`, we can safely mask <|AUDIO_OUT|>.\n\n \"\"\"\n\n def __init__(\n self,\n whisper_processor: WhisperProcessor,\n audio_in_token_id,\n audio_out_token_id,\n pad_token_id,\n audio_stream_bos_id,\n audio_stream_eos_id,\n round_to=8,\n pad_left=False,\n encode_whisper_embed=True,\n return_audio_in_tokens=True,\n audio_num_codebooks=None,\n use_delay_pattern=False,\n disable_audio_codes_transform=False,\n chunk_size_seconds=30, # Maximum duration for each chunk\n add_new_bos_eos_for_long_chunk=True,\n mask_audio_out_token_label=True,\n ):\n self.whisper_processor = whisper_processor\n self.round_to = round_to\n self.pad_left = pad_left\n self.audio_in_token_id = audio_in_token_id\n self.audio_out_token_id = audio_out_token_id\n self.audio_stream_bos_id = audio_stream_bos_id\n self.audio_stream_eos_id = audio_stream_eos_id\n self.pad_token_id = pad_token_id\n self.encode_whisper_embed = encode_whisper_embed\n self.return_audio_in_tokens = return_audio_in_tokens\n self.audio_num_codebooks = audio_num_codebooks\n self.use_delay_pattern = use_delay_pattern\n if encode_whisper_embed:\n self.chunk_size_seconds = chunk_size_seconds\n self.chunk_size_samples = int(chunk_size_seconds * whisper_processor.feature_extractor.sampling_rate)\n else:\n self.chunk_size_seconds = None\n self.chunk_size_samples = None\n self.disable_audio_codes_transform = disable_audio_codes_transform\n self.add_new_bos_eos_for_long_chunk = add_new_bos_eos_for_long_chunk\n self.mask_audio_out_token_label = mask_audio_out_token_label\n\n def _process_and_duplicate_audio_tokens(\n self, input_ids: torch.Tensor, audio_idx: int, wv: torch.Tensor, sr: int, labels: Optional[torch.Tensor] = None\n ) -> Tuple[torch.Tensor, torch.Tensor, int]:\n \"\"\"Process long audio and duplicate corresponding audio tokens.\n\n Args:\n input_ids: Input token ids\n audio_idx: Index of the audio token in the sequence\n wv: Audio waveform\n sr: Sample rate\n labels: Optional label ids to be duplicated alongside input ids\n\n Returns:\n Tuple of:\n - New input ids with duplicated audio tokens\n - New label ids (if labels were provided) or None\n - Number of chunks created\n \"\"\"\n # Calculate number of chunks needed\n total_samples = len(wv)\n num_chunks = math.ceil(total_samples / self.chunk_size_samples)\n\n if num_chunks <= 1:\n return input_ids, labels, 1\n\n # Get the three tokens: <|audio_bos|><|AUDIO|><|audio_eos|>\n audio_token_seq = input_ids[audio_idx - 1 : audio_idx + 2]\n # Duplicate sequence for each chunk\n duplicated_sequence = audio_token_seq.repeat(num_chunks)\n\n # Create new input_ids with duplicated tokens\n new_input_ids = torch.cat([input_ids[: audio_idx - 1], duplicated_sequence, input_ids[audio_idx + 2 :]])\n\n # If labels are provided, duplicate them as well\n new_labels = None\n if labels is not None:\n label_seq = labels[audio_idx - 1 : audio_idx + 2]\n duplicated_labels = label_seq.repeat(num_chunks)\n new_labels = torch.cat([labels[: audio_idx - 1], duplicated_labels, labels[audio_idx + 2 :]])\n\n return new_input_ids, new_labels, num_chunks\n\n def __call__(self, batch: List[ChatMLDatasetSample]):\n \"\"\"Collate the input data with support for long audio processing.\"\"\"\n\n label_ids = None\n label_audio_ids = None\n if all([ele.label_ids is None for ele in batch]):\n return_labels = False\n else:\n return_labels = True\n\n if self.encode_whisper_embed:\n # Process each sample in the batch to handle long audio\n # TODO(?) The implementation here can be optimized.\n processed_batch = []\n for i in range(len(batch)):\n sample = batch[i]\n audio_in_mask = sample.input_ids == self.audio_in_token_id\n audio_in_indices = torch.where(audio_in_mask)[0]\n audio_out_mask = sample.input_ids == self.audio_out_token_id\n\n # Process each audio token and duplicate if needed\n modified_input_ids = sample.input_ids\n modified_labels = sample.label_ids if return_labels else None\n modified_waveforms_concat = []\n modified_waveforms_start = []\n modified_sample_rate = []\n offset = 0 # Track position changes from duplicating tokens\n curr_wv_offset = 0\n\n # Process input audio tokens\n for idx, audio_idx in enumerate(audio_in_indices):\n # Get the audio for this token\n wv, sr = sample.get_wv(idx) # Use idx since we want the original audio index\n if sr != self.whisper_processor.feature_extractor.sampling_rate:\n resampled_wv = librosa.resample(\n wv.cpu().numpy(),\n orig_sr=sr,\n target_sr=self.whisper_processor.feature_extractor.sampling_rate,\n )\n else:\n resampled_wv = wv.cpu().numpy()\n wv = torch.tensor(resampled_wv, device=wv.device)\n sr = self.whisper_processor.feature_extractor.sampling_rate\n\n # Process and duplicate tokens if necessary\n token_pos = audio_idx + offset\n modified_input_ids, modified_labels, num_chunks = self._process_and_duplicate_audio_tokens(\n modified_input_ids, token_pos, wv, sr, modified_labels\n )\n\n # Update audio data\n for chunk_idx in range(num_chunks):\n chunk_start = chunk_idx * self.chunk_size_samples\n chunk_end = min((chunk_idx + 1) * self.chunk_size_samples, len(wv))\n chunk_wv = wv[chunk_start:chunk_end]\n modified_waveforms_concat.append(chunk_wv)\n modified_waveforms_start.append(curr_wv_offset)\n curr_wv_offset += len(chunk_wv)\n modified_sample_rate.append(sr)\n\n # Update offset for next iteration\n offset += (num_chunks - 1) * 3 # Each new chunk adds 3 more tokens\n\n # Create new sample with modified tokens and audio data\n processed_sample = ChatMLDatasetSample(\n input_ids=modified_input_ids,\n label_ids=modified_labels if return_labels else sample.label_ids,\n audio_ids_concat=sample.audio_ids_concat,\n audio_ids_start=sample.audio_ids_start,\n audio_waveforms_concat=torch.cat(modified_waveforms_concat)\n if modified_waveforms_concat\n else sample.audio_waveforms_concat,\n audio_waveforms_start=torch.tensor(modified_waveforms_start, dtype=torch.long)\n if modified_waveforms_start\n else sample.audio_waveforms_start,\n audio_sample_rate=torch.tensor(modified_sample_rate)\n if modified_sample_rate\n else sample.audio_sample_rate,\n audio_speaker_indices=torch.tensor([]),\n # FIXME(sxjscience): The logic here is not correct for audio_label_ids_concat.\n audio_label_ids_concat=sample.audio_label_ids_concat,\n )\n # audio_in_chunk_len = len(torch.where(modified_input_ids == self.audio_in_token_id)[0])\n # assert audio_in_chunk_len == processed_sample.num_audios(), f\"Mismatch: audio_in_chunk_len={audio_in_chunk_len}, processed_sample.num_audios()={processed_sample.num_audios()}\"\n processed_batch.append(processed_sample)\n else:\n processed_batch = batch\n\n # Get the max sequence length based on processed batch\n max_seq_length = _ceil_to_nearest(max([len(sample.input_ids) for sample in processed_batch]), self.round_to)\n\n # Get the ids for audio-in and audio-out for each batch\n audio_in_wv_l = []\n audio_in_ids_l = []\n audio_out_ids_l = []\n audio_out_ids_group_loc_l = []\n audio_in_label_ids_l = None\n audio_out_label_ids_l = None\n reward_l = []\n\n if return_labels:\n audio_out_no_train_flag = [] # Whether the audio-out data should be trained on or not.\n\n # Process the audio inputs and outputs\n for i in range(len(processed_batch)):\n audio_in_mask = processed_batch[i].input_ids == self.audio_in_token_id\n audio_out_mask = processed_batch[i].input_ids == self.audio_out_token_id\n audio_ids = torch.ones_like(processed_batch[i].input_ids)\n audio_ids[audio_in_mask ^ audio_out_mask] = torch.cumsum(audio_ids[audio_in_mask ^ audio_out_mask], 0) - 1\n audio_in_ids = audio_ids[audio_in_mask]\n audio_out_ids = audio_ids[audio_out_mask]\n\n if return_labels:\n audio_out_no_train_flag.append(processed_batch[i].label_ids[audio_out_mask] < 0)\n if self.mask_audio_out_token_label:\n processed_batch[i].label_ids[audio_out_mask] = -100\n\n # Process audio inputs\n if self.return_audio_in_tokens:\n audio_in_ids_l.extend(\n [processed_batch[i].get_audio_codes(idx)[: self.audio_num_codebooks, :] for idx in audio_in_ids]\n )\n if processed_batch[i].audio_label_ids_concat is not None:\n if audio_in_label_ids_l is None:\n audio_in_label_ids_l = []\n audio_in_label_ids_l.extend(\n [\n processed_batch[i].get_audio_codes_labels(idx)[: self.audio_num_codebooks, :]\n for idx in audio_in_ids\n ]\n )\n\n audio_out_ids_l.extend(\n [processed_batch[i].get_audio_codes(idx)[: self.audio_num_codebooks, :] for idx in audio_out_ids]\n )\n audio_out_ids_group_loc_l.append(i)\n if processed_batch[i].reward is not None:\n reward_l.append(processed_batch[i].reward)\n\n if processed_batch[i].audio_label_ids_concat is not None:\n if audio_out_label_ids_l is None:\n audio_out_label_ids_l = []\n audio_out_label_ids_l.extend(\n [\n processed_batch[i].get_audio_codes_labels(idx)[: self.audio_num_codebooks, :]\n for idx in audio_out_ids\n ]\n )\n\n if self.encode_whisper_embed:\n for idx in audio_in_ids:\n wv, sr = processed_batch[i].get_wv(idx)\n resampled_wv = wv.cpu().numpy()\n # Split long audio into chunks\n total_samples = len(resampled_wv)\n for chunk_start in range(0, total_samples, self.chunk_size_samples):\n chunk_end = min(chunk_start + self.chunk_size_samples, total_samples)\n chunk = resampled_wv[chunk_start:chunk_end]\n audio_in_wv_l.append(chunk)\n # assert len(audio_in_wv_l) == processed_batch[i].num_audios(), \\\n # f\"Assertion failed: Mismatch in number of audios. \" \\\n # f\"Expected {processed_batch[i].num_audios()}, but got {len(audio_in_wv_l)} at index {i}.\"\n\n if return_labels:\n audio_out_no_train_flag = torch.cat(audio_out_no_train_flag, dim=0)\n\n # Process all audio features\n if len(audio_in_wv_l) > 0:\n feature_ret = self.whisper_processor.feature_extractor(\n audio_in_wv_l,\n sampling_rate=self.whisper_processor.feature_extractor.sampling_rate,\n return_attention_mask=True,\n padding=\"max_length\",\n )\n audio_features = torch.from_numpy(feature_ret[\"input_features\"])\n audio_feature_attention_mask = torch.from_numpy(feature_ret[\"attention_mask\"])\n else:\n if self.encode_whisper_embed:\n audio_features = torch.zeros(\n (\n 0,\n self.whisper_processor.feature_extractor.feature_size,\n self.whisper_processor.feature_extractor.nb_max_frames,\n ),\n dtype=torch.float32,\n )\n audio_feature_attention_mask = torch.zeros(\n (0, self.whisper_processor.feature_extractor.nb_max_frames), dtype=torch.int32\n )\n else:\n audio_features = None\n audio_feature_attention_mask = None\n\n # Process audio input tokens\n if len(audio_in_ids_l) > 0:\n # Append audio-stream-bos and eos tokens\n new_audio_in_ids_l = []\n for ele in audio_in_ids_l:\n if self.disable_audio_codes_transform:\n # Do not add audio-stream-bos or eos tokens.\n # This may indicate that the sample comes from ConstantLengthDatasetWithBuffer.\n audio_codes = ele\n else:\n audio_codes = torch.cat(\n [\n torch.full((ele.shape[0], 1), self.audio_stream_bos_id, dtype=torch.long),\n ele,\n torch.full((ele.shape[0], 1), self.audio_stream_eos_id, dtype=torch.long),\n ],\n dim=1,\n )\n if self.use_delay_pattern:\n audio_codes = build_delay_pattern_mask(\n audio_codes.unsqueeze(0),\n bos_token_id=self.audio_stream_bos_id,\n pad_token_id=self.audio_stream_eos_id,\n )[0].squeeze(0)\n new_audio_in_ids_l.append(audio_codes)\n audio_in_ids = torch.cat(new_audio_in_ids_l, dim=1).long()\n audio_in_ids_start = torch.cumsum(\n torch.tensor([0] + [audio_codes.shape[1] for audio_codes in new_audio_in_ids_l[:-1]]), dim=0\n )\n else:\n audio_in_ids = torch.zeros((0, 0), dtype=torch.long)\n audio_in_ids_start = torch.zeros(0, dtype=torch.long)\n\n # Process audio output tokens\n audio_out_ids_start_group_loc = None\n if len(audio_out_ids_l) > 0:\n new_audio_out_ids_l = []\n label_audio_ids_l = []\n for idx, ele in enumerate(audio_out_ids_l):\n if self.disable_audio_codes_transform:\n # Do not add audio-stream-bos or eos tokens.\n # This may indicate that the sample comes from ConstantLengthDatasetWithBuffer.\n audio_codes = ele\n if return_labels:\n label_audio_ids = audio_out_label_ids_l[idx]\n else:\n audio_codes = torch.cat(\n [\n torch.full((ele.shape[0], 1), self.audio_stream_bos_id, dtype=torch.long),\n ele,\n torch.full((ele.shape[0], 1), self.audio_stream_eos_id, dtype=torch.long),\n ],\n dim=1,\n )\n if return_labels:\n label_audio_ids = torch.cat(\n [\n torch.full((ele.shape[0], 1), -100, dtype=torch.long),\n ele,\n torch.full((ele.shape[0], 1), self.audio_stream_eos_id, dtype=torch.long),\n ],\n dim=1,\n )\n if self.use_delay_pattern:\n audio_codes = build_delay_pattern_mask(\n audio_codes.unsqueeze(0),\n bos_token_id=self.audio_stream_bos_id,\n pad_token_id=self.audio_stream_eos_id,\n )[0].squeeze(0)\n if return_labels:\n label_audio_ids = build_delay_pattern_mask(\n label_audio_ids.unsqueeze(0),\n bos_token_id=-100,\n pad_token_id=-100,\n )[0].squeeze(0)\n new_audio_out_ids_l.append(audio_codes)\n\n if return_labels:\n if audio_out_no_train_flag[idx]:\n label_audio_ids[:] = -100\n label_audio_ids_l.append(label_audio_ids)\n\n audio_out_ids = torch.cat(new_audio_out_ids_l, dim=1).long()\n if return_labels:\n label_audio_ids = torch.cat(label_audio_ids_l, dim=1).long()\n audio_out_ids_start = torch.cumsum(\n torch.tensor([0] + [audio_codes.shape[1] for audio_codes in new_audio_out_ids_l[:-1]]), dim=0\n )\n audio_out_ids_start_group_loc = torch.tensor(audio_out_ids_group_loc_l, dtype=torch.long)\n else:\n audio_out_ids = torch.zeros((0, 0), dtype=torch.long)\n audio_out_ids_start = torch.zeros(0, dtype=torch.long)\n if return_labels:\n label_audio_ids = torch.zeros((0, 0), dtype=torch.long)\n\n reward = torch.tensor(reward_l, dtype=torch.float32)\n\n # Handle padding for input ids and attention mask\n if self.pad_left:\n input_ids = torch.stack(\n [\n F.pad(ele.input_ids, (max_seq_length - len(ele.input_ids), 0), value=self.pad_token_id)\n for ele in processed_batch\n ]\n )\n if return_labels:\n label_ids = torch.stack(\n [\n F.pad(ele.label_ids, (max_seq_length - len(ele.label_ids), 0), value=-100)\n for ele in processed_batch\n ]\n )\n attention_mask = torch.stack(\n [\n F.pad(torch.ones_like(ele.input_ids), (max_seq_length - len(ele.input_ids), 0), value=0)\n for ele in processed_batch\n ]\n )\n else:\n input_ids = torch.stack(\n [\n F.pad(ele.input_ids, (0, max_seq_length - len(ele.input_ids)), value=self.pad_token_id)\n for ele in processed_batch\n ]\n )\n if return_labels:\n label_ids = torch.stack(\n [\n F.pad(ele.label_ids, (0, max_seq_length - len(ele.label_ids)), value=-100)\n for ele in processed_batch\n ]\n )\n attention_mask = torch.stack(\n [\n F.pad(torch.ones_like(ele.input_ids), (0, max_seq_length - len(ele.input_ids)), value=0)\n for ele in processed_batch\n ]\n )\n\n if not self.return_audio_in_tokens:\n audio_in_ids = None\n audio_in_ids_start = None\n\n # Apply audio_num_codebooks limit if specified\n if self.audio_num_codebooks is not None:\n if audio_in_ids is not None:\n audio_in_ids = audio_in_ids[: self.audio_num_codebooks]\n if audio_out_ids is not None:\n audio_out_ids = audio_out_ids[: self.audio_num_codebooks]\n if label_audio_ids is not None:\n label_audio_ids = label_audio_ids[: self.audio_num_codebooks]\n\n return HiggsAudioBatchInput(\n input_ids=input_ids,\n attention_mask=attention_mask,\n audio_features=audio_features,\n audio_feature_attention_mask=audio_feature_attention_mask,\n audio_out_ids=audio_out_ids,\n audio_out_ids_start=audio_out_ids_start,\n audio_out_ids_start_group_loc=audio_out_ids_start_group_loc,\n audio_in_ids=audio_in_ids,\n audio_in_ids_start=audio_in_ids_start,\n label_ids=label_ids,\n label_audio_ids=label_audio_ids,\n reward=reward,\n )\n"
},
{
"path": "boson_multimodal/data_types.py",
"content": "\"\"\"Basic data types for multimodal ChatML format.\"\"\"\n\nfrom dataclasses import dataclass\nfrom typing import Dict, List, Optional, Union\n\n\n@dataclass\nclass AudioContent:\n audio_url: str\n # Base64 encoded audio bytes\n raw_audio: Optional[str] = None\n offset: Optional[float] = None\n duration: Optional[float] = None\n row_id: Optional[int] = None\n type: str = \"audio\"\n\n\n@dataclass\nclass TextContent:\n text: str\n type: str = \"text\"\n\n\n@dataclass\nclass Message:\n role: str\n content: Union[str, AudioContent, TextContent, List[Union[str, AudioContent, TextContent]]]\n recipient: Optional[str] = None\n\n\n@dataclass\nclass ChatMLSample:\n \"\"\"Dataclass to hold multimodal ChatML data.\"\"\"\n\n messages: List[Message]\n start_index: Optional[int] = None # We will mask the messages[:start_index] when finetuning the LLM.\n misc: Optional[Dict] = None\n speaker: Optional[str] = None\n"
},
{
"path": "boson_multimodal/dataset/__init__.py",
"content": ""
},
{
"path": "boson_multimodal/dataset/chatml_dataset.py",
"content": "import dacite\nimport pandas as pd\nimport torch\nimport json\n\nimport numpy as np\nimport multiprocessing as mp\n\nfrom dataclasses import dataclass, fields\nfrom abc import ABC, abstractmethod\nfrom typing import Union, List, Dict, Optional\n\nfrom ..data_types import ChatMLSample, TextContent, AudioContent\nfrom ..constants import AUDIO_IN_TOKEN, AUDIO_OUT_TOKEN\n\nfrom loguru import logger\n\n# Whisper processor, 30 sec -> 3000 features\n# Then we divide 4 in the audio towker, we decrease 3000 features to 750, which gives 25 Hz\nWHISPER_EMBED_NUM_HIDDEN_STATE_PER_SEC = 25\n\n\n@dataclass\nclass ChatMLDatasetSample:\n input_ids: torch.LongTensor # Shape (seq_len,): The input text tokens.\n label_ids: torch.LongTensor # Shape (seq_len,): The label ids.\n audio_ids_concat: torch.LongTensor # Shape (num_codebooks, audio_seq_len): The audio tokens that are concatenated.\n # Here `audio_seq_len` is the length of the concatenated audio tokens.`\n audio_ids_start: (\n torch.LongTensor\n ) # Shape (num_audios,): The start index of each audio token in the concatenated audio tokens.\n audio_waveforms_concat: (\n torch.Tensor\n ) # Shape (total_wv_length,): The concatenated audio waveforms for audio-in features.\n audio_waveforms_start: (\n torch.LongTensor\n ) # Shape (num_audios,): The start index of each audio waveform in the concatenated audio waveforms.\n audio_sample_rate: torch.Tensor # Shape (num_audios,): The sampling rate of the audio waveforms.\n audio_speaker_indices: (\n torch.LongTensor\n ) # Shape (num_audios,) -1 means unknown speaker: The speaker indices for each audio.\n audio_label_ids_concat: Optional[torch.LongTensor] = (\n None # Shape (num_codebooks, audio_seq_len): The audio tokens that are concatenated.\n )\n # Here `audio_seq_len` is the length of the concatenated audio tokens.`\n reward: Optional[float] = None\n\n def num_audios(self):\n return max(len(self.audio_waveforms_start), len(self.audio_ids_start))\n\n def get_audio_codes(self, idx):\n code_start = self.audio_ids_start[idx]\n if idx < len(self.audio_ids_start) - 1:\n code_end = self.audio_ids_start[idx + 1]\n else:\n code_end = self.audio_ids_concat.shape[-1]\n\n return self.audio_ids_concat[:, code_start:code_end]\n\n def get_audio_codes_labels(self, idx):\n if self.audio_label_ids_concat is None:\n return None\n code_start = self.audio_ids_start[idx]\n if idx < len(self.audio_ids_start) - 1:\n code_end = self.audio_ids_start[idx + 1]\n else:\n code_end = self.audio_ids_concat.shape[-1]\n\n return self.audio_label_ids_concat[:, code_start:code_end]\n\n def get_wv(self, idx):\n wv_start = self.audio_waveforms_start[idx]\n sr = self.audio_sample_rate[idx]\n if idx < len(self.audio_waveforms_start) - 1:\n wv_end = self.audio_waveforms_start[idx + 1]\n else:\n wv_end = self.audio_waveforms_concat.shape[-1]\n return self.audio_waveforms_concat[wv_start:wv_end], sr\n\n def cal_num_tokens(\n self,\n encode_whisper_embed: bool = True,\n encode_audio_in_tokens: bool = False,\n encode_audio_out_tokens: bool = True,\n audio_in_token_id: int = 128015,\n audio_out_token_id: int = 128016,\n ) -> int:\n # we firstly exclude <|AUDIO|> and <|AUDIO_OUT|> because we do late merging and replace those position with actual audio features and audio token ids\n # It's assumed that we always have audio_ids when audio_waveforms are there (but not vice-versa)\n num_tokens = len(self.input_ids) - len(self.audio_ids_start)\n\n if encode_whisper_embed and len(self.audio_waveforms_concat) > 0:\n audio_lengths = torch.diff(self.audio_waveforms_start)\n if len(audio_lengths):\n # Sum before calling .item()\n num_tokens += (\n (\n np.ceil(WHISPER_EMBED_NUM_HIDDEN_STATE_PER_SEC * audio_lengths / self.audio_sample_rate[:-1])\n ).sum()\n ).item()\n # add the last audio's token estimation\n num_tokens += (\n np.ceil(\n WHISPER_EMBED_NUM_HIDDEN_STATE_PER_SEC\n * (self.audio_waveforms_concat.shape[0] - self.audio_waveforms_start[-1])\n / self.audio_sample_rate[-1]\n )\n ).item()\n\n if self.audio_ids_concat.size(1) > 0:\n audio_io_ids = self.input_ids[\n (self.input_ids == audio_in_token_id) | (self.input_ids == audio_out_token_id)\n ]\n audio_io_id_lengths = torch.concat(\n [\n torch.diff(self.audio_ids_start),\n torch.tensor([self.audio_ids_concat.shape[-1] - self.audio_ids_start[-1]]),\n ]\n )\n if encode_audio_in_tokens:\n num_tokens += torch.sum(audio_io_id_lengths[audio_io_ids == audio_in_token_id]).item()\n\n if encode_audio_out_tokens:\n num_tokens += torch.sum(audio_io_id_lengths[audio_io_ids == audio_out_token_id]).item()\n\n return int(num_tokens)\n\n @classmethod\n def merge(\n cls,\n samples: List[\"ChatMLDatasetSample\"],\n eos_token_id: int,\n ignore_index: int,\n padding_size: Optional[int] = None,\n ) -> \"ChatMLDatasetSample\":\n \"\"\"Merges a list of ChatMLDatasetSample instances, inserting eos_token_id and ignore_index between them, and adjusting offsets for audio_ids_start and audio_waveforms_start.\n\n Args:\n samples (List[ChatMLDatasetSample]): List of samples to merge.\n eos_token_id (int): Tokens to be inserted into input_ids between samples.\n ignore_index (int): Default label for padding.\n padding_size (Optional[int]): If provided, pad the sequence to with this length.\n\n Returns:\n ChatMLDatasetSample: Merged and potentially padded sample.\n \"\"\"\n if not samples:\n logger.fatal(\"The samples list is empty and cannot be merged.\")\n raise ValueError(\"The samples list is empty and cannot be merged.\")\n\n # Initialize empty lists for concatenation\n input_ids_list = []\n label_ids_list = []\n audio_ids_concat_list = []\n audio_ids_start_list = []\n audio_waveforms_concat_list = []\n audio_waveforms_start_list = []\n audio_sample_rate_list = []\n audio_speaker_indices_list = []\n\n # Track offsets\n audio_ids_offset = 0\n audio_waveforms_offset = 0\n\n for sample in samples:\n # Add input_ids and label_ids with padding\n if input_ids_list:\n input_ids_list.append(torch.tensor([eos_token_id], dtype=torch.long))\n label_ids_list.append(torch.tensor([ignore_index], dtype=torch.long))\n input_ids_list.append(sample.input_ids)\n label_ids_list.append(sample.label_ids)\n\n # Add audio_ids_concat and handle empty audio ids\n if sample.audio_ids_concat.size(1) > 0:\n audio_ids_concat_list.append(sample.audio_ids_concat)\n\n # Offset and add audio_ids_start\n audio_ids_start_list.append(sample.audio_ids_start + audio_ids_offset)\n audio_ids_offset += sample.audio_ids_concat.size(\n 1\n ) # (num_codebooks, seq_len): Update offset by audio_seq_len\n\n # Add audio_waveforms_concat\n if sample.audio_waveforms_concat.size(0) > 0:\n # Check dimensions of the audio waveform to ensure consistency\n if (\n audio_waveforms_concat_list\n and sample.audio_waveforms_concat.dim() != audio_waveforms_concat_list[0].dim()\n ):\n logger.warning(\n f\"Skipping audio waveform with inconsistent dimensions: expected {audio_waveforms_concat_list[0].dim()}D, got {sample.audio_waveforms_concat.dim()}D\"\n )\n continue\n\n audio_waveforms_concat_list.append(sample.audio_waveforms_concat)\n audio_waveforms_start_list.append(sample.audio_waveforms_start + audio_waveforms_offset)\n audio_waveforms_offset += sample.audio_waveforms_concat.size(0)\n\n # Add audio_sample_rate and audio_speaker_indices\n audio_sample_rate_list.append(sample.audio_sample_rate)\n\n audio_speaker_indices_list.append(sample.audio_speaker_indices)\n\n # Concatenate all tensors\n input_ids = torch.cat(input_ids_list, dim=0)\n label_ids = torch.cat(label_ids_list, dim=0)\n\n # Apply padding if padding_size is specified\n if padding_size is not None and padding_size > 0:\n input_ids = torch.cat([input_ids, torch.full((padding_size,), eos_token_id, dtype=torch.long)], dim=0)\n label_ids = torch.cat([label_ids, torch.full((padding_size,), ignore_index, dtype=torch.long)], dim=0)\n\n # Safely concatenate audio tensors with proper error handling\n try:\n audio_ids_concat = torch.cat(audio_ids_concat_list, dim=1) if audio_ids_concat_list else torch.tensor([[]])\n audio_ids_start = torch.cat(audio_ids_start_list, dim=0) if audio_ids_start_list else torch.tensor([])\n\n # Check for dimensional consistency in audio waveforms\n if audio_waveforms_concat_list:\n dims = [t.dim() for t in audio_waveforms_concat_list]\n if not all(d == dims[0] for d in dims):\n # If dimensions don't match, log warning and filter out the problematic tensors\n logger.warning(\n f\"Inconsistent dimensions in audio waveforms: {dims}. Filtering to keep only consistent ones.\"\n )\n expected_dim = max(set(dims), key=dims.count) # Most common dimension\n audio_waveforms_concat_list = [t for t in audio_waveforms_concat_list if t.dim() == expected_dim]\n\n # Recalculate audio_waveforms_start with the filtered list\n if audio_waveforms_concat_list:\n audio_waveforms_offset = 0\n audio_waveforms_start_list = []\n for waveform in audio_waveforms_concat_list:\n audio_waveforms_start_list.append(torch.tensor([audio_waveforms_offset]))\n audio_waveforms_offset += waveform.size(0)\n\n audio_waveforms_concat = (\n torch.cat(audio_waveforms_concat_list, dim=0) if audio_waveforms_concat_list else torch.tensor([])\n )\n audio_waveforms_start = (\n torch.cat(audio_waveforms_start_list, dim=0) if audio_waveforms_start_list else torch.tensor([])\n )\n audio_sample_rate = (\n torch.cat(audio_sample_rate_list, dim=0) if audio_sample_rate_list else torch.tensor([])\n )\n audio_speaker_indices = (\n torch.cat(audio_speaker_indices_list, dim=0) if audio_speaker_indices_list else torch.tensor([])\n )\n\n except RuntimeError as e:\n logger.error(f\"Error during tensor concatenation: {str(e)}\")\n logger.warning(\"Falling back to empty audio tensors\")\n # Fall back to empty tensors\n audio_ids_concat = torch.tensor([[]])\n audio_ids_start = torch.tensor([])\n audio_waveforms_concat = torch.tensor([])\n audio_waveforms_start = torch.tensor([])\n audio_sample_rate = torch.tensor([])\n audio_speaker_indices = torch.tensor([])\n\n # Create the merged sample\n merged_sample = cls(\n input_ids=input_ids,\n label_ids=label_ids,\n audio_ids_concat=audio_ids_concat,\n audio_ids_start=audio_ids_start,\n audio_waveforms_concat=audio_waveforms_concat,\n audio_waveforms_start=audio_waveforms_start,\n audio_sample_rate=audio_sample_rate,\n audio_speaker_indices=audio_speaker_indices,\n )\n\n return merged_sample\n\n\n@dataclass\nclass RankedChatMLDatasetSampleTuple:\n samples: List[ChatMLDatasetSample]\n scores: List[float]\n\n def max_score_sample(self) -> ChatMLDatasetSample:\n idx = self.scores.index(max(self.scores))\n self.samples[idx].reward = self.scores[idx]\n return self.samples[idx]\n\n def min_score_sample(self) -> ChatMLDatasetSample:\n idx = self.scores.index(min(self.scores))\n self.samples[idx].reward = self.scores[idx]\n return self.samples[idx]\n\n\n@dataclass\nclass ChatMLDatasetStorageSample:\n input_tokens: torch.LongTensor\n label_tokens: torch.LongTensor\n audio_bytes_cache_dir_index: int\n audio_codes_cache_dir_index: int\n audio_bytes_indices: torch.LongTensor\n audio_codes_indices: torch.LongTensor\n speaker_indices: torch.LongTensor\n file_index: int\n original_sample_index: int\n\n\n# TODO(sxjscience): We need to revist the logic about parsing speaker ids.\n# Currently, we assume that the speaker id is stored at the \"misc\" field in ChatMLSample.\ndef prepare_chatml_sample(sample: Union[ChatMLSample, Dict], tokenizer):\n \"\"\"Preprocess the ChatML sample to get the tokens for the text part.\n\n Args:\n sample (ChatMLSample): The ChatML sample to preprocess.\n tokenizer: The tokenizer to use for encoding the text.\n\n \"\"\"\n\n try:\n if not isinstance(sample, ChatMLSample):\n # Handle all fields that could be NaN\n if \"speaker\" in sample and pd.isna(sample[\"speaker\"]):\n sample[\"speaker\"] = None\n if \"start_index\" in sample and pd.isna(sample[\"start_index\"]):\n sample[\"start_index\"] = None\n if \"content\" in sample and pd.isna(sample[\"content\"]):\n sample[\"content\"] = \"\"\n\n # Convert any other potential NaN values in nested structures\n def convert_nan_to_none(obj):\n import numpy as np\n\n if isinstance(obj, (pd.Series, np.ndarray)):\n return obj.tolist()\n elif pd.api.types.is_scalar(obj) and pd.isna(obj):\n return None\n elif isinstance(obj, dict):\n return {k: convert_nan_to_none(v) for k, v in obj.items()}\n elif isinstance(obj, (list, tuple)): # Fixed: Handle both list and tuple\n return [convert_nan_to_none(item) for item in obj]\n return obj\n\n # Clean the sample data\n clean_sample = convert_nan_to_none(sample)\n\n val_keys = []\n for field in fields(ChatMLSample):\n if field.name in clean_sample:\n val_keys.append(field.name)\n clean_sample = {k: clean_sample[k] for k in val_keys}\n\n try:\n sample = dacite.from_dict(\n data_class=ChatMLSample, data=clean_sample, config=dacite.Config(strict=True, check_types=True)\n )\n except Exception as e:\n print(f\"Failed to convert to ChatMLSample: {e}\")\n print(f\"Clean sample: {json.dumps(clean_sample, indent=2)}\")\n return None, None, None, None\n\n input_tokens = []\n label_tokens = []\n audio_contents = []\n speaker_id = None\n if sample.speaker is not None:\n speaker_id = sample.speaker\n elif sample.misc is not None:\n if \"speaker\" in sample.misc:\n speaker_id = sample.misc[\"speaker\"]\n\n total_m = len(sample.messages)\n for turn_id, message in enumerate(sample.messages):\n role = message.role\n recipient = message.recipient\n content = message.content\n content_l = []\n\n if isinstance(content, str):\n content_l.append(TextContent(text=content))\n elif isinstance(content, TextContent):\n content_l.append(content)\n elif isinstance(content, AudioContent):\n content_l.append(content)\n elif isinstance(content, list):\n for ele in content:\n if isinstance(ele, str):\n content_l.append(TextContent(text=ele))\n else:\n content_l.append(ele)\n if turn_id == 0:\n prefix = f\"<|begin_of_text|><|start_header_id|>{role}<|end_header_id|>\\n\\n\"\n else:\n prefix = f\"<|start_header_id|>{role}<|end_header_id|>\\n\\n\"\n eot_postfix = \"<|eot_id|>\"\n eom_postfix = \"<|eom_id|>\"\n\n prefix_tokens = tokenizer.encode(prefix, add_special_tokens=False)\n input_tokens.extend(prefix_tokens)\n label_tokens.extend([-100 for _ in prefix_tokens])\n\n if recipient:\n assert role == \"assistant\", \"Recipient is only available for assistant role.\"\n recipient_tokens = tokenizer.encode(f\"{recipient}<|recipient|>\", add_special_tokens=False)\n input_tokens.extend(recipient_tokens)\n label_tokens.extend(recipient_tokens)\n\n for content in content_l:\n if content.type == \"text\":\n text_tokens = tokenizer.encode(content.text, add_special_tokens=False)\n input_tokens.extend(text_tokens)\n if role == \"assistant\" and (sample.start_index is None or turn_id >= sample.start_index):\n label_tokens.extend(text_tokens)\n else:\n label_tokens.extend([-100 for _ in text_tokens])\n\n elif content.type == \"audio\":\n # Generate the text-part of the audio tokens\n audio_contents.append(content)\n if role == \"user\" or role == \"system\":\n # Add the text tokens\n text_tokens = tokenizer.encode(\n f\"<|audio_bos|><|AUDIO|><|audio_eos|>\",\n add_special_tokens=False,\n )\n input_tokens.extend(text_tokens)\n label_tokens.extend([-100 for _ in text_tokens])\n elif role == \"assistant\":\n # Add the text tokens for audio-out part.\n text_tokens = tokenizer.encode(\n f\"<|audio_out_bos|><|AUDIO_OUT|><|audio_eos|>\",\n add_special_tokens=False,\n )\n input_tokens.extend(text_tokens)\n if sample.start_index is None or turn_id >= sample.start_index:\n label_tokens.extend(text_tokens)\n else:\n label_tokens.extend([-100 for _ in text_tokens])\n next_id = turn_id + 1\n if role == \"assistant\" and next_id != total_m and sample.messages[next_id].role == \"assistant\":\n postfix_tokens = tokenizer.encode(eom_postfix, add_special_tokens=False)\n input_tokens.extend(postfix_tokens)\n else:\n postfix_tokens = tokenizer.encode(eot_postfix, add_special_tokens=False)\n input_tokens.extend(postfix_tokens)\n if role == \"assistant\" and (sample.start_index is None or turn_id >= sample.start_index):\n label_tokens.extend(postfix_tokens)\n else:\n label_tokens.extend([-100 for _ in postfix_tokens])\n\n return input_tokens, label_tokens, audio_contents, speaker_id\n\n except Exception as e:\n print(f\"Error in prepare_chatml_sample: {str(e)}\")\n print(f\"Sample data: {json.dumps(sample, indent=2)}\")\n return None, None, None, None\n\n\ndef extract_generation_prompt_from_input_tokens(input_tokens, tokenizer):\n \"\"\"Extract the generation prompt and reference answer from the input tokens.\n\n For example:\n\n Input Text = '<|begin_of_text|><|start_header_id|>user<|end_header_id|>\\n\\n\n What words do you hear from the provided audio? Write it down for me.<|audio_bos|><|AUDIO|><|audio_eos|><|eot_id|>\n <|start_header_id|>assistant<|end_header_id|>\\n\\nAt first they went by quick, too quick to even get.<|eot_id|>'\n\n -->\n\n Prompt = '<|begin_of_text|><|start_header_id|>user<|end_header_id|>\\n\\n\n What words do you hear from the provided audio? Write it down for me.<|audio_bos|><|AUDIO|><|audio_eos|><|eot_id|>\n <|start_header_id|>assistant<|end_header_id|>\\n\\n',\n Reference = 'At first they went by quick, too quick to even get.'\n\n Args:\n input_tokens: The input tokens.\n audio_contents: The audio contents.\n tokenizer: The tokenizer to use for decoding the text.\n\n Returns:\n prompt_tokens: The tokens for the prompt.\n reference_answer: The reference answer.\n num_audios_in_reference: The number of audios in the reference answer.\n\n \"\"\"\n input_text = tokenizer.decode(input_tokens)\n generation_prefix = \"<|start_header_id|>assistant<|end_header_id|>\\n\\n\"\n postfix = \"<|eot_id|>\"\n assert generation_prefix in input_text\n generation_prompt_end_loc = input_text.rfind(generation_prefix) + len(generation_prefix)\n generation_prompt = input_text[:generation_prompt_end_loc]\n reference_answer = input_text[generation_prompt_end_loc : input_text.find(postfix, generation_prompt_end_loc)]\n num_audios_in_reference = reference_answer.count(AUDIO_IN_TOKEN) + reference_answer.count(AUDIO_OUT_TOKEN)\n return tokenizer.encode(generation_prompt, add_special_tokens=False), reference_answer, num_audios_in_reference\n\n\ndef prepare_chatml_dataframe_single_process(df, tokenizer):\n \"\"\"Prepare the ChatML DataFrame.\"\"\"\n ret = []\n for _, row in df.iterrows():\n input_tokens, label_tokens, audio_contents, speaker_id = prepare_chatml_sample(row.to_dict(), tokenizer)\n ret.append((input_tokens, label_tokens, audio_contents, speaker_id))\n return ret\n\n\ndef prepare_chatml_dataframe(df, tokenizer, num_process=16):\n if num_process is None:\n return prepare_chatml_dataframe_single_process(df, tokenizer)\n else:\n num_process = max(min(len(df) // 1000, num_process), 1)\n workloads = np.array_split(df, num_process)\n with mp.Pool(num_process) as pool:\n ret = pool.starmap(\n prepare_chatml_dataframe_single_process, [(workload, tokenizer) for workload in workloads]\n )\n return sum(ret, [])\n\n\nclass DatasetInterface(ABC):\n @abstractmethod\n def __getitem__(self, idx) -> Union[\"ChatMLDatasetSample\", \"RankedChatMLDatasetSampleTuple\"]:\n \"\"\"Retrieve a dataset sample by index.\"\"\"\n raise NotImplementedError\n\n\nclass IterableDatasetInterface(ABC):\n @abstractmethod\n def __iter__(self) -> Union[\"ChatMLDatasetSample\", \"RankedChatMLDatasetSampleTuple\"]:\n \"\"\"Retrieve a sample by iterating through the dataset.\"\"\"\n raise NotImplementedError\n\n\n@dataclass\nclass DatasetInfo:\n dataset_type: str\n group_type: Optional[str] = None\n mask_text: Optional[bool] = None # Whether to mask the text tokens for pretraining samples.\n"
},
{
"path": "boson_multimodal/model/__init__.py",
"content": ""
},
{
"path": "boson_multimodal/model/higgs_audio/__init__.py",
"content": "from transformers import AutoConfig, AutoModel\n\nfrom .configuration_higgs_audio import HiggsAudioConfig, HiggsAudioEncoderConfig\nfrom .modeling_higgs_audio import HiggsAudioModel\n\n\nAutoConfig.register(\"higgs_audio_encoder\", HiggsAudioEncoderConfig)\nAutoConfig.register(\"higgs_audio\", HiggsAudioConfig)\nAutoModel.register(HiggsAudioConfig, HiggsAudioModel)\n"
},
{
"path": "boson_multimodal/model/higgs_audio/audio_head.py",
"content": "\"\"\"Projector that maps hidden states from the LLM component to multimodal logits.\"\"\"\n\nimport torch\nfrom torch import nn\n\nfrom dataclasses import dataclass\nfrom typing import Optional, Tuple\n\nfrom .common import HiggsAudioPreTrainedModel\nfrom .configuration_higgs_audio import HiggsAudioConfig\n\n\n@dataclass\nclass HiggsAudioDecoderLayerOutput:\n logits: torch.FloatTensor\n audio_logits: torch.FloatTensor\n attentions: Optional[Tuple[torch.FloatTensor, ...]] = None\n past_key_values: Optional[Tuple[Tuple[torch.FloatTensor]]] = None\n\n\nclass HiggsAudioDecoderProjector(HiggsAudioPreTrainedModel):\n \"\"\"Projection layers that map hidden states from the LLM component to audio / text logits.\n\n We support two type of audio head:\n - Basic Audio Head:\n Directly map the hidden states to audio logits for all the codebooks.\n \"\"\"\n\n def __init__(self, config: HiggsAudioConfig, layer_idx: Optional[int] = None):\n super().__init__(config)\n self.text_lm_head = nn.Linear(config.text_config.hidden_size, config.text_config.vocab_size, bias=False)\n self.audio_lm_head = nn.Linear(\n config.text_config.hidden_size, config.audio_num_codebooks * (config.audio_codebook_size + 2), bias=False\n )\n\n # Initialize weights and apply final processing\n self.post_init()\n\n def forward(\n self,\n hidden_states,\n audio_out_mask,\n label_audio_ids=None,\n attention_mask=None,\n position_ids=None,\n past_key_values=None,\n use_cache=None,\n output_attentions=None,\n output_hidden_states=None,\n output_audio_hidden_states=False,\n cache_position=None,\n ):\n \"\"\"\n Args:\n hidden_states (`torch.Tensor` of shape `(batch_size, seq_len, hidden_size)`):\n Hidden states from the LLM component\n audio_out_mask (`torch.Tensor` of shape `(batch_size, seq_len)`):\n Mask for identifying the audio out tokens.\n label_audio_ids (`torch.Tensor` of shape `(num_codebooks, num_audio_out_tokens)`):\n Label tokens for the audio-out part. This is used for calculating the logits if RQ-Transformer is used.\n attention_mask (`torch.Tensor` of shape `(batch_size, seq_len)`):\n Mask to avoid performing attention on padding token indices\n position_ids (`torch.Tensor` of shape `(batch_size, seq_len)`):\n Position ids for the input tokens\n\n Returns:\n logits (`torch.Tensor` of shape `(batch_size, seq_len, vocab_size)`):\n Logits for text tokens\n audio_logits (`torch.Tensor` of shape `(num_audio_out_tokens, audio_num_codebooks * audio_codebook_size)`):\n Logits for audio tokens. We ensure `num_text_tokens + num_audio_tokens == batch_size * seq_len`\n \"\"\"\n logits = self.text_lm_head(hidden_states)\n\n all_hidden_states = () if output_hidden_states else None\n all_self_attns = () if output_attentions else None\n next_decoder_cache = None\n\n if self.config.audio_decoder_proj_num_layers > 0:\n # create position embeddings to be shared across the decoder layers\n position_embeddings = self.rotary_emb(hidden_states, position_ids)\n for decoder_layer in self.transformer_layers:\n if output_hidden_states:\n all_hidden_states += (hidden_states,)\n\n if self.gradient_checkpointing and self.training:\n layer_outputs = self._gradient_checkpointing_func(\n decoder_layer.__call__,\n hidden_states,\n attention_mask,\n position_ids,\n past_key_values,\n output_attentions,\n use_cache,\n cache_position,\n position_embeddings,\n )\n else:\n layer_outputs = decoder_layer(\n hidden_states,\n attention_mask=attention_mask,\n position_ids=position_ids,\n past_key_value=past_key_values,\n output_attentions=output_attentions,\n use_cache=use_cache,\n cache_position=cache_position,\n position_embeddings=position_embeddings,\n )\n hidden_states = layer_outputs[0]\n hidden_states = self.norm(hidden_states)\n\n if output_hidden_states:\n all_hidden_states += (hidden_states,)\n\n if output_attentions:\n all_self_attns += (layer_outputs[1],)\n\n if use_cache:\n next_decoder_cache = layer_outputs[2 if output_attentions else 1]\n\n next_cache = next_decoder_cache if use_cache else None\n\n audio_logits = self.audio_lm_head(hidden_states[audio_out_mask])\n\n if output_audio_hidden_states:\n audio_hidden_states = hidden_states[audio_out_mask]\n else:\n audio_hidden_states = None\n\n return logits, audio_logits, all_self_attns, all_hidden_states, audio_hidden_states, next_cache\n"
},
{
"path": "boson_multimodal/model/higgs_audio/common.py",
"content": "from torch import nn\n\nfrom transformers.modeling_utils import PreTrainedModel\n\nfrom .configuration_higgs_audio import HiggsAudioConfig\n\n\nclass HiggsAudioPreTrainedModel(PreTrainedModel):\n config_class = HiggsAudioConfig\n base_model_prefix = \"model\"\n supports_gradient_checkpointing = True\n _no_split_modules = []\n _skip_keys_device_placement = \"past_key_values\"\n _supports_flash_attn_2 = True\n _supports_sdpa = True\n\n def _init_weights(self, module):\n std = self.config.init_std if hasattr(self.config, \"init_std\") else self.config.audio_encoder_config.init_std\n\n if isinstance(module, (nn.Linear, nn.Conv1d)):\n module.weight.data.normal_(mean=0.0, std=std)\n if module.bias is not None:\n module.bias.data.zero_()\n elif isinstance(module, nn.Embedding):\n module.weight.data.normal_(mean=0.0, std=std)\n if module.padding_idx is not None:\n module.weight.data[module.padding_idx].zero_()\n"
},
{
"path": "boson_multimodal/model/higgs_audio/configuration_higgs_audio.py",
"content": "from transformers.configuration_utils import PretrainedConfig\nfrom transformers.models.auto import CONFIG_MAPPING\n\n\nclass HiggsAudioEncoderConfig(PretrainedConfig):\n \"\"\"Configuration of the Audio encoder in Higgs-Audio.\"\"\"\n\n model_type = \"higgs_audio_encoder\"\n\n def __init__(\n self,\n num_mel_bins=128,\n encoder_layers=32,\n encoder_attention_heads=20,\n encoder_ffn_dim=5120,\n encoder_layerdrop=0.0,\n d_model=1280,\n dropout=0.0,\n attention_dropout=0.0,\n activation_function=\"gelu\",\n activation_dropout=0.0,\n scale_embedding=False,\n init_std=0.02,\n max_source_positions=1500,\n pad_token_id=128001,\n **kwargs,\n ):\n super().__init__(**kwargs)\n\n self.num_mel_bins = num_mel_bins\n self.d_model = d_model\n self.encoder_layers = encoder_layers\n self.encoder_attention_heads = encoder_attention_heads\n self.encoder_ffn_dim = encoder_ffn_dim\n self.dropout = dropout\n self.attention_dropout = attention_dropout\n self.activation_function = activation_function\n self.activation_dropout = activation_dropout\n self.encoder_layerdrop = encoder_layerdrop\n self.num_hidden_layers = encoder_layers\n self.init_std = init_std\n self.scale_embedding = scale_embedding # scale factor will be sqrt(d_model) if True\n self.max_source_positions = max_source_positions\n self.pad_token_id = pad_token_id\n\n\nclass HiggsAudioConfig(PretrainedConfig):\n r\"\"\"\n This is the configuration class for the HiggsAudioModel.\n\n Args:\n text_config (`Union[AutoConfig, dict]`):\n The config object or dictionary of the text backbone.\n audio_encoder_config (`Union[AutoConfig, dict]`):\n The config object or dictionary of the whisper encoder.\n The audio encoder will be bidirectional and will be only available for audio understanding.\n audio_tokenizer_config\n The config object or dictionary of the audio tokenizer.\n audio_adapter_type\n The type of audio adapter to use. We support two types of adapter:\n - stack:\n We stack additional Transformer layers after the main LLM backbone for audio generation.\n - dual_ffn:\n For selected part of the LLM backbone, we replace the text FFN with a dual FFN architecture\n that contains an additional audio FFN. The audio FFN will be triggered when the location is marked for audio tokens.\n - dual_ffn_fast_forward:\n We pick a few layers in the LLM backbone to plug-in the audio FFN. For the remaining layers,\n the audio hidden states will be directly fast-forward to the next layer.\n This reduces the computational cost for audio generation.\n audio_embed_avg (`bool`, *optional*, defaults to False):\n Whether to average the audio embeddings before sending them to the text attention layer.\n audio_ffn_hidden_size\n The hidden size of the audio feedforward network in dual-path FFN\n audio_ffn_intermediate_size\n The intermediate size of the audio feedforward network in dual-path FFN\n audio_dual_ffn_layers\n The layers in the LLM backbone to plug-in the dual FFN layer (mixture of audio FFN and text FFN).\n audio_decoder_proj_num_attention (`int`, *optional*, defaults to 0):\n The number of attention heads in the audio decoder projection layer.\n use_delay_pattern (`bool`, *optional*, defaults to False):\n Whether to use delay pattern in the audio decoder.\n skip_audio_tower (`bool`, *optional*, defaults to False):\n Whether to skip the audio tower in the audio encoder.\n use_audio_out_embed_projector (`bool`, *optional*, defaults to False):\n Whether to use an embedding projector to map audio out embeddings.\n use_audio_out_self_attention (`bool`, *optional*, defaults to False):\n Whether to use self-attention to aggregate information from audio-tokens before sending to the text attention layer.\n audio_num_codebooks (`int`, *optional*, defaults to 12):\n The number of codebooks in RVQGAN.\n audio_codebook_size (`int`, *optional*, defaults to 1024):\n The size of each codebook in RVQGAN.\n audio_stream_bos_id\n The id of the bos in the audio stream\n audio_stream_eos_id\n The id of the eos in the audio stream\n audio_bos_token (`str`, *optional*, defaults to \"<|audio_bos|>\"):\n The special `<|audio_bos|>` token. In Higgs-Audio, it is mapped to 128011,\n which is the index of `<|reserved_special_token_3|>` in Llama-3.1-8B-Instruct's tokenizer.\n audio_eos_token (`str`, *optional*, defaults to \"<|audio_eos|>\"):\n The special `<|audio_eos|>` token. We use 128012 as the default value,\n which is the index of `<|reserved_special_token_4|>` in Llama-3.1-8B-Instruct's tokenizer.\n audio_out_bos_token (`str`, *optional*, defaults to \"<|audio_out_bos|>\"):\n The special `<|audio_out_bos|>` token. We use 128013 as the default value,\n which is the index of `<|reserved_special_token_5|>` in Llama-3.1-8B-Instruct's tokenizer.\n audio_token (`str`, *optional*, defaults to \"<|AUDIO|>\"):\n The special `<|AUDIO|>` token. We use 128015 as the default value,\n which is the index of `<|reserved_special_token_7|>` in Llama-3.1-8B-Instruct's tokenizer.\n This token indicates that the location should be filled in with whisper features.\n audio_out_token (`str`, *optional*, defaults to \"<|AUDIO_OUT|>\"):\n The special `<|AUDIO_OUT|>` token. We use 128016 as the default value,\n which is the index of `<|reserved_special_token_8|>` in Llama-3.1-8B-Instruct's tokenizer.\n This token indicates that the location should be filled in with audio tokens extracted via audio tokenizer.\n \"\"\"\n\n model_type = \"higgs_audio\"\n is_composition = True\n\n def __init__(\n self,\n text_config=None,\n audio_encoder_config=None,\n audio_tokenizer_config=None,\n audio_adapter_type=\"stack\",\n audio_embed_avg=False,\n audio_ffn_hidden_size=4096,\n audio_ffn_intermediate_size=14336,\n audio_dual_ffn_layers=None,\n audio_decoder_proj_num_layers=0,\n encode_whisper_embed=True,\n encode_audio_in_tokens=False,\n use_delay_pattern=False,\n skip_audio_tower=False,\n use_audio_out_embed_projector=False,\n use_audio_out_self_attention=False,\n use_rq_transformer=False,\n rq_transformer_hidden_size=None,\n rq_transformer_intermediate_size=None,\n rq_transformer_num_attention_heads=None,\n rq_transformer_num_key_value_heads=None,\n rq_transformer_num_hidden_layers=3,\n audio_num_codebooks=12,\n audio_codebook_size=1024,\n audio_stream_bos_id=1024,\n audio_stream_eos_id=1025,\n audio_bos_token=\"<|audio_bos|>\",\n audio_eos_token=\"<|audio_eos|>\",\n audio_out_bos_token=\"<|audio_out_bos|>\",\n audio_in_token=\"<|AUDIO|>\",\n audio_out_token=\"<|AUDIO_OUT|>\",\n audio_in_token_idx=128015,\n audio_out_token_idx=128016,\n pad_token_id=128001,\n audio_out_bos_token_id=128013,\n audio_eos_token_id=128012,\n **kwargs,\n ):\n if isinstance(audio_encoder_config, dict):\n audio_encoder_config[\"model_type\"] = (\n audio_encoder_config[\"model_type\"] if \"model_type\" in audio_encoder_config else \"higgs_audio_encoder\"\n )\n audio_encoder_config = CONFIG_MAPPING[audio_encoder_config[\"model_type\"]](**audio_encoder_config)\n elif audio_encoder_config is None:\n audio_encoder_config = HiggsAudioEncoderConfig()\n\n if isinstance(text_config, dict):\n text_config[\"model_type\"] = text_config[\"model_type\"] if \"model_type\" in text_config else \"llama\"\n text_config = CONFIG_MAPPING[text_config[\"model_type\"]](**text_config)\n elif text_config is None:\n text_config = CONFIG_MAPPING[\"llama\"]()\n\n assert audio_adapter_type in [\n \"stack\",\n \"dual_ffn\",\n \"dual_ffn_fast_forward\",\n ], f\"Invalid audio adapter type: {audio_adapter_type}\"\n if audio_adapter_type.startswith(\"dual_ffn\"):\n assert audio_dual_ffn_layers is not None, (\n \"audio_dual_ffn_layers must be specified when using dual_ffn adapter.\"\n )\n self.text_config = text_config\n self.audio_encoder_config = audio_encoder_config\n self.audio_tokenizer_config = audio_tokenizer_config\n self.audio_adapter_type = audio_adapter_type\n self.audio_embed_avg = audio_embed_avg\n self.audio_ffn_hidden_size = audio_ffn_hidden_size\n self.audio_ffn_intermediate_size = audio_ffn_intermediate_size\n self.audio_dual_ffn_layers = audio_dual_ffn_layers\n self.audio_decoder_proj_num_layers = audio_decoder_proj_num_layers\n self.encode_whisper_embed = encode_whisper_embed\n self.encode_audio_in_tokens = encode_audio_in_tokens\n self.use_delay_pattern = use_delay_pattern\n self.skip_audio_tower = skip_audio_tower\n self.use_audio_out_embed_projector = use_audio_out_embed_projector\n self.use_audio_out_self_attention = use_audio_out_self_attention\n\n self.use_rq_transformer = use_rq_transformer\n\n if self.use_rq_transformer:\n assert not self.use_delay_pattern, \"Delay pattern is not supported if you turned on RQ-Transformer!\"\n self.rq_transformer_hidden_size = rq_transformer_hidden_size\n self.rq_transformer_intermediate_size = rq_transformer_intermediate_size\n self.rq_transformer_num_attention_heads = rq_transformer_num_attention_heads\n self.rq_transformer_num_key_value_heads = rq_transformer_num_key_value_heads\n self.rq_transformer_num_hidden_layers = rq_transformer_num_hidden_layers\n\n if use_rq_transformer:\n # For RQ-Transformer, we set the hidden_size to the same as the text model's hidden size if it is not specified.\n if self.rq_transformer_hidden_size is None:\n self.rq_transformer_hidden_size = text_config.hidden_size\n assert self.rq_transformer_hidden_size % 128 == 0\n if self.rq_transformer_intermediate_size is None:\n self.rq_transformer_intermediate_size = text_config.intermediate_size\n if self.rq_transformer_num_attention_heads is None:\n self.rq_transformer_num_attention_heads = self.rq_transformer_hidden_size // 128\n if self.rq_transformer_num_key_value_heads is None:\n self.rq_transformer_num_key_value_heads = self.rq_transformer_hidden_size // 128 // 4\n assert self.rq_transformer_hidden_size % self.rq_transformer_num_attention_heads == 0\n assert self.rq_transformer_hidden_size % self.rq_transformer_num_key_value_heads == 0\n\n self.audio_num_codebooks = audio_num_codebooks\n self.audio_codebook_size = audio_codebook_size\n self.audio_bos_token = audio_bos_token\n self.audio_eos_token = audio_eos_token\n self.audio_out_bos_token = audio_out_bos_token\n self.audio_in_token = audio_in_token\n self.audio_out_token = audio_out_token\n self.audio_in_token_idx = audio_in_token_idx\n self.audio_out_token_idx = audio_out_token_idx\n self.audio_stream_bos_id = audio_stream_bos_id\n self.audio_stream_eos_id = audio_stream_eos_id\n self.audio_out_bos_token_id = audio_out_bos_token_id\n self.audio_eos_token_id = audio_eos_token_id\n\n super().__init__(**kwargs)\n self.pad_token_id = pad_token_id\n"
},
{
"path": "boson_multimodal/model/higgs_audio/cuda_graph_runner.py",
"content": "import torch\nimport torch.nn as nn\nfrom typing import Optional, List, Dict, Tuple, Union\nimport gc\n\nfrom transformers.cache_utils import Cache\n\n\n_NUM_WARMUP_ITERS = 2\n\n\nclass CUDAGraphRunner(nn.Module):\n def __init__(self, model):\n super().__init__()\n self.model = model\n\n self.input_buffers: Dict[str, torch.Tensor] = {}\n self.output_buffers: Dict[str, torch.Tensor] = {}\n\n self._graph: Optional[torch.cuda.CUDAGraph] = None\n\n @property\n def graph(self):\n assert self._graph is not None\n return self._graph\n\n def capture(\n self,\n hidden_states: torch.Tensor,\n causal_mask: torch.Tensor,\n position_ids: torch.Tensor,\n audio_discrete_codes_mask: torch.Tensor,\n cache_position: torch.Tensor,\n past_key_values: Union[Cache, List[torch.FloatTensor]],\n use_cache: bool,\n audio_attention_mask: torch.Tensor,\n fast_forward_attention_mask: torch.Tensor,\n output_attentions: bool,\n output_hidden_states: bool,\n is_decoding_audio_token: Optional[bool] = None,\n is_using_cuda_graph: Optional[bool] = False,\n stream: torch.cuda.Stream = None,\n memory_pool: Optional[Tuple[int, int]] = None,\n ):\n assert self._graph is None\n # Run warmup iterations\n for _ in range(_NUM_WARMUP_ITERS):\n self.model(\n hidden_states=hidden_states,\n causal_mask=causal_mask,\n position_ids=position_ids,\n audio_discrete_codes_mask=audio_discrete_codes_mask,\n cache_position=cache_position,\n past_key_values=past_key_values,\n use_cache=use_cache,\n audio_attention_mask=audio_attention_mask,\n fast_forward_attention_mask=fast_forward_attention_mask,\n output_attentions=output_attentions,\n output_hidden_states=output_hidden_states,\n is_decoding_audio_token=is_decoding_audio_token,\n is_using_cuda_graph=is_using_cuda_graph,\n )\n\n torch.cuda.synchronize()\n\n # Capture the graph\n self._graph = torch.cuda.CUDAGraph()\n with torch.cuda.graph(self._graph, pool=memory_pool, stream=stream):\n out_hidden_states, all_hidden_states, all_self_attns = self.model(\n hidden_states=hidden_states,\n causal_mask=causal_mask,\n position_ids=position_ids,\n audio_discrete_codes_mask=audio_discrete_codes_mask,\n cache_position=cache_position,\n past_key_values=past_key_values,\n use_cache=use_cache,\n audio_attention_mask=audio_attention_mask,\n fast_forward_attention_mask=fast_forward_attention_mask,\n output_attentions=output_attentions,\n output_hidden_states=output_hidden_states,\n is_decoding_audio_token=is_decoding_audio_token,\n is_using_cuda_graph=is_using_cuda_graph,\n )\n # hidden_states_out = torch.ops._C.weak_ref_tensor(outputs[0])\n # del outputs\n gc.collect()\n torch.cuda.synchronize()\n\n # Save input and output buffers\n self.input_buffers = {\n \"hidden_states\": hidden_states,\n \"causal_mask\": causal_mask,\n \"position_ids\": position_ids,\n \"audio_discrete_codes_mask\": audio_discrete_codes_mask,\n \"cache_position\": cache_position,\n \"past_key_values\": past_key_values,\n \"audio_attention_mask\": audio_attention_mask,\n \"fast_forward_attention_mask\": fast_forward_attention_mask,\n }\n self.output_buffers = {\n \"hidden_states\": out_hidden_states,\n \"all_hidden_states\": all_hidden_states,\n \"all_self_attns\": all_self_attns,\n }\n\n def forward(\n self,\n hidden_states: torch.Tensor,\n causal_mask: torch.Tensor,\n position_ids: torch.Tensor,\n audio_discrete_codes_mask: torch.Tensor,\n cache_position: torch.Tensor,\n audio_attention_mask: torch.Tensor,\n fast_forward_attention_mask: torch.Tensor,\n **kwargs,\n ) -> torch.Tensor:\n # Copy input tensors to buffers\n self.input_buffers[\"hidden_states\"].copy_(hidden_states, non_blocking=True)\n self.input_buffers[\"causal_mask\"].copy_(causal_mask, non_blocking=True)\n self.input_buffers[\"position_ids\"].copy_(position_ids, non_blocking=True)\n self.input_buffers[\"audio_discrete_codes_mask\"].copy_(audio_discrete_codes_mask, non_blocking=True)\n self.input_buffers[\"cache_position\"].copy_(cache_position, non_blocking=True)\n self.input_buffers[\"audio_attention_mask\"].copy_(audio_attention_mask, non_blocking=True)\n self.input_buffers[\"fast_forward_attention_mask\"].copy_(fast_forward_attention_mask, non_blocking=True)\n\n # Run the captured graph\n self.graph.replay()\n\n return self.output_buffers[\"hidden_states\"], None, None\n"
},
{
"path": "boson_multimodal/model/higgs_audio/custom_modules.py",
"content": "import torch\nimport torch.nn as nn\n\n\nclass PartiallyFrozenEmbedding(nn.Module):\n \"\"\"Split an existing `nn.Embedding` module that splits the embedding into:\n\n - A frozen embedding for indices [0..freeze_until_idx].\n - A trainable embedding for indices [freeze_until_idx+1..vocab_size-1].\n\n This should work with both Zero-2 and Zero-3 seamlessly\n \"\"\"\n\n def __init__(self, original_embedding: nn.Embedding, freeze_until_idx: int):\n \"\"\"\n :param original_embedding: An instance of nn.Embedding (the original embedding layer).\n :param freeze_until_idx: The index up to which the embedding is frozen (excluding). The freeze_until_idx is not frozen.\n \"\"\"\n super().__init__()\n self.freeze_until_idx = freeze_until_idx\n self.original_vocab_size = original_embedding.num_embeddings\n self.embedding_dim = original_embedding.embedding_dim\n\n # Split the original embedding into frozen and trainable parts\n self.embedding_frozen = nn.Embedding(\n freeze_until_idx,\n self.embedding_dim,\n dtype=original_embedding.weight.dtype,\n device=original_embedding.weight.device,\n )\n self.embedding_trainable = nn.Embedding(\n self.original_vocab_size - freeze_until_idx,\n self.embedding_dim,\n dtype=original_embedding.weight.dtype,\n device=original_embedding.weight.device,\n )\n\n # Copy weights from the original embedding into the frozen and trainable parts\n with torch.no_grad():\n self.embedding_frozen.weight.copy_(original_embedding.weight[:freeze_until_idx])\n self.embedding_trainable.weight.copy_(original_embedding.weight[freeze_until_idx:])\n\n # Freeze the frozen embedding\n self.embedding_frozen.weight.requires_grad = False\n\n def forward(self, input_ids: torch.Tensor) -> torch.Tensor:\n \"\"\"\n Forward pass for the split embedding wrapper.\n :param input_ids: Tensor of shape [batch_size, seq_len] with indices in [0..original_vocab_size-1].\n \"\"\"\n # Masks to separate frozen and trainable indices\n # (bsz, seq_len)\n mask_frozen = input_ids < self.freeze_until_idx\n mask_trainable = ~mask_frozen\n\n # Output tensor for embedding results\n batch_size, seq_len = input_ids.shape\n embeddings = torch.zeros(\n batch_size,\n seq_len,\n self.embedding_dim,\n device=input_ids.device,\n dtype=self.embedding_frozen.weight.dtype,\n )\n\n # Handle frozen embedding\n if mask_frozen.any():\n frozen_ids = input_ids[mask_frozen]\n frozen_emb = self.embedding_frozen(frozen_ids)\n embeddings[mask_frozen] = frozen_emb\n\n # Handle trainable embedding\n if mask_trainable.any():\n # Adjust trainable IDs to the local index space of the trainable embedding\n trainable_ids = input_ids[mask_trainable] - (self.freeze_until_idx)\n trainable_emb = self.embedding_trainable(trainable_ids)\n embeddings[mask_trainable] = trainable_emb\n\n return embeddings\n\n def to_unsplit(self) -> nn.Embedding:\n unsplit_embedding = nn.Embedding(\n self.original_vocab_size,\n self.embedding_dim,\n dtype=self.embedding_frozen.weight.dtype,\n device=self.embedding_frozen.weight.device,\n )\n\n with torch.no_grad():\n unsplit_embedding.weight[: self.freeze_until_idx].copy_(self.embedding_frozen.weight)\n unsplit_embedding.weight[self.freeze_until_idx :].copy_(self.embedding_trainable.weight)\n\n return unsplit_embedding\n\n\nclass PartiallyFrozenLinear(nn.Module):\n \"\"\"A wrapper around nn.Linear to partially freeze part of the weight matrix.\"\"\"\n\n def __init__(self, original_linear: nn.Linear, freeze_until_idx: int):\n \"\"\"\n :param original_linear: The original nn.Linear layer.\n :param freeze_until_idx: The index up to which the rows of the weight matrix are frozen.\n \"\"\"\n super().__init__()\n assert original_linear.bias is None, \"Currently only support linear module without bias\"\n\n self.freeze_until_idx = freeze_until_idx\n self.input_dim = original_linear.in_features\n self.output_dim = original_linear.out_features\n\n # Create frozen and trainable linear layers\n self.linear_frozen = nn.Linear(\n self.input_dim,\n freeze_until_idx,\n bias=False,\n dtype=original_linear.weight.dtype,\n device=original_linear.weight.device,\n )\n self.linear_trainable = nn.Linear(\n self.input_dim,\n self.output_dim - freeze_until_idx,\n bias=False,\n dtype=original_linear.weight.dtype,\n device=original_linear.weight.device,\n )\n\n # Copy weights from the original linear layer\n with torch.no_grad():\n self.linear_frozen.weight.copy_(original_linear.weight[:freeze_until_idx])\n self.linear_trainable.weight.copy_(original_linear.weight[freeze_until_idx:])\n\n # Freeze the frozen linear layer\n self.linear_frozen.weight.requires_grad = False\n\n def forward(self, input_tensor):\n # input_tensor: (bsz, seq_len, hidden_state_dim)\n frozen_output = self.linear_frozen(input_tensor)\n trainable_output = self.linear_trainable(input_tensor)\n return torch.cat((frozen_output, trainable_output), dim=-1)\n\n def to_unsplit(self) -> nn.Linear:\n unsplit_linear = nn.Linear(\n self.input_dim,\n self.output_dim,\n bias=False,\n dtype=self.linear_frozen.weight.dtype,\n device=self.linear_frozen.weight.device,\n )\n\n # Copy weights from the frozen and trainable layers into the unsplit linear layer\n with torch.no_grad():\n unsplit_linear.weight[: self.freeze_until_idx].copy_(self.linear_frozen.weight)\n unsplit_linear.weight[self.freeze_until_idx :].copy_(self.linear_trainable.weight)\n\n return unsplit_linear\n"
},
{
"path": "boson_multimodal/model/higgs_audio/modeling_higgs_audio.py",
"content": "\"\"\"Higgs-Audio is an end-to-end multimodal model with the capability to understand and generate text / audio.\"\"\"\n\nimport torch\nimport torch.nn as nn\nimport math\nimport glob\nimport functools\nimport os\nfrom collections import defaultdict, OrderedDict\nfrom dataclasses import dataclass\nfrom enum import Enum\nfrom safetensors.torch import load_file\nfrom typing import Optional, Tuple, Union, List, Dict, Any\n\nfrom transformers import AutoTokenizer\nfrom transformers.modeling_outputs import BaseModelOutput\nfrom transformers.models.whisper.modeling_whisper import WhisperEncoderLayer\nfrom transformers.models.llama.modeling_llama import (\n LlamaDecoderLayer,\n LlamaRMSNorm,\n LlamaRotaryEmbedding,\n LLAMA_ATTENTION_CLASSES,\n LlamaMLP,\n LlamaRMSNorm,\n)\nfrom transformers.modeling_attn_mask_utils import AttentionMaskConverter\nfrom transformers.cache_utils import Cache, DynamicCache, StaticCache\nfrom transformers.generation import GenerationMixin, GenerationConfig, LogitsProcessorList, StoppingCriteriaList\nfrom transformers.generation.utils import GenerateNonBeamOutput\nfrom transformers.utils import logging, ModelOutput\n\nfrom .common import HiggsAudioPreTrainedModel\nfrom .utils import (\n merge_input_ids_with_audio_features,\n count_parameters,\n)\nfrom .configuration_higgs_audio import HiggsAudioConfig, HiggsAudioEncoderConfig\nfrom .custom_modules import PartiallyFrozenLinear, PartiallyFrozenEmbedding\nfrom .cuda_graph_runner import CUDAGraphRunner\nfrom .audio_head import HiggsAudioDecoderProjector\n\nlogger = logging.get_logger(__name__)\n\n\nclass GenerationMode(Enum):\n \"\"\"Enum for different generation modes in HiggsAudio model.\"\"\"\n\n TEXT = 0 # Text generation mode\n AUDIO_INIT = 1 # Audio generation mode initialization\n AUDIO_IN_PROGRESS = 2 # Audio generation mode in progress\n\n\ndef _whisper_encoder_zero_shape_forward(whisper_encoder, *args, **kwargs):\n \"\"\"The whisper encoder does not support zero-shape tensor by default due to the following implementations\n\n key_states = self._shape(self.k_proj(current_states), -1, bsz)\n\n If `bsz` is 0, the \"-1\" dimension will be ambiguous and triggers error in the shape inference pass.\n\n See also: https://github.com/huggingface/transformers/blob/30335093276212ce74938bdfd85bfd5df31a668a/src/transformers/models/whisper/modeling_whisper.py#L306-L307\n\n This function monkey-patches all `_shape` functions in the whisper encoder's self-attention layers to ensure function supports zero-shape tensor.\n\n #FIXME!!!! This is a temporary workaround and should be removed once the upstream issue is resolved.\n\n \"\"\"\n\n global _higgs_flash_attention_forward\n\n def _patched_shape(tensor: torch.Tensor, seq_len: int, bsz: int, num_heads: int, head_dim: int):\n if seq_len == -1:\n return tensor.view(bsz, tensor.shape[1], num_heads, head_dim).transpose(1, 2).contiguous()\n else:\n return tensor.view(bsz, seq_len, num_heads, head_dim).transpose(1, 2).contiguous()\n\n def _patched_scaled_dot_product_attention(\n query, key, value, attn_mask=None, dropout_p=0.0, is_causal=False, scale=None, enable_gqa=False\n ) -> torch.Tensor:\n # IMPORTANT! Implementation here is wrong and is only for the purpose of obtaining the correct attn_weight shape\n if enable_gqa:\n key = key.repeat_interleave(query.size(-3) // key.size(-3), -3)\n value = value.repeat_interleave(query.size(-3) // value.size(-3), -3)\n\n attn_weight = query @ key.transpose(-2, -1)\n return attn_weight @ value\n\n # Apply monkey-patch\n if whisper_encoder.config._attn_implementation != \"flash_attention_2\":\n old_shape_functions = []\n for layer in whisper_encoder.layers:\n old_shape_functions.append(getattr(layer.self_attn, \"_shape\"))\n layer.self_attn._shape = functools.partial(\n _patched_shape, num_heads=layer.self_attn.num_heads, head_dim=layer.self_attn.head_dim\n )\n\n original_scaled_dot_product_attention = torch.nn.functional.scaled_dot_product_attention\n torch.nn.functional.scaled_dot_product_attention = _patched_scaled_dot_product_attention\n\n out = whisper_encoder(*args, **kwargs)\n torch.nn.functional.scaled_dot_product_attention = original_scaled_dot_product_attention\n\n # Restore the original shape functions\n if whisper_encoder.config._attn_implementation != \"flash_attention_2\":\n for layer, old_shape_function in zip(whisper_encoder.layers, old_shape_functions):\n layer.self_attn._shape = old_shape_function\n\n return out\n\n\ndef _prepare_4d_causal_attention_mask_with_cache_position(\n attention_mask: torch.Tensor,\n sequence_length: int,\n target_length: int,\n dtype: torch.dtype,\n device: torch.device,\n min_dtype: float,\n cache_position: torch.Tensor,\n batch_size: int,\n):\n \"\"\"\n Creates a causal 4D mask of shape `(batch_size, 1, query_length, key_value_length)` from a 2D mask of shape\n `(batch_size, key_value_length)`, or if the input `attention_mask` is already 4D, do nothing.\n\n Args:\n attention_mask (`torch.Tensor`):\n A 2D attention mask of shape `(batch_size, key_value_length)` or a 4D attention mask of shape `(batch_size, 1, query_length, key_value_length)`.\n sequence_length (`int`):\n The sequence length being processed.\n target_length (`int`):\n The target length: when generating with static cache, the mask should be as long as the static cache, to account for the 0 padding, the part of the cache that is not filled yet.\n dtype (`torch.dtype`):\n The dtype to use for the 4D attention mask.\n device (`torch.device`):\n The device to plcae the 4D attention mask on.\n min_dtype (`float`):\n The minimum value representable with the dtype `dtype`.\n cache_position (`torch.Tensor`):\n Indices depicting the position of the input sequence tokens in the sequence.\n batch_size (`torch.Tensor`):\n Batch size.\n \"\"\"\n if attention_mask is not None and attention_mask.dim() == 4:\n # In this case we assume that the mask comes already in inverted form and requires no inversion or slicing.\n causal_mask = attention_mask\n else:\n causal_mask = torch.full((sequence_length, target_length), fill_value=min_dtype, dtype=dtype, device=device)\n if sequence_length != 1:\n causal_mask = torch.triu(causal_mask, diagonal=1)\n causal_mask *= torch.arange(target_length, device=device) > cache_position.reshape(-1, 1)\n causal_mask = causal_mask[None, None, :, :].expand(batch_size, 1, -1, -1)\n if attention_mask is not None:\n causal_mask = causal_mask.clone() # copy to contiguous memory for in-place edit\n mask_length = attention_mask.shape[-1]\n padding_mask = causal_mask[:, :, :, :mask_length] + attention_mask[:, None, None, :]\n padding_mask = padding_mask == 0\n causal_mask[:, :, :, :mask_length] = causal_mask[:, :, :, :mask_length].masked_fill(\n padding_mask, min_dtype\n )\n\n return causal_mask\n\n\nclass HiggsAudioFeatureProjector(nn.Module):\n \"\"\"Projector that maps audio features extracted by Whisper to hidden state of the text model.\"\"\"\n\n def __init__(self, config: HiggsAudioConfig):\n super().__init__()\n self.linear = nn.Linear(config.audio_encoder_config.d_model, config.text_config.hidden_size, bias=True)\n\n def forward(self, audio_features):\n hidden_states = self.linear(audio_features)\n return hidden_states\n\n\n# Revised on top of transformers.models.qwen2_audio.modeling_qwen2_audio with Qwen2AudioEncoder --> HiggsAudioEncoder\n# The code was originally borrowed from WhisperEncoder\nclass HiggsAudioEncoder(HiggsAudioPreTrainedModel):\n \"\"\"\n Transformer encoder consisting of *config.encoder_layers* self attention layers. Each layer is a\n [`WhisperEncoderLayer`].\n\n Args:\n config: HiggsAudioEncoderConfig\n \"\"\"\n\n # Ignore copy\n config_class = HiggsAudioEncoderConfig\n main_input_name = \"input_features\"\n _no_split_modules = [\"WhisperEncoderLayer\"]\n\n def __init__(self, config: HiggsAudioEncoderConfig):\n super().__init__(config)\n self.dropout = config.dropout\n self.layerdrop = config.encoder_layerdrop\n\n embed_dim = config.d_model\n self.num_mel_bins = config.num_mel_bins\n self.padding_idx = config.pad_token_id\n self.max_source_positions = config.max_source_positions\n self.embed_scale = math.sqrt(embed_dim) if config.scale_embedding else 1.0\n\n self.conv1 = nn.Conv1d(self.num_mel_bins, embed_dim, kernel_size=3, padding=1)\n self.conv2 = nn.Conv1d(embed_dim, embed_dim, kernel_size=3, stride=2, padding=1)\n\n self.embed_positions = nn.Embedding(self.max_source_positions, embed_dim)\n self.embed_positions.requires_grad_(False)\n\n # Flash Attention 2 does not support zero shape tensor, so we have to use sdpa implementation for the Whisper component.\n self.layers = nn.ModuleList([WhisperEncoderLayer(config) for _ in range(config.encoder_layers)])\n self.layer_norm = nn.LayerNorm(config.d_model)\n # Ignore copy\n self.avg_pooler = nn.AvgPool1d(2, stride=2)\n\n self.gradient_checkpointing = False\n # Initialize weights and apply final processing\n self.post_init()\n\n def _freeze_parameters(self):\n for param in self.parameters():\n param.requires_grad = False\n self._requires_grad = False\n\n def get_input_embeddings(self) -> nn.Module:\n return self.conv1\n\n def set_input_embeddings(self, value: nn.Module):\n self.conv1 = value\n\n def forward(\n self,\n input_features,\n attention_mask=None,\n head_mask=None,\n output_attentions=None,\n output_hidden_states=None,\n return_dict=None,\n check_seq_length=True,\n ):\n r\"\"\"\n Args:\n input_features (`torch.LongTensor` of shape `(batch_size, feature_size, sequence_length)`):\n Float values of mel features extracted from the raw speech waveform. Raw speech waveform can be\n obtained by loading a `.flac` or `.wav` audio file into an array of type `List[float]` or a\n `numpy.ndarray`, *e.g.* via the soundfile library (`pip install soundfile`). To prepare the array into\n `input_features`, the [`AutoFeatureExtractor`] should be used for extracting the mel features, padding\n and conversion into a tensor of type `torch.FloatTensor`. See [`~WhisperFeatureExtractor.__call__`]\n attention_mask (`torch.Tensor`)`, *optional*):\n HiggsAudio does not support masking of the `input_features`, this argument is preserved for compatibility,\n but it is not used. By default the silence in the input log mel spectrogram are ignored.\n head_mask (`torch.Tensor` of shape `(encoder_layers, encoder_attention_heads)`, *optional*):\n Mask to nullify selected heads of the attention modules. Mask values selected in `[0, 1]`:\n\n - 1 indicates the head is **not masked**,\n - 0 indicates the head is **masked**.\n output_attentions (`bool`, *optional*):\n Whether or not to return the attentions tensors of all attention layers. See `attentions` under\n returned tensors for more detail.\n output_hidden_states (`bool`, *optional*):\n Whether or not to return the hidden states of all layers. See `hidden_states` under returned tensors\n for more detail.\n return_dict (`bool`, *optional*):\n Whether or not to return a [`~utils.ModelOutput`] instead of a plain tuple.\n \"\"\"\n\n expected_seq_length = self.config.max_source_positions * self.conv1.stride[0] * self.conv2.stride[0]\n if check_seq_length and (input_features.shape[-1] != expected_seq_length):\n raise ValueError(\n f\"HiggsAudio expects the mel input features to be of length {expected_seq_length}, but found {input_features.shape[-1]}. Make sure to pad the input mel features to {expected_seq_length}.\"\n )\n\n output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions\n output_hidden_states = (\n output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states\n )\n return_dict = return_dict if return_dict is not None else self.config.use_return_dict\n\n # Ignore copy\n input_features = input_features.to(dtype=self.conv1.weight.dtype, device=self.conv1.weight.device)\n\n inputs_embeds = nn.functional.gelu(self.conv1(input_features))\n inputs_embeds = nn.functional.gelu(self.conv2(inputs_embeds))\n\n inputs_embeds = inputs_embeds.permute(0, 2, 1)\n embed_pos = self.embed_positions.weight\n\n hidden_states = inputs_embeds + embed_pos\n hidden_states = nn.functional.dropout(hidden_states, p=self.dropout, training=self.training)\n\n encoder_states = () if output_hidden_states else None\n all_attentions = () if output_attentions else None\n\n # check if head_mask has a correct number of layers specified if desired\n if head_mask is not None:\n assert head_mask.size()[0] == (len(self.layers)), (\n f\"The head_mask should be specified for {len(self.layers)} layers, but it is for {head_mask.size()[0]}.\"\n )\n\n for idx, encoder_layer in enumerate(self.layers):\n if output_hidden_states:\n encoder_states = encoder_states + (hidden_states,)\n # add LayerDrop (see https://arxiv.org/abs/1909.11556 for description)\n to_drop = False\n if self.training:\n dropout_probability = torch.rand([])\n if dropout_probability < self.layerdrop: # skip the layer\n to_drop = True\n\n # Ignore copy\n if to_drop:\n layer_outputs = (None, None)\n else:\n if self.gradient_checkpointing and self.training:\n layer_outputs = self._gradient_checkpointing_func(\n encoder_layer.__call__,\n hidden_states,\n attention_mask,\n (head_mask[idx] if head_mask is not None else None),\n output_attentions,\n )\n else:\n layer_outputs = encoder_layer(\n hidden_states,\n attention_mask,\n layer_head_mask=(head_mask[idx] if head_mask is not None else None),\n output_attentions=output_attentions,\n )\n\n hidden_states = layer_outputs[0]\n\n if output_attentions:\n all_attentions = all_attentions + (layer_outputs[1],)\n\n # Ignore copy\n hidden_states = hidden_states.permute(0, 2, 1)\n # If the sequence length after average pooling is not divisible by the sequence parallel size, we would duplicate it across the sequence parallel ranks.\n # In this case, gradients need to be scaled up because the subsequent scaling up in the function _apply_audio_tower is skipped.\n hidden_states = self.avg_pooler(hidden_states)\n\n hidden_states = hidden_states.permute(0, 2, 1)\n\n hidden_states = self.layer_norm(hidden_states)\n\n if output_hidden_states:\n encoder_states = encoder_states + (hidden_states,)\n\n if not return_dict:\n return tuple(v for v in [hidden_states, encoder_states, all_attentions] if v is not None)\n return BaseModelOutput(\n last_hidden_state=hidden_states, hidden_states=encoder_states, attentions=all_attentions\n )\n\n # Ignore copy\n def _get_feat_extract_output_lengths(self, input_lengths: torch.LongTensor):\n \"\"\"\n Computes the output length of the convolutional layers and the output length of the audio encoder\n \"\"\"\n input_lengths = (input_lengths - 1) // 2 + 1\n output_lengths = (input_lengths - 2) // 2 + 1\n return input_lengths, output_lengths\n\n\nclass HiggsAudioDualFFNDecoderLayer(nn.Module):\n \"\"\"We implement a dual-path FFN decoder layer where the audio tokens and text tokens go through separate FFN layers.\n\n The audio and text tokens share the text-attention layer, but will be encoded with separate feedforward layers.\n In addition, the audio tokens can be configured to go through separate attention layer.\n\n Following is an illustration:\n\n t t t a a a t t t\n |\n | (shared attention layer)\n v\n h_t h_t h_t h_a h_a h_a h_t h_t h_t\n |\n | (separate text/audio hidden states)\n v\n [h_t h_t h_t h_t h_t h_t], [h_a, h_a, h_a]\n | |\n | (separate FFNs) |\n v v\n [o_t o_t o_t o_t o_t o_t], [o_a, o_a, o_a]\n |\n | (reorder)\n v\n o_t o_t o_t o_a o_a o_a o_t o_t o_t\n\n This has a few advantages:\n 1) We are able to use a smaller FFN, or even bypass the FFN for audio tokens. This accelerates the inference speed.\n 2) The Audio-FFN introduces more trainable parameters to the model.\n This should have the same effect as the mixture-of-expert layer and we may expect better performance due to parameter scaling.\n 3) We can replace the original FFN in LLMs with the dual-path FFN without changing the number of FLOPs.\n\n\n \"\"\"\n\n def __init__(\n self, config: HiggsAudioConfig, layer_idx: int, fast_forward: bool = False, use_audio_attention: bool = False\n ):\n super().__init__()\n text_config = config.text_config\n self.hidden_size = text_config.hidden_size\n self.layer_idx = layer_idx\n self.self_attn = LLAMA_ATTENTION_CLASSES[config._attn_implementation](config=text_config, layer_idx=layer_idx)\n\n self.mlp = LlamaMLP(text_config)\n\n if not fast_forward:\n if use_audio_attention:\n self.audio_attn = LLAMA_ATTENTION_CLASSES[config._attn_implementation](\n config=text_config, layer_idx=layer_idx + 1\n )\n self.audio_post_audio_attn_layer_norm = LlamaRMSNorm(\n text_config.hidden_size, eps=text_config.rms_norm_eps\n )\n\n self.audio_mlp = LlamaMLP(text_config)\n self.audio_input_layernorm = LlamaRMSNorm(text_config.hidden_size, eps=text_config.rms_norm_eps)\n self.audio_post_attention_layernorm = LlamaRMSNorm(text_config.hidden_size, eps=text_config.rms_norm_eps)\n\n self.use_audio_attention = use_audio_attention\n self.fast_forward = fast_forward\n if self.fast_forward:\n assert not self.use_audio_attention, (\n \"We cannot use audio_attention if the layer is marked as fast-forward.\"\n )\n self.input_layernorm = LlamaRMSNorm(text_config.hidden_size, eps=text_config.rms_norm_eps)\n self.post_attention_layernorm = LlamaRMSNorm(text_config.hidden_size, eps=text_config.rms_norm_eps)\n\n def forward(\n self,\n hidden_states: torch.Tensor,\n attention_mask: Optional[torch.Tensor] = None,\n audio_attention_mask: Optional[torch.Tensor] = None,\n fast_forward_attention_mask: Optional[torch.Tensor] = None,\n position_ids: Optional[torch.LongTensor] = None,\n audio_out_mask: Optional[torch.BoolTensor] = None,\n is_decoding_audio_token: Optional[bool] = None,\n past_key_value: Optional[Cache] = None,\n output_attentions: Optional[bool] = False,\n use_cache: Optional[bool] = False,\n cache_position: Optional[torch.LongTensor] = None,\n position_embeddings: Optional[Tuple[torch.Tensor, torch.Tensor]] = None, # will become mandatory in v4.46\n is_using_cuda_graph: Optional[bool] = False,\n **kwargs,\n ):\n \"\"\"\n Args:\n hidden_states (`torch.FloatTensor`): input to the layer of shape `(batch, seq_len, embed_dim)`\n attention_mask (`torch.FloatTensor`, *optional*):\n attention mask of size `(batch_size, sequence_length)` if flash attention is used or `(batch_size, 1,\n query_sequence_length, key_sequence_length)` if default attention is used.\n position_ids\n IDs of positions in the input sequence\n audio_out_mask\n Mask for identifying the audio tokens. Size (batch_size, sequence_length)\n 1 --> location contains audio_out\n 0 --> location does not contain audio_out\n\n When use_cache is True and not in torch compile mode, the audio_out_mask contains audio_out masks for\n all tokens up to the current token. That means, it has size (batch_size, sequence_length) while\n hidden_states will have size (batch_size, 1). In the torch compile mode, the audio_out_mask will have\n size (batch_size, 1).\n is_decoding_audio_token\n Used in the torch compile mode to determine if the current token is an audio token or not.\n past_key_value (`Cache`, *optional*): cached past key and value projection states. We fetch the corresponding cached key/value via the layer_idx.\n output_attentions (`bool`, *optional*):\n Whether or not to return the attentions tensors of all attention layers. See `attentions` under\n returned tensors for more detail.\n use_cache (`bool`, *optional*):\n If set to `True`, `past_key_values` key value states are returned and can be used to speed up decoding\n (see `past_key_values`).\n cache_position (`torch.LongTensor` of shape `(sequence_length)`, *optional*):\n Indices depicting the position of the input sequence tokens in the sequence\n position_embeddings (`Tuple[torch.FloatTensor, torch.FloatTensor]`, *optional*):\n Tuple containing the cosine and sine positional embeddings of shape `(batch_size, seq_len, head_dim)`,\n with `head_dim` being the embedding dimension of each attention head.\n is_using_cuda_graph (`bool`, *optional*):\n Indicates whether the model is running by cuda graph.\n kwargs (`dict`, *optional*):\n Arbitrary kwargs to be ignored, used for FSDP and other methods that injects code\n into the model\n \"\"\"\n residual = hidden_states\n target_length = hidden_states.shape[1]\n use_static_cache = isinstance(past_key_value, StaticCache)\n decode_stage = hidden_states.shape[1] == 1\n if is_using_cuda_graph:\n assert decode_stage and use_static_cache, (\n \"The CUDA graph mode should only be used in the decoding stage with static cache.\"\n )\n\n # If we are decoding an audio token and the layer is marked as fast-forward,\n # we can skip it.\n if is_decoding_audio_token and self.fast_forward:\n return (hidden_states,)\n\n has_audio_out = audio_out_mask is not None and audio_out_mask.shape[0] > 0\n\n audio_out_mask_sq = audio_out_mask\n\n if self.fast_forward and has_audio_out:\n original_hidden_states = hidden_states.clone()\n min_dtype = torch.finfo(hidden_states.dtype).min\n if attention_mask is None:\n attention_mask = ~audio_out_mask\n\n if self.self_attn.config._attn_implementation != \"flash_attention_2\":\n sequence_length = audio_out_mask.shape[1]\n attention_mask = _prepare_4d_causal_attention_mask_with_cache_position(\n attention_mask=attention_mask,\n sequence_length=sequence_length,\n target_length=sequence_length,\n dtype=hidden_states.dtype,\n min_dtype=min_dtype,\n device=hidden_states.device,\n cache_position=cache_position,\n batch_size=hidden_states.shape[0],\n )\n if use_cache:\n attention_mask = attention_mask[:, :, -target_length:, :]\n elif len(attention_mask.shape) == 2:\n # Attention mask has shape (batch_size, sequence_length)\n # We should be using flash attention 2\n attention_mask = attention_mask * ~audio_out_mask\n elif len(attention_mask.shape) == 4:\n # When using static cache, the attention mask was already preprocessed in the previous layer\n if use_static_cache:\n attention_mask = fast_forward_attention_mask\n else:\n if use_cache:\n # Attention mask has shape (batch_size, 1, query_length, key_length)\n # In addition, the attention mask should be inverted, that means \"1\" (attend_to) --> \"0\", and \"0\" --> minimal dtype value.\n attention_mask = attention_mask.masked_fill(\n audio_out_mask[:, -target_length:].reshape(audio_out_mask.shape[0], 1, target_length, 1)\n | audio_out_mask.reshape(audio_out_mask.shape[0], 1, 1, audio_out_mask.shape[1]),\n min_dtype,\n )\n else:\n attention_mask = attention_mask.masked_fill(\n audio_out_mask.reshape(audio_out_mask.shape[0], 1, audio_out_mask.shape[1], 1)\n | audio_out_mask.reshape(audio_out_mask.shape[0], 1, 1, audio_out_mask.shape[1]),\n min_dtype,\n )\n else:\n raise NotImplementedError(f\"Unsupported attention_mask format, attention_mask={attention_mask}\")\n\n if (\n self.self_attn.config._attn_implementation == \"sdpa\"\n and attention_mask is not None\n and attention_mask.device.type == \"cuda\"\n and not output_attentions\n ):\n # Attend to all tokens in fully masked rows in the causal_mask, for example the relevant first rows when\n # using left padding. This is required by F.scaled_dot_product_attention memory-efficient attention path.\n # Details: https://github.com/pytorch/pytorch/issues/110213\n attention_mask = AttentionMaskConverter._unmask_unattended(attention_mask, min_dtype)\n\n if has_audio_out and not self.fast_forward:\n # Apply separate layernorm layers for audio tokens and text tokens\n if use_cache:\n hidden_states = torch.where(\n audio_out_mask_sq[:, -target_length:].unsqueeze(-1),\n self.audio_input_layernorm(hidden_states),\n self.input_layernorm(hidden_states),\n )\n else:\n hidden_states = torch.where(\n audio_out_mask_sq.unsqueeze(-1),\n self.audio_input_layernorm(hidden_states),\n self.input_layernorm(hidden_states),\n )\n else:\n hidden_states = self.input_layernorm(hidden_states)\n\n # Audio Attention\n if self.use_audio_attention and has_audio_out:\n if use_static_cache:\n assert audio_attention_mask is not None, (\n \"audio_attention_mask should not be None when using static cache.\"\n )\n\n if audio_attention_mask is None:\n no_audio_out_mask = (~audio_out_mask)[:, -target_length:].reshape(\n audio_out_mask.shape[0], 1, target_length, 1\n ) | (~audio_out_mask).reshape(audio_out_mask.shape[0], 1, 1, audio_out_mask.shape[1])\n min_dtype = torch.finfo(hidden_states.dtype).min\n\n if attention_mask is None:\n audio_attention_mask = audio_out_mask\n\n if self.audio_attn.config._attn_implementation != \"flash_attention_2\":\n sequence_length = audio_out_mask.shape[1]\n audio_attention_mask = _prepare_4d_causal_attention_mask_with_cache_position(\n attention_mask=audio_attention_mask,\n sequence_length=sequence_length,\n target_length=sequence_length,\n dtype=hidden_states.dtype,\n min_dtype=min_dtype,\n device=hidden_states.device,\n cache_position=cache_position,\n batch_size=hidden_states.shape[0],\n )\n if use_cache:\n audio_attention_mask = audio_attention_mask[:, :, -target_length:, :]\n audio_attention_mask = audio_attention_mask.masked_fill(no_audio_out_mask, min_dtype)\n elif len(attention_mask.shape) == 2:\n # Attention mask has shape (batch_size, sequence_length)\n audio_attention_mask = attention_mask * audio_out_mask\n elif len(attention_mask.shape) == 4:\n # Attention mask has shape (batch_size, 1, query_length, key_length)\n # In addition, the attention mask should be inverted. This means \"1\" (attend_to) --> \"0\", and \"0\" --> minimal dtype value.\n audio_attention_mask = attention_mask.masked_fill(no_audio_out_mask, min_dtype)\n else:\n raise NotImplementedError(f\"Unsupported attention_mask format, attention_mask={attention_mask}\")\n\n if (\n self.audio_attn.config._attn_implementation == \"sdpa\"\n and audio_attention_mask is not None\n and audio_attention_mask.device.type == \"cuda\"\n and not output_attentions\n ):\n # Attend to all tokens in fully masked rows in the causal_mask, for example the relevant first rows when\n # using left padding. This is required by F.scaled_dot_product_attention memory-efficient attention path.\n # Details: https://github.com/pytorch/pytorch/issues/110213\n audio_attention_mask = AttentionMaskConverter._unmask_unattended(audio_attention_mask, min_dtype)\n\n audio_attention_mask = audio_attention_mask.contiguous()\n\n audio_hidden_states, audio_self_attn_weights, audio_present_key_value = self.audio_attn(\n hidden_states=hidden_states,\n attention_mask=audio_attention_mask,\n position_ids=position_ids,\n past_key_value=past_key_value,\n output_attentions=output_attentions,\n use_cache=use_cache,\n cache_position=cache_position,\n position_embeddings=position_embeddings,\n **kwargs,\n )\n audio_hidden_states = residual + audio_hidden_states\n if use_cache:\n residual = torch.where(\n audio_out_mask_sq[:, -target_length:].unsqueeze(-1), audio_hidden_states, residual\n )\n else:\n residual = torch.where(audio_out_mask_sq.unsqueeze(-1), audio_hidden_states, residual)\n audio_hidden_states = self.audio_post_audio_attn_layer_norm(audio_hidden_states)\n if use_cache:\n hidden_states = torch.where(\n audio_out_mask_sq[:, -target_length:].unsqueeze(-1), audio_hidden_states, hidden_states\n )\n else:\n hidden_states = torch.where(audio_out_mask_sq.unsqueeze(-1), audio_hidden_states, hidden_states)\n\n # Text Attention\n hidden_states, self_attn_weights, present_key_value = self.self_attn(\n hidden_states=hidden_states,\n attention_mask=attention_mask,\n position_ids=position_ids,\n past_key_value=past_key_value,\n output_attentions=output_attentions,\n use_cache=use_cache,\n cache_position=cache_position,\n position_embeddings=position_embeddings,\n **kwargs,\n )\n hidden_states = residual + hidden_states\n\n # Apply Dual-path FFN\n residual = hidden_states\n\n if has_audio_out and not self.fast_forward:\n if use_cache:\n real_audio_out_mask = audio_out_mask_sq[:, -target_length:]\n else:\n real_audio_out_mask = audio_out_mask_sq\n\n # Make whole graph in decode stage\n if decode_stage and is_using_cuda_graph:\n assert is_decoding_audio_token is not None, (\n \"is_decoding_audio_token should be present in the decoding stage.\"\n )\n if is_decoding_audio_token:\n hidden_states = self.audio_post_attention_layernorm(hidden_states)\n hidden_states = self.audio_mlp(hidden_states)\n else:\n hidden_states = self.post_attention_layernorm(hidden_states)\n hidden_states = self.mlp(hidden_states)\n residual = residual + hidden_states\n else:\n text_hidden_states = self.post_attention_layernorm(hidden_states[~real_audio_out_mask])\n audio_hidden_states = self.audio_post_attention_layernorm(hidden_states[real_audio_out_mask])\n\n text_hidden_states = self.mlp(text_hidden_states)\n residual[~real_audio_out_mask] += text_hidden_states\n\n audio_hidden_states = self.audio_mlp(audio_hidden_states)\n residual[real_audio_out_mask] += audio_hidden_states\n\n hidden_states = residual\n else:\n hidden_states = self.post_attention_layernorm(hidden_states)\n hidden_states = self.mlp(hidden_states)\n hidden_states = residual + hidden_states\n\n if self.fast_forward and has_audio_out:\n if use_cache:\n hidden_states = torch.where(\n audio_out_mask_sq[:, -target_length:].unsqueeze(-1), original_hidden_states, hidden_states\n )\n else:\n hidden_states = torch.where(audio_out_mask_sq.unsqueeze(-1), original_hidden_states, hidden_states)\n\n outputs = (hidden_states,)\n\n if output_attentions:\n if self.use_audio_attention:\n # The returned attn weights have shape (batch_size, num_heads + num_audio_attn_heads, seq_length, seq_length)\n outputs += (torch.concat([self_attn_weights, audio_self_attn_weights], dim=1),)\n else:\n # The returned attn weights have shape (batch_size, num_heads, seq_length, seq_length)\n outputs += (self_attn_weights,)\n\n if use_cache:\n outputs += (present_key_value,)\n\n return outputs\n\n\n@dataclass\nclass HiggsAudioModelOutputWithPast(ModelOutput):\n loss: Optional[torch.FloatTensor] = None\n llm_loss: Optional[torch.FloatTensor] = None\n audio_loss: Optional[torch.FloatTensor] = None\n codebook_losses: Optional[torch.FloatTensor] = None\n logits: Optional[torch.FloatTensor] = None\n expanded_input_ids: Optional[torch.LongTensor] = None\n expanded_labels: Optional[torch.LongTensor] = None\n audio_in_mask: Optional[torch.BoolTensor] = None\n audio_in_discrete_codes_mask: Optional[torch.BoolTensor] = None\n audio_out_mask: Optional[torch.BoolTensor] = None\n attention_mask: Optional[torch.BoolTensor] = None\n audio_logits: Optional[torch.FloatTensor] = None\n past_key_values: Optional[Cache] = None\n hidden_states: Optional[Tuple[torch.FloatTensor, ...]] = None\n audio_hidden_states: Optional[Tuple[torch.FloatTensor, ...]] = None\n attentions: Optional[Tuple[torch.FloatTensor, ...]] = None\n\n\n@dataclass\nclass HiggsAudioGenerationOutput(ModelOutput):\n \"\"\"\n Outputs of HiggsAudio generation models, when using non-beam methods.\n\n Args:\n sequences (`torch.LongTensor` of shape `(batch_size, sequence_length)`):\n The generated sequences. The second dimension (sequence_length) is either equal to `max_length` or shorter\n if all batches finished early due to the `eos_token_id`.\n audio_sequences (`tuple(torch.LongTensor)` *optional*):\n The generated discrete audio codes. These codes can be used to fill-in related locations of <|AUDIO_OUT|> at input sequences.\n scores (`tuple(torch.FloatTensor)` *optional*, returned when `output_scores=True`):\n Processed prediction scores of the language modeling head (scores for each vocabulary token before SoftMax)\n at each generation step. Tuple of `torch.FloatTensor` with up to `max_new_tokens` elements (one element for\n each generated token).\n If the generated token is a text token, the tensor will have shape `(batch_size, config.vocab_size)`.\n If the generated token is an audio token, the tensor will have shape `(config.audio_num_codebooks, self.audio_codebook_size)`\n logits (`tuple(torch.FloatTensor)` *optional*, returned when `output_logits=True`):\n Unprocessed prediction scores of the language modeling head or the audio head (scores for each vocabulary token before SoftMax)\n at each generation step. Tuple of `torch.FloatTensor` with up to `max_new_tokens` elements (one element for\n each generated token).\n If the generated token is a text token, the tensor will have shape `(batch_size, config.vocab_size)`.\n If the generated token is an audio token, the tensor will have shape `(config.audio_num_codebooks, self.audio_codebook_size)`\n attentions (`tuple(tuple(torch.FloatTensor))`, *optional*, returned when `output_attentions=True`):\n Tuple (one element for each generated token) of tuples (one element for each layer of the decoder) of\n `torch.FloatTensor` of shape `(batch_size, num_heads, generated_length, sequence_length)`.\n hidden_states (`tuple(tuple(torch.FloatTensor))`, *optional*, returned when `output_hidden_states=True`):\n Tuple (one element for each generated token) of tuples (one element for each layer of the decoder) of\n `torch.FloatTensor` of shape `(batch_size, generated_length, hidden_size)`.\n past_key_values (`tuple(tuple(torch.FloatTensor)))`, *optional*, returned when `use_cache=True`):\n Returns the model cache, used to speed up decoding. Different models have a different cache format, check\n the model's documentation. Usually, a [`~cache_utils.Cache`] instance.\n \"\"\"\n\n sequences: torch.LongTensor = None\n audio_sequences: Optional[List[torch.LongTensor]] = None\n scores: Optional[Tuple[torch.FloatTensor]] = None\n logits: Optional[Tuple[torch.FloatTensor]] = None\n attentions: Optional[Tuple[Tuple[torch.FloatTensor]]] = None\n hidden_states: Optional[Tuple[Tuple[torch.FloatTensor]]] = None\n past_key_values: Optional[Tuple[Tuple[Tuple[torch.FloatTensor]]]] = None\n\n\nclass HiggsAudioModel(HiggsAudioPreTrainedModel, GenerationMixin):\n \"\"\"Higgs-Audio is an end-to-end multimodal model with the capability to understand and generate text / audio.\n\n Consider the following example for mixed text/audio understanding / generation:\n\n - input_tokens: ![](\"../figures/higgs_audio_v2_architecture_combined.png\")
\n\n\nOur model is built on top of [Llama-3.2-3B](https://huggingface.co/meta-llama/Llama-3.2-3B). To enhance the model’s ability to process audio tokens, we incorporate the \"DualFFN\" architecture as an audio adapter. DualFFN acts as an audio-specific expert, boosting the LLM's performance with minimal computational overhead. Our implementation preserves 91% of the original LLM’s training speed with the inclusion of DualFFN.\n\nSince our audio tokenizer is based on Residual Vector-Quantization (RVQ) and contains multiple codebooks, we adopt the [delay pattern](https://proceedings.neurips.cc/paper_files/paper/2023/file/94b472a1842cd7c56dcb125fb2765fbd-Paper-Conference.pdf) to enable simultaneous code generation across codebooks while supporting streaming.\n\n![](\"../figures/higgs_audio_v2_open_source_delay_pattern.png\")
\n\n\n## DualFFN Performance Ablation Study\n\nTo assess the effectiveness of DualFFN, we trained two smaller models based on LLaMA-3.1-1B: one incorporating DualFFN and one without. Both models were trained for 250K steps with a learning rate of 5e-4 on a subset of the AudioVerse dataset. We evaluated their performance on SeedTTS-Eval, with the results presented in the figures below. The model equipped with DualFFN consistently outperforms its counterpart in terms of word error rate (WER) and speaker similarity.\n\n- SeedTTS-EN\n\n![](\"../figures/dual_ffn_comparison_seed_tts_en_wer.png\")
![](\"../figures/dual_ffn_comparison_seed_tts_en_sim.png\")
\n\n- SeedTTS-ZH\n\n![](\"../figures/dual_ffn_comparison_seed_tts_zh_wer.png\")
![](\"../figures/dual_ffn_comparison_seed_tts_zh_sim.png\")
\n\nWe may notice that the model with DualFFN consistently outperforms the model without DualFFN in terms of word-error-rate (WER) and speaker similarity.\n"
},
{
"path": "tech_blogs/TOKENIZER_BLOG.md",
"content": "# Higgs Audio Tokenizer\n\nIn this work, we introduce a new discretized audio tokenizer that runs at just **25 frames per second** while keeping—or even improving—audio quality compared to tokenizers with twice the bitrate. Our model is the first to train on **24 kHz data** covering speech, music, and sound events in one unified system. It also uses a simple non-diffusion encoder/decoder for fast, batch inference.\n\n\n\n## Basics of Audio Quantization\n\nAn audio signal sampled at $f_s$ Hz is first split into frames by an encoder with hop size $M$, giving a frame rate $f_r = \\frac{f_s}{M}\\quad\\text{(frames/s)}.$\nTwo common quantizers are:\n\n- **Residual Vector Quantization (RVQ)**: $N_q$ cascaded layers with codebook size $N_{cb}$ each. When $N_{cb}=1$, it reduces to single-vector quantization. \n- **Finite Scalar Quantization (FSQ)**: A single layer ($N_q=1$) with codebook size $N_{cb}$. \n\nIf every combination of codewords is a token, the vocabulary size is $N_{cb}^{N_q}$, and each token needs $N_q\\log_2 N_{cb}$ bits. The overall bitrate (bits/s, BPS) is simply $f_r \\times N_q \\log_2 N_{cb}.$ \nWe aim to push this bitrate as low as possible without hurting audio fidelity.\n\n## What Makes Ours Better\n\n- **Low Frame Rate**: At 25 fps, our tokenizer halves the frame rate of many baselines when still maintaining high audio quality. \n- **Unified 24 kHz Training**: We mix speech, music, and sound-event clips in one model, capturing both semantic and acoustic details, hugely facilitating the training of audio language models.\n- **Fast Inference**: By avoiding diffusion steps, our encoder/decoder processes batches quickly, making it practical for real-time or large-scale tasks.\n\n\n## Data and Evaluation Metrics\n\nWe test on four subsets, available [here](https://huggingface.co/datasets/bosonai/AudioTokenBench):\n\n- **Speech, Music, Sound Event**: Includes 1,000 clips for each category, with each clip lasting 10 seconds. Clips are randomly sampled from [DAPS](https://ccrma.stanford.edu/~gautham/Site/daps.html) (Speech), [MUSDB](https://sigsep.github.io/datasets/musdb.html) (Music), and [AudioSet](https://research.google.com/audioset/index.html) (Sound Event).\n\n- **Audiophile**: Contains 150 clips, each 30 seconds long, curated from eleven high-fidelity test discs. The clips feature both music and sound events, selected for audio quality evaluation.\n\nWe measure:\n\n- **Acoustic Quality**: STFT distance between the original and reconstructed audio. \n- **Semantic Integrity**: Semantic preservation of the original audio using [SeedTTS](https://arxiv.org/abs/2406.02430)[15] dataset on English and Chinese. \n- **Aesthetics**: SOTA unified model-based quality assessment, [Meta Audiobox Aesthetics](https://github.com/facebookresearch/audiobox-aesthetics)[8], for Content Enjoyment (CE), Content Usefulness (CU) .\n\n\nWe compare our tokenizer with a wide range of baselines, from tokenizers mainly built for better acoustic reconstruction and compression rate, to those focused on semantic integrity, and to tokenizers used in existing large audio language models. We also compare with tokenizers that are pretrained specifically on speech or on music.\n\n\nThe tables below summarize the tokenizers evaluated. As shown, our tokenizer achieves a well-rounded balance of efficiency, semantic fidelity, and acoustic quality.\n\n### Accoustic Evaluation\n\nWe use the STFT metric here for simplicity. The baselines are ordered chronologically, grouped by whether semantic distillation (SD) is applied.Despite DAC’s top acoustic quality at 12× the bitrate, our tokenizer leads all other baselines.\n\n\n| Tokenizer | 💬 | 🎵 | 🥁 | SD | $f_s$ | $f_r$ | BPS* (k) ↓ | Speech ↓ | Sound Event ↓ | Music ↓ | Audiophile ↓ |\n|-----------|----|----|----|----|-------|-------|--------------------------|----------|----------------|--------|--------------|\n| [Encodec](https://huggingface.co/facebook/encodec_24khz)[3] | ✓ | ✓ | ✓ | | 24 | 75 | 24 | 1.96 | 2.65 | 2.52 | 2.30 |\n| [DAC](https://huggingface.co/hance-ai/descript-audio-codec-24khz)[2] | ✓ | ✓ | ✓ | | 24 | 75 | 24 | **1.13** | **1.45** | **1.34** | **1.62** |\n| [SNAC-24k](https://huggingface.co/hubertsiuzdak/snac_24khz)[6] | ✓ | | | | 24 | (12, 23, 47) | 0.98 | 1.92 | 2.69 | 2.54 | 2.52 |\n| [SNAC-44.1k](https://huggingface.co/hubertsiuzdak/snac_44khz)[6] | | ✓ | ✓ | | 44.1 | (14, 29, 57, 115) | 2.6 | 1.83 | 2.25 | 2.05 | 2.00 |\n| [WavTokenizer](https://huggingface.co/novateur/WavTokenizer-medium-music-audio-75token/blob/main/wavtokenizer_medium_music_audio_320_24k_v2.ckpt)[7] | | ✓ | ✓ | | 24 | 75 | 0.9 | 1.93 | 2.44 | 2.17 | 2.15 |\n| [WavTokenizer (Speech)](https://huggingface.co/novateur/WavTokenizer-large-speech-75token/tree/main)[7] | ✓ | | | | 24 | 75 | 0.9 | 1.78 | 2.47 | 2.42 | 2.47 |\n| [MuCodec](https://huggingface.co/haoheliu/audioldm_48k/tree/main)[11] | | ✓ | | | 48 | 25 | 0.35 | 2.87 | 3.69 | 3.36 | 2.97 |\n| [FlowDec-75m](https://github.com/facebookresearch/FlowDec?tab=readme-ov-file)[12] | ✓ | ✓ | ✓ | | 48 | 75 | 7.5 | 1.73 | 2.14 | 2.01 | 2.03 |\n| [FlowDec-25s](https://github.com/facebookresearch/FlowDec?tab=readme-ov-file)[12] | ✓ | ✓ | ✓ | | 48 | 25 | 4 | 1.94 | 2.42 | 2.25 | 2.33 |\n| [SpeechTokenizer](https://huggingface.co/fnlp/SpeechTokenizer/tree/main/speechtokenizer_hubert_avg)[14] | ✓ | | | ✓ | 16 | 50 | 4 | 3.21 | 3.58 | 3.65 | 3.69 |\n| [SemantiCodec](https://huggingface.co/haoheliu/SemantiCodec/tree/main/semanticodec_tokenrate_100)[5] | ✓ | ✓ | ✓ | ✓ | 16 | 50 | 1.4 | 3.05 | 3.28 | 3.24 | 3.18 |\n| [Mimi](https://huggingface.co/docs/transformers/en/model_doc/mimi)[13] | ✓ | | | ✓ | 24 | 12.5 | 4.4 | 1.77 | 2.40 | 2.30 | 2.15 |\n| [XCodec](https://huggingface.co/ZhenYe234/xcodec/blob/main/config_hubert_general.yaml)[1] | ✓ | ✓ | ✓ | ✓ | 16 | 50 | 4 | 2.95 | 3.16 | 3.00 | 3.03 |\n| [CosyVoice 2](https://huggingface.co/FunAudioLLM/CosyVoice2-0.5B)[13] | ✓ | | | ✓ | 16 | 25 | -** | 2.30 | 3.30 | 3.14 | 3.25 |\n| [XCodec2](https://huggingface.co/HKUST-Audio/xcodec2/blob/main/ckpt/epoch%3D4-step%3D1400000.ckpt)[9] | ✓ | | | ✓ | 16 | 50 | 0.8 | 3.06 | 3.72 | 3.62 | 3.64 |\n| [XY](https://huggingface.co/fnlp/XY_Tokenizer_TTSD_V0/tree/main)[10] | ✓ | | | ✓ | 24 | 12.5 | 1 | 1.89 | 2.51 | 2.40 | 2.26 |\n| Ours | ✓ | ✓ | ✓ | ✓ | 24 | 25 | 2 | **1.62** | **2.03** | **1.85** | **1.80** |\n\n\n\n\\* Bits-per-second is calculated according to the checkpoint the author provided.\n\n\\*\\* CosyVoice 2 uses the continuous feature as the conditioning, we include it for completeness.\n\n\n### Semantic Evaluation\nHere we only compare with tokenizers that are trained with semantic distillation.\n[SeedTTS](https://github.com/BytedanceSpeech/seed-tts-eval) is a dataset includes prompt/target audio and texts. We reconstructed the target audio, and use the word error rate (WER) and speaker similarity (SIM) metrics to evaluate the semantic integrity. SIM is calculated by the similarity between the prompt audio and reconstructed targeted audio with [WavLM-large](https://drive.google.com/file/d/1-aE1NfzpRCLxA4GUxX9ITI3F9LlbtEGP/view) as the embedding model. \n\nThe following table shows that our tokenizer achieves comparable performance to tokenizers that 2.2x the bitrate of our model.\n\n| Model | BPS (k) | en WER ↓ | en SIM ↑ | zh WER ↓ | zh SIM ↑ |\n|------------------|---------|------------|------------|------------|------------|\n| [SpeechTokenizer](https://huggingface.co/fnlp/SpeechTokenizer/tree/main/speechtokenizer_hubert_avg) | 4 | 2.82 | 0.63 | 2.04 | 0.65 |\n| [SemantiCodec](https://huggingface.co/haoheliu/SemantiCodec/tree/main/semanticodec_tokenrate_100) | 1.4 | 3.46 | 0.56 | 2.18 | 0.60 |\n| [Mimi](https://huggingface.co/docs/transformers/en/model_doc/mimi) | 4.4 | **2.35** | **0.70** | **1.48** | **0.72** |\n| [XCodec](https://huggingface.co/ZhenYe234/xcodec/blob/main/config_hubert_general.yaml) | 4.0 | 2.68 | 0.63 | 1.66 | 0.66 |\n| [CosyVoice 2](https://huggingface.co/FunAudioLLM/CosyVoice2-0.5B) | - | 3.17 | 0.65 | 2.11 | 0.70 |\n| [XCodec2](https://huggingface.co/HKUST-Audio/xcodec2/blob/main/ckpt/epoch%3D4-step%3D1400000.ckpt) | 0.8 | 2.74 | 0.62 | 1.91 | 0.67 |\n| [XY-MOSS-TTSD](https://huggingface.co/fnlp/XY_Tokenizer_TTSD_V0/tree/main) | 1.0 | 2.72 | 0.61 | 1.58 | 0.67 |\n| Ours | 2.0 | 2.52 | 0.67 | **1.48** | 0.71 |\n\n\n\n### Audiobox Aesthetics Evaluation\n\nThis model based evaluation[8] further demonstrates the superiority of our tokenizer. CU is the Content Usefulness and CE is the Content Enjoyment. Each term is rated on a scale of 1-10. Notably, our tokenizer performs best on the Audiophile set—demonstrating a clear advantage when the original audio quality is high.\n\n\n| Model | BPS (k) | Music CE ↑ | Music CU ↑ | Sound Event CE ↑ | Sound Event CU ↑ | Speech CE ↑ | Speech CU ↑ | Audiophile CE ↑ | Audiophile CU ↑ |\n|------------------|---------|--------------|--------------|--------------------|--------------------|---------------|---------------|--------------------|--------------------|\n| Origin | - | 6.20 | 7.10 | 4.47 | 5.64 | 5.03 | 4.87 | 7.17 | 7.65 |\n| [SpeechTokenizer](https://huggingface.co/fnlp/SpeechTokenizer/tree/main/speechtokenizer_hubert_avg) | 4.0 | 3.55 | 5.22 | 3.03 | 4.50 | 4.68 | 4.58 | 3.59 | 5.07 |\n| [SemantiCodec](https://huggingface.co/haoheliu/SemantiCodec/tree/main/semanticodec_tokenrate_100) | 1.4 | 6.01 | 6.83 | 4.22 | 5.30 | 4.28 | 4.12 | 6.97 | 7.43 |\n| [Mimi](https://huggingface.co/docs/transformers/en/model_doc/mimi) | 4.4 | 6.01 | 6.83 | 4.26 | 5.35 | 4.87 | 4.72 | 6.80 | 7.29 |\n| [XCodec](https://huggingface.co/ZhenYe234/xcodec/blob/main/config_hubert_general.yaml) | 4.0 | **6.30** | **7.10** | **4.43** | 5.45 | **4.96** | **4.79** | 7.06 | 7.49 |\n| [CosyVoice 2](https://huggingface.co/FunAudioLLM/CosyVoice2-0.5B) | - | 5.21 | 6.14 | 4.08 | 4.73 | **4.91** | **4.75** | 5.97 | 6.56 |\n| [XCodec2](https://huggingface.co/HKUST-Audio/xcodec2/blob/main/ckpt/epoch%3D4-step%3D1400000.ckpt) | 0.8 | 4.38 | 5.66 | 3.43 | 4.63 | **4.93** | **4.78** | 4.56 | 5.46 |\n| [XY-MOSS-TTSD](https://huggingface.co/fnlp/XY_Tokenizer_TTSD_V0/tree/main) | 1.0 | 5.77 | 6.80 | 4.23 | 5.34 | 4.88 | 4.72 | 6.95 | 7.48 |\n| Ours | 2.0 | **6.35** | **7.15** | **4.47** | **5.51** | 4.90 | 4.70 | **7.21** | **7.66** |\n\n\n\nNote that since some tokenizers are trained on 16 kHz data, we upsample their audio outputs to 24 kHz before computing metrics. Different upsampling methods may cause slight variations (e.g., 4.36 vs. 4.43 for XCodec Sound Event CE). We report the best results we could obtain and highlight any results within 0.05 of the best one. \n\n\n\n\n\n\n\n\n\n\n\n\n## Reference \n[1] [Ye, Zhen, et al. \"Codec does matter: Exploring the semantic shortcoming of codec for audio language model.\" Proceedings of the AAAI Conference on Artificial Intelligence. Vol. 39. No. 24. 2025.](https://arxiv.org/abs/2408.17175)\n\n[2] [Kumar, Rithesh, et al. \"High-fidelity audio compression with improved rvqgan.\" Advances in Neural Information Processing Systems 36 (2023): 27980-27993.](https://dl.acm.org/doi/10.5555/3666122.3667336)\n\n[3] [Défossez, Alexandre, et al. \"High fidelity neural audio compression.\" arXiv preprint arXiv:2210.13438 (2022).](https://arxiv.org/abs/2210.13438)\n\n[4] [Défossez, Alexandre, et al. \"Moshi: a speech-text foundation model for real-time dialogue.\" arXiv preprint arXiv:2410.00037 (2024).](https://arxiv.org/abs/2410.00037)\n\n[5] [Liu, Haohe, et al. \"Semanticodec: An ultra low bitrate semantic audio codec for general sound.\" IEEE Journal of Selected Topics in Signal Processing (2024).](https://ieeexplore.ieee.org/document/10768970)\n\n[6] [Siuzdak, Hubert, Florian Grötschla, and Luca A. Lanzendörfer. \"Snac: Multi-scale neural audio codec.\" arXiv preprint arXiv:2410.14411 (2024).](https://arxiv.org/abs/2410.14411)\n\n[7] [Ji, Shengpeng, et al. \"Wavtokenizer: an efficient acoustic discrete codec tokenizer for audio language modeling.\" arXiv preprint arXiv:2408.16532 (2024).](https://arxiv.org/abs/2408.16532)\n\n[8] [Tjandra, Andros, et al. \"Meta audiobox aesthetics: Unified automatic quality assessment for speech, music, and sound.\" arXiv preprint arXiv:2502.05139 (2025).](https://arxiv.org/abs/2502.05139)\n\n[9] [Ye, Zhen, et al. \"Llasa: Scaling Train-Time and Inference-Time Compute for Llama-based Speech Synthesis.\" arXiv preprint arXiv:2502.04128 (2025).](https://arxiv.org/abs/2502.04128)\n\n[10] [Gong, Yitian, et al. \"XY-Tokenizer: Mitigating the Semantic-Acoustic Conflict in Low-Bitrate Speech Codecs.\" arXiv preprint arXiv:2506.23325 (2025).](https://arxiv.org/abs/2506.23325)\n\n[11] [Xu, Yaoxun, et al. \"MuCodec: Ultra Low-Bitrate Music Codec.\" arXiv preprint arXiv:2409.13216 (2024).](https://arxiv.org/abs/2409.13216)\n\n[12] [Welker, Simon, et al. \"FlowDec: A flow-based full-band general audio codec with high perceptual quality.\" arXiv preprint arXiv:2503.01485 (2025).](https://arxiv.org/abs/2503.01485)\n\n[13] [Du, Zhihao, et al. \"Cosyvoice 2: Scalable streaming speech synthesis with large language models.\" arXiv preprint arXiv:2412.10117 (2024).](https://arxiv.org/abs/2412.10117)\n\n[14] [Zhang, Xin, et al. \"Speechtokenizer: Unified speech tokenizer for speech large language models.\" arXiv preprint arXiv:2308.16692 (2023).](https://arxiv.org/abs/2308.16692)\n\n[15] [Anastassiou, Philip, et al. \"Seed-tts: A family of high-quality versatile speech generation models.\" arXiv preprint arXiv:2406.02430 (2024).](https://arxiv.org/abs/2406.02430)\n"
}
]