rlx_moshi/
backend.rs

1//! Unified Moshi LM backend — eager CPU (ndarray) or native RLX graphs.
2
3use crate::checkpoint::MoshiCheckpoint;
4use crate::config::{GenerateConfig, MoshiVariant};
5use crate::generate::GenerateState;
6use crate::lm::LmModel;
7use crate::rlx_gen::{RlxGenerateState, RlxLm};
8use crate::sampling::LogitsProcessor;
9use crate::weights::{open_lm, open_lm_from_weights};
10use anyhow::{Result, ensure};
11use rlx_runtime::Device;
12use std::path::Path;
13
14/// How to load and run the Moshi temporal + DepFormer stack.
15pub enum MoshiLm {
16    /// Eager CPU reference (ndarray).
17    Eager(LmModel),
18    /// Native RLX graphs (CPU or GPU).
19    Rlx(RlxLm),
20}
21
22/// Per-generation autoregressive state.
23pub enum MoshiGenState {
24    Eager(GenerateState),
25    Rlx(RlxGenerateState),
26}
27
28impl MoshiLm {
29    pub fn open(
30        model_dir: &Path,
31        variant: MoshiVariant,
32        checkpoint: MoshiCheckpoint,
33        device: Device,
34    ) -> Result<Self> {
35        // Native RLX graph backend: explicit opt-in, or any non-CPU device.
36        if std::env::var_os("RLX_MOSHI_NATIVE").is_some() || device != Device::Cpu {
37            return Ok(Self::Rlx(RlxLm::open(
38                model_dir, variant, checkpoint, device,
39            )?));
40        }
41
42        // CPU eager (ndarray) path. All checkpoint formats load candle-free.
43        let cfg = variant.lm_config();
44        let weights_path = checkpoint.lm_weights_path(model_dir);
45        if checkpoint.is_mlx() {
46            let weights =
47                crate::mlx_weights::load_eager_weight_map(&weights_path, checkpoint, &cfg)?;
48            return Ok(Self::Eager(open_lm_from_weights(cfg, weights)?));
49        }
50        let lm = if checkpoint.is_gguf() {
51            let weights = crate::gguf::load_gguf_weight_map(&weights_path, &cfg)?;
52            open_lm_from_weights(cfg, weights)?
53        } else {
54            open_lm(model_dir, cfg)?
55        };
56        Ok(Self::Eager(lm))
57    }
58
59    pub fn config(&self) -> &crate::config::LmConfig {
60        match self {
61            Self::Eager(m) => m.config(),
62            Self::Rlx(m) => m.config(),
63        }
64    }
65
66    pub fn reset_state(&mut self) {
67        match self {
68            Self::Eager(m) => m.reset_state(),
69            Self::Rlx(_) => {}
70        }
71    }
72
73    pub fn new_gen_state(
74        &self,
75        max_steps: usize,
76        text_lp: LogitsProcessor,
77        audio_lp: LogitsProcessor,
78        gen_cfg: GenerateConfig,
79    ) -> Result<MoshiGenState> {
80        match self {
81            Self::Eager(_) => Ok(MoshiGenState::Eager(GenerateState::new(
82                max_steps, text_lp, audio_lp, gen_cfg,
83            ))),
84            Self::Rlx(_) => Ok(MoshiGenState::Rlx(RlxGenerateState::new(
85                max_steps, text_lp, audio_lp, gen_cfg,
86            ))),
87        }
88    }
89}
90
91impl MoshiGenState {
92    pub fn config(&self) -> &GenerateConfig {
93        match self {
94            Self::Eager(s) => s.config(),
95            Self::Rlx(s) => s.config(),
96        }
97    }
98
99    pub fn step_idx(&self) -> usize {
100        match self {
101            Self::Eager(s) => s.step_idx(),
102            Self::Rlx(s) => s.step_idx(),
103        }
104    }
105
106    pub fn text_tokens(&self) -> &[u32] {
107        match self {
108            Self::Eager(s) => s.text_tokens(),
109            Self::Rlx(s) => s.text_tokens(),
110        }
111    }
112
113    pub fn step(&mut self, lm: &mut MoshiLm, text_token: u32, input_audio: &[u32]) -> Result<u32> {
114        match (self, lm) {
115            (Self::Eager(s), MoshiLm::Eager(m)) => s.step(m, text_token, input_audio),
116            (Self::Rlx(s), MoshiLm::Rlx(m)) => s.step(m, text_token, input_audio),
117            #[allow(unreachable_patterns)]
118            _ => {
119                ensure!(false, "Moshi LM / generation backend mismatch");
120                Ok(0)
121            }
122        }
123    }
124
125    pub fn last_audio_frame(&self) -> Option<Vec<u32>> {
126        match self {
127            Self::Eager(s) => s.last_audio_frame(),
128            Self::Rlx(s) => s.last_audio_frame(),
129        }
130    }
131
132    pub fn reset(
133        &mut self,
134        lm: &mut MoshiLm,
135        max_steps: usize,
136        text_lp: LogitsProcessor,
137        audio_lp: LogitsProcessor,
138    ) -> Result<()> {
139        let _ = (max_steps, text_lp, audio_lp);
140        match (self, lm) {
141            (Self::Eager(s), MoshiLm::Eager(m)) => {
142                s.reset(m);
143                Ok(())
144            }
145            (Self::Rlx(s), MoshiLm::Rlx(_m)) => {
146                // lps are kept; just clear KV cache + token buffers.
147                s.reset();
148                Ok(())
149            }
150            #[allow(unreachable_patterns)]
151            _ => anyhow::bail!("Moshi LM / generation backend mismatch"),
152        }
153    }
154}
155
156/// Pick the LM device: honor CPU, keep an available GPU, else fall back to CPU.
157pub fn resolve_lm_device(requested: Device, _checkpoint: MoshiCheckpoint) -> Device {
158    if requested == Device::Cpu {
159        return Device::Cpu;
160    }
161    if rlx_runtime::is_available(requested) {
162        requested
163    } else {
164        Device::Cpu
165    }
166}
rlx_moshi/backend.rs

rlx_moshi/
backend.rs
