rlx_dac/
codec.rs

1use crate::audio::{load_wav_mono, write_wav_mono};
2use crate::build::{build_decoder, build_encoder};
3use crate::codes::DacCodes;
4use crate::config::DacConfig;
5use crate::graph::CodecGraph;
6use crate::layers::{Decoder, Encoder};
7use crate::quantize::{ResidualVectorQuantizer, build_quantizer};
8use crate::weights::WeightStore;
9use anyhow::{Result, ensure};
10use ndarray::{Array2, ArrayView2};
11use rlx_runtime::Device;
12use std::cell::RefCell;
13use std::collections::HashMap;
14use std::path::{Path, PathBuf};
15use std::time::Instant;
16
17pub const SAMPLE_RATE_24KHZ: u32 = 24_000;
18pub const SAMPLE_RATE_16KHZ: u32 = 16_000;
19pub const SAMPLE_RATE_44KHZ: u32 = 44_100;
20
21#[derive(Debug, Clone)]
22pub struct RoundtripStats {
23    pub encode_ms: f64,
24    pub decode_ms: f64,
25    pub num_frames: usize,
26    pub pcm_samples: usize,
27}
28
29pub struct DacCodec {
30    model_dir: PathBuf,
31    cfg: DacConfig,
32    device: Device,
33    encoder: Encoder,
34    quantizer: ResidualVectorQuantizer,
35    decoder: Decoder,
36    // Compiled encoder/decoder graphs, cached by input length. Empty (and
37    // unused) when `device == Cpu`, which runs the exact ndarray path.
38    enc_graphs: RefCell<HashMap<usize, CodecGraph>>,
39    dec_graphs: RefCell<HashMap<(usize, usize), CodecGraph>>,
40}
41
42impl DacCodec {
43    pub fn open(model_dir: impl AsRef<Path>) -> Result<Self> {
44        Self::open_on(model_dir, Device::Cpu)
45    }
46
47    pub fn open_on(model_dir: impl AsRef<Path>, device: Device) -> Result<Self> {
48        let model_dir = model_dir.as_ref().to_path_buf();
49        let cfg = if model_dir.join("config.json").is_file() {
50            DacConfig::load(&model_dir.join("config.json"))?
51        } else {
52            DacConfig::config_24khz()
53        };
54        let weights = model_dir.join("model.safetensors");
55        ensure!(
56            weights.is_file(),
57            "missing {} — run fetch or set RLX_DAC_DIR",
58            weights.display()
59        );
60        let store = WeightStore::open(&weights)?;
61        Ok(Self {
62            model_dir,
63            cfg: cfg.clone(),
64            device,
65            encoder: build_encoder(&cfg, &store)?,
66            quantizer: build_quantizer(&cfg, &store)?,
67            decoder: build_decoder(&cfg, &store)?,
68            enc_graphs: RefCell::new(HashMap::new()),
69            dec_graphs: RefCell::new(HashMap::new()),
70        })
71    }
72
73    /// Run the encoder on the selected backend. CPU uses the exact ndarray
74    /// reference; every other backend runs the compiled rlx graph.
75    fn encode_z(&self, input: ArrayView2<f32>) -> Result<Array2<f32>> {
76        if self.device == Device::Cpu && std::env::var_os("RLX_DAC_CPU_GRAPH").is_none() {
77            return Ok(self.encoder.forward(input));
78        }
79        let in_len = input.shape()[1];
80        let mut cache = self.enc_graphs.borrow_mut();
81        let graph = match cache.get_mut(&in_len) {
82            Some(g) => g,
83            None => {
84                let g = CodecGraph::encoder(self.device, &self.encoder, in_len)?;
85                cache.entry(in_len).or_insert(g)
86            }
87        };
88        let flat: Vec<f32> = input.iter().copied().collect();
89        graph.run(&flat)
90    }
91
92    /// Run the decoder on the selected backend (see [`Self::encode_z`]).
93    fn decode_z(&self, z_q: ArrayView2<f32>) -> Result<Array2<f32>> {
94        if self.device == Device::Cpu && std::env::var_os("RLX_DAC_CPU_GRAPH").is_none() {
95            return Ok(self.decoder.forward(z_q));
96        }
97        let (in_c, in_t) = (z_q.shape()[0], z_q.shape()[1]);
98        let mut cache = self.dec_graphs.borrow_mut();
99        let graph = match cache.get_mut(&(in_c, in_t)) {
100            Some(g) => g,
101            None => {
102                let g = CodecGraph::decoder(self.device, &self.decoder, in_c, in_t)?;
103                cache.entry((in_c, in_t)).or_insert(g)
104            }
105        };
106        let flat: Vec<f32> = z_q.iter().copied().collect();
107        graph.run(&flat)
108    }
109
110    pub fn config(&self) -> &DacConfig {
111        &self.cfg
112    }
113
114    pub fn model_dir(&self) -> &Path {
115        &self.model_dir
116    }
117
118    pub fn device(&self) -> Device {
119        self.device
120    }
121
122    pub fn sample_rate(&self) -> u32 {
123        self.cfg.sample_rate
124    }
125
126    pub fn hop_length(&self) -> usize {
127        self.cfg.hop_length
128    }
129
130    fn preprocess(&self, pcm: &[f32]) -> Vec<f32> {
131        let len = pcm.len();
132        let target = len.div_ceil(self.cfg.hop_length) * self.cfg.hop_length;
133        let mut out = pcm.to_vec();
134        out.resize(target, 0.0);
135        out
136    }
137
138    /// Encode mono PCM at the model sample rate.
139    pub fn encode_pcm(&self, pcm: &[f32], num_quantizers: Option<usize>) -> Result<DacCodes> {
140        ensure!(!pcm.is_empty(), "empty PCM");
141        let padded = self.preprocess(pcm);
142        let mut input = Array2::<f32>::zeros((1, padded.len()));
143        for (i, &s) in padded.iter().enumerate() {
144            input[[0, i]] = s;
145        }
146        let z = self.encode_z(input.view())?;
147        let nq = num_quantizers.unwrap_or(self.cfg.n_codebooks);
148        let (_z_q, code_rows) = self.quantizer.encode(z.view(), nq);
149        let t = code_rows.first().map(|c| c.len()).unwrap_or(0);
150        let mut frames = Vec::with_capacity(t);
151        for ti in 0..t {
152            let mut row = Vec::with_capacity(nq);
153            for qi in 0..nq {
154                row.push(code_rows[qi][ti]);
155            }
156            frames.push(row);
157        }
158        Ok(DacCodes {
159            frames,
160            num_quantizers: nq,
161        })
162    }
163
164    /// Decode codec frames → mono PCM at the model sample rate.
165    pub fn decode_codes(&self, codes: &DacCodes) -> Result<Vec<f32>> {
166        ensure!(!codes.frames.is_empty(), "empty codec frames");
167        let z_q = self.quantizer.from_codes(&codes.to_quantizer_rows());
168        let wav = self.decode_z(z_q.view())?;
169        ensure!(wav.dim().0 >= 1, "decoder produced no channels");
170        Ok(wav.row(0).to_vec())
171    }
172
173    /// Full forward pass: encode then decode, trimmed to the original PCM length.
174    pub fn forward_pcm(
175        &self,
176        pcm: &[f32],
177        num_quantizers: Option<usize>,
178    ) -> Result<(DacCodes, Vec<f32>)> {
179        let orig_len = pcm.len();
180        let codes = self.encode_pcm(pcm, num_quantizers)?;
181        let mut recon = self.decode_codes(&codes)?;
182        recon.truncate(orig_len.min(recon.len()));
183        Ok((codes, recon))
184    }
185
186    pub fn encode_wav(
187        &self,
188        wav: impl AsRef<Path>,
189        num_quantizers: Option<usize>,
190    ) -> Result<DacCodes> {
191        let pcm = load_wav_mono(wav.as_ref(), self.cfg.sample_rate)?;
192        self.encode_pcm(&pcm, num_quantizers)
193    }
194
195    pub fn decode_wav(
196        &self,
197        codes: &DacCodes,
198        out: impl AsRef<Path>,
199        trim_to_samples: Option<usize>,
200    ) -> Result<()> {
201        let mut pcm = self.decode_codes(codes)?;
202        if let Some(n) = trim_to_samples {
203            pcm.truncate(n.min(pcm.len()));
204        }
205        write_wav_mono(out.as_ref(), &pcm, self.cfg.sample_rate)
206    }
207
208    pub fn roundtrip_pcm(
209        &self,
210        pcm: &[f32],
211        num_quantizers: Option<usize>,
212    ) -> Result<(DacCodes, Vec<f32>, RoundtripStats)> {
213        let t0 = Instant::now();
214        let codes = self.encode_pcm(pcm, num_quantizers)?;
215        let num_frames = codes.num_frames();
216        let encode_ms = t0.elapsed().as_secs_f64() * 1000.0;
217        let t1 = Instant::now();
218        let mut recon = self.decode_codes(&codes)?;
219        let decode_ms = t1.elapsed().as_secs_f64() * 1000.0;
220        recon.truncate(pcm.len().min(recon.len()));
221        Ok((
222            codes,
223            recon,
224            RoundtripStats {
225                encode_ms,
226                decode_ms,
227                num_frames,
228                pcm_samples: pcm.len(),
229            },
230        ))
231    }
232
233    pub fn roundtrip_wav(
234        &self,
235        in_wav: impl AsRef<Path>,
236        out_wav: impl AsRef<Path>,
237        num_quantizers: Option<usize>,
238    ) -> Result<DacCodes> {
239        let pcm = load_wav_mono(in_wav.as_ref(), self.cfg.sample_rate)?;
240        let (codes, recon, _) = self.roundtrip_pcm(&pcm, num_quantizers)?;
241        write_wav_mono(out_wav.as_ref(), &recon, self.cfg.sample_rate)?;
242        Ok(codes)
243    }
244}
245
246/// Unified [`rlx_core::AudioCodec`] view (see rlx-mimi for the rationale).
247impl rlx_core::AudioCodec for DacCodec {
248    fn info(&self) -> rlx_core::CodecInfo {
249        let hop = self.cfg.hop_length.max(1);
250        rlx_core::CodecInfo {
251            sample_rate: self.cfg.sample_rate,
252            frame_rate: self.cfg.sample_rate as f32 / hop as f32,
253            hop_length: self.cfg.hop_length,
254            channels: 1,
255            max_quantizers: self.cfg.n_codebooks,
256            codebook_size: self.cfg.codebook_size,
257        }
258    }
259
260    fn device(&self) -> Device {
261        self.device
262    }
263
264    fn encode_pcm(&self, pcm: &[f32], num_quantizers: Option<usize>) -> Result<rlx_core::RvqCodes> {
265        let codes = DacCodec::encode_pcm(self, pcm, num_quantizers)?;
266        Ok(rlx_core::RvqCodes::new(codes.frames, codes.num_quantizers))
267    }
268
269    fn decode_codes(&self, codes: &rlx_core::RvqCodes) -> Result<Vec<f32>> {
270        let dc = DacCodes {
271            frames: codes.frames.clone(),
272            num_quantizers: codes.num_quantizers,
273        };
274        DacCodec::decode_codes(self, &dc)
275    }
276}
rlx_dac/codec.rs

rlx_dac/
codec.rs
