25 lines
1.4 KiB
Python
25 lines
1.4 KiB
Python
"""Export pinned MTPLX vision stages; oracle only, never application code.
|
|
Usage: python tools/qwen-vision-reference.py MODEL_DIR IMAGE OUTPUT_DIR
|
|
Run in the pinned MTPLX reference environment with GPU access.
|
|
"""
|
|
import sys,json,hashlib,pathlib
|
|
import mlx.core as mx
|
|
import numpy as np
|
|
from mtplx.vision import load_vision_tower
|
|
from mtplx.vision.processing import decode_image,preprocess_images
|
|
root=pathlib.Path(sys.argv[1]); image=pathlib.Path(sys.argv[2]); out=pathlib.Path(sys.argv[3]);out.mkdir(exist_ok=True)
|
|
def save(name,x):
|
|
mx.eval(x); a=np.asarray(x.astype(mx.float32));a.tofile(out/(name+'.f32')); print(json.dumps(dict(stage=name,shape=list(x.shape),dtype=str(x.dtype),min=float(a.min()),max=float(a.max()),sha256=hashlib.sha256(a.tobytes()).hexdigest())),flush=True)
|
|
pixels,grids=preprocess_images([decode_image(image.read_bytes())],json.loads((root/'preprocessor_config.json').read_text()))
|
|
(out/'grid.json').write_text(json.dumps(grids[0]));save('pixels',pixels); print(json.dumps(dict(grids=grids)),flush=True)
|
|
tower=load_vision_tower(root)
|
|
h=tower.patch_embed(pixels.astype(tower.patch_embed.proj.weight.dtype));save('patch',h)
|
|
p=tower.fast_pos_embed_interpolate(grids);save('position',p)
|
|
h=h+p;r=tower.rot_pos_emb(grids);save('rotary',r)
|
|
for i,b in enumerate(tower.blocks):
|
|
h=b(h,[],r)
|
|
if i in [0,26]:save('block'+str(i),h)
|
|
else:mx.eval(h)
|
|
print(json.dumps(dict(block=i)),flush=True)
|
|
h=tower.merger(h);save('embeddings',h)
|