File size: 1,777 Bytes
47d8ad6 | 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 | """Convert-only fp16 export of DenseTrackStep (real SAM3.1 weights) and save.
Splits convert from verify so the RAM-heavy predict loop doesn't OOM the
convert on a 16GB host. Verify separately with verify_coreml_lean.py.
"""
import time
import torch
import common
import dense_wrapper as dw
def main():
cache = torch.load("eager_cache.pt", weights_only=False)
wrapper, model = dw.build_wrapper()
inputs = dw.frame_inputs(model, cache, 3) # mid frame exercises cond-tpos path
common.hide_triton_stub()
with torch.no_grad():
try:
ep = torch.export.export(wrapper, inputs)
except Exception as e:
print(f"strict export failed ({type(e).__name__}); retry strict=False")
ep = torch.export.export(wrapper, inputs, strict=False)
ep = ep.run_decompositions({})
print("torch.export OK")
import coremltools as ct
from coremltools.converters.mil.frontend.torch.torch_op_registry import (
register_torch_op,
)
from coremltools.converters.mil.frontend.torch.ops import _get_inputs
from coremltools.converters.mil.mil import Builder as mb
@register_torch_op(torch_alias=["where.scalarother"])
def where_scalarother(context, node):
cond, a, b = _get_inputs(context=context, node=node, expected=3)
context.add(mb.select(cond=cond, a=a, b=b), node.name)
t0 = time.time()
mlmodel = ct.convert(
ep,
minimum_deployment_target=ct.target.iOS18,
compute_units=ct.ComputeUnit.CPU_ONLY, # convert/load only; deploy sets CU
)
print(f"fp16 convert OK in {time.time()-t0:.1f}s")
mlmodel.save("dense_sam3_trackstep.mlpackage")
print("saved dense_sam3_trackstep.mlpackage")
if __name__ == "__main__":
main()
|