jhcodec_1.4m / README.md
junhyeok lee
JHCodec-1.4M checkpoint and model card
24910ef
|
Raw
History Blame Contribute Delete
5.57 kB
---
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