qts/model/
loader.rs

1use std::ffi::{CStr, CString};
2use std::fs::File;
3use std::io::{Read, Seek, SeekFrom};
4use std::path::{Path, PathBuf};
5use std::ptr::NonNull;
6
7use qts_ggml::sys;
8
9use super::paths::ModelPaths;
10use crate::Qwen3TtsError;
11
12pub struct GgufFile {
13    ctx: NonNull<sys::gguf_context>,
14    meta_ctx: Option<NonNull<sys::ggml_context>>,
15    path: PathBuf,
16}
17
18#[derive(Debug, Clone)]
19pub struct GgufTensorInfo {
20    pub name: String,
21    pub ty: sys::ggml_type,
22    pub dims: Vec<usize>,
23    pub offset: usize,
24    pub size: usize,
25}
26
27impl GgufFile {
28    pub fn open(path: impl AsRef<Path>) -> Result<Self, Qwen3TtsError> {
29        let path = path.as_ref().to_path_buf();
30        let path_str = path.to_str().ok_or(Qwen3TtsError::InvalidPath)?;
31        let path_c = CString::new(path_str)?;
32
33        let mut meta_ctx = std::ptr::null_mut();
34        let params = sys::gguf_init_params {
35            no_alloc: true,
36            ctx: &mut meta_ctx,
37        };
38
39        let ctx = unsafe { sys::gguf_init_from_file(path_c.as_ptr(), params) };
40        let ctx = NonNull::new(ctx).ok_or_else(|| Qwen3TtsError::InvalidGguf(path.clone()))?;
41
42        Ok(Self {
43            ctx,
44            meta_ctx: NonNull::new(meta_ctx),
45            path,
46        })
47    }
48
49    #[must_use]
50    pub fn ctx_ptr(&self) -> *mut sys::gguf_context {
51        self.ctx.as_ptr()
52    }
53
54    #[must_use]
55    pub fn path(&self) -> &Path {
56        &self.path
57    }
58
59    #[must_use]
60    pub fn meta_ctx_ptr(&self) -> Option<*mut sys::ggml_context> {
61        self.meta_ctx.map(NonNull::as_ptr)
62    }
63
64    #[must_use]
65    pub fn n_tensors(&self) -> i64 {
66        unsafe { sys::gguf_get_n_tensors(self.ctx.as_ptr()) }
67    }
68
69    #[must_use]
70    pub fn data_offset(&self) -> usize {
71        unsafe { sys::gguf_get_data_offset(self.ctx.as_ptr()) }
72    }
73
74    #[must_use]
75    pub fn key_index(&self, key: &str) -> Option<i64> {
76        let key = CString::new(key).ok()?;
77        let idx = unsafe { sys::gguf_find_key(self.ctx.as_ptr(), key.as_ptr()) };
78        (idx >= 0).then_some(idx)
79    }
80
81    #[must_use]
82    pub fn get_u32(&self, key: &str, default: u32) -> u32 {
83        self.key_index(key)
84            .map(|idx| unsafe { sys::gguf_get_val_u32(self.ctx.as_ptr(), idx) })
85            .unwrap_or(default)
86    }
87
88    #[must_use]
89    pub fn get_f32(&self, key: &str, default: f32) -> f32 {
90        self.key_index(key)
91            .map(|idx| unsafe { sys::gguf_get_val_f32(self.ctx.as_ptr(), idx) })
92            .unwrap_or(default)
93    }
94
95    #[must_use]
96    pub fn get_str(&self, key: &str) -> Option<String> {
97        let idx = self.key_index(key)?;
98        if unsafe { sys::gguf_get_kv_type(self.ctx.as_ptr(), idx) }
99            != sys::gguf_type_GGUF_TYPE_STRING
100        {
101            return None;
102        }
103        let value = unsafe { sys::gguf_get_val_str(self.ctx.as_ptr(), idx) };
104        if value.is_null() {
105            return None;
106        }
107        Some(
108            unsafe { CStr::from_ptr(value) }
109                .to_string_lossy()
110                .into_owned(),
111        )
112    }
113
114    #[must_use]
115    pub fn get_arr_str(&self, key: &str) -> Option<Vec<String>> {
116        let idx = self.key_index(key)?;
117        let n = unsafe { sys::gguf_get_arr_n(self.ctx.as_ptr(), idx) };
118        let mut out = Vec::with_capacity(n);
119        for i in 0..n {
120            let s = unsafe { sys::gguf_get_arr_str(self.ctx.as_ptr(), idx, i) };
121            if s.is_null() {
122                continue;
123            }
124            out.push(unsafe { CStr::from_ptr(s) }.to_string_lossy().into_owned());
125        }
126        Some(out)
127    }
128
129    pub fn tensor_info(&self, name: &str) -> Result<GgufTensorInfo, Qwen3TtsError> {
130        let meta_ctx = self
131            .meta_ctx_ptr()
132            .ok_or_else(|| Qwen3TtsError::InvalidGguf(self.path.clone()))?;
133        let name_c = CString::new(name)?;
134        let tensor = unsafe { sys::ggml_get_tensor(meta_ctx, name_c.as_ptr()) };
135        let tensor =
136            NonNull::new(tensor).ok_or_else(|| Qwen3TtsError::MissingTensor(name.into()))?;
137
138        let idx = self
139            .tensor_index(name)
140            .ok_or_else(|| Qwen3TtsError::MissingTensor(name.into()))?;
141        let n_dims = unsafe { sys::ggml_n_dims(tensor.as_ptr()) } as usize;
142        let dims = (0..n_dims)
143            .map(|dim| unsafe { (*tensor.as_ptr()).ne[dim] as usize })
144            .collect::<Vec<_>>();
145
146        Ok(GgufTensorInfo {
147            name: name.into(),
148            ty: unsafe { sys::gguf_get_tensor_type(self.ctx.as_ptr(), idx) },
149            dims,
150            offset: unsafe { sys::gguf_get_tensor_offset(self.ctx.as_ptr(), idx) },
151            size: unsafe { sys::gguf_get_tensor_size(self.ctx.as_ptr(), idx) },
152        })
153    }
154
155    pub fn read_tensor_f32(&self, name: &str) -> Result<(GgufTensorInfo, Vec<f32>), Qwen3TtsError> {
156        let info = self.tensor_info(name)?;
157        let raw = self.read_tensor_bytes_by_info(&info)?;
158
159        let element_count = info
160            .dims
161            .iter()
162            .copied()
163            .try_fold(1usize, |acc, dim| acc.checked_mul(dim))
164            .ok_or_else(|| Qwen3TtsError::InvalidTensor(name.into()))?;
165
166        let data = match info.ty {
167            sys::ggml_type_GGML_TYPE_F32 => {
168                if info.size != element_count * std::mem::size_of::<f32>() {
169                    return Err(Qwen3TtsError::InvalidTensor(name.into()));
170                }
171                raw.as_chunks::<4>()
172                    .0
173                    .iter()
174                    .map(|bytes| f32::from_le_bytes(*bytes))
175                    .collect()
176            }
177            sys::ggml_type_GGML_TYPE_F16 => {
178                if info.size != element_count * std::mem::size_of::<sys::ggml_fp16_t>() {
179                    return Err(Qwen3TtsError::InvalidTensor(name.into()));
180                }
181                let mut out = Vec::with_capacity(element_count);
182                for bytes in raw.as_chunks::<2>().0 {
183                    let value = u16::from_le_bytes(*bytes);
184                    out.push(unsafe { sys::ggml_fp16_to_fp32(value) });
185                }
186                out
187            }
188            _ => return Err(Qwen3TtsError::UnsupportedTensorType(name.into())),
189        };
190
191        Ok((info, data))
192    }
193
194    pub fn read_tensor_bytes(
195        &self,
196        name: &str,
197    ) -> Result<(GgufTensorInfo, Vec<u8>), Qwen3TtsError> {
198        let info = self.tensor_info(name)?;
199        let raw = self.read_tensor_bytes_by_info(&info)?;
200        Ok((info, raw))
201    }
202
203    fn tensor_index(&self, name: &str) -> Option<i64> {
204        for idx in 0..self.n_tensors() {
205            let tensor_name = unsafe { sys::gguf_get_tensor_name(self.ctx.as_ptr(), idx) };
206            if tensor_name.is_null() {
207                continue;
208            }
209            if unsafe { CStr::from_ptr(tensor_name) }.to_string_lossy() == name {
210                return Some(idx);
211            }
212        }
213        None
214    }
215
216    fn read_tensor_bytes_by_info(&self, info: &GgufTensorInfo) -> Result<Vec<u8>, Qwen3TtsError> {
217        let mut file = File::open(&self.path)?;
218        file.seek(SeekFrom::Start((self.data_offset() + info.offset) as u64))?;
219        let mut raw = vec![0u8; info.size];
220        file.read_exact(&mut raw)?;
221        Ok(raw)
222    }
223}
224
225impl Drop for GgufFile {
226    fn drop(&mut self) {
227        unsafe {
228            sys::gguf_free(self.ctx.as_ptr());
229            if let Some(meta_ctx) = self.meta_ctx {
230                sys::ggml_free(meta_ctx.as_ptr());
231            }
232        }
233    }
234}
235
236/// Validate that required model artifacts exist and the main checkpoint parses as GGUF.
237pub fn load_and_validate(paths: &ModelPaths) -> Result<(), Qwen3TtsError> {
238    if !paths.main_gguf.is_file() {
239        return Err(Qwen3TtsError::ModelFile(paths.main_gguf.clone()));
240    }
241    if !paths.vocoder_onnx.is_file() {
242        return Err(Qwen3TtsError::ModelFile(paths.vocoder_onnx.clone()));
243    }
244    validate_gguf_file(&paths.main_gguf)?;
245    Ok(())
246}
247
248fn validate_gguf_file(path: &Path) -> Result<(), Qwen3TtsError> {
249    GgufFile::open(path).map(|_| ())
250}
qts/model/loader.rs

qts/model/
loader.rs
