File size: 6,037 Bytes
dee7f43 | 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 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 | import os,sys,re,json,shutil
from pathlib import Path
os.environ.update(VDS_BATCH2_ROOT='/workspace/vds-litert-batch2',HF_TOKEN_FILE='/workspace/.secrets/hf_token',HF_FIREVIEWER_DEST_REPO='fireviewer/litert-models',TORCH_NUM_THREADS='8',HF_HOME='/workspace/cache/huggingface')
import torch,numpy as np,litert_torch
from ai_edge_litert.interpreter import Interpreter,OpResolverType
from huggingface_hub import HfApi
import dynamic_lowerings
root=Path('/workspace/vds-litert-batch2');sys.path.insert(0,str(root/'package/scripts'))
import convert_batch2 as worker
worker.WORK=root/'dynamic-work'
key=sys.argv[1]
spec=next((s for s in worker.all_specs() if s['key']==key),None)
if spec is None:
import previous_native
old=next(s for s in json.loads((root/'previous-models.json').read_text()) if s['key']==key)
spec={**old,'group':'general_models','route':key,'title':key,'task':'vision embeddings' if key=='dinov2' else 'pose estimation'}
worker.CONV[key]=lambda s,e,o:getattr(previous_native,key)(s,e,o,worker)
original_prefetched=worker.prefetched
worker.prefetched=lambda k:json.loads((root/'previous-prefetch.json').read_text())[k] if k==key else original_prefetched(k)
validation={}
adapter_sha256=worker.sha256(root/'dynamic_lowerings.py')
def dynamic_export(module,args,out):
global validation
module=module.cpu().eval();args=tuple(a.cpu() for a in args)
for sub in module.modules():
if hasattr(sub,'set_attn_implementation'):
try:sub.set_attn_implementation('eager')
except (ValueError,AttributeError):pass
h,w=args[0].shape[-2:]
stride=14 if spec['route'] in {'depth','dinov2'} else 16 if spec['route']=='dinov3_fireviewer' else 32
if h%stride or w%stride:stride=1
height=stride*torch.export.Dim('height_blocks',min=4,max=128) if stride>1 else torch.export.Dim('height',min=128,max=2048)
width=stride*torch.export.Dim('width_blocks',min=4,max=128) if stride>1 else torch.export.Dim('width',min=128,max=2048)
dims=[{} for a in args];dims[0]={2:height,3:width}
masks=[]
for i,a in enumerate(args[1:],1):
if a.ndim==3 and tuple(a.shape[-2:])==(h,w):dims[i]={1:height,2:width};masks.append(i)
shapes=[(h,w),(h-2*stride,w+stride),(h+stride,w-2*stride)]
def inputs(hh,ww):
aa=list(args);aa[0]=torch.from_numpy(np.random.default_rng(42).normal(size=(1,3,hh,ww)).astype(np.float32))
for i in masks:aa[i]=torch.ones((1,hh,ww),dtype=args[i].dtype)
return tuple(aa)
with torch.inference_mode():
for hh,ww in shapes:module(*inputs(hh,ww))
if os.environ.get('REUSE_ARTIFACT'):
shutil.copy2(os.environ['REUSE_ARTIFACT'],out)
else:
converted=litert_torch.convert(module,args,dynamic_shapes=tuple(dims))
out.parent.mkdir(parents=True,exist_ok=True);converted.export(str(out));del converted
interp=Interpreter(model_path=str(out),num_threads=8,experimental_op_resolver_type=OpResolverType.BUILTIN_WITHOUT_DEFAULT_DELEGATES)
details=interp.get_input_details();mapping={}
for j,d in enumerate(details):
match=re.search(r'args_(\d+)',d['name']);index=int(match.group(1)) if match else j
mapping[index]=d
if mapping[0]['shape_signature'].tolist()[2:]!=[-1,-1]:raise RuntimeError('Spatial signature was fixed by converter')
checks=[]
for hh,ww in shapes:
aa=inputs(hh,ww)
for i,a in enumerate(aa):interp.resize_tensor_input(mapping[i]['index'],tuple(a.shape),strict=True)
interp.allocate_tensors()
for i,a in enumerate(aa):interp.set_tensor(mapping[i]['index'],a.numpy())
interp.invoke()
with torch.inference_mode():ref=module(*aa)
if isinstance(ref,torch.Tensor):ref=(ref,)
refs=[v.detach().numpy() for v in ref]
errors=[]
for j,d in enumerate(interp.get_output_details()):
a=interp.get_tensor(d['index']);candidates=[r for r in refs if r.shape==a.shape]
if len(candidates)==1:r=candidates[0]
else:
match=re.search(r'output_(\d+)',d['name'])
if not match:raise RuntimeError(f'Ambiguous output mapping for {d["name"]} {a.shape}')
r=refs[int(match.group(1))]
if r.shape!=a.shape:raise RuntimeError('Output signature mismatch');np.testing.assert_allclose(a,r,rtol=1e-3,atol=1e-3)
if not np.isfinite(a).all():raise RuntimeError('Nonfinite output')
errors.append(float(np.max(np.abs(a-r))))
checks.append({'image_shape':[1,3,hh,ww],'max_abs_errors':errors})
print('DYNAMIC_PARITY_OK',key,hh,ww,errors,flush=True)
validation={'dynamic_validated':True,'checks':checks,'atol':1e-3,'rtol':1e-3,'alignment':stride,'source_resolution_resampled':False,'runtime_requirement':'LiteRT builtin CPU; disable default XNNPACK delegates for these dynamic graphs','android_tested':False,'accuracy_evaluated':False,'conversion_adapter_sha256':adapter_sha256}
return {**worker.inspect_tflite(out),**validation}
base_publish=worker.publish
def publish(api,s,e,out,arts,report):
if not validation.get('dynamic_validated'):raise RuntimeError('Dynamic validation gate did not pass')
report.update(validation)
return base_publish(api,s,e,out,arts,report)
worker.direct_export=dynamic_export;worker.publish=publish
if spec['route']=='rfdetr_medium':
from rfdetr.export._litert.exporter import LiteRTExporter
def patched_convert(self,litert,wrapped_model,input_tensors,output_file):
dynamic_export(wrapped_model,(input_tensors,),Path(output_file))
LiteRTExporter._convert_and_save=patched_convert
if key=='fireviewer_rfdetr_medium_v110' and os.environ.get('RF_NATIVE')=='1':
import rfdetr_native
worker.CONV['rfdetr_medium']=lambda s,e,o:rfdetr_native.convert(s,e,o,worker)
if key=='fireviewer_dfine_m_strict_v1':
import dfine_native
worker.CONV['onnx_existing']=lambda s,e,o:dfine_native.convert(s,e,o,worker)
worker.process(HfApi(token=worker.token()),spec)
|