rlx_kittentts/
assets.rs

1// RLX — versatile ML compiler + runtime.
2// Copyright (C) 2026 Eugene Hauptmann, Nataliya Kosmyna.
3//
4// This program is free software: you can redistribute it and/or modify
5// it under the terms of the GNU General Public License as published by
6// the Free Software Foundation, version 3.
7//
8// This program is distributed in the hope that it will be useful,
9// but WITHOUT ANY WARRANTY; without even the implied warranty of
10// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
11// GNU General Public License for more details.
12//
13// You should have received a copy of the GNU General Public License
14// along with this program. If not, see <https://www.gnu.org/licenses/>.
15
16//! Model directory discovery and path layout.
17
18use std::path::{Path, PathBuf};
19
20use anyhow::{Context, Result, bail};
21
22use crate::config::{DEFAULT_HF_REPO, ModelConfig};
23
24/// Default local checkout from `just fetch-kittentts`.
25pub const DEFAULT_LOCAL_DIR: &str = ".cache/kittentts-mini-0.8";
26
27/// Resolved ONNX + voices (+ optional native weights) for one checkpoint directory.
28#[derive(Debug, Clone)]
29pub struct ModelLayout {
30    pub dir: PathBuf,
31    pub config: ModelConfig,
32    pub onnx: PathBuf,
33    pub voices: PathBuf,
34    pub native_weights: Option<PathBuf>,
35}
36
37impl ModelLayout {
38    pub fn resolve(model_dir: &Path) -> Result<Self> {
39        let dir = model_dir
40            .canonicalize()
41            .unwrap_or_else(|_| model_dir.to_path_buf());
42        let config = ModelConfig::load_from_dir(&dir)?;
43        let onnx = dir.join(&config.model_file);
44        let voices = dir.join(&config.voices);
45        let native_weights = find_native_weights(&dir);
46        if !onnx.is_file() && native_weights.is_none() {
47            bail!(
48                "ONNX model missing: {}\n\
49                 Fetch weights: `just fetch-kittentts` or set RLX_KITTENTTS_DIR",
50                onnx.display()
51            );
52        }
53        if !voices.is_file() {
54            bail!("voices NPZ missing: {}", voices.display());
55        }
56        Ok(Self {
57            native_weights,
58            dir,
59            config,
60            onnx,
61            voices,
62        })
63    }
64
65    /// Voice keys from the NPZ plus friendly alias names from config.
66    pub fn voice_names(&self) -> Result<Vec<String>> {
67        let raw = crate::load_npz(&self.voices)
68            .with_context(|| format!("load voices {}", self.voices.display()))?;
69        let mut names: Vec<String> = raw.into_keys().collect();
70        for alias in self.config.voice_aliases.keys() {
71            if !names.iter().any(|n| n == alias) {
72                names.push(alias.clone());
73            }
74        }
75        names.sort();
76        Ok(names)
77    }
78}
79
80/// Best-effort model directory for CLI and tests.
81pub fn default_model_dir() -> Result<PathBuf> {
82    for key in ["RLX_KITTENTTS_DIR", "KITTENTTS_MODEL_DIR"] {
83        if let Ok(raw) = std::env::var(key) {
84            let p = PathBuf::from(&raw);
85            if layout_exists(&p) {
86                return Ok(p);
87            }
88        }
89    }
90
91    let local = PathBuf::from(DEFAULT_LOCAL_DIR);
92    if layout_exists(&local) {
93        return Ok(local);
94    }
95
96    #[cfg(feature = "hf-download")]
97    if let Ok(p) = hf_snapshot_dir(DEFAULT_HF_REPO) {
98        return Ok(p);
99    }
100
101    if let Some(p) = model_dir_from_bundle_manifest() {
102        return Ok(p);
103    }
104
105    if let Some(p) = hf_hub_cache_snapshot(DEFAULT_HF_REPO) {
106        return Ok(p);
107    }
108
109    bail!(
110        "KittenTTS weights not found.\n\
111         Quick start:\n\
112           just fetch-kittentts\n\
113           just kittentts-demo\n\
114         Or set RLX_KITTENTTS_DIR to a directory containing config.json, \
115         the ONNX file, and voices.npz."
116    )
117}
118
119pub fn layout_exists(dir: &Path) -> bool {
120    dir.join("config.json").is_file()
121}
122
123fn home_dir() -> PathBuf {
124    std::env::var("HOME")
125        .or_else(|_| std::env::var("USERPROFILE"))
126        .map(PathBuf::from)
127        .unwrap_or_else(|_| PathBuf::from("."))
128}
129
130pub fn hf_hub_root() -> PathBuf {
131    if let Ok(h) = std::env::var("HF_HOME") {
132        return PathBuf::from(h).join("hub");
133    }
134    if let Ok(h) = std::env::var("HUGGINGFACE_HUB_CACHE") {
135        return PathBuf::from(h);
136    }
137    home_dir().join(".cache").join("huggingface").join("hub")
138}
139
140/// Locate a cached HF snapshot without downloading (no `hf-hub` dependency).
141fn hf_hub_cache_snapshot(repo_id: &str) -> Option<PathBuf> {
142    let cache_name = format!("models--{}", repo_id.replace('/', "--"));
143    let snapshots = hf_hub_root().join(cache_name).join("snapshots");
144    let mut candidates: Vec<PathBuf> = std::fs::read_dir(&snapshots)
145        .ok()
146        .into_iter()
147        .flatten()
148        .flatten()
149        .map(|e| e.path())
150        .filter(|snap| layout_exists(snap))
151        .collect();
152    candidates.sort();
153    candidates.into_iter().last()
154}
155
156fn model_dir_from_bundle_manifest() -> Option<PathBuf> {
157    let manifest = Path::new(env!("CARGO_MANIFEST_DIR"))
158        .join("../kitten_tts_mini_rlx/weights/rlx_bundle/manifest.json");
159    let data = std::fs::read_to_string(&manifest).ok()?;
160    let v: serde_json::Value = serde_json::from_str(&data).ok()?;
161    let onnx = v.get("source_onnx")?.as_str()?;
162    let dir = PathBuf::from(onnx).parent()?.to_path_buf();
163    layout_exists(&dir).then_some(dir)
164}
165
166/// RLX ONNX bundle colocated under `weights_dir/rlx_bundle`.
167pub fn find_rlx_bundle_colocated(weights_dir: &Path) -> Option<PathBuf> {
168    let weights_dir = weights_dir
169        .canonicalize()
170        .unwrap_or_else(|_| weights_dir.to_path_buf());
171    let in_dir = weights_dir.join("rlx_bundle");
172    if in_dir.join("graph.json").is_file() {
173        return in_dir.canonicalize().ok().or(Some(in_dir));
174    }
175    None
176}
177
178/// RLX ONNX bundle (`graph.json` + weights) under a native weights directory.
179pub fn find_rlx_bundle(weights_dir: &Path) -> Option<PathBuf> {
180    if let Some(dir) = find_rlx_bundle_colocated(weights_dir) {
181        return Some(dir);
182    }
183    if let Ok(raw) =
184        std::env::var("RLX_ONNX_BUNDLE").or_else(|_| std::env::var("KITTEN_RLX_BUNDLE"))
185    {
186        let p = PathBuf::from(raw);
187        if p.join("graph.json").is_file() {
188            return p.canonicalize().ok().or(Some(p));
189        }
190    }
191    None
192}
193
194fn has_native_weight_file(dir: &Path) -> bool {
195    dir.join("model.safetensors").is_file() || dir.join("model.gguf").is_file()
196}
197
198/// Workspace decomposed weights (`crates/kitten_tts_mini_rlx/weights`).
199pub fn default_native_weights_dir() -> Option<PathBuf> {
200    if let Ok(raw) = std::env::var("KITTEN_RLX_WEIGHTS") {
201        let p = PathBuf::from(raw);
202        if has_native_weight_file(&p) || find_rlx_bundle_colocated(&p).is_some() {
203            return Some(p);
204        }
205    }
206    let sibling = Path::new(env!("CARGO_MANIFEST_DIR")).join("../kitten_tts_mini_rlx/weights");
207    let sibling = sibling.canonicalize().unwrap_or(sibling);
208    if has_native_weight_file(&sibling) || find_rlx_bundle_colocated(&sibling).is_some() {
209        return Some(sibling);
210    }
211    None
212}
213
214/// Decomposed RLX weights (`model.safetensors` or `model.gguf`), if present.
215pub fn find_native_weights(model_dir: &Path) -> Option<PathBuf> {
216    if let Ok(raw) = std::env::var("KITTEN_RLX_WEIGHTS") {
217        let p = PathBuf::from(raw);
218        let p = p.canonicalize().unwrap_or(p);
219        if has_native_weight_file(&p) || find_rlx_bundle_colocated(&p).is_some() {
220            return Some(p);
221        }
222    }
223    let model_dir = model_dir
224        .canonicalize()
225        .unwrap_or_else(|_| model_dir.to_path_buf());
226    if has_native_weight_file(&model_dir) || find_rlx_bundle_colocated(&model_dir).is_some() {
227        return Some(model_dir);
228    }
229    default_native_weights_dir()
230}
231
232#[cfg(feature = "hf-download")]
233pub fn hf_snapshot_dir(repo_id: &str) -> Result<PathBuf> {
234    let api = hf_hub::api::sync::ApiBuilder::new()
235        .with_cache_dir(hf_hub_root())
236        .build()
237        .context("hf_hub ApiBuilder")?;
238    let repo = api.model(normalize_repo_id(repo_id));
239    let config = repo.get("config.json").with_context(|| {
240        format!(
241            "locate {repo_id} in Hugging Face cache under {}\n\
242             Download once: `just fetch-kittentts`",
243            hf_hub_root().display()
244        )
245    })?;
246    config
247        .parent()
248        .map(Path::to_path_buf)
249        .context("config.json has no parent (snapshot dir)")
250}
251
252#[cfg(not(feature = "hf-download"))]
253pub fn hf_snapshot_dir(_repo_id: &str) -> Result<PathBuf> {
254    bail!("rebuild with `--features hf-download` on rlx-kittentts")
255}
256
257pub fn normalize_repo_id(repo_id: &str) -> String {
258    if repo_id.contains('/') {
259        repo_id.to_string()
260    } else {
261        format!("KittenML/{repo_id}")
262    }
263}
rlx_kittentts/assets.rs

rlx_kittentts/
assets.rs
