| """Source-byte, strict-state, native-feature/token and codec-path verification.""" |
| import argparse,bisect,gc,hashlib,json,pathlib,struct,sys,time,io |
| import torch |
| import numpy as np |
| from PIL import Image |
| from component import build_model,load_model,tokenizer_view |
| class SourceRanges: |
| def __init__(self,root,mapping): |
| self.shards={};self.root=root;self.dcp=any('chunks' in e for e in mapping.values()) |
| if self.dcp:return |
| for shard in {e['source_shard'] for e in mapping.values()}: |
| parts=sorted((int(p.name.rsplit('_',1)[1]),p) for p in (root/'.download').glob(shard.replace('/','__')+'_*') if not p.name.endswith('.part')) |
| assert parts |
| self.shards[shard]=(parts,parts[0][0]-min(e['data_offsets'][0] for e in mapping.values() if e['source_shard']==shard)) |
| def read(self,e): |
| if self.dcp: |
| t=torch.empty(e['shape'],dtype=DTYPES[e['dtype']]);coverage=torch.zeros(e['shape'][0],dtype=torch.bool) |
| for c in e['chunks']: |
| raw=(self.root/'.download'/c['cache_file']).read_bytes();assert hashlib.sha256(raw).hexdigest()==c['serialized_sha256'] |
| part=torch.load(io.BytesIO(raw),map_location='cpu',weights_only=True);assert list(part.shape)==c['tensor_sizes'] and part.dtype==t.dtype |
| a=c['tensor_offsets'][0];b=a+c['tensor_sizes'][0];assert not coverage[a:b].any();coverage[a:b]=True |
| assert c['tensor_offsets'][1:]==[0]*(t.ndim-1) and c['tensor_sizes'][1:]==list(t.shape)[1:] |
| t[tuple(slice(a,a+b) for a,b in zip(c['tensor_offsets'],c['tensor_sizes']))]=part |
| assert coverage.all();return bytearray(t.view(torch.uint8).numpy().tobytes()) |
| parts,base=self.shards[e['source_shard']];starts=[p[0] for p in parts];a,b=e['data_offsets'];a+=base;b+=base;result=bytearray() |
| while a<b: |
| start,p=parts[bisect.bisect_right(starts,a)-1] |
| with p.open('rb') as f: |
| f.seek(a-start);data=f.read(min(b-a,p.stat().st_size-a+start)) |
| if not data:raise RuntimeError('Source gap') |
| result.extend(data);a+=len(data) |
| return result |
| DTYPES={'BF16':torch.bfloat16,'F32':torch.float32,'F16':torch.float16,'I64':torch.int64,'I32':torch.int32,'BOOL':torch.bool} |
| def tree(x): |
| if isinstance(x,torch.Tensor):return x.detach().cpu().clone() |
| if isinstance(x,dict):return {k:tree(v) for k,v in x.items()} |
| if isinstance(x,(list,tuple)):return [tree(v) for v in x] |
| return x |
|
|
| def compare(a,b): |
| if isinstance(a,torch.Tensor): |
| assert a.shape==b.shape |
| if not a.is_floating_point():assert torch.equal(a,b);return 0. |
| assert torch.isfinite(a).all() and torch.isfinite(b).all() |
| torch.testing.assert_close(a,b,rtol=2e-4,atol=2e-4) |
| return (a.float()-b.float()).abs().max().item() |
| if isinstance(a,dict):return max([compare(a[k],b[k]) for k in a] or [0.]) |
| if isinstance(a,(list,tuple)): |
| assert len(a)==len(b);return max([compare(x,y) for x,y in zip(a,b)] or [0.]) |
| assert a==b;return 0. |
|
|
| def image(h,w): |
| y,x=np.mgrid[:h,:w];return Image.fromarray(np.stack([(x*7+y*3)%256,(x*2+y*11)%256,(x*13+y*5)%256],-1).astype('uint8')) |
| def run_case(m,root,case): |
| kind=json.loads((root/'config.json').read_text()).get('component',root.name) |
| if kind=='pixelumm_pixel_io': |
| px,tubes,hs,t=case['input'];return {'image_und':m.pixel_patch_embed(px),'image_gen':m.pixel_patch_embed_mot_gen(px),'video_gen':m.pixel_video_patch_embed_mot_gen(tubes),'image_head':m.unpatchify_head(hs),'video_head':m.video_unpatchify_head(hs),'time_mlp_inactive_in_release':m.time_embedder(t)} |
| if kind=='sensenova_pixel_io': |
| px,g,hs,t=case['input'];return {'understanding':m.encode_pixels(px,g),'generation':m.encode_pixels(px,g,True),'pixel_head':m.decode_hidden(hs,(2,3)),'time':m.fm_modules['timestep_embedder'](t),'noise':m.fm_modules['noise_scale_embedder'](t)} |
| if kind=='mingtok': |
| x=case['input'];o=m(x);pix=m.forward_pixel_decoder(o['x_norm_patchtokens']);sem=m.forward_feature_decoder_wo_cache(o['latent']*m.scaling_factor+m.mean)['x_norm_patchtokens'] |
| return {'latent':o['latent'],'semantic':o['x_norm_patchtokens'],'reconstruction':pix,'semantic_from_latent':sem} |
| if kind=='uniar_bsq': |
| px,g=case['input'];bits=m(px,grid_thw=g,bsq_only=True,direct_indices=True) |
| joined=torch.cat([bits[0]]+bits[1],dim=-1) |
| assert torch.all((joined==0)|(joined==1)) |
| features=m.bsq_indices_to_features(joined,multiscale_idx='all',with_merger=True) |
| direct=m(px,grid_thw=g) |
| compare(features,direct) |
| return {'bits':joined,'decoded_features':features,'direct_features':direct} |
| t=tokenizer_view(m,root);ims=case['input'];inp=t.image_processor(ims);hs=t.visual(inp['pixel_values'],inp['image_grid_thw']);ids=t._quantize(hs,inp['image_grid_thw']) |
| actual=t.encode_batch(ims) |
| assert actual==[x.tolist() for x in ids.split(inp['image_grid_thw'].prod(-1).tolist())] |
| assert ((ids>=0)&(ids<t.codebook_size)).all() |
| |
| codes=torch.nn.functional.normalize(t.vqmodel.quantize.embedding.weight,dim=-1)[ids] |
| inverse=torch.nn.functional.linear(codes,t.vqmodel.post_quant_conv.weight[:,:,0,0],t.vqmodel.post_quant_conv.bias) |
| return {'hidden':hs,'ids':ids,'inverse_vq_features':inverse,'api_token_ids':actual} |
|
|
| def main(root): |
| torch.set_num_threads(1);torch.manual_seed(17);start=time.time();mapping=json.loads((root/'tensor_mapping.json').read_text());ranges=SourceRanges(root,mapping) |
| with (root/'model.safetensors').open('rb') as f: |
| n=struct.unpack('<Q',f.read(8))[0];header=json.loads(f.read(n));base=8+n |
| for k,e in mapping.items(): |
| h=header[k];assert h['shape']==e['shape'] and h['dtype']==e['dtype'];a,b=h['data_offsets'];f.seek(base+a);assert f.read(b-a)==ranges.read(e),k |
| print('ALL SOURCE BYTES MATCH',flush=True) |
| |
| ref=build_model(root,attention='eager');sd=ref.state_dict();assert set(sd)==set(mapping),(set(sd)-set(mapping),set(mapping)-set(sd)) |
| with torch.no_grad(): |
| for k,e in mapping.items():sd[k].copy_(torch.frombuffer(ranges.read(e),dtype=DTYPES[e['dtype']]).reshape(e['shape'])) |
| del sd;gc.collect();kind=json.loads((root/'config.json').read_text()).get('component',root.name) |
| if kind=='pixelumm_pixel_io':cases=[{'name':name,'input':(torch.randn(n,768),torch.randn(n,3072),torch.randn(n,4096),torch.tensor([0.,.5,1.]))} for name,n in [('single_patch',1),('six_patches',6),('multi_frame_tubes',24)]] |
| elif kind=='sensenova_pixel_io': |
| cases=[] |
| for name,grid in [('square',[[4,4]]),('rectangle',[[4,6]]),('packed_two_images',[[4,4],[4,6]])]: |
| g=torch.tensor(grid);cases.append({'name':name,'input':(torch.randn(int(g.prod(-1).sum()),768),g,torch.randn(1,6,4096),torch.tensor([0.,.5,1.]))}) |
| elif kind=='mingtok':cases=[{'name':f'square_{s}','input':torch.linspace(-1,1,3*s*s).reshape(1,3,s,s)} for s in [64,128]] |
| elif kind=='uniar_bsq': |
| vc=ref.config;cases=[] |
| for name,grid in [('square',[[1,4,4]]),('rectangle',[[1,4,6]]),('packed_two_images',[[1,4,4],[1,4,6]])]: |
| g=torch.tensor(grid);px=torch.randn(int(g.prod(-1).sum()),3*vc.temporal_patch_size*vc.patch_size**2,dtype=torch.bfloat16);cases.append({'name':name,'input':(px,g)}) |
| else:cases=[{'name':'square_rgb','input':[image(64,64)]},{'name':'rectangle_rgb','input':[image(64,96)]},{'name':'packed_two_images','input':[image(64,64),image(64,96)]}] |
| results=[] |
| with torch.inference_mode(): |
| for c in cases:results.append(tree(run_case(ref,root,c)));print('reference '+c['name'],flush=True) |
| del ref;gc.collect();exported=load_model(root);checks=[] |
| with torch.inference_mode(): |
| for c,r in zip(cases,results): |
| o=tree(run_case(exported,root,c));err=compare(o,r);checks.append({'case':c['name'],'max_abs_error':err,'passed':True});print('PASS '+c['name']+' '+str(err),flush=True) |
| |
| if kind=='pixelumm_pixel_io':smoke={'name':'256_image_and_video_tube_io','input':(torch.randn(256,768),torch.randn(256,3072),torch.randn(256,4096),torch.tensor([.5]))} |
| elif kind=='sensenova_pixel_io':smoke={'name':'256_native_pixel_io','input':(torch.randn(256,768),torch.tensor([[16,16]]),torch.randn(1,6,4096),torch.tensor([.5]))} |
| elif kind=='mingtok':smoke={'name':'512_pixel_reconstruction','input':torch.from_numpy(np.asarray(image(512,512)).copy()).permute(2,0,1)[None].float()/127.5-1} |
| elif kind=='llada_image_tokenizer':smoke={'name':'256_rgb_tokenization','input':[image(256,256)]} |
| else: |
| g=torch.tensor([[1,28,28]]);smoke={'name':'448_native_patch_encoding','input':(torch.randn(784,1536,dtype=torch.bfloat16),g)} |
| with torch.inference_mode():out=tree(run_case(exported,root,smoke)) |
| compare(out,out) |
| shapes={k:list(v.shape) for k,v in out.items() if isinstance(v,torch.Tensor)} |
| if kind=='mingtok': |
| arr=((out['reconstruction'][0].permute(1,2,0)+1)*127.5).round().clamp(0,255).byte().numpy();Image.fromarray(arr).save(root/'reconstruction_smoke.png');image(512,512).save(root/'input_smoke.png') |
| report={'source_tensor_bytes_equal':True,'strict_load':True,'reference':'Pinned upstream native classes with all original source-range tensors loaded directly. Exported explicit loader uses the assembled checkpoint. Eager reference vs SDPA export where supported.','checks':checks,'smoke':{'case':smoke['name'],'shapes':shapes,'passed':True},'dtype':'BF16' if kind=='uniar_bsq' else 'FP32','torch':torch.__version__,'transformers':__import__('transformers').__version__,'tolerance':{'rtol':2e-4,'atol':2e-4},'seed':17,'elapsed_seconds':round(time.time()-start,2),'scope':'Component encoding/quantization/features/reconstruction as available; no language backbone, full text-to-image generation or downstream accuracy evaluation.'} |
| (root/'validation.json').write_text(json.dumps(report,indent=2)+'\n');print('VALIDATED '+root.name,flush=True) |
| if __name__=='__main__': |
| p=argparse.ArgumentParser();p.add_argument('root',type=pathlib.Path);main(p.parse_args().root) |
|
|
