refactor: share tensor writer across ggml converters

convert_edgetam_to_ggml.py

@@ -19,14 +19,12 @@
 import struct
 import sys
 import os
-import numpy as np
+from ggml_writer import KEEP_F32, write_tensor
 
 # ── Constants ──────────────────────────────────────────────────────────────────
 
 MAGIC   = 0x73616D32   # "sam2" — reuse SAM2 magic, dispatch on backbone_type
 VERSION = 1
-FTYPE_F32 = 0
-FTYPE_F16 = 1
 
 # EdgeTAM hyperparameters (from edgetam.yaml)
 EDGETAM_HPARAMS = {
@@ -472,48 +470,6 @@ def write_header(fout, ftype, n_tensors, hp):
     fout.write(struct.pack("<i", hp["mem_attn_ca_k_size"]))
 
 
-def write_tensor(fout, name, data, ftype):
-    """Write one tensor record with 32-byte aligned data."""
-    n_dims = len(data.shape)
-    name_bytes = name.encode("utf-8")
-
-    # 1D tensors, embeddings, positions → always f32
-    use_f16 = (ftype == FTYPE_F16 and n_dims >= 2
-               and "embed" not in name
-               and "pos_embed" not in name
-               and "tpos" not in name
-               and "pe_gaussian" not in name
-               and "token" not in name
-               and "no_obj" not in name
-               and "no_mem" not in name
-               and "gamma" not in name
-               and "latents" not in name)
-
-    dtype_id = FTYPE_F16 if use_f16 else FTYPE_F32
-
-    if use_f16:
-        data = data.astype(np.float16)
-    else:
-        data = data.astype(np.float32)
-
-    fout.write(struct.pack("<i", n_dims))
-    fout.write(struct.pack("<i", len(name_bytes)))
-    fout.write(struct.pack("<i", dtype_id))
-
-    # ggml expects dimensions in reverse order (column-major)
-    for dim in reversed(data.shape):
-        fout.write(struct.pack("<i", dim))
-
-    fout.write(name_bytes)
-
-    # Pad to 32-byte alignment
-    pos = fout.tell()
-    pad = (32 - pos % 32) % 32
-    fout.write(b"\x00" * pad)
-
-    fout.write(data.tobytes())
-
-
 # ── Main ─────────────────────────────────────────────────────────────────────
 
 def main():
@@ -590,7 +546,7 @@ def main():
 
         for name in sorted(all_tensors.keys()):
             data = all_tensors[name]
-            write_tensor(fout, name, data, args.ftype)
+            write_tensor(fout, name, data, args.ftype, KEEP_F32 + ("latents",))
 
     file_size = os.path.getsize(args.output)
     print(f"Done. {len(all_tensors)} tensors, {file_size / (1024*1024):.1f} MB")

convert_sam2_to_ggml.py

@@ -12,14 +12,12 @@
 import struct
 import sys
 import os
-import numpy as np
+from ggml_writer import write_tensor
 
 # ── Constants ──────────────────────────────────────────────────────────────────
 
 MAGIC   = 0x73616D32   # "sam2"
 VERSION = 1
-FTYPE_F32 = 0
-FTYPE_F16 = 1
 
 # ── Default hyperparameters (SAM2.1 HieraL) ─────────────────────────────────
 
@@ -277,47 +275,6 @@ def write_header(fout, ftype: int, n_tensors: int, hparams: dict):
         fout.write(struct.pack("<i", hp[flag]))
 
 
-def write_tensor(fout, name: str, data: np.ndarray, ftype: int):
-    """Write one tensor record with 32-byte aligned data."""
-    n_dims = len(data.shape)
-    name_bytes = name.encode("utf-8")
-
-    # 1D tensors, embeddings, positions → always f32
-    use_f16 = (ftype == FTYPE_F16 and n_dims >= 2
-               and "embed" not in name
-               and "pos_embed" not in name
-               and "tpos" not in name
-               and "pe_gaussian" not in name
-               and "token" not in name
-               and "no_obj" not in name
-               and "no_mem" not in name
-               and "gamma" not in name)
-
-    dtype_id = FTYPE_F16 if use_f16 else FTYPE_F32
-
-    if use_f16:
-        data = data.astype(np.float16)
-    else:
-        data = data.astype(np.float32)
-
-    fout.write(struct.pack("<i", n_dims))
-    fout.write(struct.pack("<i", len(name_bytes)))
-    fout.write(struct.pack("<i", dtype_id))
-
-    # ggml expects dimensions in reverse order (column-major)
-    for dim in reversed(data.shape):
-        fout.write(struct.pack("<i", dim))
-
-    fout.write(name_bytes)
-
-    # Pad to 32-byte alignment
-    pos = fout.tell()
-    pad = (32 - pos % 32) % 32
-    fout.write(b"\x00" * pad)
-
-    fout.write(data.tobytes())
-
-
 # ── Config parsing ───────────────────────────────────────────────────────────
 
 def detect_variant_from_checkpoint(state_dict: dict) -> str:

convert_sam3_to_ggml.py

@@ -14,13 +14,12 @@
 import os
 import re
 import numpy as np
+from ggml_writer import KEEP_F32, write_tensor
 
 # ── Constants ──────────────────────────────────────────────────────────────────
 
 MAGIC   = 0x73616D33   # "sam3"
 VERSION = 3
-FTYPE_F32 = 0
-FTYPE_F16 = 1
 
 # ── Hyperparameter defaults ───────────────────────────────────────────────────
 
@@ -211,50 +210,6 @@ def write_header(fout, ftype: int, n_tensors: int, visual_only: bool = False):
         fout.write(struct.pack("<i", val))
 
 
-def write_tensor(fout, name: str, data: np.ndarray, ftype: int):
-    """Write one tensor record with 32-byte aligned data."""
-    n_dims = len(data.shape)
-    name_bytes = name.encode("utf-8")
-
-    # Determine storage dtype
-    # 1D tensors, embeddings, and positions → always f32
-    use_f16 = (ftype == FTYPE_F16 and n_dims >= 2
-               and "embed" not in name
-               and "pos_embed" not in name
-               and "tpos" not in name
-               and "pe_gaussian" not in name
-               and "freqs_cis" not in name
-               and "token" not in name
-               and "no_obj" not in name
-               and "no_mem" not in name
-               and "gamma" not in name)
-
-    dtype_id = FTYPE_F16 if use_f16 else FTYPE_F32
-
-    if use_f16:
-        data = data.astype(np.float16)
-    else:
-        data = data.astype(np.float32)
-
-    # Write: n_dims, name_len, dtype, shape (reversed), name, padding, data
-    fout.write(struct.pack("<i", n_dims))
-    fout.write(struct.pack("<i", len(name_bytes)))
-    fout.write(struct.pack("<i", dtype_id))
-
-    # ggml expects dimensions in reverse order (column-major)
-    for dim in reversed(data.shape):
-        fout.write(struct.pack("<i", dim))
-
-    fout.write(name_bytes)
-
-    # Pad to 32-byte alignment
-    pos = fout.tell()
-    pad = (32 - pos % 32) % 32
-    fout.write(b"\x00" * pad)
-
-    fout.write(data.tobytes())
-
-
 # ── Tokenizer embedding ──────────────────────────────────────────────────────
 
 TOK_MAGIC = 0x746F6B00   # "tok\0"
@@ -393,7 +348,7 @@ def main():
         write_header(fout, args.ftype, len(renamed), visual_only=args.visual_only)
 
         for i, (name, data) in enumerate(renamed.items()):
-            write_tensor(fout, name, data, args.ftype)
+            write_tensor(fout, name, data, args.ftype, KEEP_F32 + ("freqs_cis",))
             if (i + 1) % 100 == 0 or i == len(renamed) - 1:
                 print(f"  [{i+1}/{len(renamed)}] {name}  {list(data.shape)}")
 

ggml_writer.py

@@ -0,0 +1,21 @@
+"""Tensor record writer shared by the convert_*_to_ggml.py scripts."""
+
+import struct
+import numpy as np
+
+FTYPE_F32 = 0
+FTYPE_F16 = 1
+
+# Names containing any of these (embeddings, positions, tokens, ...) stay f32.
+KEEP_F32 = ("embed", "tpos", "pe_gaussian", "token", "no_obj", "no_mem", "gamma")
+
+
+def write_tensor(fout, name, data, ftype, keep_f32=KEEP_F32):
+    """Write one tensor record with 32-byte aligned data. 1D tensors always stay f32."""
+    use_f16 = ftype == FTYPE_F16 and data.ndim >= 2 and not any(s in name for s in keep_f32)
+    name_bytes = name.encode("utf-8")
+    fout.write(struct.pack("<3i", data.ndim, len(name_bytes), FTYPE_F16 if use_f16 else FTYPE_F32))
+    fout.write(struct.pack(f"<{data.ndim}i", *reversed(data.shape)))  # ggml dims are reversed
+    fout.write(name_bytes)
+    fout.write(b"\x00" * (-fout.tell() % 32))
+    fout.write(data.astype(np.float16 if use_f16 else np.float32).tobytes())