Skip to main content

qts/model/
paths.rs

1use std::path::{Path, PathBuf};
2
3/// Paths to the main GGUF plus vocoder artifacts.
4#[derive(Debug, Clone)]
5pub struct ModelPaths {
6    /// Main TTS checkpoint (e.g. `qwen3-tts-0.6b-f16.gguf`).
7    pub main_gguf: PathBuf,
8    /// Vocoder artifact exported for ONNX Runtime.
9    pub vocoder_onnx: PathBuf,
10}
11
12impl ModelPaths {
13    pub fn new(main_gguf: PathBuf, vocoder_onnx: PathBuf) -> Self {
14        Self {
15            main_gguf,
16            vocoder_onnx,
17        }
18    }
19
20    /// Resolve under a single directory using the conventional filenames.
21    pub fn from_model_dir(dir: impl AsRef<Path>) -> Self {
22        let dir = dir.as_ref();
23        Self {
24            main_gguf: choose_model_file(
25                dir,
26                &[
27                    "qwen3-tts-0.6b-f16.gguf",
28                    "qwen3-tts-0.6b-q8_0.gguf",
29                    "qwen3-tts-0.6b-q6_k.gguf",
30                    "qwen3-tts-0.6b-q5_k.gguf",
31                    "qwen3-tts-0.6b-q4_k.gguf",
32                ],
33            ),
34            vocoder_onnx: dir.join("qwen3-tts-vocoder.onnx"),
35        }
36    }
37
38    pub fn main_exists(&self) -> bool {
39        self.main_gguf.is_file()
40    }
41
42    pub fn vocoder_exists(&self) -> bool {
43        self.vocoder_onnx.is_file()
44    }
45}
46
47/// Separate mode-specific 1.7B QTS exports. All three GGUF paths may name one bundled file.
48#[derive(Debug, Clone)]
49pub struct ModeModelPaths {
50    pub talker_gguf: PathBuf,
51    pub code_predictor_gguf: PathBuf,
52    pub tokenizer_gguf: PathBuf,
53    pub vocoder_onnx: PathBuf,
54}
55
56impl ModeModelPaths {
57    pub fn from_model_dir(dir: impl AsRef<Path>, mode: crate::QtsMode) -> Self {
58        let dir = dir.as_ref();
59        let stem = format!("qwen3-tts-1.7b-{}", mode.as_str());
60        let choose = |component: &str| {
61            let candidates = ["f16", "q8_0", "q6_k", "q5_k", "q4_k"]
62                .map(|quant| format!("{stem}-{component}-{quant}.gguf"));
63            let refs = candidates.iter().map(String::as_str).collect::<Vec<_>>();
64            choose_model_file(dir, &refs)
65        };
66        let bundled = ["f16", "q8_0", "q6_k", "q5_k", "q4_k"]
67            .map(|quant| dir.join(format!("{stem}-{quant}.gguf")))
68            .into_iter()
69            .find(|path| path.is_file());
70        Self {
71            talker_gguf: bundled.clone().unwrap_or_else(|| choose("talker")),
72            code_predictor_gguf: bundled.clone().unwrap_or_else(|| choose("code-predictor")),
73            tokenizer_gguf: bundled.unwrap_or_else(|| dir.join(format!("{stem}-tokenizer.gguf"))),
74            vocoder_onnx: dir.join("qwen3-tts-vocoder.onnx"),
75        }
76    }
77}
78
79fn choose_model_file(dir: &Path, candidates: &[&str]) -> PathBuf {
80    for candidate in candidates {
81        let path = dir.join(candidate);
82        if path.is_file() {
83            return path;
84        }
85    }
86
87    dir.join(candidates[0])
88}
89
90#[cfg(test)]
91mod tests {
92    use super::*;
93
94    #[test]
95    fn from_model_dir_names() {
96        let p = ModelPaths::from_model_dir("/models");
97        assert!(p.main_gguf.ends_with("qwen3-tts-0.6b-f16.gguf"));
98        assert!(p.vocoder_onnx.ends_with("qwen3-tts-vocoder.onnx"));
99    }
100
101    #[test]
102    fn from_model_dir_prefers_existing_main_quantized_files() {
103        let dir = std::env::temp_dir().join("qwen_tts_native_paths_quantized");
104        let _ = std::fs::remove_dir_all(&dir);
105        std::fs::create_dir_all(&dir).unwrap();
106        std::fs::write(dir.join("qwen3-tts-0.6b-q4_k.gguf"), b"").unwrap();
107
108        let p = ModelPaths::from_model_dir(&dir);
109        assert!(p.main_gguf.ends_with("qwen3-tts-0.6b-q4_k.gguf"));
110        assert!(p.vocoder_onnx.ends_with("qwen3-tts-vocoder.onnx"));
111
112        let _ = std::fs::remove_dir_all(dir);
113    }
114    #[test]
115    fn mode_specific_1_7b_paths_share_vocoder() {
116        for mode in [
117            crate::QtsMode::Base,
118            crate::QtsMode::CustomVoice,
119            crate::QtsMode::VoiceDesign,
120        ] {
121            let paths = ModeModelPaths::from_model_dir("/missing-qts-models", mode);
122            let stem = format!("qwen3-tts-1.7b-{}", mode.as_str());
123            assert!(paths
124                .talker_gguf
125                .ends_with(format!("{stem}-talker-f16.gguf")));
126            assert!(paths
127                .code_predictor_gguf
128                .ends_with(format!("{stem}-code-predictor-f16.gguf")));
129            assert!(paths
130                .tokenizer_gguf
131                .ends_with(format!("{stem}-tokenizer.gguf")));
132            assert!(paths.vocoder_onnx.ends_with("qwen3-tts-vocoder.onnx"));
133        }
134    }
135    #[test]
136    fn mode_paths_resolve_split_quantization_or_bundled_export() {
137        let fixture = crate::model::test_fixture::TestGguf::write(&[], &[], &[]);
138        let dir = fixture.0.with_extension("models");
139        std::fs::create_dir(&dir).unwrap();
140        std::fs::write(dir.join("qwen3-tts-1.7b-customvoice-talker-q4_k.gguf"), b"").unwrap();
141        std::fs::write(
142            dir.join("qwen3-tts-1.7b-customvoice-code-predictor-q8_0.gguf"),
143            [],
144        )
145        .unwrap();
146        let paths = ModeModelPaths::from_model_dir(&dir, crate::QtsMode::CustomVoice);
147        assert!(paths
148            .talker_gguf
149            .ends_with("qwen3-tts-1.7b-customvoice-talker-q4_k.gguf"));
150        assert!(paths
151            .code_predictor_gguf
152            .ends_with("qwen3-tts-1.7b-customvoice-code-predictor-q8_0.gguf"));
153        std::fs::write(dir.join("qwen3-tts-1.7b-customvoice-q8_0.gguf"), b"").unwrap();
154        let paths = ModeModelPaths::from_model_dir(&dir, crate::QtsMode::CustomVoice);
155        assert!(paths
156            .talker_gguf
157            .ends_with("qwen3-tts-1.7b-customvoice-q8_0.gguf"));
158        assert_eq!(paths.talker_gguf, paths.code_predictor_gguf);
159        assert_eq!(paths.talker_gguf, paths.tokenizer_gguf);
160        std::fs::remove_dir_all(dir).unwrap();
161    }
162}