Official Implementation of JHCodec

Reconstruct! Don't Encode: Self-Supervised Representation Reconstruction Loss for High-Intelligibility and Low-Latency Streaming Neural Audio Codec

JHCodec is a pure Transformer decoder based neural audio codec with residual vector quantization. It shows state-of-the-art performance with minimal latency.

New Checkpoint

JHCodec-1.4M (jhcodec/jhcodec_1.4m, jhcodec_mimi_1400000.pt) is now the default checkpoint, trained for 1.4M steps. It is a drop-in replacement for the 1M release — same architecture, same state_dict keys, same 20 ms / 320-sample framing.

python jhcodec/inference.py --from_hf \
    --input_file /path/to/input.wav --output_file /path/to/output.wav

--from_hf downloads it automatically. To stay on the previous 1M checkpoint, pass --repo_id jhcodec/jhcodec (or load_pretrained_jhcodec(repo_id='jhcodec/jhcodec')). See Official Checkpoints for the full list.

Overview

This repository contains the implementation for training and inference neural audio codecs with end-to-end training capabilities. The codec supports:

• Multiple RVQ architectures (DAC, MIMI)
• Supports end-to-end training leveraging (distilled) w2v-bert-2.0 semantic features
• SSRR and non-SSRR variants

TODO

• Revise Readme
• Upload checkpoint
• Upload to HuggingFace
• Auto-download from HuggingFace
• Upload to PyPI (probably after the review)
• Make non-anonymous (after the review)

Installation

pip install -e .

Requirements

• Python >= 3.10 (required for using the X | None union type syntax in type hints; see PEP 604), or manually remove this syntax if using an older Python version
• PyTorch/TorchAudio with CUDA support: tested with torch==2.6.0+cu124 and torch==2.9.1+cu128
• omegaconf==2.3.0: for configuration management
• Flash-Attention: For fast train/inference. We tested with flash-attn==2.7.4.post1 and flash-attn==2.8.3.
• HF transformers: Required only for running baselines and w2v-bert2.0. JHCodec inference has no dependency on it.
• huggingface_hub: Required only if you want to use --from_hf to auto-download official checkpoints/configs.
• MLX: Apple Silicon only, and only for jhcodec/model/codec_mlx.py. Not needed for the CUDA or CPU paths.
• A CUDA toolkit with nvcc plus ninja (pip install ninja): only for the fused CUDA step kernels, which are JIT-compiled on first use. Without them the code falls back to the unfused path automatically.

We have provided a Shell Script to help set up the environment. PLEASE DO NOT RUN It Directly. INSTEAD, REVIEW THE SCRIPT AND MODIFY IT AS NEEDED FOR YOUR SYSTEM.

OUR MODEL REQUIRES ONLY THE MINIMUM DEPENDENCIES LISTED ABOVE.

Fixed: CPU reconstruction quality. Earlier releases produced degraded audio when run without CUDA. The cause was the pure-torch rotary fallback: it built the rotated tensor but never wrote it back, so on any non-CUDA device the model ran with no positional information at all. The CUDA path was never affected — it uses the Triton kernel, which was always correct. On a 16 kHz test clip at 8 codebooks, the fix moves CPU round-trip SNR from 4.6 dB to 12.8 dB (waveform correlation 0.853 to 0.975).

For training

• PhaseAug
• alias-free-torch

To install both required libraries, run:

pip install omegaconf==2.3.0
pip install alias-free-torch==0.0.6 phaseaug

Flash-Attention should be installed carefully. Please read the official README.

Official Checkpoints

Project Structure

codec_paper/
├── jhcodec/                      # Main package
│   ├── model/                    # Model implementations
│   │   ├── codec.py              # Main codec models (JHCodec, JHCodecMimi)
│   │   ├── sw2v.py               # streaming wav2vec encoder
│   │   ├── attention.py          # Transformer decoder stack, RoPE attention, KV cache
│   │   ├── discriminator.py      # Discriminator for adversarial training
│   │   ├── vq.py                 # Vector quantization modules
│   │   ├── attention_cudagraph.py # CUDA-graph capture for the transformer stack
│   │   ├── codec_cudagraph.py    # CUDA-graph streaming codec (JHCodecMimiCudaGraph)
│   │   ├── sw2v_cudagraph.py     # CUDA-graph streaming sw2v encoders
│   │   └── codec_mlx.py          # MLX port of JHCodecMimi inference (Apple Silicon)
│   ├── kernel/                   # Custom kernels
│   │   ├── rotary_kernel.py      # Rotary positional embedding's kernel, adopted from FlashAttn
│   │   ├── vq_kernel.py          # Vector quantization kernel (Triton)
│   │   ├── fused_step.py         # JIT loader for the fused CUDA step kernels
│   │   └── cuda/                 # Raw CUDA sources for the per-frame decode step
│   │       ├── fused_step_kernels.cu
│   │       └── fused_step_binding.cpp
│   ├── loss/                     # Custom loss functions
│   │   └── multiscalemelspec.py  # Implements MultiScaleMelSpectrogramLoss used for perceptual audio training
│   ├── train_codec_e2e_w2v.py    # End-to-end training script
│   ├── decode_eval.py            # Decoding and evaluation script
│   └── dataloader.py             # Data loading utilities
├── config/                       # Configuration files
│   ├── config_dac_norecon.json   # without SED without reconstruction
│   ├── config_dac_recon.json     # without SED with reconstruction
│   ├── config_mimi_norecon.json  # with SED without reconstruction
│   └── config_mimi_recon.json    # with SED with reconstruction
└── setup.py

Training

Training Command:

Data Preparation Using the main of jhcodec/dataloader.py

The main block of dataloader.py demonstrates how to construct and inspect an AudioDataset:

from jhcodec.dataloader import AudioDataset, collate_fn
from torch.utils.data import DataLoader

dataset = AudioDataset(
    audio_dir='./data',                  # Path to your data
    sample_rate=16000,
    segment_duration=10.24,
    training=True,
    init_dataset=False,                  # Use True to scan files initially (slow), or False to load from cache
    cache_dir='cache_dir/dataloader/v9', # location of the cache
    use_mel=False,                       # Set True to return also Mel features
)

Notes:

• Initial dataset caching may take a while; once done, restart with init_dataset=False for faster loading.
• Requires all dependencies (see top part of jhcodec/dataloader.py).
• You can add a custom dataset by modifying the dictionary at the top of dataloader.py.

For DAC with reconstruction:

python jhcodec/train_codec_e2e_w2v.py \
    --experiment_name paper/dac_recon \
    --config config/config_dac_recon.json \
    --resume # if resume

For MIMI with reconstruction:

python jhcodec/train_codec_e2e_w2v.py \
    --experiment_name paper/mimi_recon \
    --config config/config_mimi_recon.json \
    --resume

Available Configurations:

• config_dac_norecon.json - DAC without reconstruction
• config_dac_recon.json - DAC with reconstruction
• config_mimi_norecon.json - MIMI without reconstruction
• config_mimi_recon.json - MIMI with reconstruction (main)

Training Parameters

Key training parameters (configurable in JSON config files):

• learning_rate: 1e-4
• batch_size: 42
• num_epochs: 100
• warmup_steps: 1000
• discriminator_start_steps: 10000
• Loss weights for reconstruction, VQ, commit, feature matching, and adversarial losses

Decoding

Single Files

python jhcodec/inference.py \
    --from_hf \
    --input_file /path/to/input.wav \
    --output_file /path/to/output.wav \
    --num_codebooks 8 \
    --device 'cuda'

--from_hf downloads jhcodec/jhcodec_1.4m by default. Pass --repo_id jhcodec/jhcodec to use the 1M model instead.

Single Files (load from a local checkpoint)

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'

Multiple Files

python jhcodec/decode_eval.py \
    --config config/config_dac_norecon.json \
    --checkpoint /path/to/checkpoint_300000.pt \
    --name jhcodec_dac_norecon \
    --glob_pattern "/path/to/audio/*.wav" \
    --out_dir "out_dir" \
    --hierarchy 4

Arguments:

• --config: Path to configuration file
• --checkpoint: Path to model checkpoint
• --name: Model name for output directory
• --glob_pattern: Glob pattern for input audio files
• --hierarchy: Depth of quantization hierarchy (default: 4)
• --out_dir: Output directory

Supported Datasets

The decoding script supports various audio datasets:

• LibriSpeech: /data/LibriSpeech/test-other/*/*/*.flac
• TITW: /data/titw/titw_hard/test/*.wav
• MLS: /data/MLS/mls_*/test/audio/*/*/*.flac

Use in Python

See jhcodec/inference.py for the full script. 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.

Offline (whole utterance at once)

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)

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.

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:

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()

CUDA Graph Streaming

For per-frame streaming, the Python/launch overhead of many tiny kernels dominates the step time. jhcodec/model/codec_cudagraph.py captures the whole per-frame pipeline into a single CUDA graph and replays it, which removes that overhead.

JHCodecMimiCudaGraph subclasses JHCodecMimi and only swaps the transformer stacks for graph-capable equivalents, so the state_dict is identical — an existing checkpoint loads unchanged, and the eager encode / decode methods still work as a reference.

Requirements and constraints:

• CUDA + flash-attn (CudaGraphStream asserts on this; there is no CPU path).
• Strict order: construct → load checkpoint → .to(device) → .eval() → make_*_stream. RVQMimi.train(False) (i.e. .eval()) calls register_up_vq(), which bakes the current codebook weights through the up linears into lookup-only VQ modules. Loading a checkpoint after .eval() leaves those tables holding the randomly initialized weights, and load_state_dict(..., strict=True) will also fail on the extra up_vqs.* / semantic_up_vq.* keys the model now has. The module itself warns about this: DO NOT USE EVAL BEFORE LOADING CHECKPOINT. Capture must then come after .eval(), since the graph traces whichever decode path is active. If you ever call .train() again the tables are dropped (up_vqs = None) and any captured stream is stale — re-eval() and re-capture.
• Each captured graph is specialized to one (batch, n_codebooks, max_frames). Call make_*_stream again to specialize differently.
• The KV cache is bounded by max_frames; exceeding it raises. Call stream.reset() to start a new utterance.
• step() returns the static output buffer, overwritten on the next step(). Pass clone=True if you keep the result around.

import torch
import omegaconf
import jhcodec.utils as utils
from jhcodec.model.codec_cudagraph import JHCodecMimiCudaGraph

DEVICE = torch.device('cuda')
config = omegaconf.OmegaConf.load('config/config_mimi_recon.json')

# 1. construct
model = JHCodecMimiCudaGraph(config.model, training=False)
# 2. load -- same checkpoint as JHCodecMimi, the state_dict keys are identical.
#    This MUST happen before .eval().
utils.load_checkpoint(model, None, None, 'jhcodec_mimi_1400000.pt', strict_model=True)
# 3. move, then 4. eval() -- registers the fused up+VQ tables from the loaded weights
model = model.to(DEVICE).eval()

FRAME = model.feature_size        # 320 samples
N_CB = model.num_codebooks
B, MAX_FRAMES = 1, 1600           # 1600 frames = 32 s at 20 ms/frame

# 5. capture, only now that the model is loaded and in eval mode
# encode: step(x=[B, FRAME]) -> indices [B, 1, N_CB]
enc_stream = model.make_encode_stream(B, MAX_FRAMES, N_CB, DEVICE)
# decode: step(indices=[B, 1, N_CB]) -> audio [B, FRAME]
dec_stream = model.make_decode_stream(B, MAX_FRAMES, N_CB, DEVICE)

x = torch.randn(B, 200 * FRAME, device=DEVICE)   # your streaming source

indices = []
for t in range(x.shape[1] // FRAME):
    indices.append(enc_stream.step(x=x[:, t * FRAME:(t + 1) * FRAME], clone=True))

chunks = []
for frame_indices in indices:
    chunks.append(dec_stream.step(indices=frame_indices, clone=True))
decoded = torch.cat(chunks, dim=1)               # [B, T]

enc_stream.reset()   # rewind both streams before the next utterance
dec_stream.reset()

To use an already-loaded eager model (e.g. from load_pretrained_jhcodec, which returns a plain JHCodecMimi), copy the weights across — the keys match:

from jhcodec.utils import load_pretrained_jhcodec

eager = load_pretrained_jhcodec(repo_id='jhcodec/jhcodec_1.4m')
model = JHCodecMimiCudaGraph(config.model, training=False)
model.load_state_dict(eager.state_dict())
model = model.to(DEVICE).eval()

The SW2V encoders have the same interface via AudioEncoderCudaGraph / AudioEncoderPhonemicCudaGraph in jhcodec/model/sw2v_cudagraph.py — make_encode_stream(batch, max_frames, device), then step(x=[B, in_features]) returning [B, 1, rvq_dim].

Running the modules directly checks graph output against the eager frame-by-frame path and prints the per-step speedup:

python jhcodec/model/attention_cudagraph.py   # decoder stack: graph vs eager max diff
python jhcodec/model/codec_cudagraph.py       # full codec: index mismatch, audio diff, latency
python jhcodec/model/sw2v_cudagraph.py        # sw2v encoder

Fused CUDA Step Kernels

On top of the graph capture, the per-frame decode step runs through a small raw-CUDA extension in jhcodec/kernel/cuda/, loaded lazily by jhcodec/kernel/fused_step.py. It is compiled on first use with torch's cpp_extension.load (cached under ~/.cache/torch_extensions), so it needs nvcc and ninja available at runtime. Nothing about this is required: if the extension cannot be built the loader logs a warning and the unfused graph path is used instead, with identical outputs.

What is fused:

• fused_attn_step — rotary + KV-append + windowed attention in one kernel, replacing the triton rotary, the casts/copies, and FA2's splitkv+combine for a 1-token query.
• bf16 warp-per-row GEMV with fused epilogues (SwiGLU, residual-add + norm), replacing the fp32 cuBLAS GEMVs that dominated the step (the weight read is the bottleneck at M=1).
• fused_rvq_encode / fused_rvq_decode — the whole semantic + residual codebook loop as one persistent block per token, instead of ~60 tiny kernels. RVQMimi.encode uses this in eager mode too.
• LayerNorm and RMSNorm, with the norm affine and the residual-branch scales folded into the neighbouring projections at capture time (exact, done in fp32).

Attention spans are routed per layer at capture: the custom attention tail is linear in the key count, so layers whose worst-case span exceeds JHCODEC_FUSED_ATTN_MAX_KEYS fall back to FA2's splitkv kvcache kernel while keeping the GEMV/fold/add+LN fusions.

Measured on an idle H200 with config_mimi_recon.json, batch 1: the decode step goes from 655 us (graph baseline) to 279 us — 2.3x over the captured graph and ~9x over eager — with reconstruction SNR/L1 unchanged on the trained checkpoint.

inference_optimization_path.md records the full path, the per-stage measurements, and the negative results.

Apple Silicon (MLX)

jhcodec/model/codec_mlx.py is an MLX port of JHCodecMimi inference for Apple Silicon. It mirrors the torch model operation for operation, so a torch checkpoint loads unchanged and produces the same indices and audio — there is no separate MLX checkpoint. Inference only; training stays on the torch path. Both the whole-utterance and the frame-by-frame paths are implemented.

import mlx.core as mx
import omegaconf
from jhcodec.model.codec_mlx import JHCodecMimiMLX

config = omegaconf.OmegaConf.load('config/config_mimi_recon.json')
model = JHCodecMimiMLX.from_torch_checkpoint('jhcodec_mimi_1400000.pt', config.model)
# or JHCodecMimiMLX.from_torch_state_dict(state_dict, config.model)

indices = model.encode(audio)          # [B, T // 320, n_codebooks], mx.int32
audio_out = model.decode(indices)      # [B, T]

model.fuse()                           # optional, exact: folds norm affine + branch scales
model.quantize(bits=8)                 # optional, ~2x faster per frame; must come after fuse()

state = model.decoder_state()          # frame by frame
frame = model.decode_step(indices[:, t:t + 1], state)      # [B, 320]

• fuse() folds the LayerNorm affine and the residual-branch scales into the neighbouring matmuls. Exact — the same arithmetic, done once on the weights rather than every frame.
• quantize(bits=8) quantizes only the large projections and routes them through mx.quantized_matmul, roughly halving the per-frame cost. The RVQ codebooks, the small latent_dim=16 down/up projections, and every bias / norm parameter are left in full precision, since quantization error there lands straight on the nearest-code decision. It must run after fuse(), which rewrites the weights it packs.
• Measured on an M3 Pro, batch 1, per 20 ms frame: ~6 ms encode and ~6 ms decode. fuse() and the ring buffer are exactness/memory wins rather than speedups — a single-frame step is limited by GEMV efficiency at M=1, not by the elementwise ops around it.

Running the module directly checks the MLX output against torch and the streaming path against the whole-utterance path:

python jhcodec/model/codec_mlx.py

Configuration

Configuration files are JSON-based and include:

• Model Architecture: Encoder/decoder layers, attention heads, embedding dimensions
• Vector Quantization: Codebook size, number of codebooks, embedding dimensions
• Training: Learning rate, batch size, loss weights, discriminator settings
• Data: Sample rate, segment duration, data directories
• Logging: Checkpoint intervals, tensorboard settings

Example configuration structure:

{
    "model": {
        "encoder": {...},
        "decoder": {...},
        "rvq": {
            "type": "dac",
            "num_codebooks": 8,
            "codebook_size": 1024
        }
    },
    "training": {...},
    "loss": {...},
    "data": {...}
}

Citation

If you find this work useful, please cite:

@article{lee2026reconstruct,
  title={Reconstruct! Don't Encode: Self-Supervised Representation Reconstruction Loss for High-Intelligibility and Low-Latency Streaming Neural Audio Codec},
  author={Lee, Junhyeok and He, Xiluo and Lee, Jihwan and Wang, Helin and Narayanan, Shrikanth and Thebaud, Thomas and Moro-Velazquez, Laureano and Villalba, Jes{\'u}s and Dehak, Najim},
  journal={arXiv preprint arXiv:2603.05887},
  year={2026}
}

Main Contact

Anonymous. Contact: jhcodec843@gmail.com Submitted to Interspeech 2026

References

• EnCodec
• DAC
• Mimi
• AudioCraft
• HiFiGAN
• PhaseAug
• AliasFreeTorch
• FlashAttention
• MLX

License

MIT License