JHCodec-1.4M checkpoint and model card

.gitattributes

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

README.md

@@ -0,0 +1,147 @@
+---
+license: mit
+pipeline_tag: audio-to-audio
+tags:
+- audio
+- audio-codec
+- speech
+- representation-learning
+---
+
+# Model Card for JHCodec
+
+JHCodec is a pure Transformer decoder-based neural audio codec with residual vector quantization (RVQ). It achieves state-of-the-art performance with minimal latency and high intelligibility through self-supervised representation reconstruction (SSRR) loss.
+
+- **Paper:** [Reconstruct! Don't Encode: Self-Supervised Representation Reconstruction Loss for High-Intelligibility and Low-Latency Streaming Neural Audio Codec](https://huggingface.co/papers/2603.05887)
+- **GitHub Repository:** [https://github.com/jhcodec843/jhcodec](https://github.com/jhcodec843/jhcodec)
+- **Demo:** [https://jhcodec843.github.io/jhcodec/](https://jhcodec843.github.io/jhcodec/)
+- **License:** MIT
+
+## Model Details
+
+This checkpoint corresponds to the **JHCodec-1.4M** model variant (`jhcodec_mimi_1400000.pt`). JHCodec uses a self-supervised representation reconstruction loss to improve codec training, enhancing intelligibility by reconstructing distilled self-supervised representations from codec outputs. It features a zero-lookahead architecture designed for real-time streaming deployment.
+
+The model operates on 16 kHz mono audio in frames of `FRAME_SIZE = 320` samples (20 ms), so the input length must be a multiple of 320.
+
+### Requirements
+
+- Python >= 3.10
+- PyTorch/TorchAudio with CUDA support (tested with `torch==2.6.0+cu124` and `torch==2.9.1+cu128`)
+- [omegaconf==2.3.0](https://omegaconf.readthedocs.io/en/2.3_branch/)
+- [Flash-Attention](https://github.com/Dao-AILab/flash-attention) (required for the reported performance; tested with `flash-attn==2.7.4.post1` and `flash-attn==2.8.3`)
+- [huggingface_hub](https://huggingface.co/docs/huggingface_hub/index) — only if you auto-download the official checkpoint
+
+**Note: Running on CPU currently leads to degraded reconstruction quality.**
+
+## Usage
+
+### Inference via CLI
+
+Download this checkpoint and point `--checkpoint` at it (`--from_hf` fetches the 1M
+variant from `jhcodec/jhcodec`):
+
+```bash
+python jhcodec/inference.py \
+    --config config/config_mimi_recon.json \
+    --checkpoint jhcodec_mimi_1400000.pt \
+    --input_file /path/to/input.wav \
+    --output_file /path/to/output.wav \
+    --num_codebooks 8 \
+    --device 'cuda'
+```
+
+### Use in Python (offline, whole utterance at once)
+
+```python
+import torch
+import torch.nn.functional as F
+import torchaudio
+from jhcodec.utils import load_pretrained_jhcodec
+
+DEVICE = 'cuda'
+SAMPLE_RATE = 16000
+FRAME_SIZE = 320       # 20 ms hop; input length must be a multiple of this
+NUM_CODEBOOKS = 8      # <= config.model.rvq.num_codebooks
+
+codec = load_pretrained_jhcodec(repo_id='jhcodec/jhcodec_1.4m').to(DEVICE).eval()
+
+x, sr = torchaudio.load('input.wav')
+if sr != SAMPLE_RATE:
+    x = torchaudio.transforms.Resample(sr, SAMPLE_RATE)(x)
+x = x[0, :].view(1, -1).to(DEVICE)               # [1, T], mono
+if x.shape[1] % FRAME_SIZE != 0:
+    x = F.pad(x, (0, FRAME_SIZE - x.shape[1] % FRAME_SIZE))
+
+# encode/decode are already decorated with @torch.no_grad()
+n_codebooks = torch.tensor([NUM_CODEBOOKS], device=DEVICE)
+indices, _ = codec.encode(x, n_codebooks, inference_cache=None)       # [1, T//320, NUM_CODEBOOKS]
+decoded, _ = codec.decode(indices, n_codebooks, inference_cache=None) # [1, T]
+
+torchaudio.save('output.wav', decoded.detach().cpu(), SAMPLE_RATE)
+```
+
+### Use in Python (streaming, frame by frame)
+
+Pass the returned `inference_cache` back in on every call. The encoder and the decoder each keep their own cache, so use two separate variables and start both at `None`.
+
+```python
+encoder_cache = None
+indices = []
+for i in range(0, x.shape[1], FRAME_SIZE):
+    frame_indices, encoder_cache = codec.encode(
+        x[:, i:i + FRAME_SIZE], n_codebooks, inference_cache=encoder_cache)
+    indices.append(frame_indices)                 # each [1, 1, NUM_CODEBOOKS]
+
+decoder_cache = None
+chunks = []
+for frame_indices in indices:
+    audio_chunk, decoder_cache = codec.decode(
+        frame_indices, n_codebooks, inference_cache=decoder_cache)
+    chunks.append(audio_chunk)                    # each [1, 320]
+decoded = torch.cat(chunks, dim=1)                # [1, T]
+```
+
+To load a local checkpoint instead of the Hugging Face one:
+
+```python
+import omegaconf
+import jhcodec.utils as utils
+from jhcodec.model.codec import JHCodecMimi
+
+config = omegaconf.OmegaConf.load('config/config_mimi_recon.json')
+codec = JHCodecMimi(config.model, training=False)
+utils.load_checkpoint(codec, None, None, 'jhcodec_mimi_1400000.pt', strict_model=True)
+codec = codec.to(DEVICE).eval()
+```
+
+For CUDA-graph per-frame streaming (`JHCodecMimiCudaGraph`, whose `state_dict` is identical to `JHCodecMimi`), see the [GitHub repository README](https://github.com/jhcodec843/jhcodec).
+
+## Intended Use
+
+- Real-time low-latency audio codecs for speech-to-speech models
+- Research into neural codecs and generative modeling
+- Serving as a neural front-end for speech recognition or synthesis pipelines
+- Compressing large audio datasets
+
+### Out-of-Scope Use
+
+- Any malicious, deceptive, or privacy-violating applications
+
+## Training Details
+
+Please refer to the GitHub repository README.
+
+## Citation
+
+```bibtex
+@article{jhcodec2026,
+  title={Reconstruct! Don't Encode: Self-Supervised Representation Reconstruction Loss for High-Intelligibility and Low-Latency Streaming Neural Audio Codec},
+  author={Anonymous},
+  journal={arXiv preprint arXiv:2603.05887},
+  year={2026}
+}
+```
+
+## Authors
+
+Anonymous, Submitted to Interspeech 2026

config.json

@@ -0,0 +1,139 @@
+{
+    "model": {
+        "mlp_in":{
+            "in_features": 320,
+            "hidden_features": 768,
+            "out_features": 1024,
+            "compute_dtype": "float32"
+        },
+        "mlp_out":{
+            "in_features": 1024,
+            "hidden_features": 768,
+            "out_features": 320,
+            "compute_dtype": "float32"
+        },
+        "decoder": {
+            "n_layers": 8,
+            "n_embd": 1024,
+            "n_hidden": 4096,
+            "n_heads": 16,
+            "head_dim": 64,
+            "compute_dtype": "float32",
+            "window_size": 15,
+            "dropout_rate": 0.1,
+            "drop_path_rate": 0.1
+        },
+        "encoder": {
+            "n_layers": 8,
+            "n_embd": 1024,
+            "n_hidden": 4096,
+            "n_heads": 16,
+            "head_dim": 64,
+            "compute_dtype": "float32",
+            "window_size": 15,
+            "dropout_rate": 0.1,
+            "drop_path_rate": 0.1
+        },
+        "rvq": {
+            "type": "mimi",
+            "num_codebooks": 8,
+            "codebook_size": 1024,
+            "embedding_dim": 1024,
+            "latent_dim": 16,
+            "updown_linears": true,
+            "codebook_weight_dtype": "float32"
+        },
+        "training": {
+            "quantizer_dropout": 0.5,
+            "encoder_mask_rate": 0.1,
+            "decoder_mask_rate": 0.1,
+            "noise_augmentation": 0.1
+        }
+    },
+    "w2v":{
+        "mlp_in":{
+            "in_features": 320,
+            "hidden_features": 768,
+            "out_features": 1024,
+            "compute_dtype": "float32"
+        },
+        "encoder": {
+            "n_layers": 8,
+            "n_embd": 1024,
+            "n_hidden": 4096,
+            "n_heads": 16,
+            "head_dim": 64,
+            "compute_dtype": "float32",
+            "window_size": 15,
+            "dropout_rate": 0.1
+        },
+        "rvq": {
+            "num_codebooks": 8,
+            "codebook_size": 1024,
+            "embedding_dim": 1024,
+            "latent_dim": 1024,
+            "updown_linears": false,
+            "codebook_weight_dtype": "float32"
+        },
+        "training": {
+            "noise_masking": 0.1
+        }
+    },
+    "training":{
+        "resume": false,
+        "loss_type": "l1",
+        "strict_model": false,
+        "load_discriminator": true,
+        "learning_rate": 1e-4,
+        "weight_decay": 1e-2,
+        "discriminator_start_steps": 10000,
+        "discriminator_segment_duration": 2.56,
+        "masking_stop_steps": 100000,
+        "apply_apa": false,
+        "warmup_steps": 1000,
+        "min_lr": 1e-6,
+        "num_epochs": 100,
+        "use_continuous": 0.1,
+        "max_grad_norm": 1000.0,
+        "batch_size": 42,
+        "gradient_accumulation_steps": 1,
+        "num_workers": 4,
+        "use_phaseaug": true,
+        "init_dataset": false,
+        "profile": false,
+        "verbose_grad_norm": false,
+        "verbose_norm_threshold_max": 5.0,
+        "verbose_norm_threshold_min": 0.001,
+        "verbose_paramter_norm": false,
+        "use_discriminator": true,
+        "encoder_checkpoint": null,
+        "rect_checkpoint": null,
+        "decoder_checkpoint": null,
+        "codec_checkpoint": null,
+        "sw2v_checkpoint": "/data/jhcodec/sw2v/cossim/checkpoints/checkpoint_60000.pt"
+    },
+    "loss":{
+        "recon_loss_weight": 0.1,
+        "vq_loss_weight": 1.0,
+        "commit_loss_weight": 0.1,
+        "fm_loss_weight": 1.0,
+        "adv_loss_weight": 1.0,
+        "w2v_loss_weight": 1.0,
+        "semantic_loss_weight": 1.0
+    },
+    "data": {
+        "audio_dir": "/data",
+        "sample_rate": 16000,
+        "segment_duration": 10.24,
+        "cache_dir": "/data/dataloader/v9" 
+    },
+    "logging": {
+        "log_interval": 100,
+        "save_interval": 500,
+        "eval_interval": 3000,
+        "experiment_dir": "/data/jhcodec/{experiment_name}",
+        "checkpoint_dir": "/data/jhcodec/{experiment_name}/checkpoints",
+        "tensorboard_dir": "/data/jhcodec/{experiment_name}/tensorboard",
+        "n_samples": 3
+      }
+}

jhcodec_mimi_1400000.pt

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