vllm_omni.model_executor.models.audex.checkpoint ¶

Audex checkpoint layout preparation.

The nvidia/Nemotron-Labs-Audex-2B repo deduplicates weights: the checkpoint_folder_audiogen (and checkpoint_folder_textonly) index files reference safetensors shards that physically live only under checkpoint_folder_full. The official repo ships a prepare script that symlinks them; this replicates that step automatically so users can pass the repo root without any manual preparation.

XCODEC1_DEFAULT_REPO `module-attribute` ¶

XCODEC1_DEFAULT_REPO = (
    "hf-audio/xcodec-hubert-general-balanced"
)

logger `module-attribute` ¶

logger = init_logger(__name__)

ensure_audex_snapshot ¶

ensure_audex_snapshot(
    model: str, profile: str = "tts"
) -> str

Resolve model to a local repo-root directory, downloading if needed.

Local paths pass through untouched. For an HF repo id, download (or reuse from cache) exactly the subset the requesting stage's profile needs, so per-stage model_subdir joins land on a real snapshot instead of the repo-id string on a fresh cache. Falls back to a cached snapshot when offline.

ensure_audiogen_weights ¶

ensure_audiogen_weights(model_dir: str) -> None

Link index-referenced shards missing from model_dir from the full checkpoint.

No-op when the shards are already present (or the layout is unexpected); raises only if a referenced shard exists nowhere.

ensure_xcodec1_snapshot ¶

ensure_xcodec1_snapshot(model: str | None) -> str

Resolve the external XCodec1 checkpoint (TTA waveform decoder).

Accepts a local directory or an HF repo id; None/empty falls back to the official hf-audio/xcodec-hubert-general-balanced.