zeroae/calliope-snac-4b-base-4k

Calliope SNAC 4B Base (4K)

Stage-1 multilingual SNAC prior for the Calliope text-to-speech project — a continued-pretrain of `nvidia/Nemotron-H-4B-Base-8K` with the vocabulary augmented by 12,288 SNAC codec tokens and a slot router that enforces the codec's C·M·F·F·M·F·F frame pattern at audio-mode positions.

This is the HuggingFace safetensors version, loadable via AutoModelForCausalLM.from_pretrained(..., trust_remote_code=True). The Megatron-Bridge FSDP DCP format is at the sibling repo `zeroae/calliope-snac-4b-base-4k.megatron` (private) for Bridge-based continued training.

What this is and isn't. This is a pretrained prior, not a finished TTS system. Training mixed text-only and SNAC-audio-only documents — the cross-modal text→SNAC bridge is a separate stage-2 finetune objective and was not learned here. Use this checkpoint as the starting point for a TTS finetune, not as an end-to-end speech model.

Quick start: text generation

import torch
from transformers import AutoModelForCausalLM, AutoTokenizer

REPO = "zeroae/calliope-snac-4b-base-4k"

tokenizer = AutoTokenizer.from_pretrained(REPO, trust_remote_code=True)
model = AutoModelForCausalLM.from_pretrained(
    REPO,
    dtype=torch.bfloat16,
    device_map="cuda",
    low_cpu_mem_usage=True,
    trust_remote_code=True,
)

# Text-only generation (the text path is preserved at near-base quality).
# The slot router masks all SNAC tokens to -inf in text mode, so text
# generation is unaffected by the augmented vocab.
ids = tokenizer("In multilingual TTS, prosody", return_tensors="pt").input_ids.to("cuda")
out = model.generate(ids, max_new_tokens=32, do_sample=False)
print(tokenizer.decode(out[0]))

The trust_remote_code=True flag pulls modeling_nemotron_h_augmented.py from this repo, which wraps the base NemotronH-4B with the slot-router logits mask (text mode masks all SNAC tokens, audio mode masks all text — enforced per-position).

End-to-end: generate SNAC frames → decode to audio

This pretrained prior has no text→SNAC bridge (see disclaimer above); the example below shows the unconditional end-to-end pipeline that the slot router makes work: prompt with the [SNAC] marker, generate tokens (which the slot router constrains to the C·M·F·F·M·F·F frame pattern), parse them back into the three SNAC codebooks, and decode to a waveform via the upstream `hubertsiuzdak/snac_24khz` codec.

Expect the audio to be a babble/noise — the model is sampling unconditionally from its learned audio distribution; no text guides the content. The point is to demonstrate the mechanics; quality requires a stage-2 finetune that learns the text→SNAC bridge.

import torch
from transformers import AutoModelForCausalLM, AutoTokenizer
from snac import SNAC          # pip install snac
import torchaudio

REPO = "zeroae/calliope-snac-4b-base-4k"

# --- 1. Load LM ---------------------------------------------------------
tok = AutoTokenizer.from_pretrained(REPO, trust_remote_code=True)
model = AutoModelForCausalLM.from_pretrained(
    REPO, dtype=torch.bfloat16, device_map="cuda",
    low_cpu_mem_usage=True, trust_remote_code=True,
).eval()

# Vocab layout (from augmented.yaml, also visible in the repo)
SNAC_OPEN, SNAC_CLOSE = 100, 101
C_BASE, M_BASE, F_BASE = 131072, 135168, 139264   # start of each codebook range
N_FRAMES = 50                                      # ~4 s at SNAC-24kHz's coarse rate
N_TOKENS = N_FRAMES * 7                            # 7 tokens / frame (C,M,F,F,M,F,F)

# --- 2. Generate inside an [SNAC] ... span ------------------------------
# The slot router (modeling_nemotron_h_augmented.py) carries its
# (in_slot_mode, slot_counter) state across forward calls via
# self._slot_router_state, so KV caching just works: prefill computes
# routing from initial state, subsequent forwards advance from the
# cached final state. No special flags needed.
prompt = torch.tensor([[tok.bos_token_id, SNAC_OPEN]], device="cuda")
with torch.no_grad():
    out = model.generate(
        prompt,
        max_new_tokens=N_TOKENS,
        do_sample=True, temperature=0.8, top_p=0.95,
    )

# --- 3. Parse the C/M/F/F/M/F/F frames back into codebook indices --------
gen = out[0, prompt.shape[1]:].tolist()
gen = gen[: (len(gen) // 7) * 7]                    # truncate to whole frames
c_codes, m_codes, f_codes = [], [], []
for i in range(0, len(gen), 7):
    frame = gen[i:i + 7]
    c_codes.append(frame[0] - C_BASE)               # slot 0: C
    m_codes.append(frame[1] - M_BASE)               # slot 1: M
    f_codes.append(frame[2] - F_BASE)               # slot 2: F
    f_codes.append(frame[3] - F_BASE)               # slot 3: F
    m_codes.append(frame[4] - M_BASE)               # slot 4: M
    f_codes.append(frame[5] - F_BASE)               # slot 5: F
    f_codes.append(frame[6] - F_BASE)               # slot 6: F

# Sanity-check the slot router did its job (codes within [0, 4096))
assert all(0 <= c < 4096 for c in c_codes), "C codes out of range — slot router off?"
assert all(0 <= m < 4096 for m in m_codes), "M codes out of range"
assert all(0 <= f < 4096 for f in f_codes), "F codes out of range"

# --- 4. Decode the three codebooks to a 24 kHz waveform -----------------
codec = SNAC.from_pretrained("hubertsiuzdak/snac_24khz").eval().to("cuda")
codes = [
    torch.tensor([c_codes], dtype=torch.long, device="cuda"),   # [B=1, N_FRAMES]
    torch.tensor([m_codes], dtype=torch.long, device="cuda"),   # [B=1, 2*N_FRAMES]
    torch.tensor([f_codes], dtype=torch.long, device="cuda"),   # [B=1, 4*N_FRAMES]
]
with torch.no_grad():
    audio = codec.decode(codes)                                 # [1, 1, num_samples]

# --- 5. Save ------------------------------------------------------------
torchaudio.save("calliope_unconditional.wav", audio.squeeze(0).cpu(), sample_rate=24000)
print(f"saved {audio.shape[-1] / 24000:.2f} s of audio  "
      f"({len(c_codes)} frames, {len(c_codes) + len(m_codes) + len(f_codes)} codes)")

Dependencies: pip install snac torchaudio in addition to transformers torch. Wall-clock for 50 frames (~4 s of audio): a few seconds on a GB10 with KV caching on (the default).

Token-budget rule of thumb: SNAC-24kHz's coarse rate is ~12 Hz, so one frame ≈ 83 ms of audio. To pre-allocate max_new_tokens for a given duration:

N_TOKENS = int(seconds * 12) * 7      # 7 tokens per frame

Why this demo's audio sounds bad (and that's expected)

The model has never seen text + [SNAC]…[/SNAC] parallel sequences — only text-only documents and SNAC-only documents, mixed at the batch level. Unconditional sampling from the SNAC distribution produces something codec-plausible (the slot router guarantees the bit-stream is structurally valid, and the codec can always decode), but it has no semantic content. It's the analogue of letting a language model generate without a prompt — you get gibberish that has the shape of the training distribution. A stage-2 TTS finetune on text → SNAC parallel data is what makes this conditional and intelligible.

Architecture

The SNAC frame layout is [C, M, F, F, M, F, F] — 7 tokens per coarse frame, one coarse (C) → two mid (M) → four fine (F) — matching SNAC-24kHz's 1:2:4 residual-quantizer hierarchy.

Training summary

Final validation losses (iter 75,000)

Held-out per-source dev splits, blended at the same per-phase weights as training:

The coarse-to-fine ordering C < M < F is preserved at every eval over the run, consistent with SNAC's residual hierarchy. Text PPL ~10 is approximately base NemotronH-Base quality (the 30% text rehearsal anchor held throughout) — the augmented vocab did not catastrophically forget the base model's text capability.

For context, the random-baseline PPL over each 4096-token SNAC range is 4096; the iter-0 augmented-baseline PPL was ~7200 (slightly above random because the newly-added rows perturbed the softmax). Final range_C PPL 66 ≈ 62× better than random on coarse codes.

Data composition

Bin total: 601,910 samples × 4096 tokens = 2.466 B tokens. 44 sources pooled across 10 languages, blended via per-source disjoint phase windows + without-replacement sampling (provably no document seen by more than one phase or twice within a phase).

Realized per-phase quality mix (% of phase, train split):

Source corpora (all encoded into Megatron .bin/.idx format via the project's format-snac → format-phases pipeline):

• —Text rehearsal: FineWeb-EN (sample-10BT) + FineWeb-2 (deu/spa/fra/ita/por/kor/jpn/rus/cmn)
• —Bulk-noisy audio: Emilia-YODAS (EN/ZH/DE/FR/JA/KO), Emilia-zh, GigaSpeech-XL (audiobook/podcast/youtube), VoxPopuli (DE/FR/IT/PL/ES)
• —Clean read: LibriTTS-R (EN), MLS (DE/FR/ES/IT/PT)
• —Studio + crowd: HiFi-TTS, VCTK, AISHELL-1, AISHELL-3, CommonVoice-17 (EN/DE/FR/ES/IT/PT/JA/KO/RU)

Intended use

• —Starting point for a stage-2 TTS finetune on parallel text → SNAC data. The cross-modal bridge is not present in this checkpoint; supervised finetuning on text-aligned SNAC sequences is what lights it up.
• —Multilingual SNAC perplexity benchmarking across the 10 NemotronH-supported languages.
• —Acoustic embedding extraction — pool residual stream activations over a SNAC sequence for downstream classification (language ID, speaker family, audio quality scoring).
• —Audio-only continuation / infilling — given partial SNAC, generate plausible continuation. Distribution-in, distribution-out.

Out of scope and limitations

• —Not a TTS system. Stage-1 mixed text-only and SNAC-only documents. There is no learned bridge between text and audio in this checkpoint. Prompting it with text and expecting SNAC output (or vice versa) will not work cleanly without a stage-2 finetune.
• —No speaker conditioning. Speaker tokens / voice control are deferred to the downstream TTS finetune by design.
• —4K context, not 8K. Architecturally the base supports 8K; the augmented model's slot router was never exercised on sequences > 4K. Use 4K and below for now.
• —Languages outside NemotronH's supported 10 were dropped during data design — do not expect quality on e.g. Polish, Indonesian, Vietnamese, Thai, Arabic.
• —HiFi-TTS was in the training mix. If your downstream evaluation uses HiFi-TTS speakers as "held-out studio voices," this prior has already seen them — the strict version of the held-out-voice gate cannot be measured on this checkpoint. (Future Calliope stage-1 versions hold HiFi-TTS out entirely; see the v3 plan link below.)
• —Convergence ceiling. The loss plateaued well above the original optimistic target (lm 3.0-3.5 hoped, lm 3.94 reached). The conservative forecast fit at iter-20k called the final values to within 0.02 nats — the model is on its forecast trajectory, just lower-quality than initial intuition suggested. Diagnostically traced to LR/optimizer regime (high decoupledlr perturbing already-converged text rows, body-learning bottleneck), not to data quantity. A v3 design is planned that addresses these (WSD LR schedule, lower decoupledlr, mean-init confirmed already in use, marginal data blending).

Format details

This repository ships the model in standard HuggingFace safetensors format. Files:

The Megatron-Bridge FSDP DCP format of the same trained weights is at the sibling private repo `zeroae/calliope-snac-4b-base-4k.megatron` — use that one if you want to continue pretraining via Bridge.

Provenance & references

• —W&B run: ynbrpbmx (project calliope, exp stage1-phase-full-v2)
• —Original training plan: `docs/superpowers/plans/2026-05-02-snac-stage1-multilingual-prior.md`
• —Predictions vs forecast vs actuals: `2026-05-07-snac-stage1-phase-end-predictions.md`, `2026-05-10-snac-stage1-forecast.md`
• —v3 design (next iteration — WSD schedule, marginal blend, HiFi-TTS holdout): `2026-05-16-snac-stage1-v3.md`
• —Base model: `nvidia/Nemotron-H-4B-Base-8K`
• —Codec: `hubertsiuzdak/snac_24khz`

License

This checkpoint is a derivative of nvidia/Nemotron-H-4B-Base-8K. The base model's license terms apply to redistribution and use of these weights. Refer to the linked base model for the authoritative license; this repo does not extend or restrict those terms.

The augmented modeling code (modeling_nemotron_h_augmented.py, augmented.yaml, augmentation specs) is © Zero A.E., LLC and licensed for research use under terms TBD — please contact the org before commercial use.