Prepare Audio8-TTS-Realistic-0.6B private release

Complete inference package, model card, multilingual evaluation, and verified SGLang Omni compatibility.

.gitattributes

@@ -33,3 +33,4 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 *.zip filter=lfs diff=lfs merge=lfs -text
 *.zst filter=lfs diff=lfs merge=lfs -text
 *tfevents* filter=lfs diff=lfs merge=lfs -text
+tokenizer.json filter=lfs diff=lfs merge=lfs -text

LICENSE

@@ -0,0 +1,201 @@
+                                 Apache License
+                           Version 2.0, January 2004
+                        http://www.apache.org/licenses/
+
+   TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION
+
+   1. Definitions.
+
+      "License" shall mean the terms and conditions for use, reproduction,
+      and distribution as defined by Sections 1 through 9 of this document.
+
+      "Licensor" shall mean the copyright owner or entity authorized by
+      the copyright owner that is granting the License.
+
+      "Legal Entity" shall mean the union of the acting entity and all
+      other entities that control, are controlled by, or are under common
+      control with that entity. For the purposes of this definition,
+      "control" means (i) the power, direct or indirect, to cause the
+      direction or management of such entity, whether by contract or
+      otherwise, or (ii) ownership of fifty percent (50%) or more of the
+      outstanding shares, or (iii) beneficial ownership of such entity.
+
+      "You" (or "Your") shall mean an individual or Legal Entity
+      exercising permissions granted by this License.
+
+      "Source" form shall mean the preferred form for making modifications,
+      including but not limited to software source code, documentation
+      source, and configuration files.
+
+      "Object" form shall mean any form resulting from mechanical
+      transformation or translation of a Source form, including but
+      not limited to compiled object code, generated documentation,
+      and conversions to other media types.
+
+      "Work" shall mean the work of authorship, whether in Source or
+      Object form, made available under the License, as indicated by a
+      copyright notice that is included in or attached to the work
+      (an example is provided in the Appendix below).
+
+      "Derivative Works" shall mean any work, whether in Source or Object
+      form, that is based on (or derived from) the Work and for which the
+      editorial revisions, annotations, elaborations, or other modifications
+      represent, as a whole, an original work of authorship. For the purposes
+      of this License, Derivative Works shall not include works that remain
+      separable from, or merely link (or bind by name) to the interfaces of,
+      the Work and Derivative Works thereof.
+
+      "Contribution" shall mean any work of authorship, including
+      the original version of the Work and any modifications or additions
+      to that Work or Derivative Works thereof, that is intentionally
+      submitted to Licensor for inclusion in the Work by the copyright owner
+      or by an individual or Legal Entity authorized to submit on behalf of
+      the copyright owner. For the purposes of this definition, "submitted"
+      means any form of electronic, verbal, or written communication sent
+      to the Licensor or its representatives, including but not limited to
+      communication on electronic mailing lists, source code control systems,
+      and issue tracking systems that are managed by, or on behalf of, the
+      Licensor for the purpose of discussing and improving the Work, but
+      excluding communication that is conspicuously marked or otherwise
+      designated in writing by the copyright owner as "Not a Contribution."
+
+      "Contributor" shall mean Licensor and any individual or Legal Entity
+      on behalf of whom a Contribution has been received by Licensor and
+      subsequently incorporated within the Work.
+
+   2. Grant of Copyright License. Subject to the terms and conditions of
+      this License, each Contributor hereby grants to You a perpetual,
+      worldwide, non-exclusive, no-charge, royalty-free, irrevocable
+      copyright license to reproduce, prepare Derivative Works of,
+      publicly display, publicly perform, sublicense, and distribute the
+      Work and such Derivative Works in Source or Object form.
+
+   3. Grant of Patent License. Subject to the terms and conditions of
+      this License, each Contributor hereby grants to You a perpetual,
+      worldwide, non-exclusive, no-charge, royalty-free, irrevocable
+      (except as stated in this section) patent license to make, have made,
+      use, offer to sell, sell, import, and otherwise transfer the Work,
+      where such license applies only to those patent claims licensable
+      by such Contributor that are necessarily infringed by their
+      Contribution(s) alone or by combination of their Contribution(s)
+      with the Work to which such Contribution(s) was submitted. If You
+      institute patent litigation against any entity (including a
+      cross-claim or counterclaim in a lawsuit) alleging that the Work
+      or a Contribution incorporated within the Work constitutes direct
+      or contributory patent infringement, then any patent licenses
+      granted to You under this License for that Work shall terminate
+      as of the date such litigation is filed.
+
+   4. Redistribution. You may reproduce and distribute copies of the
+      Work or Derivative Works thereof in any medium, with or without
+      modifications, and in Source or Object form, provided that You
+      meet the following conditions:
+
+      (a) You must give any other recipients of the Work or
+          Derivative Works a copy of this License; and
+
+      (b) You must cause any modified files to carry prominent notices
+          stating that You changed the files; and
+
+      (c) You must retain, in the Source form of any Derivative Works
+          that You distribute, all copyright, patent, trademark, and
+          attribution notices from the Source form of the Work,
+          excluding those notices that do not pertain to any part of
+          the Derivative Works; and
+
+      (d) If the Work includes a "NOTICE" text file as part of its
+          distribution, then any Derivative Works that You distribute must
+          include a readable copy of the attribution notices contained
+          within such NOTICE file, excluding those notices that do not
+          pertain to any part of the Derivative Works, in at least one
+          of the following places: within a NOTICE text file distributed
+          as part of the Derivative Works; within the Source form or
+          documentation, if provided along with the Derivative Works; or,
+          within a display generated by the Derivative Works, if and
+          wherever such third-party notices normally appear. The contents
+          of the NOTICE file are for informational purposes only and
+          do not modify the License. You may add Your own attribution
+          notices within Derivative Works that You distribute, alongside
+          or as an addendum to the NOTICE text from the Work, provided
+          that such additional attribution notices cannot be construed
+          as modifying the License.
+
+      You may add Your own copyright statement to Your modifications and
+      may provide additional or different license terms and conditions
+      for use, reproduction, or distribution of Your modifications, or
+      for any such Derivative Works as a whole, provided Your use,
+      reproduction, and distribution of the Work otherwise complies with
+      the conditions stated in this License.
+
+   5. Submission of Contributions. Unless You explicitly state otherwise,
+      any Contribution intentionally submitted for inclusion in the Work
+      by You to the Licensor shall be under the terms and conditions of
+      this License, without any additional terms or conditions.
+      Notwithstanding the above, nothing herein shall supersede or modify
+      the terms of any separate license agreement you may have executed
+      with Licensor regarding such Contributions.
+
+   6. Trademarks. This License does not grant permission to use the trade
+      names, trademarks, service marks, or product names of the Licensor,
+      except as required for reasonable and customary use in describing the
+      origin of the Work and reproducing the content of the NOTICE file.
+
+   7. Disclaimer of Warranty. Unless required by applicable law or
+      agreed to in writing, Licensor provides the Work (and each
+      Contributor provides its Contributions) on an "AS IS" BASIS,
+      WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or
+      implied, including, without limitation, any warranties or conditions
+      of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A
+      PARTICULAR PURPOSE. You are solely responsible for determining the
+      appropriateness of using or redistributing the Work and assume any
+      risks associated with Your exercise of permissions under this License.
+
+   8. Limitation of Liability. In no event and under no legal theory,
+      whether in tort (including negligence), contract, or otherwise,
+      unless required by applicable law (such as deliberate and grossly
+      negligent acts) or agreed to in writing, shall any Contributor be
+      liable to You for damages, including any direct, indirect, special,
+      incidental, or consequential damages of any character arising as a
+      result of this License or out of the use or inability to use the
+      Work (including but not limited to damages for loss of goodwill,
+      work stoppage, computer failure or malfunction, or any and all
+      other commercial damages or losses), even if such Contributor
+      has been advised of the possibility of such damages.
+
+   9. Accepting Warranty or Additional Liability. While redistributing
+      the Work or Derivative Works thereof, You may choose to offer,
+      and charge a fee for, acceptance of support, warranty, indemnity,
+      or other liability obligations and/or rights consistent with this
+      License. However, in accepting such obligations, You may act only
+      on Your own behalf and on Your sole responsibility, not on behalf
+      of any other Contributor, and only if You agree to indemnify,
+      defend, and hold each Contributor harmless for any liability
+      incurred by, or claims asserted against, such Contributor by reason
+      of your accepting any such warranty or additional liability.
+
+   END OF TERMS AND CONDITIONS
+
+   APPENDIX: How to apply the Apache License to your work.
+
+      To apply the Apache License to your work, attach the following
+      boilerplate notice, with the fields enclosed by brackets "[]"
+      replaced with your own identifying information. (Don't include
+      the brackets!)  The text should be enclosed in the appropriate
+      comment syntax for the file format. We also recommend that a
+      file or class name and description of purpose be included on the
+      same "printed page" as the copyright notice for easier
+      identification within third-party archives.
+
+   Copyright [yyyy] [name of copyright owner]
+
+   Licensed under the Apache License, Version 2.0 (the "License");
+   you may not use this file except in compliance with the License.
+   You may obtain a copy of the License at
+
+       http://www.apache.org/licenses/LICENSE-2.0
+
+   Unless required by applicable law or agreed to in writing, software
+   distributed under the License is distributed on an "AS IS" BASIS,
+   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+   See the License for the specific language governing permissions and
+   limitations under the License.

NOTICE

@@ -0,0 +1,4 @@
+audio8_tts
+
+The model architecture is inspired by the DualAR design used in Fish Audio
+S2 Pro.

README.md

@@ -1,3 +1,306 @@
 ---
 license: apache-2.0
+language:
+  - zh
+  - en
+  - ja
+  - ko
+  - de
+  - fr
+  - es
+  - it
+  - ru
+  - ar
+  - fi
+  - hi
+  - id
+  - ms
+  - nl
+  - "no"
+  - pl
+  - pt
+  - sv
+  - th
+  - tr
+  - vi
+  - yue
+library_name: transformers
+pipeline_tag: text-to-speech
+base_model: Edge0/Audio8-TTS-Preview-0.6b
+tags:
+  - audio
+  - text-to-speech
+  - tts
+  - voice-cloning
+  - zero-shot
+  - multilingual
 ---
+
+<div align="center">
+
+<img src="./audio8-banner.jpeg" alt="Audio8" width="760">
+
+# Audio8-TTS-Realistic-0.6B
+
+**More realistic speech. More languages. Still 0.6B.**
+
+[![GitHub](https://img.shields.io/badge/GitHub-Audio8__TTS-black?style=for-the-badge&logo=github)](https://github.com/Edge0-AI/Audio8_TTS)
+[![License](https://img.shields.io/badge/License-Apache%202.0-blue?style=for-the-badge)](./LICENSE)
+![Demo — coming soon](https://img.shields.io/badge/Demo-Coming%20soon-lightgrey?style=for-the-badge&logo=githubpages)
+
+</div>
+
+**Audio8-TTS-Realistic-0.6B** builds on [Audio8-TTS-Preview-0.6b](https://huggingface.co/Edge0/Audio8-TTS-Preview-0.6b), delivering **substantially improved speech realism** while retaining its compact **0.6B** architecture and zero-shot voice cloning capabilities.
+
+This release supports **22 languages plus Cantonese**, adding **12 languages** to Preview's coverage. It uses the same Transformers inference interface as Preview. This repository includes the complete model checkpoint, its 44.1 kHz neural audio codec, tokenizer, processor, and custom Transformers code.
+
+## Highlights
+
+- **Improved realism:** more natural prosody and expressive delivery, building on Preview.
+- **Broader multilingual coverage:** 22 languages plus Cantonese, with 12 newly supported languages.
+- **Compact model:** 601 million main-model parameters, with the audio codec bundled separately.
+- **Flexible generation:** zero-shot voice cloning from a reference recording and transcript, or speech generation without a reference.
+- **GPU serving:** an SGLang Omni adapter with paged attention, dynamic batching, and an OpenAI-compatible speech API.
+
+## Supported Languages
+
+| Coverage | Languages |
+|---|---|
+| Retained from Preview | Chinese, Dutch, English, French, German, Italian, Japanese, Korean, Polish, Spanish |
+| Newly supported | Arabic, Finnish, Hindi, Indonesian, Malay, Norwegian, Portuguese, Russian, Swedish, Thai, Turkish, Vietnamese |
+| Also retained | Cantonese |
+
+Language quality varies; see the per-language evaluation below. Benchmark coverage includes some additional languages and is distinct from the supported-language list.
+
+## Model Details
+
+Audio8 TTS uses a DualAR architecture inspired by [Fish Audio S2 Pro](https://github.com/fishaudio/fish-speech). The slow AR transformer predicts one semantic token for each audio frame. The fast AR transformer predicts the frame's codec codebooks, conditioned on the slow hidden state and preceding codebooks.
+
+| Component | Configuration |
+|---|---|
+| Main model | 601,159,424 parameters, excluding the codec |
+| Slow AR | 24 layers, width 896, 14 attention heads, 2 KV heads |
+| Fast AR | 4 layers, width 896, 14 attention heads, 2 KV heads |
+| Acoustic tokens | 10 codebooks, 4,096 entries per codebook |
+| Codec | 44.1 kHz, 2,048 samples per model frame (~21.5 frames/s) |
+| Context | Up to 2,048 packed text/audio positions |
+
+The bundled codec handles both reference-audio encoding and waveform decoding, so no additional codec checkpoint is required. The **0.6B** size refers to the main DualAR model; the codec is excluded from that parameter count.
+
+## Installation
+
+Python 3.10 or newer and a CUDA-capable GPU are recommended.
+
+```bash
+pip install "torch>=2.5.0" "torchaudio>=2.5.0" \
+  "transformers>=4.57.0,<5" "soundfile>=0.12" "safetensors>=0.4"
+```
+
+The Transformers release package was smoke-tested with PyTorch 2.8.0 and Transformers 4.57.3 on an NVIDIA H20 using BF16. Checks covered model loading, English and Chinese generation with and without reference audio, and generation in Arabic, Cantonese, Portuguese, and Russian. These are runtime checks, separate from the benchmark results below.
+
+## Usage
+
+The model uses custom Transformers code. Review the files in this repository, then load it with `trust_remote_code=True`.
+
+### Zero-shot voice cloning
+
+Use a clean reference recording. Its transcript must match the spoken content exactly.
+
+```python
+import soundfile as sf
+import torch
+from transformers import AutoModel, AutoProcessor
+
+model_id = "Edge0/Audio8-TTS-Realistic-0.6b"
+device = "cuda" if torch.cuda.is_available() else "cpu"
+dtype = torch.bfloat16 if device == "cuda" else torch.float32
+
+processor = AutoProcessor.from_pretrained(model_id, trust_remote_code=True)
+model = AutoModel.from_pretrained(
+    model_id,
+    trust_remote_code=True,
+    dtype=dtype,
+).eval().to(device)
+
+inputs = processor(
+    text=["Welcome to Audio8 TTS. It is wonderful to hear from you."],
+    reference_audio=["reference.wav"],
+    reference_text=["The exact transcript of the reference recording."],
+    return_tensors="pt",
+)
+inputs = {name: value.to(device) for name, value in inputs.items()}
+
+with torch.inference_mode():
+    output = model.generate(
+        **inputs,
+        max_new_tokens=1024,
+        temperature=0.8,
+        top_p=0.95,
+        top_k=50,
+        do_sample=True,
+        return_dict_in_generate=True,
+    )
+    waveforms, waveform_lengths = model.decode_audio(output.codes)
+
+audio = waveforms[0, : int(waveform_lengths[0])].float().cpu().numpy()
+sf.write("output.wav", audio, model.config.codec_sample_rate)
+```
+
+### Generation without a reference
+
+Replace the processor call above with:
+
+```python
+inputs = processor(
+    text=["This utterance does not use a reference voice."],
+    return_tensors="pt",
+)
+```
+
+Then move the inputs to the model device and run the same generation and decoding steps. For command-line inference, batching, and supervised fine-tuning, see the [Audio8 TTS repository](https://github.com/Edge0-AI/Audio8_TTS), using `Edge0/Audio8-TTS-Realistic-0.6b` as the model path.
+
+## High-throughput serving: SGLang Omni
+
+The [SGLang Omni adapter](https://github.com/Edge0-AI/Audio8_TTS/tree/master/sglang_omni) provides GPU serving for Audio8 TTS. It is installed as an independent model plugin and does not overwrite SGLang Omni core files.
+
+| Capability | Support |
+|---|---|
+| Attention and batching | SGLang paged attention and dynamic batching |
+| DualAR execution | Slow AR serving with a fixed KV cache for the Fast AR codebook decoder |
+| Voice cloning | Reference-audio encoding and waveform decoding |
+| API | OpenAI-compatible `/v1/audio/speech` service |
+
+The adapter targets SGLang Omni revision `68a572348837f7b004857b4b07993c20ade4c017`, with PyTorch 2.9.1, SGLang 0.5.8, and Transformers 4.57.1. Follow the [deployment instructions](https://github.com/Edge0-AI/Audio8_TTS#sglang-omni-serving) for installation, configuration, and API examples, and set the model path to `Edge0/Audio8-TTS-Realistic-0.6b` or its downloaded snapshot. The adapter supports generation without a reference and zero-shot voice cloning.
+
+Use the **0.6B** configuration, `sglang_omni/configs/audio8_tts_0_6b.yaml`. This Realistic checkpoint passed non-streaming `/v1/audio/speech` smoke tests both with and without a reference on an NVIDIA H20, using the versions above and the adapter from Audio8_TTS commit `07e40f5d0b03fc473635ef378654bfb581027ac3`. These checks used BF16 with Torch compilation and CUDA graphs disabled; they validate compatibility, not throughput or streaming performance.
+
+## Evaluation
+
+Error rates are percentages (**lower is better**). Speaker similarity (SIM) is on a **0–1 scale** (**higher is better**). These metrics measure intelligibility and speaker similarity; they do not directly measure perceived speech realism.
+
+### Benchmark summary
+
+Results for **Audio8-TTS-Realistic (0.6B)**:
+
+| Benchmark | Coverage | Mean error % ↓ | Mean SIM ↑ |
+|---|---|---:|---:|
+| Seed-TTS | English / Chinese | **1.25** | **0.704** |
+| CV3 | 9 core languages | **4.86** | **0.661** |
+| MiniMax | 24 evaluated languages | **5.89** | **0.754** |
+
+All three summary error rates are unweighted language macro averages. Seed-TTS averages English WER and Chinese CER, excluding hard Chinese. CV3 averages Chinese, English, Japanese, Korean, German, French, Spanish, Italian, and Russian, excluding hard subsets and the 13 additional languages. MiniMax averages all 24 evaluated languages, including languages outside the supported-language list.
+
+### Reference comparison
+
+The following error rates come from the sources used in our technical report. Different evaluation pipelines and text normalization make this a reference comparison rather than a strictly matched ranking. A dash means that a result is unavailable in the comparison, not that the model lacks support.
+
+| Model | Parameters (B) | Seed-TTS EN/ZH error ↓ | CV3 (9) error ↓ | MiniMax (24) error ↓ | Result source |
+|---|---:|---:|---:|---:|---|
+| **Audio8-TTS-Realistic** | **0.6** | **1.25** | **4.86** | **5.89** | Our evaluation |
+| Qwen3-TTS-12Hz-0.6B | 0.6 | 1.12 | — | — | Qwen report |
+| Qwen3-TTS-1.7B | 1.7 | 1.01 | 7.73 | 27.41 | Qwen report / Boson evaluation |
+| Higgs Audio v3 | 4 | 1.11 | 4.41 | 2.74 | Boson evaluation |
+| Higgs Audio v2 | 4.7 | 1.17 | 4.49 | — | Our reevaluation |
+| Fish S2 Pro | 4.6 | 1.32 | 4.49 | 4.40 | Our reevaluation / Fish report |
+| MOSS-TTS-v1.5 | 8.5 | 1.52 | 6.11 | 3.78 | VoxCPM compilation / Boson evaluation |
+| CosyVoice3-1.5B | 1.5 | 1.67 | 6.89 | — | VoxCPM compilation |
+| VoxCPM2 | 2.3 | 1.41 | 5.35 | — | VoxCPM compilation |
+
+Sources: [Qwen3-TTS report](https://arxiv.org/abs/2601.15621), [Boson AI Higgs Audio v3 benchmarks](https://huggingface.co/bosonai/higgs-audio-v3-tts-4b), [Fish Audio S2 report](https://arxiv.org/abs/2603.08823), and the [VoxCPM repository](https://github.com/OpenBMB/VoxCPM). Fish S2 Pro was reevaluated on Seed-TTS and CV3 because its official evaluation uses a different normalizer; Higgs Audio v2 was evaluated locally. Cross-project SIM values are omitted here because speaker encoders and evaluation conventions differ.
+
+Parameter conventions also differ: Audio8's 0.6B excludes its separate codec; Qwen and CosyVoice use release labels; Higgs Audio v3 uses its nominal **4B** release label. The remaining sizes are rounded checkpoint counts used in our report (VoxCPM2 excludes its separate AudioVAE). They should not be interpreted as identically counted end-to-end system sizes.
+
+### Seed-TTS detail
+
+| Split | Metric | Error % ↓ | SIM ↑ |
+|---|---|---:|---:|
+| English | WER | 1.464 | 0.663 |
+| Chinese | CER | 1.029 | 0.745 |
+| Chinese, hard | CER | 12.797 | 0.702 |
+
+### Multilingual detail
+
+<details>
+<summary>CV3: 22 languages and hard subsets</summary>
+
+| Language | Metric | Error % ↓ | SIM ↑ |
+|---|---|---:|---:|
+| Chinese | CER | 1.814 | 0.6898 |
+| English | WER | 2.907 | 0.6295 |
+| Japanese | CER | 7.511 | 0.6659 |
+| Korean | CER | 4.801 | 0.6962 |
+| German | WER | 3.200 | 0.6669 |
+| French | WER | 7.938 | 0.6355 |
+| Spanish | WER | 3.389 | 0.6680 |
+| Italian | WER | 4.933 | 0.6498 |
+| Russian | WER | 7.258 | 0.6451 |
+| Arabic | WER | 21.098 | 0.6122 |
+| Finnish | WER | 32.279 | 0.7028 |
+| Hindi | WER | 21.714 | 0.7245 |
+| Indonesian | WER | 13.743 | 0.6682 |
+| Malay | WER | 17.798 | 0.6710 |
+| Dutch | WER | 10.873 | 0.6918 |
+| Norwegian | WER | 17.961 | 0.6905 |
+| Polish | WER | 19.233 | 0.6738 |
+| Portuguese | WER | 6.871 | 0.6772 |
+| Swedish | WER | 19.252 | 0.7104 |
+| Thai | WER | 48.574 | 0.6691 |
+| Turkish | WER | 12.567 | 0.6496 |
+| Vietnamese | WER | 28.111 | 0.6790 |
+
+Hard Chinese CER is **10.665%** and hard English WER is **4.736%**. These hard subsets are excluded from the 9-language summary. Each of the 22 regular language rows contains 500 evaluation items. Chinese, Japanese, and Korean use CER; the other regular language rows use WER.
+
+</details>
+
+<details>
+<summary>MiniMax: all 24 evaluated languages</summary>
+
+| Language | Metric | Error % ↓ | SIM ↑ |
+|---|---|---:|---:|
+| Arabic | CER | 2.403 | 0.7040 |
+| Cantonese | CER | 31.349 | 0.7710 |
+| Chinese | CER | 1.112 | 0.7950 |
+| Czech † | WER | 8.628 | 0.7330 |
+| Dutch | WER | 1.715 | 0.7050 |
+| English | WER | 2.161 | 0.7740 |
+| Finnish | WER | 5.072 | 0.7910 |
+| French | WER | 3.198 | 0.6660 |
+| German | WER | 0.898 | 0.7510 |
+| Greek † | CER | 4.181 | 0.7910 |
+| Hindi | CER | 9.209 | 0.8170 |
+| Indonesian | WER | 2.128 | 0.7270 |
+| Italian | WER | 1.958 | 0.7370 |
+| Japanese | CER | 5.142 | 0.7890 |
+| Korean | CER | 3.731 | 0.7850 |
+| Polish | WER | 1.957 | 0.7890 |
+| Portuguese | WER | 2.234 | 0.7860 |
+| Romanian † | WER | 17.425 | 0.7330 |
+| Russian | WER | 3.071 | 0.7490 |
+| Spanish | WER | 1.350 | 0.7450 |
+| Thai | CER | 6.674 | 0.7760 |
+| Turkish | WER | 1.560 | 0.7650 |
+| Ukrainian † | WER | 16.343 | 0.6940 |
+| Vietnamese | CER | 7.966 | 0.7120 |
+
+Each language has 100 evaluation items. Rows marked † are evaluated for completeness and are outside this release's supported-language list. The metric column preserves the evaluation suite's language-specific WER/CER choice. Cantonese is included in the 24-language average.
+
+</details>
+
+### Compared with Preview
+
+In the same archived MiniMax evaluation pipeline, the 24-language mean error rate decreased from **29.637%** for Preview to **5.894%** for Realistic, while mean SIM increased from **0.6932** to **0.7535**. This is an average over the full benchmark, including languages beyond the supported list; improvements are not uniform across every language and split. For example, ordinary Seed-TTS Chinese CER was **0.943% → 1.029%** in that same evaluation series. Preview values in its original public model card were measured in an earlier evaluation and should not be mixed with this comparison.
+
+The main release goal is improved speech realism alongside expanded language coverage. WER/CER and SIM alone do not establish a human preference score.
+
+## Limitations and Responsible Use
+
+- Quality varies by language and evaluation set. The detailed results include difficult cases, including Cantonese and some newly supported languages; the support list is not a claim of equal quality across languages.
+- Noisy reference audio or mismatched transcripts can reduce stability and speaker similarity. The reference audio, transcript, and target text share a finite context budget; split long passages into manageable utterances.
+- Speech may contain pronunciation errors, repetitions, omissions, or early stopping. Listen to generated audio before using it in a finished product.
+- Obtain permission to clone a voice and disclose synthetic audio where appropriate.
+
+## License and Acknowledgements
+
+The code and model weights are released under the [Apache License 2.0](./LICENSE). See [NOTICE](./NOTICE) for attribution details.
+
+We thank the Fish Audio team for publishing the DualAR architecture used in Fish Audio S2 Pro.

codec.pth

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

config.json

@@ -0,0 +1,64 @@
+{
+  "architectures": [
+    "ArkttsModel"
+  ],
+  "attention_o_bias": false,
+  "attention_qk_norm": false,
+  "attention_qkv_bias": true,
+  "auto_map": {
+    "AutoConfig": "configuration_arktts.ArkttsConfig",
+    "AutoModel": "modeling_arktts.ArkttsModel",
+    "AutoModelForCausalLM": "modeling_arktts.ArkttsModel"
+  },
+  "codebook_size": 4096,
+  "codec_filename": "codec.pth",
+  "codec_frame_size": 2048,
+  "codec_post_intermediate_size": 1216,
+  "codec_post_n_head": 16,
+  "codec_post_n_layer": 8,
+  "codec_post_n_local_heads": 8,
+  "codec_sample_rate": 44100,
+  "dim": 896,
+  "dropout": 0.0,
+  "dtype": "bfloat16",
+  "eos_token_id": 151645,
+  "fast_attention_o_bias": false,
+  "fast_attention_qk_norm": false,
+  "fast_attention_qkv_bias": false,
+  "fast_dim": 896,
+  "fast_head_dim": 64,
+  "fast_intermediate_size": 4864,
+  "fast_n_head": 14,
+  "fast_n_local_heads": 2,
+  "head_dim": 64,
+  "hidden_size": 896,
+  "initializer_range": 0.02,
+  "intermediate_size": 4864,
+  "max_position_embeddings": 2048,
+  "max_seq_len": 2048,
+  "model_type": "arktts",
+  "n_fast_layer": 4,
+  "n_head": 14,
+  "n_layer": 24,
+  "n_local_heads": 2,
+  "norm_eps": 1e-06,
+  "norm_fastlayer_input": true,
+  "num_attention_heads": 14,
+  "num_codebooks": 10,
+  "num_hidden_layers": 24,
+  "num_key_value_heads": 2,
+  "pad_token_id": 151643,
+  "ras_temperature": 1.0,
+  "ras_top_p": 0.9,
+  "ras_window_size": 10,
+  "rms_norm_eps": 1e-06,
+  "rope_base": 1000000.0,
+  "rope_theta": 1000000.0,
+  "scale_codebook_embeddings": false,
+  "semantic_begin_id": 151678,
+  "semantic_end_id": 155773,
+  "transformers_version": "4.57.3",
+  "use_gradient_checkpointing": false,
+  "vocab_size": 155776,
+  "tie_word_embeddings": true
+}

configuration_arktts.py

@@ -0,0 +1,114 @@
+from __future__ import annotations
+
+from transformers import PretrainedConfig
+
+
+class ArkttsConfig(PretrainedConfig):
+    model_type = "arktts"
+
+    def __init__(
+        self,
+        vocab_size: int = 155776,
+        dim: int = 896,
+        n_layer: int = 24,
+        n_head: int = 14,
+        n_local_heads: int = 2,
+        head_dim: int = 64,
+        intermediate_size: int = 4864,
+        max_seq_len: int = 2048,
+        rope_base: float = 1_000_000,
+        norm_eps: float = 1e-6,
+        dropout: float = 0.0,
+        attention_qkv_bias: bool = True,
+        attention_qk_norm: bool = False,
+        attention_o_bias: bool = False,
+        tie_word_embeddings: bool = True,
+        codebook_size: int = 4096,
+        num_codebooks: int = 10,
+        semantic_begin_id: int = 151678,
+        semantic_end_id: int = 155773,
+        n_fast_layer: int = 4,
+        fast_dim: int = 896,
+        fast_n_head: int = 14,
+        fast_n_local_heads: int = 2,
+        fast_head_dim: int = 64,
+        fast_intermediate_size: int = 4864,
+        fast_attention_qkv_bias: bool = False,
+        fast_attention_qk_norm: bool = False,
+        fast_attention_o_bias: bool = False,
+        norm_fastlayer_input: bool = True,
+        initializer_range: float = 0.02,
+        use_gradient_checkpointing: bool = False,
+        codec_filename: str = "codec.pth",
+        codec_sample_rate: int = 44100,
+        codec_frame_size: int = 2048,
+        codec_post_n_layer: int = 8,
+        codec_post_n_head: int = 16,
+        codec_post_n_local_heads: int = 8,
+        codec_post_intermediate_size: int = 1216,
+        ras_window_size: int = 10,
+        ras_temperature: float = 1.0,
+        ras_top_p: float = 0.9,
+        eos_token_id: int = 151645,
+        pad_token_id: int = 151643,
+        **kwargs,
+    ):
+        kwargs.pop("model_type", None)
+        kwargs.pop("audio_embed_dim", None)
+        kwargs.pop("is_reward_model", None)
+        super().__init__(
+            eos_token_id=eos_token_id,
+            pad_token_id=pad_token_id,
+            tie_word_embeddings=tie_word_embeddings,
+            **kwargs,
+        )
+        self.vocab_size = int(vocab_size)
+        self.dim = int(dim)
+        self.hidden_size = self.dim
+        self.n_layer = int(n_layer)
+        self.num_hidden_layers = self.n_layer
+        self.n_head = int(n_head)
+        self.num_attention_heads = self.n_head
+        self.n_local_heads = int(n_local_heads)
+        self.num_key_value_heads = self.n_local_heads
+        self.head_dim = int(head_dim)
+        self.intermediate_size = int(intermediate_size)
+        self.max_seq_len = int(max_seq_len)
+        self.max_position_embeddings = self.max_seq_len
+        self.rope_base = float(rope_base)
+        self.rope_theta = self.rope_base
+        self.norm_eps = float(norm_eps)
+        self.rms_norm_eps = self.norm_eps
+        self.dropout = float(dropout)
+        self.attention_qkv_bias = bool(attention_qkv_bias)
+        self.attention_qk_norm = bool(attention_qk_norm)
+        self.attention_o_bias = bool(attention_o_bias)
+        self.codebook_size = int(codebook_size)
+        self.num_codebooks = int(num_codebooks)
+        self.semantic_begin_id = int(semantic_begin_id)
+        self.semantic_end_id = int(semantic_end_id)
+        self.n_fast_layer = int(n_fast_layer)
+        self.fast_dim = int(fast_dim)
+        self.fast_n_head = int(fast_n_head)
+        self.fast_n_local_heads = int(fast_n_local_heads)
+        self.fast_head_dim = int(fast_head_dim)
+        self.fast_intermediate_size = int(fast_intermediate_size)
+        self.fast_attention_qkv_bias = bool(fast_attention_qkv_bias)
+        self.fast_attention_qk_norm = bool(fast_attention_qk_norm)
+        self.fast_attention_o_bias = bool(fast_attention_o_bias)
+        self.norm_fastlayer_input = bool(norm_fastlayer_input)
+        self.initializer_range = float(initializer_range)
+        self.use_gradient_checkpointing = bool(use_gradient_checkpointing)
+        self.codec_filename = str(codec_filename)
+        self.codec_sample_rate = int(codec_sample_rate)
+        self.codec_frame_size = int(codec_frame_size)
+        self.codec_post_n_layer = int(codec_post_n_layer)
+        self.codec_post_n_head = int(codec_post_n_head)
+        self.codec_post_n_local_heads = int(codec_post_n_local_heads)
+        self.codec_post_intermediate_size = int(codec_post_intermediate_size)
+        self.ras_window_size = int(ras_window_size)
+        self.ras_temperature = float(ras_temperature)
+        self.ras_top_p = float(ras_top_p)
+
+
+__all__ = ["ArkttsConfig"]

generation_config.json

@@ -0,0 +1,10 @@
+{
+  "do_sample": true,
+  "eos_token_id": 151645,
+  "max_new_tokens": 512,
+  "pad_token_id": 151643,
+  "temperature": 0.7,
+  "top_p": 0.9,
+  "transformers_version": "4.57.3",
+  "top_k": 50
+}

model.safetensors

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

modeling_arktts.py

@@ -0,0 +1,831 @@
+from __future__ import annotations
+
+import math
+import shutil
+from dataclasses import dataclass
+from pathlib import Path
+from typing import Optional
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from torch import Tensor
+from torch.nn.attention import SDPBackend, sdpa_kernel
+from transformers import GenerationConfig, PreTrainedModel
+from transformers.generation import (
+    LogitsProcessor,
+    LogitsProcessorList,
+    StoppingCriteriaList,
+)
+from transformers.modeling_outputs import ModelOutput
+from transformers.utils.hub import cached_file
+
+from .configuration_arktts import ArkttsConfig
+
+
+@dataclass
+class ArkttsModelOutput(ModelOutput):
+    logits: Optional[Tensor] = None
+    hidden_states: Optional[Tensor] = None
+    codebook_logits: Optional[Tensor] = None
+
+
+@dataclass
+class ArkttsGenerateOutput(ModelOutput):
+    codes: Optional[Tensor] = None
+    code_lengths: Optional[Tensor] = None
+    finished: Optional[Tensor] = None
+
+
+class ArkttsSemanticLogitsProcessor(LogitsProcessor):
+    def __init__(self, semantic_begin_id: int, semantic_end_id: int, eos_token_id: int):
+        self.semantic_begin_id = int(semantic_begin_id)
+        self.semantic_end_id = int(semantic_end_id)
+        self.eos_token_id = int(eos_token_id)
+
+    def __call__(self, input_ids: Tensor, scores: Tensor) -> Tensor:
+        filtered = torch.full_like(scores, float("-inf"))
+        filtered[:, self.semantic_begin_id : self.semantic_end_id + 1] = scores[
+            :, self.semantic_begin_id : self.semantic_end_id + 1
+        ]
+        filtered[:, self.eos_token_id] = scores[:, self.eos_token_id]
+        return filtered
+
+
+class ArkttsLegacyTopKTopPLogitsProcessor(LogitsProcessor):
+    """Matches the candidate filtering order used by the original inference code."""
+
+    def __init__(self, top_k: int, top_p: float):
+        if top_k <= 0:
+            raise ValueError("top_k must be positive")
+        if not 0.0 < top_p <= 1.0:
+            raise ValueError("top_p must be in (0, 1]")
+        self.top_k = int(top_k)
+        self.top_p = float(top_p)
+
+    def __call__(self, input_ids: Tensor, scores: Tensor) -> Tensor:
+        sorted_scores, sorted_indices = torch.sort(scores, descending=True, dim=-1)
+        cumulative = torch.cumsum(torch.softmax(sorted_scores, dim=-1), dim=-1)
+        positions = torch.arange(sorted_scores.shape[-1], device=scores.device)
+        threshold = torch.tensor(self.top_p, dtype=cumulative.dtype, device=cumulative.device)
+        remove_sorted = (cumulative > threshold) | (positions >= self.top_k)
+        remove_sorted[..., 0] = False
+        remove = torch.zeros_like(remove_sorted).scatter(1, sorted_indices, remove_sorted)
+        return scores.masked_fill(remove, float("-inf"))
+
+
+class ArkttsKVCache(nn.Module):
+    def __init__(
+        self,
+        batch_size: int,
+        max_length: int,
+        heads: int,
+        head_dim: int,
+        dtype,
+        return_full: bool = False,
+    ):
+        super().__init__()
+        shape = (batch_size, heads, max_length, head_dim)
+        self.register_buffer("keys", torch.zeros(shape, dtype=dtype), persistent=False)
+        self.register_buffer("values", torch.zeros(shape, dtype=dtype), persistent=False)
+        self.return_full = bool(return_full)
+
+    def update(self, cache_position: Tensor, keys: Tensor, values: Tensor):
+        self.keys[:, :, cache_position] = keys
+        self.values[:, :, cache_position] = values
+        end = self.keys.shape[2] if self.return_full else int(cache_position[-1].item()) + 1
+        return self.keys[:, :, :end], self.values[:, :, :end]
+
+
+class ArkttsRMSNorm(nn.Module):
+    def __init__(self, dim: int, eps: float):
+        super().__init__()
+        self.eps = float(eps)
+        self.weight = nn.Parameter(torch.ones(dim))
+
+    def forward(self, x: Tensor) -> Tensor:
+        normalized = x.float() * torch.rsqrt(x.float().pow(2).mean(-1, keepdim=True) + self.eps)
+        return normalized.to(x.dtype) * self.weight
+
+
+def _precompute_rope(length: int, head_dim: int, base: float) -> Tensor:
+    frequencies = 1.0 / (
+        base ** (torch.arange(0, head_dim, 2).float()[: head_dim // 2] / head_dim)
+    )
+    phases = torch.outer(torch.arange(length), frequencies)
+    values = torch.polar(torch.ones_like(phases), phases)
+    return torch.stack((values.real, values.imag), dim=-1).to(torch.bfloat16)
+
+
+def _apply_rope(x: Tensor, rope: Tensor) -> Tensor:
+    shaped = x.float().reshape(*x.shape[:-1], -1, 2)
+    if rope.ndim == 3:
+        rope = rope[None, :, None]
+    elif rope.ndim == 4:
+        rope = rope[:, :, None]
+    else:
+        raise ValueError(f"Unexpected RoPE shape: {tuple(rope.shape)}")
+    output = torch.stack(
+        (
+            shaped[..., 0] * rope[..., 0] - shaped[..., 1] * rope[..., 1],
+            shaped[..., 1] * rope[..., 0] + shaped[..., 0] * rope[..., 1],
+        ),
+        dim=-1,
+    )
+    return output.flatten(3).to(x.dtype)
+
+
+class ArkttsAttention(nn.Module):
+    def __init__(
+        self,
+        dim: int,
+        n_head: int,
+        n_local_heads: int,
+        head_dim: int,
+        qkv_bias: bool,
+        output_bias: bool,
+        qk_norm: bool,
+        norm_eps: float,
+        dropout: float,
+        use_sdpa: bool,
+    ):
+        super().__init__()
+        total = (n_head + 2 * n_local_heads) * head_dim
+        self.wqkv = nn.Linear(dim, total, bias=qkv_bias)
+        self.wo = nn.Linear(n_head * head_dim, dim, bias=output_bias)
+        self.n_head = int(n_head)
+        self.n_local_heads = int(n_local_heads)
+        self.head_dim = int(head_dim)
+        self.dropout = float(dropout)
+        self.use_sdpa = bool(use_sdpa)
+        self.qk_norm = bool(qk_norm)
+        if self.qk_norm:
+            self.q_norm = ArkttsRMSNorm(head_dim, norm_eps)
+            self.k_norm = ArkttsRMSNorm(head_dim, norm_eps)
+        self.kv_cache: Optional[ArkttsKVCache] = None
+
+    def forward(
+        self,
+        x: Tensor,
+        rope: Tensor,
+        attention_mask: Optional[Tensor],
+        cache_position: Optional[Tensor] = None,
+    ) -> Tensor:
+        batch, length, _ = x.shape
+        query_size = self.n_head * self.head_dim
+        kv_size = self.n_local_heads * self.head_dim
+        query, key, value = self.wqkv(x).split((query_size, kv_size, kv_size), dim=-1)
+        query = query.view(batch, length, self.n_head, self.head_dim)
+        key = key.view(batch, length, self.n_local_heads, self.head_dim)
+        value = value.view(batch, length, self.n_local_heads, self.head_dim)
+        if self.qk_norm:
+            query = self.q_norm(query)
+            key = self.k_norm(key)
+        query = _apply_rope(query, rope).transpose(1, 2)
+        key = _apply_rope(key, rope).transpose(1, 2)
+        value = value.transpose(1, 2)
+        if self.kv_cache is not None:
+            if cache_position is None:
+                raise ValueError("cache_position is required when KV cache is enabled")
+            key, value = self.kv_cache.update(cache_position, key, value)
+        repeats = self.n_head // self.n_local_heads
+        key = key.repeat_interleave(repeats, dim=1)
+        value = value.repeat_interleave(repeats, dim=1)
+        if self.use_sdpa:
+            output = F.scaled_dot_product_attention(
+                query,
+                key,
+                value,
+                attn_mask=attention_mask,
+                dropout_p=self.dropout if self.training else 0.0,
+            )
+        else:
+            scores = query @ key.transpose(-2, -1) / math.sqrt(self.head_dim)
+            if attention_mask is not None:
+                scores = scores.masked_fill(~attention_mask, float("-inf"))
+            probabilities = torch.softmax(scores, dim=-1)
+            if self.training and self.dropout:
+                probabilities = F.dropout(probabilities, p=self.dropout)
+            output = probabilities @ value
+        output = output.transpose(1, 2).contiguous().view(batch, length, query_size)
+        return self.wo(output)
+
+
+class ArkttsFeedForward(nn.Module):
+    def __init__(self, dim: int, intermediate_size: int):
+        super().__init__()
+        self.w1 = nn.Linear(dim, intermediate_size, bias=False)
+        self.w2 = nn.Linear(intermediate_size, dim, bias=False)
+        self.w3 = nn.Linear(dim, intermediate_size, bias=False)
+
+    def forward(self, x: Tensor) -> Tensor:
+        return self.w2(F.silu(self.w1(x)) * self.w3(x))
+
+
+class ArkttsTransformerBlock(nn.Module):
+    def __init__(
+        self,
+        dim: int,
+        intermediate_size: int,
+        n_head: int,
+        n_local_heads: int,
+        head_dim: int,
+        qkv_bias: bool,
+        output_bias: bool,
+        qk_norm: bool,
+        norm_eps: float,
+        dropout: float,
+        use_sdpa: bool,
+    ):
+        super().__init__()
+        self.attention = ArkttsAttention(
+            dim, n_head, n_local_heads, head_dim, qkv_bias, output_bias,
+            qk_norm, norm_eps, dropout, use_sdpa,
+        )
+        self.feed_forward = ArkttsFeedForward(dim, intermediate_size)
+        self.ffn_norm = ArkttsRMSNorm(dim, norm_eps)
+        self.attention_norm = ArkttsRMSNorm(dim, norm_eps)
+
+    def forward(self, x, rope, attention_mask, cache_position=None):
+        hidden = x + self.attention(self.attention_norm(x), rope, attention_mask, cache_position)
+        return hidden + self.feed_forward(self.ffn_norm(hidden))
+
+
+class ArkttsModel(PreTrainedModel):
+    config_class = ArkttsConfig
+    base_model_prefix = ""
+    main_input_name = "input_ids"
+    _no_split_modules = ["ArkttsTransformerBlock"]
+    _supports_sdpa = True
+
+    def __init__(self, config: ArkttsConfig):
+        super().__init__(config)
+        self.embeddings = nn.Embedding(config.vocab_size, config.dim)
+        self.codebook_embeddings = nn.Embedding(config.codebook_size * config.num_codebooks, config.dim)
+        self.layers = nn.ModuleList([
+            ArkttsTransformerBlock(
+                config.dim, config.intermediate_size, config.n_head, config.n_local_heads,
+                config.head_dim, config.attention_qkv_bias, config.attention_o_bias,
+                config.attention_qk_norm, config.norm_eps, config.dropout, True,
+            )
+            for _ in range(config.n_layer)
+        ])
+        self.norm = ArkttsRMSNorm(config.dim, config.norm_eps)
+        self.fast_project_in = (
+            nn.Linear(config.dim, config.fast_dim)
+            if config.fast_dim != config.dim else nn.Identity()
+        )
+        self.fast_embeddings = nn.Embedding(config.codebook_size, config.fast_dim)
+        self.fast_layers = nn.ModuleList([
+            ArkttsTransformerBlock(
+                config.fast_dim, config.fast_intermediate_size, config.fast_n_head,
+                config.fast_n_local_heads, config.fast_head_dim,
+                config.fast_attention_qkv_bias, config.fast_attention_o_bias,
+                config.fast_attention_qk_norm, config.norm_eps, config.dropout, False,
+            )
+            for _ in range(config.n_fast_layer)
+        ])
+        self.fast_norm = ArkttsRMSNorm(config.fast_dim, config.norm_eps)
+        self.fast_output = nn.Linear(config.fast_dim, config.codebook_size, bias=False)
+        self.register_buffer(
+            "freqs_cis", _precompute_rope(config.max_seq_len, config.head_dim, config.rope_base),
+            persistent=False,
+        )
+        self.register_buffer(
+            "fast_freqs_cis",
+            _precompute_rope(config.num_codebooks, config.fast_head_dim, config.rope_base),
+            persistent=False,
+        )
+        self.__dict__["_arktts_codec"] = None
+        self.post_init()
+
+    def _init_weights(self, module):
+        if isinstance(module, (nn.Linear, nn.Embedding)):
+            module.weight.data.normal_(mean=0.0, std=self.config.initializer_range)
+            if isinstance(module, nn.Linear) and module.bias is not None:
+                module.bias.data.zero_()
+
+    def get_input_embeddings(self):
+        return self.embeddings
+
+    def set_input_embeddings(self, value):
+        self.embeddings = value
+
+    def get_output_embeddings(self):
+        return None
+
+    def _embed(self, input_ids: Tensor) -> Tensor:
+        codebook_embeds = []
+        for index in range(self.config.num_codebooks):
+            codebook_embeds.append(
+                self.codebook_embeddings(input_ids[:, index + 1] + index * self.config.codebook_size)
+            )
+        codebook_sum = torch.stack(codebook_embeds, dim=1).sum(dim=1)
+        semantic = (input_ids[:, 0] >= self.config.semantic_begin_id) & (
+            input_ids[:, 0] <= self.config.semantic_end_id
+        )
+        codebook_sum = torch.where(semantic.unsqueeze(-1), codebook_sum, 0.0)
+        return self.embeddings(input_ids[:, 0]) + codebook_sum
+
+    @staticmethod
+    def _causal_mask(attention_mask: Tensor, query_positions: Tensor, key_length: int) -> Tensor:
+        if attention_mask.shape[1] < key_length:
+            attention_mask = F.pad(
+                attention_mask,
+                (0, key_length - attention_mask.shape[1]),
+                value=0,
+            )
+        key_positions = torch.arange(key_length, device=attention_mask.device)
+        causal = key_positions[None, :] <= query_positions[:, None]
+        return causal[None, None] & attention_mask[:, None, None, :key_length].bool()
+
+    def forward(
+        self,
+        input_ids: Tensor,
+        attention_mask: Optional[Tensor] = None,
+        labels: Optional[Tensor] = None,
+        output_hidden_states: Optional[bool] = None,
+        return_dict: Optional[bool] = None,
+        **kwargs,
+    ):
+        del labels, output_hidden_states, kwargs
+        if input_ids.ndim != 3 or input_ids.shape[1] != self.config.num_codebooks + 1:
+            raise ValueError(
+                f"input_ids must have shape [B, {self.config.num_codebooks + 1}, T]"
+            )
+        batch, _, length = input_ids.shape
+        if attention_mask is None:
+            attention_mask = torch.ones((batch, length), dtype=torch.long, device=input_ids.device)
+        position_ids = attention_mask.long().cumsum(-1).sub(1).clamp_min(0)
+        rope = self.freqs_cis[position_ids]
+        mask = self._causal_mask(
+            attention_mask,
+            torch.arange(length, device=input_ids.device),
+            length,
+        )
+        hidden = self._embed(input_ids)
+        for layer in self.layers:
+            hidden = layer(hidden, rope, mask)
+        normalized = self.norm(hidden)
+        logits = F.linear(normalized, self.embeddings.weight)
+        output = ArkttsModelOutput(logits=logits, hidden_states=hidden)
+        if return_dict is False:
+            return (logits, hidden)
+        return output
+
+    def _setup_generation_caches(self, batch_size: int, max_length: int, dtype):
+        del max_length
+        for layer in self.layers:
+            layer.attention.kv_cache = ArkttsKVCache(
+                batch_size,
+                self.config.max_seq_len,
+                self.config.n_local_heads,
+                self.config.head_dim,
+                dtype,
+                return_full=True,
+            ).to(self.device)
+        for layer in self.fast_layers:
+            layer.attention.kv_cache = ArkttsKVCache(
+                batch_size, self.config.num_codebooks, self.config.fast_n_local_heads,
+                self.config.fast_head_dim, dtype, return_full=True,
+            ).to(self.device)
+
+    def _slow_step(
+        self,
+        input_ids: Tensor,
+        cache_position: Tensor,
+        position_ids: Tensor,
+        attention_mask: Tensor,
+    ) -> tuple[Tensor, Tensor]:
+        hidden = self._embed(input_ids)
+        rope = self.freqs_cis[position_ids]
+        mask = self._causal_mask(
+            attention_mask,
+            cache_position,
+            self.config.max_seq_len,
+        )
+        for layer in self.layers:
+            hidden = layer(hidden, rope, mask, cache_position)
+        hidden = hidden[:, -1:]
+        normalized = self.norm(hidden)
+        logits = F.linear(normalized, self.embeddings.weight)[:, -1]
+        fast_hidden = normalized if self.config.norm_fastlayer_input else hidden
+        return logits, fast_hidden
+
+    def _fast_step(self, hidden: Tensor, position: int) -> Tensor:
+        cache_position = torch.tensor([position], device=hidden.device, dtype=torch.long)
+        rope = self.fast_freqs_cis[cache_position]
+        key_mask = torch.ones(
+            (hidden.shape[0], self.config.num_codebooks),
+            device=hidden.device,
+            dtype=torch.bool,
+        )
+        mask = self._causal_mask(key_mask, cache_position, self.config.num_codebooks)
+        for layer in self.fast_layers:
+            hidden = layer(hidden, rope, mask, cache_position)
+        return self.fast_output(self.fast_norm(hidden))[:, -1]
+
+    @staticmethod
+    def _as_processor_list(value) -> LogitsProcessorList:
+        if value is None:
+            return LogitsProcessorList()
+        if isinstance(value, LogitsProcessorList):
+            return value
+        return LogitsProcessorList(value)
+
+    @staticmethod
+    def _sample(scores: Tensor, generator=None) -> Tensor:
+        probabilities = torch.softmax(scores, dim=-1)
+        random = torch.rand(
+            probabilities.shape,
+            dtype=probabilities.dtype,
+            device=probabilities.device,
+            generator=generator,
+        )
+        noise = -torch.log(random)
+        return torch.argmax(probabilities / noise, dim=-1)
+
+    def _processed_scores(
+        self,
+        input_ids: Tensor,
+        scores: Tensor,
+        processors: LogitsProcessorList,
+        top_k: int,
+        top_p: float,
+        temperature: float,
+    ) -> Tensor:
+        scores = processors(input_ids, scores)
+        scores = ArkttsLegacyTopKTopPLogitsProcessor(top_k, top_p)(input_ids, scores)
+        temperature_value = torch.tensor(
+            temperature, dtype=scores.dtype, device=scores.device
+        ).clamp_min(1e-5)
+        return scores / temperature_value
+
+    def _sample_semantic(
+        self,
+        history: Tensor,
+        logits: Tensor,
+        custom_processors: LogitsProcessorList,
+        top_k: int,
+        top_p: float,
+        temperature: float,
+        previous: Optional[Tensor],
+        do_sample: bool,
+        generator=None,
+    ) -> Tensor:
+        processors = LogitsProcessorList([
+            ArkttsSemanticLogitsProcessor(
+                self.config.semantic_begin_id,
+                self.config.semantic_end_id,
+                self.config.eos_token_id,
+            ),
+            *custom_processors,
+        ])
+        regular_scores = self._processed_scores(
+            history, logits, processors, top_k, top_p, temperature
+        )
+        if not do_sample:
+            return regular_scores.argmax(dim=-1)
+        normal = self._sample(regular_scores, generator=generator)
+        high_scores = self._processed_scores(
+            history,
+            logits,
+            processors,
+            top_k,
+            self.config.ras_top_p,
+            self.config.ras_temperature,
+        )
+        high = self._sample(high_scores, generator=generator)
+        if previous is None:
+            return normal
+        repeated = (previous == normal[:, None]).any(dim=1)
+        semantic = (normal >= self.config.semantic_begin_id) & (
+            normal <= self.config.semantic_end_id
+        )
+        return torch.where(repeated & semantic, high, normal)
+
+    def _generate_codebooks(
+        self,
+        slow_hidden: Tensor,
+        semantic: Tensor,
+        processors: LogitsProcessorList,
+        top_k: int,
+        top_p: float,
+        temperature: float,
+        do_sample: bool,
+        generator=None,
+    ) -> Tensor:
+        hidden = self.fast_project_in(slow_hidden)
+        self._fast_step(hidden, 0)
+        current = (semantic - self.config.semantic_begin_id).clamp(0, self.config.codebook_size - 1)
+        codebooks = [current]
+        fast_history = current[:, None]
+        hidden = self.fast_embeddings(current)[:, None]
+        for position in range(1, self.config.num_codebooks):
+            scores = self._fast_step(hidden, position)
+            scores = self._processed_scores(
+                fast_history, scores, processors, top_k, top_p, temperature
+            )
+            current = self._sample(scores, generator=generator) if do_sample else scores.argmax(dim=-1)
+            codebooks.append(current)
+            fast_history = torch.cat((fast_history, current[:, None]), dim=1)
+            hidden = self.fast_embeddings(current)[:, None]
+        return torch.stack(codebooks, dim=1)
+
+    def _prepare_prompt(
+        self,
+        input_ids=None,
+        attention_mask=None,
+        prefix_input_ids=None,
+        prefix_attention_mask=None,
+        suffix_input_ids=None,
+        suffix_attention_mask=None,
+        reference_codes=None,
+        reference_code_lengths=None,
+        reference_audio_values=None,
+        reference_audio_lengths=None,
+    ):
+        if input_ids is not None:
+            if input_ids.ndim != 3 or input_ids.shape[1] != self.config.num_codebooks + 1:
+                raise ValueError("Direct input_ids must have shape [B, num_codebooks + 1, T]")
+            if attention_mask is None:
+                attention_mask = torch.ones(
+                    input_ids.shape[0], input_ids.shape[-1], dtype=torch.long, device=input_ids.device
+                )
+            return input_ids.to(self.device), attention_mask.to(self.device)
+        if prefix_input_ids is None or suffix_input_ids is None:
+            raise ValueError("Processor output or direct input_ids is required")
+        prefix_input_ids = prefix_input_ids.to(self.device)
+        suffix_input_ids = suffix_input_ids.to(self.device)
+        prefix_attention_mask = prefix_attention_mask.to(self.device)
+        suffix_attention_mask = suffix_attention_mask.to(self.device)
+        if reference_audio_values is not None:
+            if reference_codes is not None:
+                raise ValueError("Provide reference audio or reference codes, not both")
+            reference_codes, reference_code_lengths = self.encode_audio(
+                reference_audio_values.to(self.device), reference_audio_lengths.to(self.device)
+            )
+        if reference_codes is not None:
+            reference_codes = reference_codes.to(self.device)
+            reference_code_lengths = reference_code_lengths.to(self.device)
+        batch_size = prefix_input_ids.shape[0]
+        rows = []
+        for batch_index in range(batch_size):
+            prefix = prefix_input_ids[batch_index, prefix_attention_mask[batch_index].bool()]
+            suffix = suffix_input_ids[batch_index, suffix_attention_mask[batch_index].bool()]
+            if reference_codes is None:
+                semantic_row = torch.cat((prefix, suffix))
+                values = torch.zeros(
+                    (self.config.num_codebooks + 1, semantic_row.numel()),
+                    dtype=torch.long,
+                    device=self.device,
+                )
+                values[0] = semantic_row
+            else:
+                length = int(reference_code_lengths[batch_index])
+                codes = reference_codes[batch_index, :, :length].long()
+                semantic_codes = codes[0] + self.config.semantic_begin_id
+                semantic_row = torch.cat((prefix, semantic_codes, suffix))
+                values = torch.zeros(
+                    (self.config.num_codebooks + 1, semantic_row.numel()),
+                    dtype=torch.long,
+                    device=self.device,
+                )
+                values[0] = semantic_row
+                values[1:, prefix.numel() : prefix.numel() + length] = codes
+            rows.append(values)
+        max_length = max(row.shape[1] for row in rows)
+        prompt = torch.zeros(
+            (batch_size, self.config.num_codebooks + 1, max_length),
+            dtype=torch.long,
+            device=self.device,
+        )
+        prompt[:, 0] = self.config.pad_token_id
+        prompt_mask = torch.zeros((batch_size, max_length), dtype=torch.long, device=self.device)
+        for batch_index, row in enumerate(rows):
+            start = max_length - row.shape[1]
+            prompt[batch_index, :, start:] = row
+            prompt_mask[batch_index, start:] = 1
+        return prompt, prompt_mask
+
+    @torch.inference_mode()
+    def generate(
+        self,
+        input_ids: Optional[Tensor] = None,
+        attention_mask: Optional[Tensor] = None,
+        prefix_input_ids: Optional[Tensor] = None,
+        prefix_attention_mask: Optional[Tensor] = None,
+        suffix_input_ids: Optional[Tensor] = None,
+        suffix_attention_mask: Optional[Tensor] = None,
+        reference_codes: Optional[Tensor] = None,
+        reference_code_lengths: Optional[Tensor] = None,
+        reference_audio_values: Optional[Tensor] = None,
+        reference_audio_lengths: Optional[Tensor] = None,
+        generation_config: Optional[GenerationConfig] = None,
+        logits_processor=None,
+        codebook_logits_processor=None,
+        stopping_criteria=None,
+        max_new_tokens: Optional[int] = None,
+        temperature: Optional[float] = None,
+        top_p: Optional[float] = None,
+        top_k: Optional[int] = None,
+        do_sample: Optional[bool] = None,
+        return_dict_in_generate: bool = False,
+        generator=None,
+        **kwargs,
+    ):
+        if kwargs:
+            raise TypeError(f"Unexpected generation arguments: {sorted(kwargs)}")
+        config = generation_config or getattr(self, "generation_config", GenerationConfig())
+        config_max_new = getattr(config, "max_new_tokens", None)
+        config_temperature = getattr(config, "temperature", None)
+        config_top_p = getattr(config, "top_p", None)
+        config_top_k = getattr(config, "top_k", None)
+        max_new_tokens = int(max_new_tokens if max_new_tokens is not None else (config_max_new or 512))
+        temperature = float(temperature if temperature is not None else (config_temperature or 0.7))
+        top_p = float(top_p if top_p is not None else (config_top_p or 0.9))
+        top_k = int(top_k if top_k is not None else (config_top_k or 50))
+        do_sample = bool(do_sample if do_sample is not None else getattr(config, "do_sample", True))
+        if max_new_tokens <= 0:
+            raise ValueError("max_new_tokens must be positive")
+        if temperature <= 0:
+            raise ValueError("temperature must be positive")
+        prompt, prompt_mask = self._prepare_prompt(
+            input_ids, attention_mask, prefix_input_ids, prefix_attention_mask,
+            suffix_input_ids, suffix_attention_mask, reference_codes,
+            reference_code_lengths, reference_audio_values, reference_audio_lengths,
+        )
+        batch_size, _, prompt_width = prompt.shape
+        if prompt_width >= self.config.max_seq_len:
+            raise ValueError(
+                f"Prompt length {prompt_width} must be smaller than {self.config.max_seq_len}"
+            )
+        max_new_tokens = min(max_new_tokens, self.config.max_seq_len - prompt_width)
+        self._setup_generation_caches(
+            batch_size, prompt_width + max_new_tokens, next(self.parameters()).dtype
+        )
+        semantic_processors = self._as_processor_list(logits_processor)
+        codebook_processors = self._as_processor_list(codebook_logits_processor)
+        criteria = stopping_criteria or StoppingCriteriaList()
+        if not isinstance(criteria, StoppingCriteriaList):
+            criteria = StoppingCriteriaList(criteria)
+        cache_position = torch.arange(prompt_width, device=self.device, dtype=torch.long)
+        position_ids = prompt_mask.cumsum(-1).sub(1).clamp_min(0)
+        logits, slow_hidden = self._slow_step(
+            prompt, cache_position, position_ids, prompt_mask
+        )
+        semantic_history = prompt[:, 0]
+        prompt_lengths = prompt_mask.sum(-1)
+        previous = None
+        finished = torch.zeros(batch_size, dtype=torch.bool, device=self.device)
+        code_lengths = torch.zeros(batch_size, dtype=torch.long, device=self.device)
+        generated_frames = []
+
+        for step in range(max_new_tokens):
+            active_before = ~finished
+            semantic = self._sample_semantic(
+                semantic_history, logits, semantic_processors, top_k, top_p,
+                temperature, previous, do_sample, generator,
+            )
+            codebooks = self._generate_codebooks(
+                slow_hidden, semantic, codebook_processors, top_k, top_p,
+                temperature, do_sample, generator,
+            )
+            emitted = active_before & (semantic != self.config.eos_token_id)
+            frame = torch.where(emitted[:, None], codebooks, -1)
+            generated_frames.append(frame)
+            code_lengths += emitted.long()
+            semantic_history = torch.cat((semantic_history, semantic[:, None]), dim=1)
+
+            if previous is None:
+                previous = torch.zeros(
+                    (batch_size, self.config.ras_window_size),
+                    dtype=torch.long,
+                    device=self.device,
+                )
+            else:
+                previous = previous.roll(-1, dims=1)
+                previous[:, -1] = semantic
+            finished |= semantic.eq(self.config.eos_token_id)
+            if criteria:
+                stopped = criteria(semantic_history, logits)
+                if not isinstance(stopped, Tensor):
+                    stopped = torch.full_like(finished, bool(stopped))
+                finished |= stopped.to(device=self.device, dtype=torch.bool)
+            if finished.all():
+                break
+
+            next_column = torch.cat((semantic[:, None], codebooks), dim=1).unsqueeze(-1)
+            new_valid = active_before.long()[:, None]
+            prompt_mask = torch.cat((prompt_mask, new_valid), dim=1)
+            physical_position = torch.tensor([prompt_width + step], device=self.device)
+            token_position = (prompt_lengths + step)[:, None]
+            with sdpa_kernel(SDPBackend.MATH):
+                logits, slow_hidden = self._slow_step(
+                    next_column, physical_position, token_position, prompt_mask
+                )
+
+        if generated_frames:
+            codes = torch.stack(generated_frames, dim=2)
+            max_valid = int(code_lengths.max().item()) if code_lengths.numel() else 0
+            codes = codes[:, :, :max_valid]
+        else:
+            codes = torch.empty(
+                (batch_size, self.config.num_codebooks, 0), dtype=torch.long, device=self.device
+            )
+        result = ArkttsGenerateOutput(codes=codes, code_lengths=code_lengths, finished=finished)
+        return result if return_dict_in_generate else codes
+
+    def _codec_path(self) -> str:
+        source = str(getattr(self.config, "_name_or_path", ""))
+        local = Path(source)
+        if local.is_dir() and (local / self.config.codec_filename).is_file():
+            return str(local / self.config.codec_filename)
+        resolved = cached_file(source, self.config.codec_filename)
+        if resolved is None:
+            raise FileNotFoundError(f"Could not resolve {self.config.codec_filename} from {source}")
+        return resolved
+
+    def save_pretrained(self, save_directory, *args, **kwargs):
+        result = super().save_pretrained(save_directory, *args, **kwargs)
+        source = Path(self._codec_path()).resolve()
+        destination = Path(save_directory) / self.config.codec_filename
+        destination.parent.mkdir(parents=True, exist_ok=True)
+        if source != destination.resolve():
+            shutil.copy2(source, destination)
+        return result
+
+    def load_codec(self, device=None, dtype=None):
+        codec = self.__dict__.get("_arktts_codec")
+        target_device = torch.device(device) if device is not None else self.device
+        target_dtype = dtype or self.dtype
+        if target_device.type == "cpu":
+            target_dtype = torch.float32
+        if codec is None:
+            from .modeling_arktts_codec import ArkttsCodec
+
+            codec = ArkttsCodec(self.config)
+            state = torch.load(self._codec_path(), map_location="cpu", weights_only=True)
+            if "state_dict" in state:
+                state = state["state_dict"]
+            if any("generator." in key for key in state):
+                state = {
+                    key.replace("generator.", ""): value
+                    for key, value in state.items() if "generator." in key
+                }
+            state = {
+                key: value
+                for key, value in state.items()
+                if not key.endswith(("freqs_cis", "causal_mask"))
+            }
+            codec.load_state_dict(state, strict=True)
+            codec.eval()
+            self.__dict__["_arktts_codec"] = codec
+        codec.to(device=target_device, dtype=target_dtype)
+        return codec
+
+    @torch.inference_mode()
+    def encode_audio(self, audio_values: Tensor, audio_lengths: Optional[Tensor] = None):
+        codec = self.load_codec(device=audio_values.device)
+        audio_values = audio_values.to(dtype=next(codec.parameters()).dtype)
+        return codec.encode(audio_values, audio_lengths)
+
+    @torch.inference_mode()
+    def decode_audio(self, codes: Tensor):
+        if codes.ndim == 2:
+            codes = codes.unsqueeze(0)
+        if codes.ndim != 3 or codes.shape[1] != self.config.num_codebooks:
+            raise ValueError(f"codes must have shape [B, {self.config.num_codebooks}, T]")
+        codec = self.load_codec(device=codes.device)
+        waveforms = []
+        lengths = []
+        for item in codes:
+            valid = (item >= 0).all(dim=0)
+            length = int(valid.sum().item())
+            if length == 0:
+                waveform = torch.empty(0, device=codes.device, dtype=torch.float32)
+            else:
+                waveform = codec.decode(item[:, :length].unsqueeze(0))[0, 0].float()
+            waveforms.append(waveform)
+            lengths.append(waveform.numel())
+        max_length = max(lengths, default=0)
+        padded = torch.zeros((len(waveforms), max_length), dtype=torch.float32, device=codes.device)
+        for index, waveform in enumerate(waveforms):
+            padded[index, : waveform.numel()] = waveform
+        return padded, torch.tensor(lengths, dtype=torch.long, device=codes.device)
+
+    @torch.inference_mode()
+    def generate_audio(self, **kwargs):
+        codes = self.generate(**kwargs)
+        waveforms, lengths = self.decode_audio(codes)
+        return waveforms, lengths, codes
+
+
+__all__ = [
+    "ArkttsConfig",
+    "ArkttsGenerateOutput",
+    "ArkttsLegacyTopKTopPLogitsProcessor",
+    "ArkttsModel",
+    "ArkttsModelOutput",
+    "ArkttsSemanticLogitsProcessor",
+]

modeling_arktts_codec.py

@@ -0,0 +1,520 @@
+from __future__ import annotations
+
+import math
+from dataclasses import dataclass
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from torch import Tensor
+from torch.nn.utils import weight_norm as legacy_weight_norm
+from torch.nn.utils.parametrizations import weight_norm
+
+
+def _rope(length: int, head_dim: int, base: float, device=None) -> Tensor:
+    frequencies = 1.0 / (
+        base ** (torch.arange(0, head_dim, 2, device=device).float() / head_dim)
+    )
+    phases = torch.outer(torch.arange(length, device=device), frequencies)
+    values = torch.polar(torch.ones_like(phases), phases)
+    return torch.stack((values.real, values.imag), dim=-1).to(torch.bfloat16)
+
+
+def _apply_rope(x: Tensor, values: Tensor) -> Tensor:
+    shaped = x.float().reshape(*x.shape[:-1], -1, 2)
+    values = values.view(1, shaped.shape[1], 1, shaped.shape[3], 2)
+    output = torch.stack(
+        (
+            shaped[..., 0] * values[..., 0] - shaped[..., 1] * values[..., 1],
+            shaped[..., 1] * values[..., 0] + shaped[..., 0] * values[..., 1],
+        ),
+        dim=-1,
+    )
+    return output.flatten(3).to(x.dtype)
+
+
+class ArkttsCodecRMSNorm(nn.Module):
+    def __init__(self, dim: int, eps: float = 1e-5):
+        super().__init__()
+        self.eps = float(eps)
+        self.weight = nn.Parameter(torch.ones(dim))
+
+    def forward(self, x):
+        output = x.float() * torch.rsqrt(x.float().pow(2).mean(-1, keepdim=True) + self.eps)
+        return output.to(x.dtype) * self.weight
+
+
+class ArkttsCodecLayerScale(nn.Module):
+    def __init__(self, dim: int, init_values: float = 1e-2, inplace: bool = False):
+        super().__init__()
+        self.inplace = bool(inplace)
+        self.gamma = nn.Parameter(init_values * torch.ones(dim))
+
+    def forward(self, x):
+        return x.mul_(self.gamma) if self.inplace else x * self.gamma
+
+
+@dataclass
+class ArkttsCodecTransformerConfig:
+    n_layer: int
+    n_head: int
+    dim: int
+    intermediate_size: int
+    n_local_heads: int = -1
+    head_dim: int = 64
+    rope_base: float = 10000
+    norm_eps: float = 1e-5
+    dropout_rate: float = 0.1
+    attn_dropout_rate: float = 0.1
+    channels_first: bool = True
+
+    def __post_init__(self):
+        if self.n_local_heads == -1:
+            self.n_local_heads = self.n_head
+
+
+class ArkttsCodecAttention(nn.Module):
+    def __init__(self, config: ArkttsCodecTransformerConfig):
+        super().__init__()
+        total = (config.n_head + 2 * config.n_local_heads) * config.head_dim
+        self.wqkv = nn.Linear(config.dim, total, bias=False)
+        self.wo = nn.Linear(config.head_dim * config.n_head, config.dim, bias=False)
+        self.n_head = config.n_head
+        self.n_local_heads = config.n_local_heads
+        self.head_dim = config.head_dim
+        self.attn_dropout_rate = config.attn_dropout_rate
+
+    def forward(self, x, rope_values, mask):
+        batch, length, _ = x.shape
+        query_size = self.n_head * self.head_dim
+        kv_size = self.n_local_heads * self.head_dim
+        query, key, value = self.wqkv(x).split((query_size, kv_size, kv_size), dim=-1)
+        query = query.view(batch, length, self.n_head, self.head_dim)
+        key = key.view(batch, length, self.n_local_heads, self.head_dim)
+        value = value.view(batch, length, self.n_local_heads, self.head_dim)
+        query = _apply_rope(query, rope_values).transpose(1, 2)
+        key = _apply_rope(key, rope_values).transpose(1, 2)
+        value = value.transpose(1, 2)
+        repeat = self.n_head // self.n_local_heads
+        key = key.repeat_interleave(repeat, dim=1)
+        value = value.repeat_interleave(repeat, dim=1)
+        output = F.scaled_dot_product_attention(
+            query,
+            key,
+            value,
+            attn_mask=mask,
+            dropout_p=self.attn_dropout_rate if self.training else 0.0,
+        )
+        output = output.transpose(1, 2).contiguous().view(batch, length, query_size)
+        return self.wo(output)
+
+
+class ArkttsCodecFeedForward(nn.Module):
+    def __init__(self, config: ArkttsCodecTransformerConfig):
+        super().__init__()
+        self.w1 = nn.Linear(config.dim, config.intermediate_size, bias=False)
+        self.w3 = nn.Linear(config.dim, config.intermediate_size, bias=False)
+        self.w2 = nn.Linear(config.intermediate_size, config.dim, bias=False)
+        self.dropout = nn.Dropout(config.dropout_rate)
+
+    def forward(self, x):
+        return self.w2(self.dropout(F.silu(self.w1(x)) * self.w3(x)))
+
+
+class ArkttsCodecTransformerBlock(nn.Module):
+    def __init__(self, config):
+        super().__init__()
+        self.attention = ArkttsCodecAttention(config)
+        self.feed_forward = ArkttsCodecFeedForward(config)
+        self.ffn_norm = ArkttsCodecRMSNorm(config.dim, config.norm_eps)
+        self.attention_norm = ArkttsCodecRMSNorm(config.dim, config.norm_eps)
+        self.attention_layer_scale = ArkttsCodecLayerScale(config.dim, inplace=True)
+        self.ffn_layer_scale = ArkttsCodecLayerScale(config.dim, inplace=True)
+
+    def forward(self, x, rope_values, mask):
+        hidden = x + self.attention_layer_scale(
+            self.attention(self.attention_norm(x), rope_values, mask)
+        )
+        return hidden + self.ffn_layer_scale(self.feed_forward(self.ffn_norm(hidden)))
+
+
+class ArkttsCodecWindowTransformer(nn.Module):
+    def __init__(self, config, input_dim: int, window_size: int | None, causal: bool = True):
+        super().__init__()
+        self.layers = nn.ModuleList([ArkttsCodecTransformerBlock(config) for _ in range(config.n_layer)])
+        self.norm = ArkttsCodecRMSNorm(config.dim, config.norm_eps)
+        self.window_size = window_size
+        self.causal = causal
+        self.channels_first = config.channels_first
+        self.input_proj = nn.Linear(input_dim, config.dim) if input_dim != config.dim else nn.Identity()
+        self.output_proj = nn.Linear(config.dim, input_dim) if input_dim != config.dim else nn.Identity()
+        self.look_ahead_conv = nn.Identity()
+        self.head_dim = config.head_dim
+        self.rope_base = config.rope_base
+
+    def forward(self, x, x_lens=None):
+        del x_lens
+        if self.channels_first:
+            x = x.transpose(1, 2)
+        x = self.look_ahead_conv(self.input_proj(x))
+        length = x.shape[1]
+        row = torch.arange(length, device=x.device)[:, None]
+        column = torch.arange(length, device=x.device)[None, :]
+        mask = column <= row
+        if self.window_size is not None:
+            mask &= column >= (row - self.window_size + 1).clamp_min(0)
+        mask = mask[None, None]
+        rope_values = _rope(length, self.head_dim, self.rope_base, x.device)
+        for layer in self.layers:
+            x = layer(x, rope_values, mask)
+        x = self.output_proj(self.norm(x))
+        return x.transpose(1, 2) if self.channels_first else x
+
+
+def _extra_padding(x, kernel_size: int, stride: int, padding_total: int = 0) -> int:
+    length = x.shape[-1]
+    frames = (length - kernel_size + padding_total) / stride + 1
+    ideal = (math.ceil(frames) - 1) * stride + kernel_size - padding_total
+    return ideal - length
+
+
+class ArkttsCausalConv1d(nn.Module):
+    def __init__(self, in_channels, out_channels, kernel_size, dilation=1, stride=1, groups=1):
+        super().__init__()
+        self.conv = nn.Conv1d(
+            in_channels, out_channels, kernel_size, stride=stride, dilation=dilation, groups=groups
+        )
+        self.stride = stride
+        self.kernel_size = (kernel_size - 1) * dilation + 1
+        self.padding = self.kernel_size - self.stride
+
+    def forward(self, x):
+        right = _extra_padding(x, self.kernel_size, self.stride, self.padding)
+        return self.conv(F.pad(x, (self.padding, right))).contiguous()
+
+    def apply_weight_norm(self):
+        self.conv = weight_norm(self.conv)
+        return self
+
+
+class ArkttsCausalConvTranspose1d(nn.Module):
+    def __init__(self, in_channels, out_channels, kernel_size, stride=1, dilation=1):
+        super().__init__()
+        self.conv = nn.ConvTranspose1d(
+            in_channels, out_channels, kernel_size, stride=stride, dilation=dilation
+        )
+        self.stride = stride
+        self.kernel_size = kernel_size
+
+    def forward(self, x):
+        x = self.conv(x)
+        crop = self.kernel_size - self.stride
+        return x[..., : x.shape[-1] - crop].contiguous() if crop else x.contiguous()
+
+    def apply_weight_norm(self):
+        self.conv = weight_norm(self.conv)
+        return self
+
+
+def _causal_wn_conv(*args, **kwargs):
+    return ArkttsCausalConv1d(*args, **kwargs).apply_weight_norm()
+
+
+def _causal_wn_transpose(*args, **kwargs):
+    return ArkttsCausalConvTranspose1d(*args, **kwargs).apply_weight_norm()
+
+
+@torch.jit.script
+def _arktts_snake(x: Tensor, alpha: Tensor) -> Tensor:
+    shape = x.shape
+    x = x.reshape(shape[0], shape[1], -1)
+    x = x + (alpha + 1e-9).reciprocal() * torch.sin(alpha * x).pow(2)
+    return x.reshape(shape)
+
+
+class ArkttsSnake1d(nn.Module):
+    def __init__(self, channels):
+        super().__init__()
+        self.alpha = nn.Parameter(torch.ones(1, channels, 1))
+
+    def forward(self, x):
+        return _arktts_snake(x, self.alpha)
+
+
+class ArkttsResidualUnit(nn.Module):
+    def __init__(self, dim: int, dilation: int):
+        super().__init__()
+        self.block = nn.Sequential(
+            ArkttsSnake1d(dim),
+            _causal_wn_conv(dim, dim, kernel_size=7, dilation=dilation),
+            ArkttsSnake1d(dim),
+            _causal_wn_conv(dim, dim, kernel_size=1),
+        )
+
+    def forward(self, x):
+        output = self.block(x)
+        difference = x.shape[-1] - output.shape[-1]
+        if difference > 0:
+            x = x[..., :-difference]
+        return x + output
+
+
+class ArkttsEncoderBlock(nn.Module):
+    def __init__(self, dim: int, stride: int, transformer_layers: int):
+        super().__init__()
+        modules = [
+            ArkttsResidualUnit(dim // 2, 1),
+            ArkttsResidualUnit(dim // 2, 3),
+            ArkttsResidualUnit(dim // 2, 9),
+            ArkttsSnake1d(dim // 2),
+            _causal_wn_conv(dim // 2, dim, kernel_size=2 * stride, stride=stride),
+        ]
+        if transformer_layers:
+            config = ArkttsCodecTransformerConfig(
+                n_layer=transformer_layers,
+                n_head=dim // 64,
+                dim=dim,
+                intermediate_size=dim * 3,
+            )
+            modules.append(ArkttsCodecWindowTransformer(config, dim, window_size=512))
+        else:
+            modules.append(nn.Identity())
+        self.block = nn.Sequential(*modules)
+
+    def forward(self, x):
+        return self.block(x)
+
+
+class ArkttsEncoder(nn.Module):
+    def __init__(self):
+        super().__init__()
+        dim = 64
+        modules = [_causal_wn_conv(1, dim, kernel_size=7)]
+        for stride, transformer_layers in zip((2, 4, 8, 8), (0, 0, 0, 4)):
+            dim *= 2
+            modules.append(ArkttsEncoderBlock(dim, stride, transformer_layers))
+        modules.extend((ArkttsSnake1d(dim), _causal_wn_conv(dim, 1024, kernel_size=3)))
+        self.block = nn.Sequential(*modules)
+
+    def forward(self, x):
+        return self.block(x)
+
+
+class ArkttsDecoderBlock(nn.Module):
+    def __init__(self, input_dim: int, output_dim: int, stride: int):
+        super().__init__()
+        self.block = nn.Sequential(
+            ArkttsSnake1d(input_dim),
+            _causal_wn_transpose(input_dim, output_dim, kernel_size=2 * stride, stride=stride),
+            ArkttsResidualUnit(output_dim, 1),
+            ArkttsResidualUnit(output_dim, 3),
+            ArkttsResidualUnit(output_dim, 9),
+        )
+
+    def forward(self, x):
+        return self.block(x)
+
+
+class ArkttsDecoder(nn.Module):
+    def __init__(self):
+        super().__init__()
+        channels = 1536
+        modules = [_causal_wn_conv(1024, channels, kernel_size=7)]
+        for index, stride in enumerate((8, 8, 4, 2)):
+            input_dim = channels // (2**index)
+            output_dim = channels // (2 ** (index + 1))
+            modules.append(ArkttsDecoderBlock(input_dim, output_dim, stride))
+        modules.extend((ArkttsSnake1d(output_dim), _causal_wn_conv(output_dim, 1, kernel_size=7), nn.Tanh()))
+        self.model = nn.Sequential(*modules)
+
+    def forward(self, x):
+        return self.model(x)
+
+
+class ArkttsVectorQuantizer(nn.Module):
+    def __init__(self, input_dim: int, codebook_size: int, codebook_dim: int):
+        super().__init__()
+        self.codebook_size = int(codebook_size)
+        self.codebook_dim = int(codebook_dim)
+        self.in_proj = legacy_weight_norm(nn.Conv1d(input_dim, codebook_dim, kernel_size=1))
+        self.out_proj = legacy_weight_norm(nn.Conv1d(codebook_dim, input_dim, kernel_size=1))
+        self.codebook = nn.Embedding(codebook_size, codebook_dim)
+
+    def decode_code(self, indices):
+        return F.embedding(indices, self.codebook.weight).transpose(1, 2)
+
+    def decode_latents(self, latents):
+        batch, _, length = latents.shape
+        flattened = latents.transpose(1, 2).reshape(batch * length, -1)
+        flattened = F.normalize(flattened)
+        codebook = F.normalize(self.codebook.weight)
+        distances = (
+            flattened.pow(2).sum(1, keepdim=True)
+            - 2 * flattened @ codebook.t()
+            + codebook.pow(2).sum(1, keepdim=True).t()
+        )
+        indices = (-distances).argmax(1).view(batch, length)
+        return self.decode_code(indices), indices
+
+    def forward(self, z):
+        projected = self.in_proj(z)
+        quantized, indices = self.decode_latents(projected)
+        quantized_st = projected + (quantized - projected).detach()
+        return self.out_proj(quantized_st), indices, projected
+
+
+class ArkttsResidualQuantizer(nn.Module):
+    def __init__(self, input_dim: int, n_codebooks: int, codebook_size: int, codebook_dim: int):
+        super().__init__()
+        self.n_codebooks = int(n_codebooks)
+        self.codebook_size = int(codebook_size)
+        self.quantizers = nn.ModuleList([
+            ArkttsVectorQuantizer(input_dim, codebook_size, codebook_dim)
+            for _ in range(n_codebooks)
+        ])
+
+    def forward(self, z):
+        quantized_sum = 0.0
+        residual = z
+        codes = []
+        for quantizer in self.quantizers:
+            quantized, indices, _ = quantizer(residual)
+            quantized_sum = quantized_sum + quantized
+            residual = residual - quantized
+            codes.append(indices)
+        return quantized_sum, torch.stack(codes, dim=1)
+
+    def from_codes(self, codes):
+        output = 0.0
+        for index in range(codes.shape[1]):
+            projected = self.quantizers[index].decode_code(codes[:, index])
+            output = output + self.quantizers[index].out_proj(projected)
+        return output
+
+
+class ArkttsConvNeXtBlock(nn.Module):
+    def __init__(self, dim: int):
+        super().__init__()
+        self.dwconv = ArkttsCausalConv1d(dim, dim, kernel_size=7, groups=dim)
+        self.norm = nn.LayerNorm(dim, eps=1e-6)
+        self.pwconv1 = nn.Linear(dim, 4 * dim)
+        self.act = nn.GELU()
+        self.pwconv2 = nn.Linear(4 * dim, dim)
+        self.gamma = nn.Parameter(1e-6 * torch.ones(dim))
+
+    def forward(self, x):
+        residual = x
+        x = self.dwconv(x).permute(0, 2, 1)
+        x = self.pwconv2(self.act(self.pwconv1(self.norm(x))))
+        x = (self.gamma * x).permute(0, 2, 1)
+        return residual + x
+
+
+class ArkttsDownsampleQuantizer(nn.Module):
+    def __init__(self, config=None):
+        super().__init__()
+        self.semantic_quantizer = ArkttsResidualQuantizer(1024, 1, 4096, 8)
+        self.quantizer = ArkttsResidualQuantizer(1024, 9, 1024, 8)
+        self.downsample = nn.Sequential(
+            nn.Sequential(
+                ArkttsCausalConv1d(1024, 1024, kernel_size=2, stride=2),
+                ArkttsConvNeXtBlock(1024),
+            ),
+            nn.Sequential(
+                ArkttsCausalConv1d(1024, 1024, kernel_size=2, stride=2),
+                ArkttsConvNeXtBlock(1024),
+            ),
+        )
+        self.upsample = nn.Sequential(
+            nn.Sequential(
+                ArkttsCausalConvTranspose1d(1024, 1024, kernel_size=2, stride=2),
+                ArkttsConvNeXtBlock(1024),
+            ),
+            nn.Sequential(
+                ArkttsCausalConvTranspose1d(1024, 1024, kernel_size=2, stride=2),
+                ArkttsConvNeXtBlock(1024),
+            ),
+        )
+        pre_transformer_config = ArkttsCodecTransformerConfig(
+            n_layer=8,
+            n_head=16,
+            dim=1024,
+            intermediate_size=3072,
+        )
+        post_transformer_config = ArkttsCodecTransformerConfig(
+            n_layer=int(getattr(config, "codec_post_n_layer", 8)),
+            n_head=int(getattr(config, "codec_post_n_head", 16)),
+            n_local_heads=int(getattr(config, "codec_post_n_local_heads", 8)),
+            dim=1024,
+            intermediate_size=int(getattr(config, "codec_post_intermediate_size", 1216)),
+        )
+        self.pre_module = ArkttsCodecWindowTransformer(
+            pre_transformer_config, 1024, window_size=128
+        )
+        self.post_module = ArkttsCodecWindowTransformer(
+            post_transformer_config, 1024, window_size=128
+        )
+        self.semantic_predictor_module = nn.Identity()
+
+    def forward(self, z):
+        original_length = z.shape[-1]
+        z = self.pre_module(self.downsample(z))
+        semantic, semantic_codes = self.semantic_quantizer(z)
+        residual, residual_codes = self.quantizer(z - semantic)
+        z = self.upsample(self.post_module(semantic + residual))
+        difference = original_length - z.shape[-1]
+        if difference > 0:
+            z = F.pad(z, (difference, 0))
+        elif difference < 0:
+            z = z[..., -difference:]
+        return z, torch.cat((semantic_codes, residual_codes), dim=1)
+
+    def decode(self, indices):
+        indices = indices.clone()
+        indices[:, 0].clamp_(0, self.semantic_quantizer.codebook_size - 1)
+        indices[:, 1:].clamp_(0, self.quantizer.codebook_size - 1)
+        semantic = self.semantic_quantizer.from_codes(indices[:, :1])
+        residual = self.quantizer.from_codes(indices[:, 1:])
+        return self.upsample(self.post_module(semantic + residual))
+
+
+class ArkttsCodec(nn.Module):
+    sample_rate = 44100
+    hop_length = 512
+    frame_length = 2048
+
+    def __init__(self, config=None):
+        super().__init__()
+        self.encoder = ArkttsEncoder()
+        self.quantizer = ArkttsDownsampleQuantizer(config)
+        self.decoder = ArkttsDecoder()
+
+    @torch.inference_mode()
+    def encode(self, audio, audio_lengths=None):
+        if audio.ndim == 2:
+            audio = audio[:, None]
+        if audio.ndim != 3 or audio.shape[1] != 1:
+            raise ValueError("audio must have shape [B, 1, samples]")
+        original_length = audio.shape[-1]
+        right = math.ceil(original_length / self.frame_length) * self.frame_length - original_length
+        if right:
+            audio = F.pad(audio, (0, right))
+        if audio_lengths is None:
+            audio_lengths = torch.full(
+                (audio.shape[0],), original_length, device=audio.device, dtype=torch.long
+            )
+        encoded = self.encoder(audio)
+        _, codes = self.quantizer(encoded)
+        code_lengths = torch.ceil(audio_lengths.float() / self.frame_length).long()
+        max_codes = codes.shape[-1]
+        padded = torch.full_like(codes, -1)
+        for index, length in enumerate(code_lengths.tolist()):
+            padded[index, :, : min(length, max_codes)] = codes[index, :, : min(length, max_codes)]
+        return padded, code_lengths.clamp_max(max_codes)
+
+    @torch.inference_mode()
+    def decode(self, codes):
+        return self.decoder(self.quantizer.decode(codes.long()))
+
+
+__all__ = ["ArkttsCodec"]

preprocessor_config.json

@@ -0,0 +1,10 @@
+{
+  "audio_sampling_rate": 44100,
+  "auto_map": {
+    "AutoProcessor": "processing_arktts.ArkttsProcessor"
+  },
+  "feature_extractor_type": "ArkttsProcessor",
+  "num_codebooks": 10,
+  "processor_class": "ArkttsProcessor",
+  "semantic_begin_id": 151678
+}

processing_arktts.py

@@ -0,0 +1,259 @@
+from __future__ import annotations
+
+import inspect
+import json
+import os
+import re
+from pathlib import Path
+from typing import Any, Sequence
+
+import numpy as np
+import torch
+from transformers import AutoTokenizer
+from transformers.feature_extraction_utils import BatchFeature
+from transformers.processing_utils import ProcessorMixin
+
+
+def _clean_text(text: str) -> str:
+    return " ".join(str(text).strip().split())
+
+
+def _as_list(value: Any, batch_size: int, name: str) -> list[Any]:
+    if isinstance(value, (str, Path)) or value is None or np.isscalar(value):
+        return [value] * batch_size
+    if isinstance(value, torch.Tensor) and value.ndim <= 2:
+        return [value] if batch_size == 1 else list(value)
+    if isinstance(value, np.ndarray) and value.ndim <= 2:
+        return [value] if batch_size == 1 else list(value)
+    values = list(value)
+    if len(values) != batch_size:
+        raise ValueError(f"{name} must contain {batch_size} items, got {len(values)}")
+    return values
+
+
+def _pad_1d(rows: list[torch.Tensor], pad_value: int) -> tuple[torch.Tensor, torch.Tensor]:
+    max_len = max((row.numel() for row in rows), default=0)
+    values = torch.full((len(rows), max_len), pad_value, dtype=torch.long)
+    mask = torch.zeros((len(rows), max_len), dtype=torch.long)
+    for idx, row in enumerate(rows):
+        length = row.numel()
+        values[idx, :length] = row
+        mask[idx, :length] = 1
+    return values, mask
+
+
+class ArkttsProcessor(ProcessorMixin):
+    attributes = ["tokenizer"]
+    tokenizer_class = ("PreTrainedTokenizerFast", "PreTrainedTokenizer")
+    valid_kwargs = ["num_codebooks", "semantic_begin_id", "audio_sampling_rate"]
+
+    def __init__(
+        self,
+        tokenizer,
+        num_codebooks: int = 10,
+        semantic_begin_id: int = 151678,
+        audio_sampling_rate: int = 44100,
+        **kwargs,
+    ):
+        super().__init__(tokenizer)
+        self.num_codebooks = int(num_codebooks)
+        self.semantic_begin_id = int(semantic_begin_id)
+        self.audio_sampling_rate = int(audio_sampling_rate)
+
+    @classmethod
+    def from_pretrained(cls, pretrained_model_name_or_path: str, **kwargs) -> "ArkttsProcessor":
+        trust_remote_code = bool(kwargs.pop("trust_remote_code", False))
+        shared_names = {
+            "cache_dir", "force_download", "local_files_only", "token", "revision", "subfolder"
+        }
+        shared = {key: kwargs[key] for key in list(kwargs) if key in shared_names}
+        config = {}
+        local_config = os.path.join(str(pretrained_model_name_or_path), "processor_config.json")
+        if os.path.isfile(local_config):
+            with open(local_config, "r", encoding="utf-8") as handle:
+                config = json.load(handle)
+        tokenizer = AutoTokenizer.from_pretrained(
+            pretrained_model_name_or_path,
+            use_fast=True,
+            trust_remote_code=trust_remote_code,
+            fix_mistral_regex=False,
+            **shared,
+        )
+        return cls(
+            tokenizer=tokenizer,
+            num_codebooks=config.get("num_codebooks", 10),
+            semantic_begin_id=config.get("semantic_begin_id", 151678),
+            audio_sampling_rate=config.get("audio_sampling_rate", 44100),
+        )
+
+    def _encode(self, text: str) -> torch.Tensor:
+        encode_kwargs = {"add_special_tokens": False}
+        if "allowed_special" in inspect.signature(self.tokenizer.encode).parameters:
+            encode_kwargs["allowed_special"] = "all"
+        return torch.tensor(self.tokenizer.encode(text, **encode_kwargs), dtype=torch.long)
+
+    @staticmethod
+    def _format_reference_text(text: str) -> str:
+        cleaned = _clean_text(text)
+        if re.search(r"<\|speaker:\d+\|>", cleaned):
+            return cleaned
+        return f"<|speaker:0|>{cleaned}"
+
+    def _prompt_segments(self, text: str, reference_text: str | None, has_reference: bool):
+        target = _clean_text(text)
+        if not target:
+            raise ValueError("text must not be empty")
+        def encode_parts(parts: list[str]) -> torch.Tensor:
+            return torch.cat([self._encode(part) for part in parts])
+
+        if not has_reference:
+            full = encode_parts([
+                "<|im_start|>system\n",
+                "convert the provided text to speech",
+                "<|im_end|>\n",
+                "<|im_start|>user\n",
+                target,
+                "<|im_end|>\n",
+                "<|im_start|>assistant\n<|voice|>",
+            ])
+            return full, self._encode("")
+        if not reference_text:
+            raise ValueError("reference_text is required when a reference voice is provided")
+        prefix = encode_parts([
+            "<|im_start|>system\n",
+            "convert the provided text to speech reference to the following:\n\nText:\n",
+            self._format_reference_text(reference_text),
+            "\n\nSpeech:\n",
+        ])
+        suffix = encode_parts([
+            "<|im_end|>\n",
+            "<|im_start|>user\n",
+            target,
+            "<|im_end|>\n",
+            "<|im_start|>assistant\n<|voice|>",
+        ])
+        return prefix, suffix
+
+    def _load_audio(self, value: Any, sampling_rate: int | None) -> torch.Tensor:
+        source_rate = sampling_rate
+        if isinstance(value, (str, Path)):
+            try:
+                import soundfile as sf
+            except ImportError as exc:
+                raise ImportError("soundfile is required for reference audio paths") from exc
+            array, source_rate = sf.read(str(value), dtype="float32", always_2d=True)
+            array = array.mean(axis=1)
+            audio = torch.from_numpy(np.asarray(array, dtype=np.float32))
+        else:
+            if isinstance(value, dict):
+                source_rate = value.get("sampling_rate", source_rate)
+                value = value.get("array")
+            if isinstance(value, (tuple, list)) and len(value) == 2 and np.isscalar(value[1]):
+                value, source_rate = value
+            audio = torch.as_tensor(value, dtype=torch.float32)
+            if audio.ndim == 2:
+                audio = audio.mean(dim=0)
+            if audio.ndim != 1:
+                raise ValueError(f"reference audio must be mono or channels-first, got {tuple(audio.shape)}")
+        if audio.numel() == 0:
+            raise ValueError("reference audio must not be empty")
+        if source_rate is None:
+            raise ValueError("sampling_rate is required for reference audio arrays")
+        if int(source_rate) != self.audio_sampling_rate:
+            try:
+                from torchaudio.functional import resample
+            except ImportError as exc:
+                raise ImportError("torchaudio is required to resample reference audio") from exc
+            audio = resample(audio, int(source_rate), self.audio_sampling_rate)
+        return audio.contiguous()
+
+    def __call__(
+        self,
+        text: str | Sequence[str],
+        reference_text: str | Sequence[str] | None = None,
+        reference_audio: Any = None,
+        reference_codes: Any = None,
+        sampling_rate: int | Sequence[int] | None = None,
+        return_tensors: str = "pt",
+        **kwargs,
+    ) -> BatchFeature:
+        if kwargs:
+            raise TypeError(f"Unexpected processor arguments: {sorted(kwargs)}")
+        if return_tensors != "pt":
+            raise ValueError("ArkttsProcessor currently supports return_tensors='pt' only")
+        texts = [text] if isinstance(text, str) else list(text)
+        if not texts:
+            raise ValueError("text batch must not be empty")
+        batch_size = len(texts)
+        ref_texts = _as_list(reference_text, batch_size, "reference_text")
+        if reference_audio is not None and reference_codes is not None:
+            raise ValueError("Provide reference_audio or reference_codes, not both")
+        has_reference = reference_audio is not None or reference_codes is not None
+
+        prefix_rows, suffix_rows = zip(*[
+            self._prompt_segments(item, ref_texts[idx], has_reference)
+            for idx, item in enumerate(texts)
+        ])
+        prefix_ids, prefix_mask = _pad_1d(list(prefix_rows), self.tokenizer.pad_token_id)
+        suffix_ids, suffix_mask = _pad_1d(list(suffix_rows), self.tokenizer.pad_token_id)
+        data: dict[str, torch.Tensor] = {
+            "prefix_input_ids": prefix_ids,
+            "prefix_attention_mask": prefix_mask,
+            "suffix_input_ids": suffix_ids,
+            "suffix_attention_mask": suffix_mask,
+        }
+
+        if reference_codes is not None:
+            code_items = _as_list(reference_codes, batch_size, "reference_codes")
+            loaded = []
+            for item in code_items:
+                if isinstance(item, (str, Path)):
+                    item = np.load(str(item))
+                codes = torch.as_tensor(item, dtype=torch.long)
+                if codes.ndim != 2 or codes.shape[0] != self.num_codebooks or codes.shape[1] == 0:
+                    raise ValueError(
+                        f"reference codes must have shape [{self.num_codebooks}, T>0], got {tuple(codes.shape)}"
+                    )
+                if codes.min() < 0 or codes.max() >= 4096:
+                    raise ValueError("reference codes must be in [0, 4095]")
+                loaded.append(codes)
+            max_frames = max(item.shape[1] for item in loaded)
+            padded = torch.full((batch_size, self.num_codebooks, max_frames), -1, dtype=torch.long)
+            lengths = torch.empty(batch_size, dtype=torch.long)
+            for idx, codes in enumerate(loaded):
+                lengths[idx] = codes.shape[1]
+                padded[idx, :, : codes.shape[1]] = codes
+            data["reference_codes"] = padded
+            data["reference_code_lengths"] = lengths
+
+        if reference_audio is not None:
+            audio_items = _as_list(reference_audio, batch_size, "reference_audio")
+            rate_items = _as_list(sampling_rate, batch_size, "sampling_rate")
+            loaded_audio = [self._load_audio(item, rate_items[idx]) for idx, item in enumerate(audio_items)]
+            max_samples = max(item.numel() for item in loaded_audio)
+            padded_audio = torch.zeros((batch_size, 1, max_samples), dtype=torch.float32)
+            lengths = torch.empty(batch_size, dtype=torch.long)
+            for idx, audio in enumerate(loaded_audio):
+                lengths[idx] = audio.numel()
+                padded_audio[idx, 0, : audio.numel()] = audio
+            data["reference_audio_values"] = padded_audio
+            data["reference_audio_lengths"] = lengths
+
+        return BatchFeature(data=data)
+
+    @property
+    def model_input_names(self) -> list[str]:
+        return [
+            "prefix_input_ids", "prefix_attention_mask", "suffix_input_ids",
+            "suffix_attention_mask", "reference_codes", "reference_code_lengths",
+            "reference_audio_values", "reference_audio_lengths",
+        ]
+
+    def batch_decode(self, *args, **kwargs):
+        return self.tokenizer.batch_decode(*args, **kwargs)
+
+    def decode(self, *args, **kwargs):
+        return self.tokenizer.decode(*args, **kwargs)
+
+
+__all__ = ["ArkttsProcessor"]

processor_config.json

@@ -0,0 +1,9 @@
+{
+  "audio_sampling_rate": 44100,
+  "auto_map": {
+    "AutoProcessor": "processing_arktts.ArkttsProcessor"
+  },
+  "num_codebooks": 10,
+  "processor_class": "ArkttsProcessor",
+  "semantic_begin_id": 151678
+}

tokenizer.json

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