Views
No views yet
sw2v_120000.pt). To enhance noise robustness for future applications, we incorporated noise augmentation during its training. For the checkpoint used in the paper's main experiments, see jhcodec/sw2v_60k.config.w2v.mlp_in.in_features), so the input length must be a multiple of 320. It returns one 1024-dim representation per frame.torch==2.6.0+cu124 and torch==2.9.1+cu128)flash-attn==2.7.4.post1 and flash-attn==2.8.3)1import torch
2import torch.nn.functional as F
3import torchaudio
4from jhcodec.utils import load_pretrained_sw2v
5
6DEVICE = 'cuda'
7SAMPLE_RATE = 16000
8FRAME_SIZE = 320 # 20 ms hop; input length must be a multiple of this
9
10w2v = load_pretrained_sw2v(repo_id='jhcodec/sw2v_120k').to(DEVICE).eval()
11
12x, sr = torchaudio.load('input.wav')
13if sr != SAMPLE_RATE:
14 x = torchaudio.transforms.Resample(sr, SAMPLE_RATE)(x)
15x = x[0, :].view(1, -1).to(DEVICE) # [1, T], mono
16if x.shape[1] % FRAME_SIZE != 0:
17 x = F.pad(x, (0, FRAME_SIZE - x.shape[1] % FRAME_SIZE))
18
19# encode is already decorated with @torch.no_grad()
20features, _ = w2v.encode(x, inference_cache=None) # [1, T//320, 1024]inference_cache back in on every call, starting from None.1cache = None
2feats = []
3for i in range(0, x.shape[1], FRAME_SIZE):
4 frame_feat, cache = w2v.encode(x[:, i:i + FRAME_SIZE], inference_cache=cache)
5 feats.append(frame_feat) # each [1, 1, 1024]
6features = torch.cat(feats, dim=1) # [1, T//320, 1024]1import omegaconf
2import jhcodec.utils as utils
3from jhcodec.model.sw2v import AudioEncoder
4
5config = omegaconf.OmegaConf.load('config.json')
6w2v = AudioEncoder(config.w2v, training=False)
7utils.load_checkpoint(w2v, None, None, 'sw2v_120000.pt', strict_model=True)
8w2v = w2v.to(DEVICE).eval()AudioEncoderCudaGraph in jhcodec/model/sw2v_cudagraph.py, whose state_dict is identical to AudioEncoder), see the GitHub repository README.1@article{jhcodec2026,
2 title={Reconstruct! Don't Encode: Self-Supervised Representation Reconstruction Loss for High-Intelligibility and Low-Latency Streaming Neural Audio Codec},
3 author={Anonymous},
4 journal={arXiv preprint arXiv:2603.05887},
5 year={2026}
6}