| """Export the existing Kev vision tower directly to Core ML MIL, fixed grid. |
| |
| Uses the fixed-grid/full-image-attention design in john-rocky/CoreML-LLM, |
| but reads our verified MLX weights and needs no PyTorch or model download. |
| No resampling or token reduction is performed. FP16 arithmetic differs from |
| the original BF16 tower and is explicitly evaluated, not assumed equivalent. |
| """ |
| import argparse,hashlib,json,time |
| from pathlib import Path |
| import numpy as np |
| import mlx.core as mx |
| from mlx.utils import tree_flatten |
| from mlx_vlm.models.qwen3_vl.vision import VisionModel |
| from mlx_vlm.models.qwen3_vl.config import VisionConfig |
| import coremltools as ct |
| from coremltools.converters.mil import Builder as mb |
| from coremltools.converters.mil.mil import types |
|
|
| ROOT=Path(__file__).resolve().parents[1];OUT=ROOT/'coreml_exports' |
|
|
| def main(): |
| p=argparse.ArgumentParser();p.add_argument('--grid',type=int,default=24) |
| p.add_argument('--attention',choices=['manual','sdpa'],default='manual');p.add_argument('--out');args=p.parse_args() |
| dest=OUT/(args.out or f'grid-{args.grid}-{args.attention}');dest.mkdir(parents=True,exist_ok=False) |
| config=json.loads((ROOT/'config.json').read_text())['vision_config'] |
| assert not config['deepstack_visual_indexes'] |
| checkpoint=ROOT/'model.safetensors' |
| h=hashlib.sha256() |
| with checkpoint.open('rb') as f: |
| for b in iter(lambda:f.read(8*1024**2),b''):h.update(b) |
| expected=json.loads((ROOT/'provenance.json').read_text())['sha256'];assert h.hexdigest()==expected |
| model=VisionModel(VisionConfig.from_dict(config)) |
| loaded=mx.load(str(checkpoint));model.load_weights([(k.removeprefix('vision_tower.'),v) for k,v in loaded.items() if k.startswith('vision_tower.')],strict=True) |
| del loaded;model.eval();mx.eval(model.parameters());mx.clear_cache() |
| weights={k:np.array(v.astype(mx.float32)).astype(np.float16) for k,v in tree_flatten(model.parameters())} |
| grid=mx.array([[1,args.grid,args.grid]],dtype=mx.int32);n=args.grid**2;d=config['hidden_size'];heads=config['num_heads'];hd=d//heads |
| pos=np.array(model.fast_pos_embed_interpolate(grid).astype(mx.float32)).astype(np.float16) |
| freqs=model.rot_pos_emb(grid);cos=np.array(mx.tile(mx.cos(freqs),(1,2))).astype(np.float16)[:,None,:] |
| sin=np.array(mx.tile(mx.sin(freqs),(1,2))).astype(np.float16)[:,None,:] |
| del model;mx.clear_cache() |
| (dest/'manifest.json').write_text(json.dumps({'checkpoint_sha256':expected,'grid':[1,args.grid,args.grid], |
| 'image_tokens':n//4,'config':config,'coremltools':ct.__version__,'attention':args.attention, |
| 'precision':'FP16 weights converted from original BF16; FP16 Core ML computation', |
| 'script_sha256':hashlib.sha256(Path(__file__).read_bytes()).hexdigest()},indent=2)+'\n') |
| (dest/'ane_export.py').write_bytes(Path(__file__).read_bytes()) |
| def linear(x,name):return mb.linear(x=x,weight=weights[name+'.weight'],bias=weights[name+'.bias'],name=name.replace('.','_')) |
| def norm(x,name):return mb.layer_norm(x=x,axes=[-1],gamma=weights[name+'.weight'],beta=weights[name+'.bias'],epsilon=np.float16(1e-6)) |
| def rope(x): |
| a,b=mb.split(x=x,num_splits=2,axis=-1) |
| rotated=mb.concat(values=[mb.mul(x=b,y=np.float16(-1)),a],axis=-1) |
| return mb.add(x=mb.mul(x=x,y=cos),y=mb.mul(x=rotated,y=sin)) |
| @mb.program(input_specs=[mb.TensorSpec(shape=(n,1536),dtype=types.fp16)],opset_version=ct.target.iOS18) |
| def program(pixels): |
| |
| w=weights['patch_embed.proj.weight'].transpose(0,4,1,2,3).reshape(d,1536) |
| x=mb.linear(x=pixels,weight=w,bias=weights['patch_embed.proj.bias']) |
| x=mb.add(x=x,y=pos) |
| for i in range(config['depth']): |
| prefix=f'blocks.{i}';a=norm(x,prefix+'.norm1');qkv=linear(a,prefix+'.attn.qkv') |
| qkv=mb.reshape(x=qkv,shape=[n,3,heads,hd]);q,k,v=mb.split(x=qkv,num_splits=3,axis=1) |
| q=rope(mb.squeeze(x=q,axes=[1]));k=rope(mb.squeeze(x=k,axes=[1]));v=mb.squeeze(x=v,axes=[1]) |
| q,k,v=[mb.expand_dims(x=mb.transpose(x=t,perm=[1,0,2]),axes=[0]) for t in [q,k,v]] |
| if args.attention=='sdpa':a=mb.scaled_dot_product_attention(query=q,key=k,value=v) |
| else: |
| scores=mb.matmul(x=mb.mul(x=q,y=np.float16(hd**-.5)),y=k,transpose_y=True) |
| a=mb.matmul(x=mb.softmax(x=scores,axis=-1),y=v) |
| a=mb.reshape(x=mb.transpose(x=a,perm=[0,2,1,3]),shape=[n,d]) |
| x=mb.add(x=x,y=linear(a,prefix+'.attn.proj')) |
| a=norm(x,prefix+'.norm2');a=mb.gelu(x=linear(a,prefix+'.mlp.linear_fc1'),mode='TANH_APPROXIMATION') |
| x=mb.add(x=x,y=linear(a,prefix+'.mlp.linear_fc2')) |
| x=mb.reshape(x=norm(x,'merger.norm'),shape=[n//4,4*d]) |
| x=mb.gelu(x=linear(x,'merger.linear_fc1'),mode='EXACT') |
| return mb.identity(x=linear(x,'merger.linear_fc2'),name='features') |
| print('MIL built',flush=True);t=time.perf_counter() |
| converted=ct.convert(program,convert_to='mlprogram',minimum_deployment_target=ct.target.iOS18, |
| compute_precision=ct.precision.FLOAT16,compute_units=ct.ComputeUnit.CPU_AND_NE,skip_model_load=True) |
| converted.save(str(dest/'vision.mlpackage')) |
| print(json.dumps({'saved':str(dest),'conversion_seconds':time.perf_counter()-t}),flush=True) |
|
|
| if __name__=='__main__':main() |
|
|
