rustymimi/
lib.rs

1// Copyright (c) Kyutai, all rights reserved.
2// This source code is licensed under the license found in the
3// LICENSE file in the root directory of this source tree.
4
5use pyo3::prelude::*;
6
7use ::moshi as mm;
8use mm::{candle, candle_nn, conv, mimi, seanet, transformer};
9use std::sync::{mpsc, Mutex};
10
11trait PyRes<R> {
12    #[allow(unused)]
13    fn w(self) -> PyResult<R>;
14    fn w_f<P: AsRef<std::path::Path>>(self, p: P) -> PyResult<R>;
15}
16
17impl<R, E: Into<anyhow::Error>> PyRes<R> for Result<R, E> {
18    fn w(self) -> PyResult<R> {
19        self.map_err(|e| pyo3::exceptions::PyValueError::new_err(e.into().to_string()))
20    }
21    fn w_f<P: AsRef<std::path::Path>>(self, p: P) -> PyResult<R> {
22        self.map_err(|e| {
23            let e = e.into().to_string();
24            let msg = format!("{:?}: {e}", p.as_ref());
25            pyo3::exceptions::PyValueError::new_err(msg)
26        })
27    }
28}
29
30#[macro_export]
31macro_rules! py_bail {
32    ($msg:literal $(,)?) => {
33        return Err(pyo3::exceptions::PyValueError::new_err(format!($msg)))
34    };
35    ($err:expr $(,)?) => {
36        return Err(pyo3::exceptions::PyValueError::new_err(format!($err)))
37    };
38    ($fmt:expr, $($arg:tt)*) => {
39        return Err(pyo3::exceptions::PyValueError::new_err(format!($fmt, $($arg)*)))
40    };
41}
42
43fn mimi_cfg(num_codebooks: usize, max_seq_len: Option<usize>) -> mimi::Config {
44    let seanet_cfg = seanet::Config {
45        dimension: 512,
46        channels: 1,
47        causal: true,
48        n_filters: 64,
49        n_residual_layers: 1,
50        activation: candle_nn::Activation::Elu(1.),
51        compress: 2,
52        dilation_base: 2,
53        disable_norm_outer_blocks: 0,
54        final_activation: None,
55        kernel_size: 7,
56        residual_kernel_size: 3,
57        last_kernel_size: 3,
58        lstm: 0,
59        norm: conv::Norm::WeightNorm,
60        pad_mode: conv::PadMode::Constant,
61        ratios: vec![8, 6, 5, 4],
62        true_skip: true,
63    };
64    let transformer_cfg = transformer::Config {
65        d_model: seanet_cfg.dimension,
66        num_heads: 8,
67        num_layers: 8,
68        causal: true,
69        norm_first: true,
70        bias_ff: false,
71        bias_attn: false,
72        layer_scale: Some(0.01),
73        context: 250,
74        conv_kernel_size: 5,
75        use_conv_bias: true,
76        use_conv_block: false,
77        max_period: 10000,
78        positional_embedding: transformer::PositionalEmbedding::Rope,
79        gating: None,
80        norm: mm::NormType::LayerNorm,
81
82        dim_feedforward: 2048,
83        kv_repeat: 1,
84        conv_layout: true, // see builders.py
85        cross_attention: None,
86        max_seq_len: max_seq_len.unwrap_or(8192), // the transformer works at 25hz so this is ~5 mins.
87    };
88    mimi::Config {
89        channels: 1,
90        sample_rate: 24_000.,
91        frame_rate: 12.5,
92        renormalize: true,
93        resample_method: mimi::ResampleMethod::Conv,
94        seanet: seanet_cfg,
95        transformer: transformer_cfg,
96        quantizer_n_q: num_codebooks,
97        quantizer_bins: 2048,
98        quantizer_dim: 256,
99    }
100}
101
102#[pyclass]
103struct Tokenizer {
104    mimi: mimi::Mimi,
105    device: candle::Device,
106    dtype: candle::DType,
107}
108
109#[pymethods]
110impl Tokenizer {
111    #[pyo3(signature = (path, *, num_codebooks=8, dtype="f32", max_seq_len=None))]
112    #[new]
113    fn new(
114        path: std::path::PathBuf,
115        num_codebooks: usize,
116        dtype: &str,
117        max_seq_len: Option<usize>,
118    ) -> PyResult<Self> {
119        let device = candle::Device::Cpu;
120        let dtype = match dtype {
121            "f32" => candle::DType::F32,
122            "f16" => candle::DType::F16,
123            "bf16" => candle::DType::BF16,
124            dtype => py_bail!("unsupported dtype '{dtype}'"),
125        };
126        let vb =
127            unsafe { candle_nn::VarBuilder::from_mmaped_safetensors(&[path], dtype, &device).w()? };
128        let cfg = mimi_cfg(num_codebooks, max_seq_len);
129        let mimi = mimi::Mimi::new(cfg, vb).w()?;
130        Ok(Self { mimi, device, dtype })
131    }
132
133    fn encode(&mut self, pcm_data: numpy::PyReadonlyArray3<f32>) -> PyResult<PyObject> {
134        let py = pcm_data.py();
135        let pcm_data = pcm_data.as_array();
136        let pcm_shape = pcm_data.shape().to_vec();
137        let pcm_data = match pcm_data.to_slice() {
138            None => py_bail!("input data is not contiguous"),
139            Some(data) => data,
140        };
141        let codes = py
142            .allow_threads(|| {
143                let pcm_data = candle::Tensor::from_slice(pcm_data, pcm_shape, &self.device)?
144                    .to_dtype(self.dtype)?;
145                let codes = self.mimi.encode(&pcm_data)?;
146                codes.to_vec3::<u32>()
147            })
148            .w()?;
149        let codes = numpy::PyArray3::from_vec3(py, &codes)?;
150        Ok(codes.into_any().unbind())
151    }
152
153    fn encode_step(&mut self, pcm_data: numpy::PyReadonlyArray3<f32>) -> PyResult<PyObject> {
154        let py = pcm_data.py();
155        let pcm_data = pcm_data.as_array();
156        let pcm_shape = pcm_data.shape().to_vec();
157        let pcm_data = match pcm_data.to_slice() {
158            None => py_bail!("input data is not contiguous"),
159            Some(data) => data,
160        };
161        let codes = py
162            .allow_threads(|| {
163                let pcm_data = candle::Tensor::from_slice(pcm_data, pcm_shape, &self.device)?
164                    .to_dtype(self.dtype)?;
165                let codes = self.mimi.encode_step(&pcm_data.into())?;
166                match codes.as_option() {
167                    Some(codes) => Ok::<_, candle::Error>(Some(codes.to_vec3::<u32>()?)),
168                    None => Ok(None),
169                }
170            })
171            .w()?;
172        match codes {
173            Some(codes) => {
174                let codes = numpy::PyArray3::from_vec3(py, &codes)?;
175                Ok(codes.into_any().unbind())
176            }
177            None => Ok(py.None()),
178        }
179    }
180
181    fn decode(&mut self, codes: numpy::PyReadonlyArray3<u32>, py: Python) -> PyResult<PyObject> {
182        let codes = codes.as_array();
183        let codes_shape = codes.shape().to_vec();
184        let codes = match codes.to_slice() {
185            None => py_bail!("input data is not contiguous"),
186            Some(data) => data,
187        };
188        let pcm = py
189            .allow_threads(|| {
190                let codes = candle::Tensor::from_slice(codes, codes_shape, &self.device)?;
191                let pcm = self.mimi.decode(&codes)?.to_dtype(candle::DType::F32)?;
192                pcm.to_vec3::<f32>()
193            })
194            .w()?;
195        let pcm = numpy::PyArray3::from_vec3(py, &pcm)?;
196        Ok(pcm.into_any().unbind())
197    }
198
199    fn decode_step(
200        &mut self,
201        codes: numpy::PyReadonlyArray3<u32>,
202        py: Python,
203    ) -> PyResult<PyObject> {
204        let codes = codes.as_array();
205        let codes_shape = codes.shape().to_vec();
206        let codes = match codes.to_slice() {
207            None => py_bail!("input data is not contiguous"),
208            Some(data) => data,
209        };
210        let pcm = py
211            .allow_threads(|| {
212                let codes = candle::Tensor::from_slice(codes, codes_shape, &self.device)?;
213                let pcm = self.mimi.decode_step(&codes.into())?;
214                match pcm.as_option() {
215                    Some(pcm) => {
216                        let pcm = pcm.to_dtype(candle::DType::F32)?;
217                        Ok::<_, candle::Error>(Some(pcm.to_vec3::<f32>()?))
218                    }
219                    None => Ok(None),
220                }
221            })
222            .w()?;
223        match pcm {
224            Some(pcm) => {
225                let pcm = numpy::PyArray3::from_vec3(py, &pcm)?;
226                Ok(pcm.into_any().unbind())
227            }
228            None => Ok(py.None()),
229        }
230    }
231
232    fn reset(&mut self) {
233        self.mimi.reset_state()
234    }
235}
236
237#[pyclass]
238struct StreamTokenizer {
239    #[allow(unused)]
240    dtype: candle::DType,
241    encoder_rx: Mutex<mpsc::Receiver<Vec<Vec<u32>>>>,
242    encoder_tx: mpsc::Sender<Vec<f32>>,
243    decoder_rx: Mutex<mpsc::Receiver<Vec<f32>>>,
244    decoder_tx: mpsc::Sender<Vec<Vec<u32>>>,
245}
246
247#[pymethods]
248impl StreamTokenizer {
249    #[pyo3(signature = (path, *, num_codebooks=8, dtype="f32", max_seq_len=None))]
250    #[new]
251    fn new(
252        path: std::path::PathBuf,
253        num_codebooks: usize,
254        dtype: &str,
255        max_seq_len: Option<usize>,
256    ) -> PyResult<Self> {
257        let device = candle::Device::Cpu;
258        let dtype = match dtype {
259            "f32" => candle::DType::F32,
260            "f16" => candle::DType::F16,
261            "bf16" => candle::DType::BF16,
262            dtype => py_bail!("unsupported dtype '{dtype}'"),
263        };
264        let vb =
265            unsafe { candle_nn::VarBuilder::from_mmaped_safetensors(&[path], dtype, &device).w()? };
266        let cfg = mimi_cfg(num_codebooks, max_seq_len);
267        let mut e_mimi = mimi::Mimi::new(cfg, vb).w()?;
268        let mut d_mimi = e_mimi.clone();
269        let (encoder_tx, e_rx) = mpsc::channel::<Vec<f32>>();
270        let (decoder_tx, d_rx) = mpsc::channel::<Vec<Vec<u32>>>();
271        let (d_tx, decoder_rx) = mpsc::channel::<Vec<f32>>();
272        let (e_tx, encoder_rx) = mpsc::channel::<Vec<Vec<u32>>>();
273        std::thread::spawn(move || {
274            while let Ok(pcm_data) = e_rx.recv() {
275                // Can't wait for try blocks to be a thing
276                if let Err(err) = (|| {
277                    let l = pcm_data.len();
278                    let pcm_data =
279                        candle::Tensor::from_vec(pcm_data, (1, 1, l), &candle::Device::Cpu)?
280                            .to_dtype(dtype)?;
281                    let codes = e_mimi.encode_step(&pcm_data.into())?;
282                    if let Some(codes) = codes.as_option() {
283                        let mut codes = codes.to_vec3::<u32>()?;
284                        e_tx.send(codes.remove(0))?;
285                    }
286                    Ok::<_, anyhow::Error>(())
287                })() {
288                    eprintln!("error in encoder thread {err:?}")
289                }
290            }
291        });
292        std::thread::spawn(move || {
293            while let Ok(codes) = d_rx.recv() {
294                if let Err(err) = (|| {
295                    let codes = candle::Tensor::new(codes, &candle::Device::Cpu)?.unsqueeze(2)?;
296                    let pcm_data = d_mimi.decode_step(&codes.into())?;
297                    if let Some(pcm_data) = pcm_data.as_option() {
298                        let mut pcm_data = pcm_data.to_vec3::<f32>()?;
299                        d_tx.send(pcm_data.remove(0).remove(0))?;
300                    }
301                    Ok::<_, anyhow::Error>(())
302                })() {
303                    eprintln!("error in decoder thread {err:?}")
304                }
305            }
306        });
307        Ok(Self {
308            dtype,
309            encoder_rx: Mutex::new(encoder_rx),
310            encoder_tx,
311            decoder_rx: Mutex::new(decoder_rx),
312            decoder_tx,
313        })
314    }
315
316    fn encode(&mut self, pcm_data: numpy::PyReadonlyArray1<f32>) -> PyResult<()> {
317        self.encoder_tx.send(pcm_data.as_array().to_vec()).w()?;
318        Ok(())
319    }
320
321    fn decode(&mut self, codes: numpy::PyReadonlyArray2<u32>) -> PyResult<()> {
322        let codes = codes.as_array();
323        let dims = codes.shape();
324        let codes = match codes.to_slice() {
325            None => py_bail!("input data is not contiguous"),
326            Some(data) => data.to_vec(),
327        };
328        let codes = codes.chunks_exact(dims[1]).map(|v| v.to_vec()).collect::<Vec<_>>();
329        self.decoder_tx.send(codes).w()?;
330        Ok(())
331    }
332
333    fn get_encoded(&mut self, py: Python) -> PyResult<PyObject> {
334        match self.encoder_rx.lock().unwrap().try_recv() {
335            Ok(codes) => {
336                let codes = numpy::PyArray2::from_vec2(py, &codes)?;
337                Ok(codes.into_any().unbind())
338            }
339            Err(mpsc::TryRecvError::Disconnected) => {
340                py_bail!("worker thread disconnected")
341            }
342            Err(mpsc::TryRecvError::Empty) => Ok(py.None()),
343        }
344    }
345
346    fn get_decoded(&mut self, py: Python) -> PyResult<PyObject> {
347        match self.decoder_rx.lock().unwrap().try_recv() {
348            Ok(pcm) => {
349                let pcm = numpy::PyArray1::from_vec(py, pcm);
350                Ok(pcm.into_any().unbind())
351            }
352            Err(mpsc::TryRecvError::Disconnected) => {
353                py_bail!("worker thread disconnected")
354            }
355            Err(mpsc::TryRecvError::Empty) => Ok(py.None()),
356        }
357    }
358}
359
360/// Writes an audio file using the wav format based on pcm data from a numpy array.
361///
362/// This only supports a single channel at the moment so the input array data is expected to have a
363/// single dimension.
364#[pyfunction]
365#[pyo3(signature = (filename, data, sample_rate))]
366fn write_wav(
367    filename: std::path::PathBuf,
368    data: numpy::PyReadonlyArray1<f32>,
369    sample_rate: u32,
370) -> PyResult<()> {
371    let w = std::fs::File::create(&filename).w_f(&filename)?;
372    let mut w = std::io::BufWriter::new(w);
373    let data = data.as_array().to_vec();
374    mm::wav::write_pcm_as_wav(&mut w, &data, sample_rate).w_f(&filename)?;
375    Ok(())
376}
377
378#[pymodule]
379fn rustymimi(_py: Python, m: &Bound<'_, PyModule>) -> PyResult<()> {
380    m.add_class::<Tokenizer>()?;
381    m.add_class::<StreamTokenizer>()?;
382    m.add_function(wrap_pyfunction!(write_wav, m)?)?;
383    Ok(())
384}
rustymimi/lib.rs

rustymimi/
lib.rs
