Skip to main content

rlx_mimi/
codes.rs

1/// Discrete codec output from [`crate::MimiCodec::encode_pcm`].
2#[derive(Debug, Clone)]
3pub struct MimiCodes {
4    /// Per-frame codebook indices, shape `[num_frames][num_quantizers]`.
5    pub frames: Vec<Vec<u32>>,
6    pub num_quantizers: usize,
7}
8
9impl MimiCodes {
10    pub fn num_frames(&self) -> usize {
11        self.frames.len()
12    }
13
14    /// HF layout `[num_quantizers][num_frames]` (same as `MimiModel.encode` output).
15    pub fn to_hf_layout(&self) -> Vec<Vec<u32>> {
16        let t = self.num_frames();
17        if t == 0 {
18            return Vec::new();
19        }
20        let k = self.frames[0].len();
21        let mut out = vec![vec![0u32; t]; k];
22        for (ti, row) in self.frames.iter().enumerate() {
23            for (ki, &code) in row.iter().enumerate() {
24                out[ki][ti] = code;
25            }
26        }
27        out
28    }
29
30    pub fn from_hf_layout(codes: Vec<Vec<u32>>) -> Self {
31        let k = codes.len();
32        let t = codes.first().map(|r| r.len()).unwrap_or(0);
33        let mut frames = Vec::with_capacity(t);
34        for ti in 0..t {
35            let mut row = Vec::with_capacity(k);
36            for row_k in &codes {
37                row.push(row_k[ti]);
38            }
39            frames.push(row);
40        }
41        Self {
42            frames,
43            num_quantizers: k,
44        }
45    }
46}