216 lines
9.7 KiB
Python
216 lines
9.7 KiB
Python
#!/usr/bin/env python3
|
|
"""
|
|
SNAC detokenizer -- turns `<snac_N>` "listen format" tokens back into a real
|
|
audio waveform (24kHz mono), using the same `hubertsiuzdak/snac_24khz` model
|
|
used to encode them (data_prep/laion_emotional_roleplay/tokenize_snac.py,
|
|
pipeline_pose/snac_finevideo.py).
|
|
|
|
Listen format encodes only 2 of the model's 3 hierarchical codebook levels
|
|
(base 12.5Hz level 0, and one 25Hz level 1 -- the finest 50Hz level 2 is
|
|
dropped entirely to save tokens). 3 tokens per base frame, in fixed order
|
|
(L0, L1_even, L1_odd), with these offsets added to the raw codebook index
|
|
(0-4095 each):
|
|
|
|
OFFSET_L0 = 128266 (raw + 128266)
|
|
OFFSET_L1A = 132362 (raw + 128266 + 4096)
|
|
OFFSET_L1B = 144650 (raw + 128266 + 4*4096)
|
|
|
|
Since level 2 was never encoded, this decoder reconstructs it as all-zero
|
|
codes (index 0 in each of the 4x-oversampled slots) -- `SNAC.decode()` needs
|
|
all 3 levels present, so this is the only reconstruction available without
|
|
retraining/re-deriving the missing fine detail. Expect this to sound coarser
|
|
than the original clip (level 2 carries the finest timbral detail), same
|
|
lossy tradeoff already documented for `decode_cosmos.py`'s video reconstruction.
|
|
|
|
Usage:
|
|
# From a raw list of ints (must be a multiple of 3):
|
|
python tools/decode/decode_snac.py --tokens 128266,132850,145181,... --output out.wav
|
|
|
|
# From a flattened JSONL record, pulling every <snac>...</snac> block:
|
|
python tools/decode/decode_snac.py --input-jsonl path.jsonl --record-id ID --output out.wav
|
|
"""
|
|
import argparse
|
|
import json
|
|
import re
|
|
import sys
|
|
|
|
OFFSET_L0 = 128266
|
|
OFFSET_L1A = 128266 + 4096
|
|
OFFSET_L1B = 128266 + 4 * 4096
|
|
# Speak-format (full 3-level) offsets, 2026-07-23 -- matches the Leo/Orpheus
|
|
# scheme in pipeline_video/snac_gpu.py and the corrected
|
|
# data_prep/laion_emotional_roleplay/tokenize_snac.py::encode_speak(). Group
|
|
# order per base frame: L0, L1a, L2_0, L2_1, L1b, L2_2, L2_3 (7 tokens).
|
|
OFFSET_L2 = [136458, 140554, 148746, 152842]
|
|
SAMPLE_RATE = 24000
|
|
SNAC_MODEL = "hubertsiuzdak/snac_24khz"
|
|
|
|
_SNAC_ATOMIC_RE = re.compile(r"<snac_(\d+)>")
|
|
# 2026-07-23: production wrapper is <listen>/<speak>, not <snac> -- <snac>
|
|
# kept for backward compat with any older data/samples still using it.
|
|
_SNAC_BLOCK_RE = re.compile(r"<(?:snac|listen|speak)>(.*?)</(?:snac|listen|speak)>", re.DOTALL)
|
|
|
|
|
|
def extract_snac_tokens(text: str) -> list:
|
|
"""Pull every <snac_N> id inside every <listen>...</listen> or
|
|
<speak>...</speak> (or legacy <snac>...</snac>) block, in order. Falls
|
|
back to scanning the whole text if no such wrapper is present."""
|
|
blocks = _SNAC_BLOCK_RE.findall(text)
|
|
source = " ".join(blocks) if blocks else text
|
|
return [int(x) for x in _SNAC_ATOMIC_RE.findall(source)]
|
|
|
|
|
|
def load_tokens_from_jsonl(path: str, record_id: str) -> list:
|
|
with open(path, encoding="utf-8") as f:
|
|
for line in f:
|
|
line = line.strip()
|
|
if not line:
|
|
continue
|
|
rec = json.loads(line)
|
|
rid = rec.get("video_id", rec.get("id"))
|
|
if rid == record_id:
|
|
return extract_snac_tokens(rec.get("text", ""))
|
|
raise KeyError(f"record_id={record_id!r} not found in {path}")
|
|
|
|
|
|
def decode_snac_tokens(token_ids: list, output_path: str) -> None:
|
|
"""token_ids: raw <snac_N> ids (with offsets), length must be a multiple of 3."""
|
|
if len(token_ids) % 3 != 0:
|
|
raise ValueError(f"Expected a multiple of 3 tokens (L0,L1a,L1b triplets), got {len(token_ids)}")
|
|
if not token_ids:
|
|
raise ValueError("No snac tokens to decode")
|
|
|
|
import torch
|
|
import soundfile as sf
|
|
from snac import SNAC
|
|
|
|
n0 = len(token_ids) // 3
|
|
c0 = torch.zeros(1, n0, dtype=torch.long)
|
|
c1 = torch.zeros(1, 2 * n0, dtype=torch.long)
|
|
c2 = torch.zeros(1, 4 * n0, dtype=torch.long) # level 2 was never encoded -- zero-fill
|
|
|
|
# This function assumes strict positional cycling (index%3==0 -> L0,
|
|
# ==1 -> L1a, ==2 -> L1b) matches the actual band each id belongs to --
|
|
# true for well-formed <snac>...</snac> blocks, but NOT guaranteed for a
|
|
# free-floating/unwrapped token span assembled by extract_snac_tokens()'s
|
|
# whole-text fallback (e.g. a generation that never closes </snac>,
|
|
# spliced with earlier/later snac fragments elsewhere in the text). A
|
|
# mismatched id produces an out-of-[0,4095) codebook index, which SNAC's
|
|
# embedding lookup on GPU turns into an opaque
|
|
# "CUDA error: device-side assert triggered" instead of a real error
|
|
# message (hit for real 2026-07-22, roleplay_speech/sample eval run --
|
|
# see samples/qwen3_1.7b_vla_v2_eval/2026-07-22_full_eval/SUMMARY.md).
|
|
# Validate up front so the failure is legible.
|
|
for i in range(n0):
|
|
raw_l0, raw_l1a, raw_l1b = token_ids[3 * i], token_ids[3 * i + 1], token_ids[3 * i + 2]
|
|
r0, r1a, r1b = raw_l0 - OFFSET_L0, raw_l1a - OFFSET_L1A, raw_l1b - OFFSET_L1B
|
|
for pos, (name, tok, raw) in enumerate([("L0", raw_l0, r0), ("L1a", raw_l1a, r1a), ("L1b", raw_l1b, r1b)]):
|
|
if not (0 <= raw < 4096):
|
|
raise ValueError(
|
|
f"Triplet {i} position {pos} ({name}): token <snac_{tok}> decodes to raw "
|
|
f"codebook index {raw}, outside valid [0, 4096). This id likely belongs to a "
|
|
f"different band than its position implies (offsets: L0={OFFSET_L0}, "
|
|
f"L1a={OFFSET_L1A}, L1b={OFFSET_L1B}) -- often means the input tokens aren't "
|
|
f"one clean <snac>...</snac> block (e.g. spliced fragments from an unwrapped "
|
|
f"or unclosed generation)."
|
|
)
|
|
c0[0, i] = r0
|
|
c1[0, 2 * i] = r1a
|
|
c1[0, 2 * i + 1] = r1b
|
|
|
|
device = "cuda" if torch.cuda.is_available() else "cpu"
|
|
model = SNAC.from_pretrained(SNAC_MODEL).eval().to(device)
|
|
|
|
with torch.inference_mode():
|
|
audio = model.decode([c0.to(device), c1.to(device), c2.to(device)]) # (1, 1, samples)
|
|
|
|
waveform = audio.squeeze().float().cpu().numpy()
|
|
sf.write(output_path, waveform, SAMPLE_RATE)
|
|
|
|
|
|
def decode_speak_tokens(token_ids: list, output_path: str) -> None:
|
|
"""token_ids: raw <snac_N> ids (with offsets), length must be a multiple of 7,
|
|
group order per base frame: L0, L1a, L2_0, L2_1, L1b, L2_2, L2_3 -- matches
|
|
encode_speak() in data_prep/laion_emotional_roleplay/tokenize_snac.py and
|
|
pipeline_pose/snac_finevideo.py. Unlike decode_snac_tokens() (listen-only,
|
|
zero-fills level 2), this reconstructs the REAL level-2 codes."""
|
|
if len(token_ids) % 7 != 0:
|
|
raise ValueError(f"Expected a multiple of 7 tokens (speak-format groups), got {len(token_ids)}")
|
|
if not token_ids:
|
|
raise ValueError("No snac tokens to decode")
|
|
|
|
import torch
|
|
import soundfile as sf
|
|
from snac import SNAC
|
|
|
|
n0 = len(token_ids) // 7
|
|
c0 = torch.zeros(1, n0, dtype=torch.long)
|
|
c1 = torch.zeros(1, 2 * n0, dtype=torch.long)
|
|
c2 = torch.zeros(1, 4 * n0, dtype=torch.long)
|
|
|
|
offsets = [OFFSET_L0, OFFSET_L1A, OFFSET_L2[0], OFFSET_L2[1], OFFSET_L1B, OFFSET_L2[2], OFFSET_L2[3]]
|
|
names = ["L0", "L1a", "L2_0", "L2_1", "L1b", "L2_2", "L2_3"]
|
|
for i in range(n0):
|
|
raws = []
|
|
for pos in range(7):
|
|
tok = token_ids[7 * i + pos]
|
|
raw = tok - offsets[pos]
|
|
if not (0 <= raw < 4096):
|
|
raise ValueError(
|
|
f"Group {i} position {pos} ({names[pos]}): token <snac_{tok}> decodes to raw "
|
|
f"codebook index {raw}, outside valid [0, 4096). Likely a band/position "
|
|
f"mismatch (offsets: {dict(zip(names, offsets))})."
|
|
)
|
|
raws.append(raw)
|
|
c0[0, i] = raws[0]
|
|
c1[0, 2 * i] = raws[1]
|
|
c2[0, 4 * i] = raws[2]
|
|
c2[0, 4 * i + 1] = raws[3]
|
|
c1[0, 2 * i + 1] = raws[4]
|
|
c2[0, 4 * i + 2] = raws[5]
|
|
c2[0, 4 * i + 3] = raws[6]
|
|
|
|
device = "cuda" if torch.cuda.is_available() else "cpu"
|
|
model = SNAC.from_pretrained(SNAC_MODEL).eval().to(device)
|
|
|
|
with torch.inference_mode():
|
|
audio = model.decode([c0.to(device), c1.to(device), c2.to(device)])
|
|
|
|
waveform = audio.squeeze().float().cpu().numpy()
|
|
sf.write(output_path, waveform, SAMPLE_RATE)
|
|
|
|
|
|
def main():
|
|
ap = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
|
|
ap.add_argument("--tokens", help="Comma-separated list of raw <snac_N> ids")
|
|
ap.add_argument("--input-jsonl", help="Flattened JSONL file to pull tokens from")
|
|
ap.add_argument("--record-id", help="video_id/id field to select within --input-jsonl")
|
|
ap.add_argument("--text-file", help="Plain text file containing <snac_N> tokens anywhere in it")
|
|
ap.add_argument("--format", choices=["listen", "speak"], default="listen",
|
|
help="listen = 3 tok/frame (L0+L1, zero-filled L2); speak = 7 tok/frame (real L2)")
|
|
ap.add_argument("--output", required=True)
|
|
args = ap.parse_args()
|
|
|
|
if args.tokens:
|
|
token_ids = [int(x) for x in args.tokens.split(",")]
|
|
elif args.input_jsonl and args.record_id:
|
|
token_ids = load_tokens_from_jsonl(args.input_jsonl, args.record_id)
|
|
elif args.text_file:
|
|
token_ids = extract_snac_tokens(open(args.text_file, encoding="utf-8").read())
|
|
else:
|
|
ap.error("Provide --tokens, --text-file, or (--input-jsonl and --record-id)")
|
|
|
|
group_size = 7 if args.format == "speak" else 3
|
|
print(f"Decoding {len(token_ids)} snac tokens ({args.format} format, "
|
|
f"{len(token_ids) // group_size} base frames, "
|
|
f"~{len(token_ids) / group_size / 12.5:.2f}s @ 12.5Hz base rate)...")
|
|
if args.format == "speak":
|
|
decode_speak_tokens(token_ids, args.output)
|
|
else:
|
|
decode_snac_tokens(token_ids, args.output)
|
|
print(f"Saved: {args.output}")
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|