79 lines
3.0 KiB
Python
79 lines
3.0 KiB
Python
|
#!/usr/bin/env python3
|
||
|
|
"""
|
||
|
|
SNAC audio encoder -- turns a real audio/video file into `<snac_N>` tokens
|
||
|
|
(listen format, 3 tokens/base-frame), the reverse of tools/decode/decode_snac.py.
|
||
|
|
Reuses the exact encode_listen() logic already in pipeline_pose/snac_finevideo.py
|
||
|
|
(unchanged since inception) rather than re-deriving it.
|
||
|
|
|
||
|
|
This model (vla-1.7b-qwen3-v2) only ever saw listen-format audio wrapped in
|
||
|
|
the generic <snac> tag -- NOT the newer (2026-07-23) <listen>/<speak>
|
||
|
|
convention or the speak-format L2 tokens. Output here matches that: always
|
||
|
|
listen-format, always <snac> wrapper.
|
||
|
|
|
||
|
|
Usage:
|
||
|
|
python tools/encode/encode_snac.py --input clip.wav --output tokens.txt
|
||
|
|
python tools/encode/encode_snac.py --input video.mp4 --output tokens.txt
|
||
|
|
# any format ffmpeg can read (audio extracted automatically, works on
|
||
|
|
# video files too -- just uses the audio track)
|
||
|
|
"""
|
||
|
|
import argparse
|
||
|
|
import os
|
||
|
|
import subprocess
|
||
|
|
import sys
|
||
|
|
|
||
|
|
sys.path.insert(0, os.path.join(os.path.dirname(os.path.abspath(__file__)), "..", "..", "pipeline_pose"))
|
||
|
|
|
||
|
|
SAMPLE_RATE = 24000
|
||
|
|
SNAC_MODEL = "hubertsiuzdak/snac_24khz"
|
||
|
|
|
||
|
|
|
||
|
|
def _get_ffmpeg() -> str:
|
||
|
|
import imageio_ffmpeg
|
||
|
|
return imageio_ffmpeg.get_ffmpeg_exe()
|
||
|
|
|
||
|
|
|
||
|
|
def extract_audio(input_path: str):
|
||
|
|
import numpy as np
|
||
|
|
|
||
|
|
cmd = [_get_ffmpeg(), "-y", "-i", input_path, "-vn", "-ac", "1", "-ar", str(SAMPLE_RATE), "-f", "f32le", "-"]
|
||
|
|
result = subprocess.run(cmd, capture_output=True, timeout=300)
|
||
|
|
if result.returncode != 0 or not result.stdout:
|
||
|
|
raise RuntimeError(f"ffmpeg failed to extract audio from {input_path}: {result.stderr.decode(errors='replace')[:500]}")
|
||
|
|
audio = np.frombuffer(result.stdout, dtype=np.float32).copy()
|
||
|
|
if len(audio) == 0:
|
||
|
|
raise RuntimeError(f"No audio extracted from {input_path} -- does it have an audio track?")
|
||
|
|
return audio
|
||
|
|
|
||
|
|
|
||
|
|
def encode_file(input_path: str) -> list:
|
||
|
|
import torch
|
||
|
|
from snac import SNAC
|
||
|
|
from snac_finevideo import encode_listen # reused verbatim, not re-derived
|
||
|
|
|
||
|
|
audio = extract_audio(input_path)
|
||
|
|
device = "cuda" if torch.cuda.is_available() else "cpu"
|
||
|
|
model = SNAC.from_pretrained(SNAC_MODEL).eval().to(device)
|
||
|
|
tokens = encode_listen(audio, model, device) # already "<snac_N>" strings, listen-format
|
||
|
|
return tokens
|
||
|
|
|
||
|
|
|
||
|
|
def main():
|
||
|
|
ap = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
|
||
|
|
ap.add_argument("--input", required=True, help="Audio or video file (any ffmpeg-readable format)")
|
||
|
|
ap.add_argument("--output", help="Optional: write the <snac> block to this file")
|
||
|
|
args = ap.parse_args()
|
||
|
|
|
||
|
|
tokens = encode_file(args.input)
|
||
|
|
duration_s = len(tokens) / 3 / 12.5
|
||
|
|
print(f"{len(tokens)} snac tokens ({len(tokens) // 3} base frames, ~{duration_s:.2f}s @ 12.5Hz base rate)")
|
||
|
|
block = "<snac> " + " ".join(tokens) + " </snac>"
|
||
|
|
print(block)
|
||
|
|
if args.output:
|
||
|
|
with open(args.output, "w") as f:
|
||
|
|
f.write(block)
|
||
|
|
print(f"Saved: {args.output}")
|
||
|
|
|
||
|
|
|
||
|
|
if __name__ == "__main__":
|
||
|
|
main()
|