File size: 2,907 Bytes
4aa51fc | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 | #!/usr/bin/env python3
"""WideCodec — self-contained inference (44.1 kHz neural codec, depth-20 decoder).
Encode -> decode any audio to a 44.1 kHz reconstruction. The encoder ingests mono
16 kHz; the frozen FSQ codebook (single codebook, 50 tokens/s, ~0.8 kbps) yields the
code stream; the depth-20 decoder renders it back to 44.1 kHz.
This script + the bundled `neucodec/` package are ALL you need — no other source
repo. Weights (`pytorch_model.bin`) are pulled from the HF model repo on first run.
pip install torch transformers huggingface_hub local-attention einops librosa soundfile
huggingface-cli login # if the repo is private (or export HF_TOKEN=hf_...)
python infer_widecodec.py --input my.wav --out-dir out
python infer_widecodec.py --input folder/ --out-dir out # batch a directory
"""
import argparse
import glob
import os
import sys
# make the bundled `neucodec/` package importable regardless of CWD
sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
os.environ.setdefault("HF_HUB_DISABLE_XET", "1") # plain HTTPS: Xet can hang on some networks
import librosa
import soundfile as sf
import torch
from neucodec import NeuCodec
AUDIO_EXTS = (".wav", ".mp3", ".flac", ".m4a", ".ogg", ".opus")
def main():
ap = argparse.ArgumentParser()
ap.add_argument("--input", required=True, help="audio file OR a directory of audio")
ap.add_argument("--out-dir", default="out")
ap.add_argument("--repo", default="Scicom-intl/WideCodec")
ap.add_argument("--depth", type=int, default=20, help="decoder depth (WideCodec=20)")
ap.add_argument("--device", default="cuda" if torch.cuda.is_available() else "cpu")
a = ap.parse_args()
os.makedirs(a.out_dir, exist_ok=True)
print(f"[infer] loading {a.repo} (decoder_depth={a.depth}) on {a.device} …")
model = NeuCodec._from_pretrained(
model_id=a.repo, decoder_depth=a.depth, token=os.environ.get("HF_TOKEN")
).eval().to(a.device)
sr_out = model.sample_rate
print(f"[infer] output sample_rate = {sr_out}")
if os.path.isdir(a.input):
files = sorted(f for f in glob.glob(os.path.join(a.input, "*"))
if f.lower().endswith(AUDIO_EXTS))
else:
files = [a.input]
print(f"[infer] {len(files)} file(s)")
for f in files:
wav16, _ = librosa.load(f, sr=16000, mono=True) # frozen encoder wants 16 kHz mono
x = torch.from_numpy(wav16).float().view(1, 1, -1).to(a.device)
with torch.no_grad():
codes = model.encode_code(x)
wav = model.decode_code(codes).squeeze().detach().cpu().float().numpy()
base = os.path.splitext(os.path.basename(f))[0]
out = os.path.join(a.out_dir, f"{base}_44k.wav")
sf.write(out, wav, sr_out)
print(f"[infer] {f} -> {out} ({len(wav) / sr_out:.2f}s @ {sr_out} Hz)")
if __name__ == "__main__":
main()
|