| |
| """ |
| DM-JEPA 1.5 Apple Neural Engine (ANE) Exporter. |
| Exports the 94.24M Cartesian Block Butterfly Decision Architecture to Apple Core ML (.mlpackage). |
| Optimized specifically for the 16-core Apple Neural Engine on iPhone (A16, A17 Pro, A18) and iPad/Mac (M1-M4). |
| """ |
|
|
| import os |
| import argparse |
| import numpy as np |
| import torch |
| import coremltools as ct |
| from ane_model import DMJEPA15ANEModel |
|
|
|
|
| def export_coreml( |
| checkpoint_path: str = "checkpoints/djepa_v2.0_autonomous_production.pt", |
| output_path: str = "dm_jepa_1_5_ios_ane/DMJEPA15_ANE.mlpackage", |
| seq_len: int = 64, |
| num_options: int = 4, |
| opt_len: int = 32 |
| ): |
| print("=" * 70) |
| print("🚀 Exporting DM-JEPA 1.5 to Apple Neural Engine (Core ML .mlpackage)") |
| print("=" * 70) |
|
|
| |
| print("Initializing ANE-optimized PyTorch model...") |
| model = DMJEPA15ANEModel().eval() |
|
|
| |
| print(f"Loading weights from {checkpoint_path}...") |
| ckpt = torch.load(checkpoint_path, map_location="cpu", weights_only=False) |
| sd = ckpt["model_state_dict"] if "model_state_dict" in ckpt else ckpt |
|
|
| ane_sd = {} |
| for k, v in sd.items(): |
| k_ane = k |
| k_ane = k_ane.replace("criteria_encoder.proj.0.", "criteria_encoder.proj_0.") |
| k_ane = k_ane.replace("criteria_encoder.proj.2.", "criteria_encoder.proj_2.") |
| k_ane = k_ane.replace("predictor.cell.ffn.0.", "predictor_cell.ffn_0.") |
| k_ane = k_ane.replace("predictor.cell.ffn.2.", "predictor_cell.ffn_2.") |
| k_ane = k_ane.replace("predictor.cell.", "predictor_cell.") |
| k_ane = k_ane.replace("predictor.out_proj.0.", "predictor_out_0.") |
| k_ane = k_ane.replace("predictor.out_proj.1.", "predictor_out_1.") |
| k_ane = k_ane.replace("predictor.out_proj.2.", "predictor_out_2.") |
| ane_sd[k_ane] = v |
|
|
| missing, unexpected = model.load_state_dict(ane_sd, strict=False) |
| real_missing = [m for m in missing if not m.endswith(".cos") and not m.endswith(".sin")] |
| if real_missing: |
| print(f"⚠️ Warning: Missing keys: {real_missing}") |
| else: |
| print("✅ All 181 production parameters loaded cleanly into ANE graph!") |
|
|
| |
| state_ids = torch.randint(0, 1000, (1, seq_len), dtype=torch.int32) |
| state_mask = torch.ones((1, seq_len), dtype=torch.int32) |
| opt_ids = torch.randint(0, 1000, (1, num_options, opt_len), dtype=torch.int32) |
| opt_mask = torch.ones((1, num_options, opt_len), dtype=torch.int32) |
|
|
| print("Tracing execution graph with PyTorch JIT...") |
| traced = torch.jit.trace(model, (state_ids.long(), state_mask.long(), opt_ids.long(), opt_mask.long())) |
|
|
| |
| inputs = [ |
| ct.TensorType(name="state_input_ids", shape=(1, seq_len), dtype=np.int32), |
| ct.TensorType(name="state_attention_mask", shape=(1, seq_len), dtype=np.int32), |
| ct.TensorType(name="option_input_ids", shape=(1, num_options, opt_len), dtype=np.int32), |
| ct.TensorType(name="option_attention_mask", shape=(1, num_options, opt_len), dtype=np.int32), |
| ] |
|
|
| outputs = [ |
| ct.TensorType(name="logits"), |
| ct.TensorType(name="probabilities"), |
| ct.TensorType(name="metric_energy"), |
| ] |
|
|
| print("Converting PyTorch MIL graph to Apple Neural Engine program...") |
| mlmodel = ct.convert( |
| traced, |
| inputs=inputs, |
| outputs=outputs, |
| minimum_deployment_target=ct.target.iOS16, |
| compute_precision=ct.precision.FLOAT16 |
| ) |
|
|
| |
| mlmodel.author = "DangerLabs" |
| mlmodel.license = "Apache 2.0" |
| mlmodel.short_description = "DM-JEPA 1.5 Real-Time System 1 Decision Model for Apple Neural Engine (ANE)." |
| mlmodel.version = "1.5.0" |
| |
| mlmodel.input_description["state_input_ids"] = "State / Instruction token IDs [1, 64]" |
| mlmodel.input_description["state_attention_mask"] = "State attention mask [1, 64]" |
| mlmodel.input_description["option_input_ids"] = "Candidate action token IDs [1, 4, 32]" |
| mlmodel.input_description["option_attention_mask"] = "Candidate action attention mask [1, 4, 32]" |
|
|
| mlmodel.output_description["logits"] = "Raw compatibility logits for candidate options [1, 4]" |
| mlmodel.output_description["probabilities"] = "Softmax normalized choice probabilities [1, 4]" |
| mlmodel.output_description["metric_energy"] = "Riemannian manifold metric energy [1, 4]" |
|
|
| os.makedirs(os.path.dirname(output_path), exist_ok=True) |
| print(f"Saving compiled Core ML package to {output_path}...") |
| mlmodel.save(output_path) |
| print("=" * 70) |
| print("✨ Core ML package export for Apple Neural Engine complete! ✨") |
| print(f"📦 Output: {output_path}") |
| print("=" * 70) |
|
|
|
|
| if __name__ == "__main__": |
| parser = argparse.ArgumentParser() |
| parser.add_argument("--checkpoint", default="checkpoints/djepa_v2.0_autonomous_production.pt") |
| parser.add_argument("--output", default="dm_jepa_1_5_ios_ane/DMJEPA15_ANE.mlpackage") |
| args = parser.parse_args() |
| export_coreml(args.checkpoint, args.output) |
|
|
