IndiCodec v1.0

Co-authored-by: prakhar khatri <prakhar khatri@users.noreply.huggingface.co>

.gitattributes

@@ -0,0 +1,36 @@
+*.7z filter=lfs diff=lfs merge=lfs -text
+*.arrow filter=lfs diff=lfs merge=lfs -text
+*.bin filter=lfs diff=lfs merge=lfs -text
+*.bz2 filter=lfs diff=lfs merge=lfs -text
+*.ckpt filter=lfs diff=lfs merge=lfs -text
+*.ftz filter=lfs diff=lfs merge=lfs -text
+*.gz filter=lfs diff=lfs merge=lfs -text
+*.h5 filter=lfs diff=lfs merge=lfs -text
+*.joblib filter=lfs diff=lfs merge=lfs -text
+*.lfs.* filter=lfs diff=lfs merge=lfs -text
+*.mlmodel filter=lfs diff=lfs merge=lfs -text
+*.model filter=lfs diff=lfs merge=lfs -text
+*.msgpack filter=lfs diff=lfs merge=lfs -text
+*.npy filter=lfs diff=lfs merge=lfs -text
+*.npz filter=lfs diff=lfs merge=lfs -text
+*.onnx filter=lfs diff=lfs merge=lfs -text
+*.ot filter=lfs diff=lfs merge=lfs -text
+*.parquet filter=lfs diff=lfs merge=lfs -text
+*.pb filter=lfs diff=lfs merge=lfs -text
+*.pickle filter=lfs diff=lfs merge=lfs -text
+*.pkl filter=lfs diff=lfs merge=lfs -text
+*.pt filter=lfs diff=lfs merge=lfs -text
+*.pth filter=lfs diff=lfs merge=lfs -text
+*.rar filter=lfs diff=lfs merge=lfs -text
+*.safetensors filter=lfs diff=lfs merge=lfs -text
+saved_model/**/* filter=lfs diff=lfs merge=lfs -text
+*.tar.* filter=lfs diff=lfs merge=lfs -text
+*.tar filter=lfs diff=lfs merge=lfs -text
+*.tflite filter=lfs diff=lfs merge=lfs -text
+*.tgz filter=lfs diff=lfs merge=lfs -text
+*.wasm filter=lfs diff=lfs merge=lfs -text
+*.xz 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
+assets/architecture.png filter=lfs diff=lfs merge=lfs -text

LICENSE

@@ -0,0 +1,202 @@
+
+                                 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.

README.md

@@ -0,0 +1,129 @@
+---
+license: apache-2.0
+language:
+- en
+- hi
+- bn
+- gu
+- kn
+- ml
+- mr
+- or
+- pa
+- ta
+- te
+tags:
+- audio
+- speech
+- audio-codec
+- speech-tokenizer
+- speech-language-models
+pipeline_tag: audio-to-audio
+datasets:
+- amphion/Emilia-Dataset
+- mythicinfinity/libritts
+- mythicinfinity/libritts_r
+- CSTR-Edinburgh/vctk
+- ai4bharat/Rasa
+- ai4bharat/indicvoices_r
+- ai4bharat/IndicVoices
+---
+
+# IndiCodec
+
+[Audio samples](https://codeprakhar25.github.io/IndiCodec/) · [Code on GitHub](https://github.com/codeprakhar25/IndiCodec)
+
+A 1.09 kbps neural speech codec for building speech language models and TTS. It turns speech into two streams of
+integer tokens and decodes them back to 24 kHz audio.
+
+It was trained mostly on English and 10 Indian languages, and it is strongest on Indian languages: among codecs at
+about 1 kbps, it keeps the most words in every one of the 10 we tested.
+
+## Key features
+
+- **Low bitrate:** 1.09 kbps, 75 tokens per second.
+- **Two token streams:** a content stream (25 Hz) that carries the words, and an acoustic stream (50 Hz) that carries
+  the voice and detail.
+- **Split FSQ, not RVQ:** each stream is a single FSQ grid. Content is coded first, and one acoustic grid codes what
+  remains, instead of a stack of residual codebooks. Each frame is one id per stream, with no learned codebook.
+- **16 kHz in, 24 kHz out:** any input sample rate is resampled to 16 kHz.
+- **One file:** all weights are in `model.safetensors`. Nothing else is downloaded.
+- **Commercial use allowed:** Apache 2.0. All training data is under open licences, none of them non-commercial.
+
+## Model details
+
+| | |
+|---|---|
+| Content stream | 25 Hz, FSQ 8 × 3 levels = 6,561 ids, 317 bps |
+| Acoustic stream | 50 Hz, FSQ 6 × 6 levels = 46,656 ids, 775 bps |
+| Bitrate | 1.09 kbps |
+| Parameters | 624 M in the file, 387 M of them w2v-BERT 2.0; decoding alone uses 186 M |
+| Licence | Apache 2.0 |
+
+The content stream quantises w2v-BERT 2.0 features (layer 16). The acoustic stream quantises what a CNN encoder hears
+beyond the content stream. A transformer decoder with an iSTFT head turns the sum back into audio.
+
+IndiCodec is built on [NeuCodec](https://huggingface.co/neuphonic/neucodec) by Neuphonic: it uses NeuCodec's
+architecture, and training started from its released weights. The two-stream quantiser, the encoder training and the
+training recipe are new.
+
+## Quality
+
+Speech codecs at about 1 kbps, on speakers not seen in training. Every codec is scored on the same clips.
+
+PESQ-WB (higher is better):
+
+| Test set | IndiCodec<br>1.09 kbps | Mimi, 8 codebooks<br>1.10 kbps | SNAC 24 kHz<br>0.98 kbps | NeuCodec<br>0.80 kbps |
+|---|---|---|---|---|
+| LibriTTS-R test (English) | **2.41** | 2.32 | 2.29 | 2.19 |
+| CMU-Arctic (English) | **2.30** | 2.28 | 2.17 | 2.06 |
+| Seed-TTS eval (English) | 1.99 | **2.07** | 1.82 | 1.81 |
+| FLEURS, 10 Indian languages | **1.90** | 1.84 | 1.69 | 1.69 |
+
+Words after a round trip (lower is better):
+
+| | IndiCodec | Mimi | SNAC | NeuCodec |
+|---|---|---|---|---|
+| English WER, Seed-TTS eval | **2.5 %** | 3.6 % | 4.9 % | 3.5 % |
+| English WER, CMU-Arctic | **2.3 %** | 3.2 % | 3.8 % | 2.8 % |
+| Indian-language CER, FLEURS, 10 languages | **6.3 %** | 14.3 % | 13.3 % | 8.9 % |
+| Tokens per second | 75 | 100 | 82 | 50 |
+
+English WER uses Whisper large-v3. Indian-language CER uses the IndicConformer 600M ASR; the original recordings score
+5.2 %, so IndiCodec adds 1.0 point. The 10 languages are Hindi, Bengali, Gujarati, Kannada, Malayalam, Marathi, Odia,
+Punjabi, Tamil and Telugu.
+
+## Use
+
+```bash
+hf download prakharkhatri/IndiCodec --local-dir IndiCodec && cd IndiCodec
+pip install torch==2.6.0 torchaudio==2.6.0 safetensors soundfile transformers==4.51.* vector-quantize-pytorch==1.17.8 torchtune==0.5.0 torchao==0.9.0 local-attention einops librosa
+pip install --no-deps "git+https://github.com/neuphonic/neucodec@ed3e6cd"
+```
+
+```python
+import soundfile as sf
+from indicodec import IndiCodec
+
+codec = IndiCodec.load("model.safetensors")
+wav, sr = sf.read("speech.wav")
+
+content, acoustic = codec.encode(wav, sr)    # token arrays
+audio = codec.decode(content, acoustic)      # 24 kHz
+```
+
+## Training data
+
+English: Emilia-YODAS (English part only), LibriTTS, LibriTTS-R and VCTK. Indian languages: Rasa, IndicVoices-R and
+IndicVoices from AI4Bharat. Plus about 15 h of laughs, sighs and coughs from VocalSound. All CC BY 4.0, except
+VocalSound (CC BY-SA 4.0). No audio is included in this release.
+
+## Limits
+
+- Speech only; not for music.
+- Not streaming: the encoder and the decoder both attend over the whole clip.
+
+## Licence
+
+Apache 2.0. Built on [NeuCodec](https://huggingface.co/neuphonic/neucodec) (Apache 2.0). The bundled w2v-BERT 2.0
+layers are Meta's, under MIT; its licence text is in `indicodec/w2v_bert/LICENSE`.

config.json

@@ -0,0 +1,12 @@
+{
+  "model_type": "indicodec",
+  "version": "1.0",
+  "weights": "model.safetensors",
+  "sample_rate_in": 16000,
+  "sample_rate_out": 24000,
+  "bitrate_bps": 1092,
+  "tokens_per_second": 75,
+  "content": {"hz": 25, "fsq_levels": [3, 3, 3, 3, 3, 3, 3, 3], "n_ids": 6561},
+  "acoustic": {"hz": 50, "fsq_levels": [6, 6, 6, 6, 6, 6], "n_ids": 46656},
+  "w2v_bert": {"source": "facebook/w2v-bert-2.0", "layers": 16, "bundled": true}
+}

indicodec/__init__.py

@@ -0,0 +1,3 @@
+from .infer import IndiCodec
+
+__all__ = ["IndiCodec"]

indicodec/codec_spec.json

@@ -0,0 +1,29 @@
+{
+  "name": "IndiCodec v1.0",
+  "ckpt": "model.safetensors",
+  "md5_safetensors": "7d9f2d7d04893de3706fcefcd33450ce",
+  "md5": "986a7c393c2168a00c56dfb96751a0d4",
+  "w2v_bert": {"source": "facebook/w2v-bert-2.0", "revision": "da985ba0987f70aaeb84a80f2851cfac8c697a7b", "layers": 16, "in_file": true},
+  "sample_rate_encoders": 16000,
+  "sample_rate_out": 24000,
+  "bitrate_bps": 1092,
+  "content": {
+    "levels": [3, 3, 3, 3, 3, 3, 3, 3],
+    "hz": 25,
+    "n_ids": 6561,
+    "bps": 317
+  },
+  "acoustic": {
+    "levels": [6, 6, 6, 6, 6, 6],
+    "hz": 50,
+    "n_ids": 46656,
+    "bps": 775
+  },
+  "hop_16": 320,
+  "hop_24": 480,
+  "n_fft": 1920,
+  "tail_frames": 2,
+  "level_db": null,
+  "decoder": "vocos_istft_12",
+  "amp_cuda": "bf16"
+}

indicodec/infer.py

@@ -0,0 +1,129 @@
+"""IndiCodec inference. Self-contained; the release file needs no weight download.
+
+    from indicodec import IndiCodec
+    c = IndiCodec.load("model.safetensors")      # release file; the internal .pt still loads
+    idx_s, idx_a = c.encode(wav, sr)
+    wav24 = c.decode(idx_s, idx_a, n_samples=n24)
+    wav24 = c.reconstruct(wav, sr)
+"""
+from __future__ import annotations
+
+import hashlib
+import json
+import os
+from pathlib import Path
+
+import numpy as np
+import torch
+
+_HERE = Path(__file__).resolve().parent
+_SPEC = json.loads((_HERE / "codec_spec.json").read_text())
+
+
+def md5_file(p: Path) -> str:
+    h = hashlib.md5()
+    with open(p, "rb") as f:
+        for b in iter(lambda: f.read(1 << 24), b""):
+            h.update(b)
+    return h.hexdigest()
+
+
+def n24_of(wav, sr: int) -> int:
+    return int(round(len(wav) * _SPEC["sample_rate_out"] / sr))
+
+
+def _refuse_train_flags(st: dict) -> None:
+    bad = []
+    if st.get("acoustic_continuous"):
+        bad.append("acoustic_continuous")
+    if int(st.get("acoustic_frames") or 1) != 1:
+        bad.append(f"acoustic_frames={st.get('acoustic_frames')}")
+    if st.get("rope_fix"):
+        bad.append("rope_fix")
+    if int(st.get("phase2d") or 0):
+        bad.append(f"phase2d={st.get('phase2d')}")
+    if bad:
+        raise SystemExit(f"IndiCodec infer refuses {bad}: not an IndiCodec release checkpoint")
+
+
+def read_ckpt(p: Path) -> tuple[dict, str]:
+    """md5-check a weight file against the spec and return ({"gen", "tail_frames", "levels_a", ...}, md5)."""
+    got = md5_file(p)
+    if p.suffix == ".safetensors":   # release format (2026-10-25): tensors only, no pickle
+        want = _SPEC["md5_safetensors"]
+        if got != want:
+            raise SystemExit(f"{p} md5 {got} != spec {want}")
+        from safetensors import safe_open
+        with safe_open(str(p), framework="pt", device="cpu") as f:
+            meta = f.metadata() or {}
+            gen = {k: f.get_tensor(k) for k in f.keys()}
+        st = {"gen": gen, "meta": meta, "tail_frames": int(meta.get("tail_frames", 0)),
+              "levels_a": [int(x) for x in meta.get("levels_a", "").split(",") if x]}
+    else:   # internal training checkpoint (pickle)
+        want = _SPEC["md5"]
+        if got != want:
+            raise SystemExit(f"{p} md5 {got} != spec {want}")
+        st = torch.load(p, map_location="cpu", weights_only=False)
+    return st, got
+
+
+class IndiCodec:
+    def __init__(self, model, device: str, amp: bool, ckpt_md5: str):
+        self.model = model
+        self.device = device
+        self.amp = amp
+        self.ckpt_md5 = ckpt_md5
+        self.spec = _SPEC
+
+    @classmethod
+    def load(cls, ckpt: str | Path, device: str | None = None, amp: bool | None = None) -> "IndiCodec":
+        try:
+            from .model import IndiNet
+        except ImportError:
+            from model import IndiNet  # folder-alone: infer.py + model.py side by side
+
+        p = Path(ckpt)
+        st, got = read_ckpt(p)
+        if int(st.get("tail_frames", 0)) != int(_SPEC["tail_frames"]):
+            raise SystemExit(f"ckpt tail_frames {st.get('tail_frames')} != spec {_SPEC['tail_frames']}")
+        if list(st.get("levels_a") or []) != list(_SPEC["acoustic"]["levels"]):
+            raise SystemExit(f"ckpt levels_a {st.get('levels_a')} != spec {_SPEC['acoustic']['levels']}")
+        _refuse_train_flags(st)
+        bundled = any(k.startswith("semantic_model.") for k in st["gen"])   # release file; the .pt has no w2v-BERT
+        m = IndiNet(levels_a=st.get("levels_a"), w2v_bundled=bundled)
+        m.load_gen(st["gen"])
+        m.tail_frames = int(st.get("tail_frames", 0))
+        if device is None:
+            device = "cuda" if torch.cuda.is_available() else "cpu"
+        m.to(device).eval()
+        if amp is None:
+            amp = device != "cpu"
+        return cls(m, device, amp, got)
+
+    def encode(self, wav, sr: int) -> tuple[np.ndarray, np.ndarray]:
+        try:
+            from .model import encode_wav
+        except ImportError:
+            from model import encode_wav
+        idx_s, idx_a, _ = encode_wav(self.model, wav, sr, self.device, amp=self.amp)
+        return idx_s.detach().cpu().numpy(), idx_a.detach().cpu().numpy()
+
+    @torch.no_grad()
+    def decode(self, idx_s, idx_a, n_samples: int | None = None) -> np.ndarray:
+        s = torch.as_tensor(idx_s, device=self.device)
+        a = torch.as_tensor(idx_a, device=self.device)
+        with torch.autocast("cuda", dtype=torch.bfloat16, enabled=self.amp and self.device != "cpu"):
+            y = self.model.decode_codes(s, a, mode="full").float().cpu().numpy().reshape(-1)
+        if n_samples is None:
+            return y
+        return y[:n_samples] if len(y) >= n_samples else np.pad(y, (0, n_samples - len(y)))
+
+    def reconstruct(self, wav, sr: int) -> np.ndarray:
+        return self.decode(*self.encode(wav, sr), n_samples=n24_of(wav, sr))
+
+    def deterministic(self, on: bool = True) -> None:
+        if on:
+            os.environ.setdefault("CUBLAS_WORKSPACE_CONFIG", ":4096:8")
+        torch.backends.cudnn.deterministic = on
+        torch.backends.cudnn.benchmark = not on
+        torch.use_deterministic_algorithms(on, warn_only=True)

indicodec/model.py

@@ -0,0 +1,143 @@
+"""IndiCodec inference graph. No weight download.
+
+Architecture classes still come from installed `neucodec` (Apache 2.0, commit ed3e6cd).
+The release model.safetensors holds every tensor, including the 16 w2v-BERT 2.0 (MIT) layers we use, built from the
+shipped w2v_bert/config.json. The internal .pt has no w2v-BERT, so that path still fetches it (pinned revision).
+"""
+from __future__ import annotations
+
+import math
+from pathlib import Path
+
+import numpy as np
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+SEM_LAYER = 16
+HOP16, HOP24 = 320, 480
+LEVELS_S = [3] * 8
+W2V = "facebook/w2v-bert-2.0"
+W2V_REV = "da985ba0987f70aaeb84a80f2851cfac8c697a7b"
+W2V_DIR = Path(__file__).resolve().parent / "w2v_bert"   # config (16 layers) + preprocessor config, from W2V_REV
+NEED_GEN = (
+    "fsq_s.", "fsq_a.", "backbone.", "head.",
+    "CodecEnc.", "SemanticEncoder_module.",
+)
+# training-only head (a semantic reconstruction loss); never runs at inference, dropped from the release file
+TRAIN_ONLY = ("SemanticDecoder_module.",)
+
+
+class IndiNet(nn.Module):
+    """IndiCodec encode / decode only. Refuses phase2d, rope_fix, continuous, joint FSQ."""
+
+    def __init__(self, levels_a: list[int], w2v_bundled: bool = True):
+        super().__init__()
+        from transformers import Wav2Vec2BertConfig, Wav2Vec2BertModel
+        from vector_quantize_pytorch import FSQ
+        from neucodec.codec_encoder import CodecEncoder
+        from neucodec.codec_decoder_vocos import ISTFTHead, VocosBackbone
+        from neucodec.module import SemanticEncoder
+
+        self.w2v_bundled = bool(w2v_bundled)
+        if self.w2v_bundled:   # weights come from the release file in load_gen
+            sm = Wav2Vec2BertModel(Wav2Vec2BertConfig.from_json_file(str(W2V_DIR / "config.json")))
+        else:
+            sm = Wav2Vec2BertModel.from_pretrained(W2V, revision=W2V_REV, output_hidden_states=True)
+            sm.encoder.layers = sm.encoder.layers[:SEM_LAYER]
+        if len(sm.encoder.layers) != SEM_LAYER:
+            raise SystemExit(f"w2v-BERT has {len(sm.encoder.layers)} layers, IndiCodec needs {SEM_LAYER}")
+        self.semantic_model = sm.eval().requires_grad_(False)
+        self.SemanticEncoder_module = SemanticEncoder(1024, 1024, 1024)
+        self.CodecEnc = CodecEncoder()
+        self.backbone = VocosBackbone(hidden_dim=1024, depth=12, heads=16, pos_meb_dim=64)
+        self.head = ISTFTHead(dim=1024, n_fft=HOP24 * 4, hop_length=HOP24, padding="same")
+        self.fsq_s = FSQ(levels=LEVELS_S, dim=1024)
+        self.levels_a = list(levels_a)
+        self.fsq_a = FSQ(levels=self.levels_a, dim=1024)
+        self.tail_frames = 0
+
+    def load_gen(self, gen: dict) -> None:
+        gen = {k: v for k, v in gen.items() if not k.startswith(TRAIN_ONLY)}
+        missing_pref = [p for p in NEED_GEN if not any(k.startswith(p) for k in gen)]
+        if missing_pref:
+            raise SystemExit(f"ckpt gen missing {missing_pref}; this infer needs the IndiCodec trained towers")
+        r = self.load_state_dict(gen, strict=False)
+        bad_u = [k for k in r.unexpected_keys if k.startswith("head.refine.")]
+        if bad_u:
+            raise RuntimeError(f"checkpoint has PhaseRefiner2D ({bad_u[0]}); IndiCodec infer refuses it")
+        bad_m = [m for m in r.missing_keys if self.w2v_bundled or not m.startswith("semantic_model.")]
+        if bad_m:
+            raise SystemExit(f"ckpt missing {len(bad_m)} keys, e.g. {bad_m[:5]}")
+
+    @torch.no_grad()
+    def semantic_target(self, feats: torch.Tensor) -> torch.Tensor:
+        return self.semantic_model(feats).hidden_states[SEM_LAYER].transpose(1, 2)
+
+    def towers(self, wav16: torch.Tensor, feats: torch.Tensor):
+        return self.semantic_target(feats), self.CodecEnc(wav16).transpose(1, 2)
+
+    def encode(self, wav16: torch.Tensor, feats: torch.Tensor):
+        sem, ac = self.towers(wav16, feats)
+        s_enc = self.SemanticEncoder_module(sem)
+        t = min(sem.shape[-1], ac.shape[-1], s_enc.shape[-1])
+        t -= t % 2
+        sem, ac, s_enc = sem[..., :t], ac[..., :t], s_enc[..., :t]
+        s25 = F.avg_pool1d(s_enc, 2, 2)
+        zq_s, idx_s = self.fsq_s(s25.transpose(1, 2))
+        c50 = zq_s.repeat_interleave(2, dim=1)
+        res = ac.transpose(1, 2) - c50.detach()
+        zq_a, idx_a = self.fsq_a(res)
+        return {"idx_s": idx_s, "idx_a": idx_a, "t": t}
+
+    def codes_to_latent(self, idx_s: torch.Tensor, idx_a: torch.Tensor | None) -> torch.Tensor:
+        c50 = self.fsq_s.indices_to_codes(idx_s).repeat_interleave(2, dim=1)
+        if idx_a is None:
+            return c50
+        za = self.fsq_a.indices_to_codes(idx_a)
+        t = min(c50.shape[1], za.shape[1])
+        return c50[:, :t] + za[:, :t]
+
+    def decode_latent(self, h: torch.Tensor) -> torch.Tensor:
+        return self.head(self.backbone(h))[0]
+
+    @torch.no_grad()
+    def decode_codes(self, idx_s, idx_a, mode: str = "full") -> torch.Tensor:
+        if mode == "zero_c":
+            return self.decode_latent(self.fsq_a.indices_to_codes(idx_a))
+        if mode == "zero_ac":
+            return self.decode_latent(self.codes_to_latent(idx_s, None))
+        return self.decode_latent(self.codes_to_latent(idx_s, idx_a))
+
+
+_FE = None
+
+
+def features(wav16_np_list):
+    global _FE
+    if _FE is None:
+        from transformers import AutoFeatureExtractor
+        _FE = AutoFeatureExtractor.from_pretrained(str(W2V_DIR))   # same file as W2V at W2V_REV
+    return _FE(wav16_np_list, sampling_rate=16_000, return_tensors="pt").input_features
+
+
+def to_rate(w, src_sr, tgt):
+    from scipy.signal import resample_poly
+    if int(src_sr) == int(tgt):
+        return np.ascontiguousarray(w, dtype=np.float32)
+    g = math.gcd(int(src_sr), int(tgt))
+    return np.ascontiguousarray(resample_poly(np.asarray(w, np.float64), tgt // g, src_sr // g), dtype=np.float32)
+
+
+def pad_w16(w16: np.ndarray, tail_frames: int) -> np.ndarray:
+    return np.pad(w16, (0, (-len(w16)) % HOP16 + int(tail_frames) * HOP16))
+
+
+@torch.no_grad()
+def encode_wav(model: IndiNet, wav, sr, device, amp: bool = True):
+    w16 = pad_w16(to_rate(wav, sr, 16000), getattr(model, "tail_frames", 0))
+    feats = features([w16]).to(device)
+    x = torch.from_numpy(w16)[None, None].to(device)
+    with torch.autocast("cuda", dtype=torch.bfloat16, enabled=amp and device != "cpu"):
+        e = model.encode(x, feats)
+    return e["idx_s"], e["idx_a"], e["t"]

indicodec/w2v_bert/LICENSE

@@ -0,0 +1,25 @@
+w2v-BERT 2.0 (facebook/w2v-bert-2.0, revision da985ba0987f70aaeb84a80f2851cfac8c697a7b). Its first 16 layers are
+included unchanged in IndiCodec's model.safetensors (tensors semantic_model.*). Licence: MIT, from MIT_LICENSE in
+facebookresearch/seamless_communication.
+
+MIT License
+
+Copyright (c) Meta Platforms, Inc. and affiliates.
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.

indicodec/w2v_bert/config.json

@@ -0,0 +1,82 @@
+{
+  "activation_dropout": 0.0,
+  "adapter_act": "relu",
+  "adapter_kernel_size": 3,
+  "adapter_stride": 2,
+  "add_adapter": false,
+  "apply_spec_augment": false,
+  "architectures": [
+    "Wav2Vec2BertModel"
+  ],
+  "attention_dropout": 0.0,
+  "bos_token_id": 1,
+  "classifier_proj_size": 768,
+  "codevector_dim": 768,
+  "conformer_conv_dropout": 0.1,
+  "contrastive_logits_temperature": 0.1,
+  "conv_depthwise_kernel_size": 31,
+  "ctc_loss_reduction": "sum",
+  "ctc_zero_infinity": false,
+  "diversity_loss_weight": 0.1,
+  "eos_token_id": 2,
+  "feat_proj_dropout": 0.0,
+  "feat_quantizer_dropout": 0.0,
+  "feature_projection_input_dim": 160,
+  "final_dropout": 0.1,
+  "hidden_act": "swish",
+  "hidden_dropout": 0.0,
+  "hidden_size": 1024,
+  "initializer_range": 0.02,
+  "intermediate_size": 4096,
+  "layer_norm_eps": 1e-05,
+  "layerdrop": 0.1,
+  "left_max_position_embeddings": 64,
+  "mask_feature_length": 10,
+  "mask_feature_min_masks": 0,
+  "mask_feature_prob": 0.0,
+  "mask_time_length": 10,
+  "mask_time_min_masks": 2,
+  "mask_time_prob": 0.05,
+  "max_source_positions": 5000,
+  "model_type": "wav2vec2-bert",
+  "num_adapter_layers": 1,
+  "num_attention_heads": 16,
+  "num_codevector_groups": 2,
+  "num_codevectors_per_group": 320,
+  "num_hidden_layers": 16,
+  "num_negatives": 100,
+  "output_hidden_size": 1024,
+  "pad_token_id": 0,
+  "position_embeddings_type": "relative_key",
+  "proj_codevector_dim": 768,
+  "right_max_position_embeddings": 8,
+  "rotary_embedding_base": 10000,
+  "tdnn_dilation": [
+    1,
+    2,
+    3,
+    1,
+    1
+  ],
+  "tdnn_dim": [
+    512,
+    512,
+    512,
+    512,
+    1500
+  ],
+  "tdnn_kernel": [
+    5,
+    3,
+    3,
+    1,
+    1
+  ],
+  "torch_dtype": "float32",
+  "transformers_version": "4.37.0.dev0",
+  "use_intermediate_ffn_before_adapter": false,
+  "use_weighted_layer_sum": false,
+  "vocab_size": null,
+  "xvector_output_dim": 512,
+  "output_hidden_states": true
+}

indicodec/w2v_bert/preprocessor_config.json

@@ -0,0 +1,11 @@
+{
+  "feature_extractor_type": "SeamlessM4TFeatureExtractor",
+  "feature_size": 80,
+  "num_mel_bins": 80,
+  "padding_side": "right",
+  "padding_value": 1,
+  "processor_class": "Wav2Vec2BertProcessor",
+  "return_attention_mask": true,
+  "sampling_rate": 16000,
+  "stride": 2
+}

model.safetensors

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