Add Qwen vision loading and Metal inference

tools/qwen-vision-chat-reference.py

@@ -0,0 +1,32 @@
+"""Describe a neutral image and run a no-image control through pinned MTPLX.
+Usage: test-supervisor ... --command python tools/qwen-vision-chat-reference.py MODEL_DIR IMAGE
+Uses the reference runtime only as an oracle; application inference is Rust/Metal.
+"""
+import os,sys,json,time,argparse
+from pathlib import Path
+root=Path(sys.argv[1]);image=Path(sys.argv[2])
+from mtplx.server.openai import _server_runtime_env_overrides
+os.environ.update(_server_runtime_env_overrides(argparse.Namespace(model=str(root),generation_mode='mtp'),None))
+os.environ['MTPLX_SUSTAINED_PREFILL']='1'
+os.environ['MTPLX_PREFILL_CHUNK_SIZE']='2048'
+import mlx.core as mx
+from mlx_lm.utils import load_model
+from mtplx.models import qwen4_exp as qwen
+from mtplx.runtime import _load_tokenizer_resilient,MTPLXRuntime
+from mtplx.mtp_patch import MTPContract
+from mtplx.sampling import SamplerConfig
+from mtplx.generation import generate_mtpk
+from mtplx.server.openai import ChatMessage,_encode_messages_uncached,_materialize_vision_splice
+print(json.dumps(dict(event='loading')),flush=True)
+tokenizer=_load_tokenizer_resilient(root,json.loads((root/'config.json').read_text()))
+model,_=load_model(root,lazy=False,strict=True,get_model_classes=lambda **_:(qwen.Model,qwen.ModelArgs))
+model.post_weight_load(root);assert model.attach_mtp(root)
+rt=MTPLXRuntime(model,tokenizer,root,True,MTPContract())
+for with_image in [True,False]:
+ text='Describe this image'+('<|vision_start|><|image_pad|><|vision_end|>' if with_image else '')
+ ids=_encode_messages_uncached(tokenizer,[ChatMessage(role='user',content=text)],enable_thinking=True,reasoning_effort='low',preserve_reasoning_history=True,tools=None)
+ splice=None
+ if with_image:ids,splice=_materialize_vision_splice(argparse.Namespace(args=argparse.Namespace(model=str(root))),[image.read_bytes()],ids)
+ print(json.dumps(dict(event='input',with_image=with_image,prompt='Describe this image',ids=ids)),flush=True)
+ out=generate_mtpk(rt,ids,max_tokens=1024,sampler=SamplerConfig(temperature=0.,top_p=.95,top_k=20),speculative_depth=3,seed=1,stop_token_ids=set(tokenizer.eos_token_ids),mtp_history_policy='committed',verify_strategy='batched',vision_splice=splice,capture_final_state=True,prefill_callback=lambda d:print(json.dumps(dict(event='prefill',**d)),flush=True),token_callback=lambda ids:print(json.dumps(dict(event='tokens',ids=ids)),flush=True))
+ print(json.dumps(dict(event='result',with_image=with_image,text=out.text,tokens=out.tokens,finish_reason=out.finish_reason)),flush=True)

tools/qwen-vision-reference.py

@@ -0,0 +1,24 @@
+"""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)