#!/usr/bin/env python3
"""Build a local Kimi-K2.7-Code vision tower for exo, mirroring exolabs--Kimi-K2.6-vision.

Copies the 335 vision_tower.* / mm_projector.* tensors out of the official
moonshotai/Kimi-K2.7-Code download (shards 63+64) into a single safetensors
file, with a config.json adapted from the K2.6 tower. exo's downloader falls
back to local files when the repo doesn't exist on HF, so no upload is needed.

Run with exo's venv python (has safetensors):
  ~/exo-next/.venv/bin/python ~/.exo/extract-k27-vision-tower.py
"""

import json
from pathlib import Path

from safetensors import safe_open
from safetensors.numpy import save_file

MODELS = Path("/Volumes/LLM/exo-models")
SRC = MODELS / "moonshotai--Kimi-K2.7-Code"
K26_TOWER = MODELS / "exolabs--Kimi-K2.6-vision"
DST = MODELS / "aidiffuser--Kimi-K2.7-Code-vision"

SHARDS = ["model-00063-of-000064.safetensors", "model-00064-of-000064.safetensors"]
PREFIXES = ("vision_tower.", "mm_projector.")

for shard in SHARDS:
    if not (SRC / shard).exists():
        raise SystemExit(f"missing {shard} — download not finished yet")

tensors = {}
dtypes = set()
for shard in SHARDS:
    # framework="np" keeps bf16 unsupported; use torch-free path via mlx instead
    with safe_open(str(SRC / shard), framework="pt") as f:
        for key in f.keys():
            if key.startswith(PREFIXES):
                t = f.get_tensor(key)
                dtypes.add(str(t.dtype))
                tensors[key] = t

print(f"extracted {len(tensors)} tensors, dtypes: {dtypes}")
assert len(tensors) == 335, f"expected 335 tensors, got {len(tensors)}"

DST.mkdir(exist_ok=True)

# safetensors.numpy can't write bf16; go through torch's save_file instead
from safetensors.torch import save_file as save_pt

save_pt(tensors, str(DST / "kimi_k27_vision.safetensors"))

cfg = json.loads((K26_TOWER / "config.json").read_text())
cfg["source_model"] = "moonshotai/Kimi-K2.7-Code"
cfg["description"] = (
    "Vision-only weights extracted locally from Kimi-K2.7-Code for use with MLX/exo."
)
(DST / "config.json").write_text(json.dumps(cfg, indent=1) + "\n")

size = (DST / "kimi_k27_vision.safetensors").stat().st_size
print(f"wrote {DST} ({size / 1e9:.2f} GB) — done")