Skip to main content

rlx_orpheus/decoder/
config.rs

1// RLX — versatile ML compiler + runtime.
2// Copyright (C) 2026 Eugene Hauptmann, Nataliya Kosmyna.
3
4//! SNAC 24 kHz decoder configuration (`snac_24khz_decoder_config.json`).
5
6use std::path::Path;
7
8use anyhow::{Context, Result};
9use serde::Deserialize;
10
11/// SNAC model hyper-parameters (hubertsiuzdak/snac_24khz).
12#[derive(Debug, Clone, Deserialize)]
13pub struct SnacConfig {
14    pub sampling_rate: u32,
15    pub encoder_dim: usize,
16    pub encoder_rates: Vec<usize>,
17    pub decoder_dim: usize,
18    pub decoder_rates: Vec<usize>,
19    pub attn_window_size: Option<usize>,
20    pub codebook_size: usize,
21    pub codebook_dim: usize,
22    pub vq_strides: Vec<usize>,
23    pub noise: bool,
24    pub depthwise: bool,
25    #[serde(default)]
26    pub latent_dim: Option<usize>,
27}
28
29impl SnacConfig {
30    /// Load JSON config from disk.
31    pub fn from_file(path: &Path) -> Result<Self> {
32        let text = std::fs::read_to_string(path)
33            .with_context(|| format!("read SNAC config {}", path.display()))?;
34        Self::from_json(&text).with_context(|| format!("parse SNAC config {}", path.display()))
35    }
36
37    /// Parse JSON config text.
38    pub fn from_json(text: &str) -> Result<Self> {
39        Ok(serde_json::from_str(text)?)
40    }
41
42    /// Latent channel width — explicit in export JSON or derived from encoder stack.
43    pub fn latent_dim(&self) -> usize {
44        self.latent_dim
45            .unwrap_or_else(|| self.encoder_dim * 2_usize.pow(self.encoder_rates.len() as u32))
46    }
47
48    pub fn n_codebooks(&self) -> usize {
49        self.vq_strides.len()
50    }
51}