inference.md

Inference Guide

The inference API is the RubatoCodec class (rubatocodec/model.py). It is installed by the base package:

uv sync            # inference dependencies only

Loading a model

from rubatocodec import RubatoCodec

# from a local safetensors export (preferred)
codec = RubatoCodec.from_checkpoint("path/to/model.safetensors", device="cuda")

# from the Hugging Face Hub (expects model.safetensors + config.json)
codec = RubatoCodec.from_pretrained("third-intelligence/RubatoCodec", device="cuda")

# directly from a PyTorch Lightning .ckpt (slower; loads the full checkpoint)
codec = RubatoCodec.from_checkpoint("path/to/last.ckpt", device="cuda")

By default, loading also downloads the frozen facebook/w2v-bert-2.0 semantic model from the Hugging Face Hub (used on the encode side only).

Loading for decoding only

If codes and durations are already available, load only the decoding components:

codec = RubatoCodec.from_pretrained(
    "third-intelligence/RubatoCodec", device="cuda", decoder_only=True
)
# The same option works with local .safetensors and .ckpt files:
codec = RubatoCodec.from_checkpoint(
    "path/to/model.safetensors", device="cuda", decoder_only=True
)
audio = codec.decode(enc, depth=8)  # enc was produced by a full codec instance

This skips the encoders and semantic model, supporting decode and decode_streaming but not encode.

Encoding

import soundfile as sf
import torch

data, sr = sf.read("speech.wav", dtype="float32")  # mono waveform
wav = torch.from_numpy(data)

enc = codec.encode(wav, sr=sr)  # resampled internally to codec.sr (16 kHz) if needed
enc["codes"]  # (Q=16, K) int64: quantizer indices per event; layer 0 has semantic supervision
enc["durs"]  # (K,)   int64: event duration in 50 Hz grid frames

Notes:

• Omitting sr assumes the input is already at codec.sr (16 kHz). Internal resampling can differ marginally from an external resampler.
• Encoding supports batch size 1 per call.
• Pass finite floating-point mono audio shaped (T,) or (1, T). Stereo audio must be mixed to mono before encoding.
• Encoding has no streaming mode.

Decoding

audio = codec.decode(enc)  # (1, 1, T_audio) at 44.1 kHz, full depth
audio = codec.decode(enc, depth=8)  # first 8 quantizer layers (~1.2 kbit/s)
audio = codec.decode(enc["codes"], durs=enc["durs"])  # equivalent explicit form

sf.write("recon.wav", audio[0, 0].cpu().numpy(), codec.output_sr)

depth selects how many residual-quantizer layers are used (1–16). Because the quantizer is trained with quantizer dropout, every prefix depth yields a valid reconstruction. Increasing depth retains more residual information.

The output length is exactly 882 · sum(durs) samples at 44.1 kHz (20 ms per grid frame).

Inference settings

decode(enc) uses the model's recommended settings from config.json. Both .ckpt and safetensors loading apply the same defaults when settings are absent. For the released model these are decode_window_frames=300, decode_stride_frames=100, decode_crossfade_ms=10, and a six-second attention window. Saved settings and explicit API arguments override these defaults. For example:

audio = codec.decode(enc, batch_size=16)
# Decode the entire grid without periodic window resets:
audio = codec.decode(enc, window_frames=None)

See the decode API docstring for optional decoding controls.

Use integer tensors for codes and durations, with positive durations and depth between 1 and the number of supplied code layers.

Streaming decode

decode_streaming accepts one event or a chunk of events and returns the corresponding audio incrementally. Use the same key for a stream and set finish=True on the final call to emit any remaining audio and close the session.

codes, durs = enc["codes"], enc["durs"]
K = codes.shape[1]
key = "session-0"

out = []
for i in range(K):  # e.g. events arriving one by one from a TTS language model
    chunk = codec.decode_streaming(
        codes[:, i : i + 1],
        key,
        durs=durs[i : i + 1],
        finish=(i == K - 1),
    )
    if chunk is not None:
        out.append(chunk)
audio = torch.cat(out, dim=2)

• Chunks may contain any number of events (including a flush-only call with an empty chunk and finish=True).
• Total output length is 882 · sum(durs) samples.
• depth and decoding settings must stay fixed within a session.
• One session decodes one utterance stream; concurrent sessions are separate keys.

Reconstructing a file

From a repository checkout:

uv run python -m examples.reconstruct speech.wav recon.wav --depth 8
uv run python -m examples.reconstruct speech.wav recon.wav --streaming
uv run python -m examples.reconstruct speech.wav recon.wav \
    --checkpoint path/to/model.safetensors --device cuda

The example mixes stereo input to mono, resamples through encode, and saves the result at codec.output_sr. Output includes any padding to the encoder hop.

For multiple utterances and a metadata.json accepted by the evaluation CLI, see Running evaluation.