| |
| """Replay an exported portable plan on CPU and compare its PNG to the export reference. |
| |
| For parity, use export-time PyTorch noise (seed 42) and the exact first tokenizer |
| fixture. Android's Java noise is a separate, reproducible random stream. |
| """ |
| import argparse |
| import hashlib |
| import json |
| from pathlib import Path |
| import time |
|
|
| import numpy as np |
| import torch |
| from ai_edge_litert.interpreter import Interpreter |
| from PIL import Image |
|
|
|
|
| def validate(bundle, out): |
| root = Path(bundle); m = json.loads((root/'manifest.json').read_text()) |
| assert m['format'] == 'word-image-litert-v1' and not m['compiled'] |
| for name, spec in m['files'].items(): |
| path = (root/name).resolve(); assert path.is_relative_to(root.resolve()) |
| assert path.stat().st_size == spec['bytes'] |
| assert hashlib.sha256(path.read_bytes()).hexdigest() == spec['sha256'], name |
| rng = torch.Generator().manual_seed(42) |
| values = {} |
| for e in m['externals']: |
| shape, dtype = e['shape'], np.dtype(e['dtype']) |
| if e['kind'] == 'normal': a = torch.randn(shape, generator=rng).numpy() |
| elif e['kind'] == 'constant': a = np.fromfile(root/e['path'],dtype=dtype).reshape(shape) |
| else: |
| fixture = json.loads((root/e['tokenizer']).read_text())['fixtures'][0] |
| a = np.array(fixture['ids' if e['kind']=='tokens' else 'mask'],dtype=dtype).reshape(shape) |
| values[e['tensor']] = a |
| uses = {} |
| for action in m['actions']: |
| for i in action['inputs']: uses[i['tensor']] = uses.get(i['tensor'], 0) + 1 |
| uses[m['output']] = uses.get(m['output'], 0) + 1 |
| start = time.monotonic() |
| for action in m['actions']: |
| model = Interpreter(model_path=str(root/action['model']), num_threads=8) |
| model.allocate_tensors() |
| result = model.get_signature_runner(action['signature'])(**{i['name']: values[i['tensor']] for i in action['inputs']}) |
| for o in action['outputs']: |
| a = result[o['name']].copy() |
| assert list(a.shape) == o['shape'] and np.isfinite(a).all(), (action['model'],o['name']) |
| values[o['tensor']] = a |
| for i in action['inputs']: |
| uses[i['tensor']] -= 1 |
| if uses[i['tensor']] == 0: del values[i['tensor']] |
| print(action['model'],round(time.monotonic()-start,1),flush=True) |
| del result, model |
| image = np.clip(values[m['output']][0].transpose(1,2,0)*255,0,255).round().astype(np.uint8) |
| reference = np.array(Image.open(root/'reference.png')).astype(np.float64) |
| mse = np.mean((image.astype(np.float64) - reference)**2) |
| psnr = float(10*np.log10(255**2/max(mse,1e-12))) |
| assert psnr > 40, psnr |
| Image.fromarray(image).save(out) |
| report = dict(modelId=m['modelId'], backend='desktop LiteRT CPU', psnrVsExportReference=psnr, |
| elapsedSeconds=time.monotonic()-start, actions=len(m['actions']), deviceValidated=False) |
| Path(str(out)+'.json').write_text(json.dumps(report,indent=2)+'\n') |
| print(json.dumps(report),flush=True) |
|
|
|
|
| if __name__ == '__main__': |
| p=argparse.ArgumentParser(description=__doc__); p.add_argument('bundle'); p.add_argument('--out',required=True) |
| args=p.parse_args(); validate(args.bundle,args.out) |
|
|
