rlx_moshi/
checkpoint.rs

1//! Kyutai Moshi weight checkpoint presets and HuggingFace routing.
2//!
3//! Each [`MoshiCheckpoint`] selects a quantization / file layout. Pair it with
4//! [`MoshiVoice`] (via [`crate::config::MoshiVariant`]) to resolve the correct
5//! `kyutai/moshiko-*` or `kyutai/moshika-*` repo.
6
7use crate::config::MoshiVoice;
8use std::path::{Path, PathBuf};
9
10/// Published Moshi LM weight preset (quantization + on-disk layout).
11#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
12pub enum MoshiCheckpoint {
13    /// `kyutai/moshiko-candle-bf16` — `model.safetensors` (~14 GB).
14    Bf16Safetensors,
15    /// `kyutai/moshiko-candle-q8` — `model.q8.gguf` (~8 GB).
16    Q8Gguf,
17    /// `kyutai/moshiko-mlx-q4` — `model.q4.safetensors` (~4 GB).
18    Q4MlxSafetensors,
19    /// `kyutai/moshiko-mlx-q8` — `model.q8.safetensors` (~7 GB).
20    Q8MlxSafetensors,
21    /// `kyutai/moshiko-mlx-bf16` — `model.safetensors` (~14 GB, MLX key layout).
22    MlxBf16Safetensors,
23}
24
25impl MoshiCheckpoint {
26    /// HuggingFace model repo for this preset and voice family.
27    pub fn hf_repo(self, voice: MoshiVoice) -> &'static str {
28        match (voice, self) {
29            (MoshiVoice::Moshiko, Self::Bf16Safetensors) => "kyutai/moshiko-candle-bf16",
30            (MoshiVoice::Moshiko, Self::Q8Gguf) => "kyutai/moshiko-candle-q8",
31            (MoshiVoice::Moshiko, Self::Q4MlxSafetensors) => "kyutai/moshiko-mlx-q4",
32            (MoshiVoice::Moshiko, Self::Q8MlxSafetensors) => "kyutai/moshiko-mlx-q8",
33            (MoshiVoice::Moshiko, Self::MlxBf16Safetensors) => "kyutai/moshiko-mlx-bf16",
34            (MoshiVoice::Moshika, Self::Bf16Safetensors) => "kyutai/moshika-candle-bf16",
35            (MoshiVoice::Moshika, Self::Q8Gguf) => "kyutai/moshika-candle-q8",
36            (MoshiVoice::Moshika, Self::Q4MlxSafetensors) => "kyutai/moshika-mlx-q4",
37            (MoshiVoice::Moshika, Self::Q8MlxSafetensors) => "kyutai/moshika-mlx-q8",
38            (MoshiVoice::Moshika, Self::MlxBf16Safetensors) => "kyutai/moshika-mlx-bf16",
39        }
40    }
41
42    /// Default `.cache/…` directory for a voice + preset pair (unless `RLX_MOSHI_DIR` overrides).
43    pub fn default_cache_dir(self, voice: MoshiVoice) -> PathBuf {
44        let prefix = match voice {
45            MoshiVoice::Moshiko => "moshiko",
46            MoshiVoice::Moshika => "moshika",
47        };
48        let suffix = match self {
49            Self::Bf16Safetensors => "",
50            Self::Q8Gguf => "-q8",
51            Self::Q4MlxSafetensors => "-mlx-q4",
52            Self::Q8MlxSafetensors => "-mlx-q8",
53            Self::MlxBf16Safetensors => "-mlx-bf16",
54        };
55        PathBuf::from(".cache").join(format!("{prefix}{suffix}"))
56    }
57
58    /// Primary LM weights filename inside a model directory.
59    pub fn lm_filename(self) -> &'static str {
60        match self {
61            Self::Bf16Safetensors | Self::MlxBf16Safetensors => "model.safetensors",
62            Self::Q8Gguf => "model.q8.gguf",
63            Self::Q4MlxSafetensors => "model.q4.safetensors",
64            Self::Q8MlxSafetensors => "model.q8.safetensors",
65        }
66    }
67
68    /// True when weights are Kyutai Q8 GGUF (`model.q8.gguf`).
69    pub fn is_gguf(self) -> bool {
70        matches!(self, Self::Q8Gguf)
71    }
72
73    /// MLX key layout (Q4/Q8/bf16 safetensors).
74    pub fn is_mlx(self) -> bool {
75        matches!(
76            self,
77            Self::Q4MlxSafetensors | Self::Q8MlxSafetensors | Self::MlxBf16Safetensors
78        )
79    }
80
81    /// Affine-quantized MLX weights (U32 packed + scales/biases).
82    pub fn is_mlx_quantized(self) -> bool {
83        matches!(self, Self::Q4MlxSafetensors | Self::Q8MlxSafetensors)
84    }
85
86    /// Bits per weight and group size for affine MLX quants; `(0, 0)` for full-precision MLX bf16.
87    pub fn mlx_quant_params(self) -> (u32, usize) {
88        match self {
89            Self::Q4MlxSafetensors => (4, 32),
90            Self::Q8MlxSafetensors => (8, 64),
91            _ => (0, 0),
92        }
93    }
94
95    /// Whether the Candle `moshi` GPU backend can load this preset.
96    pub fn gpu_loadable(self) -> bool {
97        self.candle_compatible() || self.is_mlx()
98    }
99
100    /// Candle-native tensor names (bf16 safetensors or Q8 GGUF).
101    pub fn candle_compatible(self) -> bool {
102        matches!(self, Self::Bf16Safetensors | Self::Q8Gguf)
103    }
104
105    /// Approximate published size class for docs / fetch hints.
106    pub fn size_class_gb(self) -> f32 {
107        match self {
108            Self::Bf16Safetensors | Self::MlxBf16Safetensors => 14.0,
109            Self::Q8Gguf => 8.0,
110            Self::Q8MlxSafetensors => 7.0,
111            Self::Q4MlxSafetensors => 4.0,
112        }
113    }
114
115    /// Parse CLI / env preset names (`bf16`, `q8`, `q4`, `q8-mlx`, `mlx-bf16`, …).
116    pub fn parse(name: &str) -> Option<Self> {
117        match name.to_ascii_lowercase().as_str() {
118            "bf16" | "f32" | "safetensors" | "default" => Some(Self::Bf16Safetensors),
119            "q8" | "q8-gguf" | "gguf" => Some(Self::Q8Gguf),
120            "q4" | "q4-mlx" | "mlx-q4" => Some(Self::Q4MlxSafetensors),
121            "q8-mlx" | "mlx-q8" => Some(Self::Q8MlxSafetensors),
122            "mlx-bf16" | "bf16-mlx" | "mlx-bf16-safetensors" => Some(Self::MlxBf16Safetensors),
123            _ => None,
124        }
125    }
126
127    /// Read `RLX_MOSHI_CHECKPOINT`, defaulting to Candle bf16 safetensors.
128    pub fn from_env_or_default() -> Self {
129        std::env::var("RLX_MOSHI_CHECKPOINT")
130            .ok()
131            .and_then(|s| Self::parse(&s))
132            .unwrap_or(Self::Bf16Safetensors)
133    }
134
135    /// Absolute path to LM weights inside `model_dir`.
136    pub fn lm_weights_path(self, model_dir: &Path) -> PathBuf {
137        model_dir.join(self.lm_filename())
138    }
139
140    /// SentencePiece model path (shared across Kyutai repos).
141    pub fn tokenizer_path(self, model_dir: &Path) -> PathBuf {
142        let _ = self;
143        model_dir.join("tokenizer_spm_32k_3.model")
144    }
145
146    /// Shared tokenizer + mimi sidecar files (same across Kyutai repos).
147    pub const SHARED_FILES: &'static [&'static str] = &[
148        "tokenizer_spm_32k_3.model",
149        "tokenizer-e351c8d8-checkpoint125.safetensors",
150    ];
151}
rlx_moshi/checkpoint.rs

rlx_moshi/
checkpoint.rs
