use std::collections::HashMap;
use std::fs::File;
use std::io::{BufWriter, Seek, Write};
use std::path::Path;
use super::config::VisionConfig;
use super::convert::VitTensor;
use super::VitConvertError;
const GGUF_MAGIC: [u8; 4] = [0x47, 0x47, 0x55, 0x46]; const GGUF_VERSION: u32 = 3;
const GGML_TYPE_F16: u32 = 1;
pub const GGML_TYPE_F32: u32 = 0;
const ALIGNMENT: u64 = 32;
#[derive(Debug, Clone)]
enum MetaValue {
String(String),
Uint32(u32),
Float32(f32),
Bool(bool),
ArrayBool(Vec<bool>),
ArrayFloat32(Vec<f32>),
}
const GGUF_TYPE_UINT32: u32 = 4;
const GGUF_TYPE_FLOAT32: u32 = 6;
const GGUF_TYPE_BOOL: u32 = 7;
const GGUF_TYPE_STRING: u32 = 8;
const GGUF_TYPE_ARRAY: u32 = 9;
pub fn write_mmproj_gguf(
output: &Path,
vision_config: &VisionConfig,
tensors: &HashMap<String, VitTensor>,
) -> Result<(), VitConvertError> {
let file = File::create(output)
.map_err(|e| VitConvertError::GgufEmit(format!("create {:?}: {}", output, e)))?;
let mut w = BufWriter::new(file);
let mut names: Vec<&String> = tensors.keys().collect();
names.sort();
let tensor_count = names.len() as u64;
let metadata = build_metadata(vision_config);
let kv_count = metadata.len() as u64;
w.write_all(&GGUF_MAGIC)?;
w.write_all(&GGUF_VERSION.to_le_bytes())?;
w.write_all(&tensor_count.to_le_bytes())?;
w.write_all(&kv_count.to_le_bytes())?;
for (key, value) in &metadata {
write_kv(&mut w, key, value)?;
}
struct Info<'a> {
name: &'a str,
shape: &'a [usize],
dtype: u32,
size: u64,
offset: u64,
}
let mut infos: Vec<Info> = Vec::with_capacity(names.len());
let mut running_offset: u64 = 0;
for name in &names {
let t = &tensors[*name];
let numel: u64 = t.shape.iter().product::<usize>() as u64;
let (gguf_dtype, bytes_per_elem) = match t.dtype {
crate::ir::DType::F32 => (GGML_TYPE_F32, 4u64),
crate::ir::DType::F16 => (GGML_TYPE_F16, 2u64),
ref other => {
return Err(VitConvertError::GgufEmit(format!(
"write_mmproj_gguf: unsupported dtype {:?} for tensor {:?} \
— only F32 (norms + biases) and F16 (weights) are emitted",
other, name
)));
}
};
let size = numel * bytes_per_elem;
infos.push(Info {
name: name.as_str(),
shape: &t.shape,
dtype: gguf_dtype,
size,
offset: running_offset,
});
running_offset = align_up(running_offset + size, ALIGNMENT);
}
for info in &infos {
write_gguf_string(&mut w, info.name)?;
w.write_all(&(info.shape.len() as u32).to_le_bytes())?;
for dim in info.shape.iter().rev() {
w.write_all(&(*dim as u64).to_le_bytes())?;
}
w.write_all(&info.dtype.to_le_bytes())?;
w.write_all(&info.offset.to_le_bytes())?;
}
let header_end = current_pos(&mut w)?;
let data_start = align_up(header_end, ALIGNMENT);
for _ in header_end..data_start {
w.write_all(&[0u8])?;
}
for (i, info) in infos.iter().enumerate() {
let t = &tensors[info.name];
w.write_all(&t.data)?;
if i + 1 < infos.len() {
let next_expected = infos[i + 1].offset;
let written = info.offset + info.size;
for _ in written..next_expected {
w.write_all(&[0u8])?;
}
}
}
w.flush()
.map_err(|e| VitConvertError::GgufEmit(format!("flush: {}", e)))?;
Ok(())
}
fn align_up(n: u64, to: u64) -> u64 {
let r = n % to;
if r == 0 {
n
} else {
n + (to - r)
}
}
fn current_pos(w: &mut BufWriter<File>) -> std::io::Result<u64> {
w.stream_position()
}
fn write_gguf_string<W: Write>(w: &mut W, s: &str) -> std::io::Result<()> {
w.write_all(&(s.len() as u64).to_le_bytes())?;
w.write_all(s.as_bytes())
}
fn write_kv<W: Write>(w: &mut W, key: &str, value: &MetaValue) -> std::io::Result<()> {
write_gguf_string(w, key)?;
match value {
MetaValue::String(s) => {
w.write_all(&GGUF_TYPE_STRING.to_le_bytes())?;
write_gguf_string(w, s)
}
MetaValue::Uint32(v) => {
w.write_all(&GGUF_TYPE_UINT32.to_le_bytes())?;
w.write_all(&v.to_le_bytes())
}
MetaValue::Float32(v) => {
w.write_all(&GGUF_TYPE_FLOAT32.to_le_bytes())?;
w.write_all(&v.to_le_bytes())
}
MetaValue::Bool(v) => {
w.write_all(&GGUF_TYPE_BOOL.to_le_bytes())?;
w.write_all(&[*v as u8])
}
MetaValue::ArrayBool(arr) => {
w.write_all(&GGUF_TYPE_ARRAY.to_le_bytes())?;
w.write_all(&GGUF_TYPE_BOOL.to_le_bytes())?;
w.write_all(&(arr.len() as u64).to_le_bytes())?;
for &b in arr {
w.write_all(&[b as u8])?;
}
Ok(())
}
MetaValue::ArrayFloat32(arr) => {
w.write_all(&GGUF_TYPE_ARRAY.to_le_bytes())?;
w.write_all(&GGUF_TYPE_FLOAT32.to_le_bytes())?;
w.write_all(&(arr.len() as u64).to_le_bytes())?;
for v in arr {
w.write_all(&v.to_le_bytes())?;
}
Ok(())
}
}
}
fn build_metadata(cfg: &VisionConfig) -> Vec<(String, MetaValue)> {
let mut kvs: Vec<(String, MetaValue)> = vec![
(
"general.architecture".into(),
MetaValue::String("clip".into()),
),
(
"general.name".into(),
MetaValue::String("hf2q-mmproj".into()),
),
("clip.has_vision_encoder".into(), MetaValue::Bool(true)),
("clip.has_text_encoder".into(), MetaValue::Bool(false)),
(
"clip.projector_type".into(),
MetaValue::String(cfg.projector_type.clone()),
),
(
"clip.vision.image_size".into(),
MetaValue::Uint32(cfg.image_size),
),
(
"clip.vision.patch_size".into(),
MetaValue::Uint32(cfg.patch_size),
),
(
"clip.vision.embedding_length".into(),
MetaValue::Uint32(cfg.hidden_size),
),
(
"clip.vision.feed_forward_length".into(),
MetaValue::Uint32(cfg.intermediate_size),
),
(
"clip.vision.attention.head_count".into(),
MetaValue::Uint32(cfg.num_attention_heads),
),
(
"clip.vision.block_count".into(),
MetaValue::Uint32(cfg.num_hidden_layers),
),
(
"clip.vision.attention.layer_norm_epsilon".into(),
MetaValue::Float32(cfg.layer_norm_eps),
),
(
"clip.vision.image_mean".into(),
MetaValue::ArrayFloat32(cfg.image_mean.to_vec()),
),
(
"clip.vision.image_std".into(),
MetaValue::ArrayFloat32(cfg.image_std.to_vec()),
),
];
if let Some(pd) = cfg.projection_dim {
kvs.push(("clip.vision.projection_dim".into(), MetaValue::Uint32(pd)));
}
if cfg.is_qwen3vl() {
kvs.push(("clip.use_gelu".into(), MetaValue::Bool(true)));
if let Some(sms) = cfg.spatial_merge_size {
kvs.push((
"clip.vision.spatial_merge_size".into(),
MetaValue::Uint32(sms),
));
}
if let Some(is_ds_bools) = cfg.build_is_deepstack_layers() {
kvs.push((
"clip.vision.is_deepstack_layers".into(),
MetaValue::ArrayBool(is_ds_bools),
));
}
}
kvs
}
#[cfg(test)]
mod tests {
use super::*;
use crate::ir::DType;
fn tiny_config() -> VisionConfig {
VisionConfig {
hidden_size: 64,
num_hidden_layers: 2,
num_attention_heads: 4,
patch_size: 4,
image_size: 16,
intermediate_size: 128,
layer_norm_eps: 1e-5,
projector_type: "mlp".into(),
projection_dim: None,
image_mean: [0.5, 0.5, 0.5],
image_std: [0.5, 0.5, 0.5],
spatial_merge_size: None,
deepstack_visual_indexes: None,
temporal_patch_size: None,
}
}
fn tiny_tensors() -> HashMap<String, VitTensor> {
let mut m = HashMap::new();
m.insert(
"v.patch_embd.weight".into(),
VitTensor {
gguf_name: "v.patch_embd.weight".into(),
shape: vec![64, 3, 4, 4],
dtype: DType::F16,
data: vec![0u8; 64 * 3 * 4 * 4 * 2],
},
);
m.insert(
"mm.0.weight".into(),
VitTensor {
gguf_name: "mm.0.weight".into(),
shape: vec![64, 64],
dtype: DType::F16,
data: vec![0u8; 64 * 64 * 2],
},
);
m
}
#[test]
fn writes_valid_gguf_magic_and_version() {
let tmp = tempfile::tempdir().unwrap();
let out = tmp.path().join("tiny.mmproj.gguf");
write_mmproj_gguf(&out, &tiny_config(), &tiny_tensors()).expect("write");
let bytes = std::fs::read(&out).expect("read");
assert!(bytes.len() > 24, "file too small");
assert_eq!(&bytes[0..4], &GGUF_MAGIC, "bad magic");
let version = u32::from_le_bytes([bytes[4], bytes[5], bytes[6], bytes[7]]);
assert_eq!(version, GGUF_VERSION);
let tensor_count = u64::from_le_bytes(bytes[8..16].try_into().unwrap());
assert_eq!(tensor_count, 2);
let kv_count = u64::from_le_bytes(bytes[16..24].try_into().unwrap());
let expected = build_metadata(&tiny_config()).len() as u64;
assert_eq!(kv_count, expected);
assert!(expected >= 14, "kv_count regression below baseline");
}
#[test]
fn align_up_works() {
assert_eq!(align_up(0, 32), 0);
assert_eq!(align_up(1, 32), 32);
assert_eq!(align_up(32, 32), 32);
assert_eq!(align_up(33, 32), 64);
}
#[test]
fn write_gguf_string_format_matches_spec() {
let mut buf = Vec::new();
write_gguf_string(&mut buf, "hello").unwrap();
assert_eq!(&buf[0..8], &5u64.to_le_bytes());
assert_eq!(&buf[8..13], b"hello");
}
#[test]
fn build_metadata_matches_mmproj_loader_keys() {
let cfg = tiny_config();
let md = build_metadata(&cfg);
let keys: Vec<&str> = md.iter().map(|(k, _)| k.as_str()).collect();
let required = [
"general.architecture",
"clip.vision.image_size",
"clip.vision.patch_size",
"clip.vision.embedding_length",
"clip.vision.feed_forward_length",
"clip.vision.attention.head_count",
"clip.vision.block_count",
];
for r in &required {
assert!(keys.contains(r), "missing required metadata key: {}", r);
}
}
#[test]
fn kv_count_in_header_matches_metadata_list() {
let tmp = tempfile::tempdir().unwrap();
let out = tmp.path().join("match-count.gguf");
write_mmproj_gguf(&out, &tiny_config(), &tiny_tensors()).unwrap();
let bytes = std::fs::read(&out).unwrap();
let kv_count = u64::from_le_bytes(bytes[16..24].try_into().unwrap());
let md = build_metadata(&tiny_config());
assert_eq!(kv_count as usize, md.len());
}
#[test]
fn empty_tensor_map_still_produces_valid_header() {
let tmp = tempfile::tempdir().unwrap();
let out = tmp.path().join("empty.gguf");
let empty_map: HashMap<String, VitTensor> = HashMap::new();
write_mmproj_gguf(&out, &tiny_config(), &empty_map).unwrap();
let bytes = std::fs::read(&out).unwrap();
assert_eq!(&bytes[0..4], &GGUF_MAGIC);
let tensor_count = u64::from_le_bytes(bytes[8..16].try_into().unwrap());
assert_eq!(tensor_count, 0);
}
#[allow(dead_code)]
fn _unused_f32_type() {
let _ = GGML_TYPE_F32;
}
fn qwen3vl_tiny_config() -> VisionConfig {
VisionConfig {
hidden_size: 64,
num_hidden_layers: 2,
num_attention_heads: 8,
patch_size: 4,
image_size: 32,
intermediate_size: 128,
layer_norm_eps: 1e-5,
projector_type: "qwen3vl_merger".into(),
projection_dim: Some(2048),
image_mean: [0.48145466, 0.4578275, 0.40821073],
image_std: [0.26862954, 0.26130258, 0.27577711],
spatial_merge_size: Some(2),
deepstack_visual_indexes: Some(vec![0, 1]),
temporal_patch_size: Some(2),
}
}
#[test]
fn wedge4f_build_metadata_emits_qwen3vl_keys_when_family_set() {
let cfg = qwen3vl_tiny_config();
let md = build_metadata(&cfg);
let keys: Vec<&str> = md.iter().map(|(k, _)| k.as_str()).collect();
for key in &[
"clip.use_gelu",
"clip.vision.spatial_merge_size",
"clip.vision.is_deepstack_layers",
"clip.vision.projection_dim",
] {
assert!(
keys.contains(key),
"Wedge-4f: build_metadata must emit Qwen3-VL key {:?} \
when cfg.is_qwen3vl(); got keys: {:?}",
key,
keys
);
}
let pt = md
.iter()
.find(|(k, _)| k == "clip.projector_type")
.map(|(_, v)| v);
match pt {
Some(MetaValue::String(s)) => {
assert_eq!(s, "qwen3vl_merger");
}
other => panic!("unexpected projector_type value: {:?}", other),
}
let is_ds = md
.iter()
.find(|(k, _)| k == "clip.vision.is_deepstack_layers")
.map(|(_, v)| v);
match is_ds {
Some(MetaValue::ArrayBool(bools)) => {
assert_eq!(bools.len(), 2, "length must equal block_count");
assert!(bools[0], "deepstack_visual_indexes[0]=0 → bools[0]=true");
assert!(bools[1], "deepstack_visual_indexes[1]=1 → bools[1]=true");
}
other => panic!("is_deepstack_layers must be ArrayBool; got {:?}", other),
}
let ug = md
.iter()
.find(|(k, _)| k == "clip.use_gelu")
.map(|(_, v)| v);
match ug {
Some(MetaValue::Bool(true)) => {}
other => panic!("use_gelu must be Bool(true); got {:?}", other),
}
}
#[test]
fn wedge4f_build_metadata_omits_qwen3vl_keys_for_clip_classic() {
let cfg = tiny_config(); let md = build_metadata(&cfg);
let keys: Vec<&str> = md.iter().map(|(k, _)| k.as_str()).collect();
for key in &[
"clip.use_gelu",
"clip.vision.spatial_merge_size",
"clip.vision.is_deepstack_layers",
] {
assert!(
!keys.contains(key),
"Wedge-4f regression: CLIP-classic build_metadata must NOT \
emit Qwen3-VL key {:?}; got keys: {:?}",
key,
keys
);
}
assert!(
!keys.contains(&"clip.vision.projection_dim"),
"projection_dim is gated on Some(_) — should be absent"
);
}
#[test]
fn wedge4f_qwen3vl_metadata_round_trips_through_gguf_reader() {
use mlx_native::gguf::{GgufFile, MetadataValue};
let tmp = tempfile::tempdir().unwrap();
let out = tmp.path().join("qwen3vl-roundtrip.mmproj.gguf");
write_mmproj_gguf(&out, &qwen3vl_tiny_config(), &tiny_tensors()).expect("write");
let gguf = GgufFile::open(&out).expect("open");
assert_eq!(
gguf.metadata_string("clip.projector_type"),
Some("qwen3vl_merger")
);
assert_eq!(gguf.metadata_u32("clip.vision.spatial_merge_size"), Some(2));
assert_eq!(gguf.metadata_u32("clip.vision.projection_dim"), Some(2048));
match gguf.metadata("clip.use_gelu") {
Some(MetadataValue::Bool(true)) => {}
other => panic!("expected Bool(true), got {:?}", other),
}
let raw = gguf
.metadata("clip.vision.is_deepstack_layers")
.expect("present");
match raw {
MetadataValue::Array(arr) => {
assert_eq!(arr.len(), 2);
let mut true_count = 0usize;
for v in arr {
if matches!(v, MetadataValue::Bool(true)) {
true_count += 1;
}
}
assert_eq!(true_count, 2, "[0,1] indexes → both true");
}
other => panic!("expected Array, got {:?}", other),
}
}
#[test]
fn wedge4f_bool_kv_writes_single_byte() {
let mut buf = Vec::new();
write_kv(&mut buf, "test.bool", &MetaValue::Bool(true)).unwrap();
let key_len = u64::from_le_bytes(buf[0..8].try_into().unwrap());
assert_eq!(key_len, 9);
assert_eq!(&buf[8..17], b"test.bool");
let type_tag = u32::from_le_bytes(buf[17..21].try_into().unwrap());
assert_eq!(type_tag, GGUF_TYPE_BOOL);
assert_eq!(buf[21], 1);
assert_eq!(buf.len(), 22, "Bool KV must be exactly 22 bytes total");
}
#[test]
fn wedge4f_array_bool_kv_layout_matches_gguf_spec() {
let mut buf = Vec::new();
write_kv(
&mut buf,
"is_deepstack",
&MetaValue::ArrayBool(vec![false, true, false, true]),
)
.unwrap();
let key_len = u64::from_le_bytes(buf[0..8].try_into().unwrap());
assert_eq!(key_len, 12);
let kv_off = 8 + 12;
let array_tag = u32::from_le_bytes(buf[kv_off..kv_off + 4].try_into().unwrap());
assert_eq!(array_tag, GGUF_TYPE_ARRAY);
let elem_tag = u32::from_le_bytes(buf[kv_off + 4..kv_off + 8].try_into().unwrap());
assert_eq!(elem_tag, GGUF_TYPE_BOOL);
let len = u64::from_le_bytes(buf[kv_off + 8..kv_off + 16].try_into().unwrap());
assert_eq!(len, 4);
assert_eq!(buf[kv_off + 16], 0);
assert_eq!(buf[kv_off + 17], 1);
assert_eq!(buf[kv_off + 18], 0);
assert_eq!(buf[kv_off + 19], 1);
}
}