1use std::fs::File;
2use std::io::BufWriter;
3use std::path::Path;
4
5use anyhow::Context;
6use indicatif::{ProgressBar, ProgressStyle};
7use safetensors::Dtype as StDtype;
8use safetensors::SafeTensors;
9
10use zsfm_gguf::{GGMLType, GGUFMetaValue, GGUFWriter};
11use zsfm_hub::ModelFiles;
12
13use crate::config::TimesFMConfig;
14use crate::tensor_map::map_tensor_name;
15
16pub struct ConvertOptions {
17 pub output_dtype: GGMLType,
18}
19
20pub fn convert(
21 model_id: &str,
22 files: &ModelFiles,
23 config: &TimesFMConfig,
24 opts: &ConvertOptions,
25 output_path: &Path,
26) -> anyhow::Result<()> {
27 let mut writer = GGUFWriter::new();
28 write_metadata(&mut writer, model_id, config);
29
30 let safetensors_path = files
31 .safetensors_shards
32 .first()
33 .context("no safetensors file downloaded")?;
34 let raw = std::fs::read(safetensors_path)
35 .with_context(|| format!("read {}", safetensors_path.display()))?;
36 let tensors = SafeTensors::deserialize(&raw).context("deserialize safetensors")?;
37
38 let total = tensors.len();
39 println!("Found {total} tensors.");
40
41 let pb = ProgressBar::new(total as u64);
42 pb.set_style(
43 ProgressStyle::with_template(
44 "{spinner:.green} [{elapsed_precise}] [{bar:40.cyan/blue}] {pos}/{len} {msg}",
45 )
46 .unwrap()
47 .progress_chars("=>-"),
48 );
49
50 let mut mapped = 0usize;
51 let mut skipped: Vec<String> = Vec::new();
52 let mut fallback_count = 0usize;
53
54 for (hf_name, tensor_view) in tensors.tensors() {
55 pb.set_message(hf_name.to_string());
56
57 let gguf_name = match map_tensor_name(&hf_name) {
58 Some(n) => n,
59 None => {
60 skipped.push(hf_name.to_string());
61 pb.inc(1);
62 continue;
63 }
64 };
65
66 let src_dtype = ggml_type_from_st(tensor_view.dtype()).with_context(|| {
67 format!(
68 "tensor {hf_name}: unsupported dtype {:?}",
69 tensor_view.dtype()
70 )
71 })?;
72
73 let raw_data = tensor_view.data();
74 let py_shape = tensor_view.shape();
75 let n_elems: usize = py_shape.iter().product();
76 let innermost = py_shape.last().copied().unwrap_or(1);
77
78 let (dst_dtype, gguf_shape, tensor_data) =
79 if opts.output_dtype == GGMLType::Q8_0 && (innermost % 32 != 0 || n_elems % 32 != 0) {
80 fallback_count += 1;
81 let data = cast_data(raw_data, src_dtype, GGMLType::F32)
82 .with_context(|| format!("tensor {hf_name}: cast failed"))?;
83 let gs = py_shape.iter().rev().map(|&d| d as u64).collect();
84 (GGMLType::F32, gs, data)
85 } else {
86 let dst = opts.output_dtype;
87 let data = cast_data(raw_data, src_dtype, dst)
88 .with_context(|| format!("tensor {hf_name}: cast failed"))?;
89 let gs = py_shape.iter().rev().map(|&d| d as u64).collect();
90 (dst, gs, data)
91 };
92
93 writer.add_tensor(gguf_name, gguf_shape, dst_dtype, tensor_data);
94 mapped += 1;
95 pb.inc(1);
96 }
97
98 pb.finish_with_message("tensors processed");
99
100 if !skipped.is_empty() {
101 eprintln!(
102 "\nWarning: {} tensor(s) skipped (unrecognised names):",
103 skipped.len()
104 );
105 for name in &skipped {
106 eprintln!(" {name}");
107 }
108 }
109 if fallback_count > 0 {
110 eprintln!(
111 "\nNote: {fallback_count} tensor(s) fell back to F32 (too small for Q8_0 blocks)."
112 );
113 }
114
115 println!("Writing {mapped} tensors to {} …", output_path.display());
116 let out_file =
117 File::create(output_path).with_context(|| format!("create {}", output_path.display()))?;
118 let mut buf = BufWriter::new(out_file);
119 writer.write_to(&mut buf)?;
120 println!("Done.");
121 Ok(())
122}
123
124fn write_metadata(writer: &mut GGUFWriter, model_id: &str, cfg: &TimesFMConfig) {
125 writer.add_metadata(
126 "general.architecture",
127 GGUFMetaValue::String("timesfm25".into()),
128 );
129 writer.add_metadata("general.name", GGUFMetaValue::String(model_id.into()));
130 writer.add_metadata(
131 "timesfm25.block_count",
132 GGUFMetaValue::Uint32(cfg.num_layers as u32),
133 );
134 writer.add_metadata(
135 "timesfm25.embedding_length",
136 GGUFMetaValue::Uint32(cfg.d_model as u32),
137 );
138 writer.add_metadata(
139 "timesfm25.feed_forward_length",
140 GGUFMetaValue::Uint32(cfg.d_ff as u32),
141 );
142 writer.add_metadata(
143 "timesfm25.attention.head_count",
144 GGUFMetaValue::Uint32(cfg.num_heads as u32),
145 );
146 writer.add_metadata(
147 "timesfm25.attention.head_dim",
148 GGUFMetaValue::Uint32(cfg.head_dim as u32),
149 );
150 writer.add_metadata(
151 "timesfm25.input_patch_len",
152 GGUFMetaValue::Uint32(cfg.input_patch_len as u32),
153 );
154 writer.add_metadata(
155 "timesfm25.output_patch_len",
156 GGUFMetaValue::Uint32(cfg.output_patch_len as u32),
157 );
158 writer.add_metadata(
159 "timesfm25.decode_index",
160 GGUFMetaValue::Uint32(cfg.decode_index as u32),
161 );
162 writer.add_metadata(
163 "timesfm25.quantile_count",
164 GGUFMetaValue::Uint32(cfg.quantiles.len() as u32),
165 );
166 writer.add_metadata(
167 "timesfm25.quantiles",
168 GGUFMetaValue::ArrayFloat32(cfg.quantiles.clone()),
169 );
170 writer.add_metadata(
171 "timesfm25.rope_theta",
172 GGUFMetaValue::Float64(cfg.rope_theta),
173 );
174 writer.add_metadata(
175 "timesfm25.rms_norm_epsilon",
176 GGUFMetaValue::Float64(cfg.rms_norm_eps),
177 );
178 writer.add_metadata(
179 "timesfm25.context_limit",
180 GGUFMetaValue::Uint32(cfg.context_limit as u32),
181 );
182}
183
184fn ggml_type_from_st(dtype: StDtype) -> anyhow::Result<GGMLType> {
185 match dtype {
186 StDtype::F32 => Ok(GGMLType::F32),
187 StDtype::F16 => Ok(GGMLType::F16),
188 StDtype::BF16 => Ok(GGMLType::BF16),
189 other => anyhow::bail!("unsupported safetensors dtype: {other:?}"),
190 }
191}
192
193fn cast_data(data: &[u8], src: GGMLType, dst: GGMLType) -> anyhow::Result<Vec<u8>> {
194 if src == dst {
195 return Ok(data.to_vec());
196 }
197 if dst == GGMLType::Q8_0 {
198 let f32_values = decode_to_f32(data, src)?;
199 return quantize_q8_0(&f32_values);
200 }
201 match (src, dst) {
202 (GGMLType::F32, GGMLType::F16) => {
203 let vals = parse_f32_le(data)?;
204 let mut out = Vec::with_capacity(vals.len() * 2);
205 for v in vals {
206 out.extend_from_slice(&f32_to_f16_bits(v).to_le_bytes());
207 }
208 Ok(out)
209 }
210 (GGMLType::F32, GGMLType::BF16) => {
211 let vals = parse_f32_le(data)?;
212 let mut out = Vec::with_capacity(vals.len() * 2);
213 for v in vals {
214 out.extend_from_slice(&((v.to_bits() >> 16) as u16).to_le_bytes());
215 }
216 Ok(out)
217 }
218 (GGMLType::F16, GGMLType::BF16) => {
219 let mut out = Vec::with_capacity(data.len());
220 for c in data.chunks_exact(2) {
221 let f32_val = f16_to_f32(u16::from_le_bytes([c[0], c[1]]));
222 out.extend_from_slice(&((f32_val.to_bits() >> 16) as u16).to_le_bytes());
223 }
224 Ok(out)
225 }
226 (GGMLType::BF16, GGMLType::F32) => {
227 let mut out = Vec::with_capacity(data.len() * 2);
228 for c in data.chunks_exact(2) {
229 let bf = u16::from_le_bytes([c[0], c[1]]);
230 out.extend_from_slice(&((bf as u32) << 16).to_le_bytes());
231 }
232 Ok(out)
233 }
234 (GGMLType::BF16, GGMLType::F16) => {
235 let mut out = Vec::with_capacity(data.len());
236 for c in data.chunks_exact(2) {
237 let bf = u16::from_le_bytes([c[0], c[1]]);
238 let f32_val = f32::from_bits((bf as u32) << 16);
239 out.extend_from_slice(&f32_to_f16_bits(f32_val).to_le_bytes());
240 }
241 Ok(out)
242 }
243 (GGMLType::F16, GGMLType::F32) => {
244 let mut out = Vec::with_capacity(data.len() * 2);
245 for c in data.chunks_exact(2) {
246 let f16 = u16::from_le_bytes([c[0], c[1]]);
247 out.extend_from_slice(&f16_to_f32(f16).to_bits().to_le_bytes());
248 }
249 Ok(out)
250 }
251 _ => anyhow::bail!("unsupported cast: {src:?} → {dst:?}"),
252 }
253}
254
255fn decode_to_f32(data: &[u8], src: GGMLType) -> anyhow::Result<Vec<f32>> {
256 match src {
257 GGMLType::F32 => parse_f32_le(data),
258 GGMLType::F16 => data
259 .chunks_exact(2)
260 .map(|c| Ok(f16_to_f32(u16::from_le_bytes([c[0], c[1]]))))
261 .collect(),
262 GGMLType::BF16 => data
263 .chunks_exact(2)
264 .map(|c| {
265 Ok(f32::from_bits(
266 (u16::from_le_bytes([c[0], c[1]]) as u32) << 16,
267 ))
268 })
269 .collect(),
270 GGMLType::Q8_0 => anyhow::bail!("Q8_0 re-quantization not supported"),
271 }
272}
273
274fn quantize_q8_0(values: &[f32]) -> anyhow::Result<Vec<u8>> {
275 const BLOCK: usize = 32;
276 if values.len() % BLOCK != 0 {
277 anyhow::bail!(
278 "Q8_0 requires elem count divisible by {BLOCK}, got {}",
279 values.len()
280 );
281 }
282 let n_blocks = values.len() / BLOCK;
283 let mut out = vec![0u8; n_blocks * 34];
284 for b in 0..n_blocks {
285 let blk = &values[b * BLOCK..(b + 1) * BLOCK];
286 let amax = blk.iter().copied().map(f32::abs).fold(0.0f32, f32::max);
287 let d = if amax == 0.0 { 0.0f32 } else { amax / 127.0 };
288 let d_inv = if d == 0.0 { 0.0f32 } else { 1.0 / d };
289 let base = b * 34;
290 out[base..base + 2].copy_from_slice(&f32_to_f16_bits(d).to_le_bytes());
291 for i in 0..BLOCK {
292 out[base + 2 + i] = (blk[i] * d_inv).round().clamp(-127.0, 127.0) as i8 as u8;
293 }
294 }
295 Ok(out)
296}
297
298fn parse_f32_le(data: &[u8]) -> anyhow::Result<Vec<f32>> {
299 if data.len() % 4 != 0 {
300 anyhow::bail!("f32 length not divisible by 4");
301 }
302 Ok(data
303 .chunks_exact(4)
304 .map(|c| f32::from_le_bytes([c[0], c[1], c[2], c[3]]))
305 .collect())
306}
307
308fn f32_to_f16_bits(v: f32) -> u16 {
309 let bits = v.to_bits();
310 let sign = ((bits >> 16) & 0x8000) as u16;
311 let exp = ((bits >> 23) & 0xFF) as i32;
312 let mantissa = bits & 0x007F_FFFF;
313 if exp == 0xFF {
314 return sign | 0x7C00 | if mantissa != 0 { 0x0200 } else { 0 };
315 }
316 let new_exp = exp - 127 + 15;
317 if new_exp >= 31 {
318 return sign | 0x7C00;
319 }
320 if new_exp <= 0 {
321 if new_exp < -10 {
322 return sign;
323 }
324 let m = (mantissa | 0x0080_0000) >> (1 - new_exp);
325 return sign | (m >> 13) as u16;
326 }
327 sign | ((new_exp as u16) << 10) | (mantissa >> 13) as u16
328}
329
330fn f16_to_f32(bits: u16) -> f32 {
331 let sign = ((bits & 0x8000) as u32) << 16;
332 let exp = ((bits >> 10) & 0x1F) as i32;
333 let mantissa = (bits & 0x03FF) as u32;
334 let f32_bits = if exp == 0 {
335 if mantissa == 0 {
336 sign
337 } else {
338 let mut m = mantissa;
339 let mut e = 0i32;
340 while m & 0x0400 == 0 {
341 m <<= 1;
342 e += 1;
343 }
344 sign | ((127 - 15 - e + 1) as u32) << 23 | (m & 0x03FF) << 13
345 }
346 } else if exp == 31 {
347 sign | 0x7F80_0000 | (mantissa << 13)
348 } else {
349 sign | ((exp + 127 - 15) as u32) << 23 | (mantissa << 13)
350 };
351 f32::from_bits(f32_bits)
352}