1use std::path::{Path, PathBuf};
2
3#[derive(Debug, Clone)]
5pub struct ModelPaths {
6 pub main_gguf: PathBuf,
8 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 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#[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}