File size: 4,224 Bytes
d3d5695 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 | import hashlib,json,time
from pathlib import Path
import torch
torch.set_num_threads(4)
from safetensors import safe_open
from component import load_model
r=Path(__file__).parent;checks=[];start=time.monotonic()
def check(name,ok,**details):
assert ok,(name,details)
checks.append({'name':name,'passed':True,**details});print(name,details,flush=True)
expected=json.loads((r/'tensor_sha256.json').read_text())
with safe_open(r/'model.safetensors',framework='pt') as f:
assert set(f.keys())==set(expected)
for k in f.keys():
v=f.get_tensor(k)
assert hashlib.sha256(v.contiguous().reshape(-1).view(torch.uint8).numpy().tobytes()).hexdigest()==expected[k],k
check('1582_original_tensor_payloads',len(expected)==1582)
check('only_active_visual_prefixes',all(k.startswith(('model.vision_tower.','model.mm_projector.','model.diffusion_connector.','model.dit.','model.vae.')) for k in expected))
m=load_model(r)
check('native_1582entry_strict_load',len(m.state_dict())==1582)
check('original_DCAE_scaling',m.model.vae.config.scaling_factor==0.41407)
with torch.inference_mode():
for size in [128,1024]:
pixels=torch.linspace(0,1,3*size*size).reshape(1,3,size,size)
features=m.model.vision_tower(pixels)
check(f'{size}px_FastViT_geometry',tuple(features.shape)==(1,(size//64)**2,3072) and torch.isfinite(features).all().item())
direct=m.model.vision_tower.vision_tower(pixels,return_image_embeddings=True)['image_embeddings'].flatten(2).transpose(1,2)
check(f'{size}px_native_features_exact',torch.equal(features,direct))
projected=m.model.mm_projector(features)
check(f'{size}px_VLM_projector',tuple(projected.shape)==(1,(size//64)**2,896) and torch.isfinite(projected).all().item())
hidden=tuple(torch.full((1,8,896),i/20) for i in range(5))
cond=m.condition(hidden)
check('MCP_preserves_sequence',tuple(cond.shape)==(1,8,2304) and torch.isfinite(cond).all().item())
changed=(hidden[0]+100,)+hidden[1:]
check('MCP_selects_final_four_layers',torch.equal(cond,m.condition(changed)))
fusion=m.model.diffusion_connector.fusion
weights=(fusion.layer_weights/fusion.temperature).softmax(0)
fused=sum(weights[i]*hidden[-4+i] for i in range(4))
direct=fusion.input_compress(fused)
for refiner in fusion.sequence_refiners:direct=refiner(direct)
direct=fusion.output_proj(direct)
check('native_weighted_fusion_exact',torch.equal(cond,direct))
for h,w in [(64,64),(64,96)]:
rgb=torch.linspace(-1,1,3*h*w).reshape(1,3,h,w)
z=m.encode_latent(rgb)
check(f'{h}x{w}_native_encode_exact',tuple(z.shape)==(1,32,h//32,w//32) and torch.equal(z,m.model.vae.encode(rgb).latent*0.41407))
decoded=m.decode_latent(z)
check(f'{h}x{w}_native_decode_exact',torch.equal(decoded,m.model.vae.decode(z/0.41407).sample))
check(f'{h}x{w}_finite_RGB',tuple(decoded.shape)==tuple(rgb.shape) and torch.isfinite(decoded).all().item(),mse=(decoded-rgb).square().mean().item())
z=torch.linspace(-1,1,32*4*4).reshape(1,32,4,4)
flow=m.model.dit(z,timestep=torch.tensor([500.]),encoder_hidden_states=cond,encoder_attention_mask=torch.ones(1,8),return_dict=False)[0]
check('complete_28layer_SANA_prediction',tuple(flow.shape)==tuple(z.shape) and torch.isfinite(flow).all().item())
report={'strict_load':True,'source_tensor_bytes_equal':True,'checks':checks,'smoke':{'passed':True,'device':'cpu','dtype':'float32'},'elapsed_seconds':time.monotonic()-start,'limitations':['Original Qwen language model is excluded; MCP/SANA smoke fixtures use synthetic VLM hidden states.','FastViT is checked at 128px and native 1024px; DC-AE reconstruction fixtures are 64px square and rectangular; no downstream quality benchmark is claimed.','Legacy base_model.model.* checkpoint entries are excluded because the pinned official inference constructor and loader use model.*; all selected active visual branches load strictly.','Native code is retained without changing sampler equations or model math. CPU float32 inference does not establish CoreML/MLX/CUDA parity.']}
(r/'validation.json').write_text(json.dumps(report,indent=2)+'\n')
print('MOBILE VALIDATED',len(checks),flush=True)
|