|
|
|
|
| """
|
| MediaTek NeuroPilot APU AOT Metadata Compilation Pipeline (LiteRT)
|
| ================================================================================
|
| Ce script automatise la compilation et l'injection de métadonnées AOT LiteRT
|
| pour les processeurs MediaTek Dimensity et leurs APU dédiées.
|
|
|
| Règles matérielles :
|
| - Dimensity 8300 / 9300 / 9400 (APU 780, 790, 890) -> W4A8 prioritaire (+ INT8 standard).
|
| - Dimensity 9000 / 9200 (APU 5.0, 680) -> Strictement INT8 pur (pas de bloc sub-byte matériel).
|
| - SESR-M5 Super-Resolution -> Strictement INT8 sur toutes les cibles.
|
| """
|
|
|
| import os
|
| import sys
|
| import argparse
|
| from pathlib import Path
|
|
|
| MEDIATEK_GENERATIONS = {
|
| "MT6983": {
|
| "name": "Dimensity 9000 / 9000+",
|
| "apu": "APU 5.0",
|
| "precision": "INT8",
|
| "devices": "OnePlus Nord 3, Vivo X80",
|
| "android": "Android 12-14"
|
| },
|
| "MT6985": {
|
| "name": "Dimensity 9200 / 9200+",
|
| "apu": "APU 680",
|
| "precision": "INT8",
|
| "devices": "Vivo X90, Xiaomi 13T Pro",
|
| "android": "Android 13-15"
|
| },
|
| "MT6897": {
|
| "name": "Dimensity 8300 / 8300-Ultra",
|
| "apu": "APU 780",
|
| "precision": "W4A8",
|
| "devices": "POCO X6 Pro",
|
| "android": "Android 14-16"
|
| },
|
| "MT6989": {
|
| "name": "Dimensity 9300 / 9300+",
|
| "apu": "APU 790",
|
| "precision": "W4A8",
|
| "devices": "Vivo X100 Pro, Pad Pro",
|
| "android": "Android 14-16"
|
| },
|
| "MT6991": {
|
| "name": "Dimensity 9400",
|
| "apu": "APU 890",
|
| "precision": "W4A8",
|
| "devices": "Vivo X200 Series, Oppo Find X8",
|
| "android": "Android 15-17"
|
| }
|
| }
|
|
|
| BASE_DIR = Path(__file__).resolve().parent
|
|
|
| def find_models_root() -> Path:
|
| env_root = os.environ.get("MODELS_ROOT")
|
| if env_root and Path(env_root).exists():
|
| return Path(env_root)
|
| if (BASE_DIR / "MobileNetV5").exists():
|
| return BASE_DIR
|
| e_drive = Path(r"E:\fine tuning\rock identifier\fine tuning rock model")
|
| if e_drive.exists():
|
| return e_drive
|
| return BASE_DIR
|
|
|
| MODELS_ROOT = find_models_root()
|
|
|
| def inject_mediatek_metadata(src_path: Path, dst_path: Path, soc_model: str, gen_info: dict, quant_type: str):
|
| if not src_path.exists():
|
| print(f" [-] Source model missing: {src_path}")
|
| return False
|
|
|
| dst_path.parent.mkdir(parents=True, exist_ok=True)
|
| with open(src_path, "rb") as f:
|
| buf = f.read()
|
|
|
| metadata_content = (
|
| f"LITERT_AOT_TARGET=MEDIATEK\n"
|
| f"SOC_MANUFACTURER=MediaTek\n"
|
| f"SOC_MODEL={soc_model}\n"
|
| f"SOC_MARKETING_NAME={gen_info.get('name', '')}\n"
|
| f"HARDWARE_ACCELERATOR={gen_info.get('apu', 'APU')}\n"
|
| f"TARGET_DEVICES={gen_info.get('devices', '')}\n"
|
| f"TARGET_OS={gen_info.get('android', 'Android 14+')}\n"
|
| f"QUANTIZATION_SCHEME={quant_type}\n"
|
| f"PRECOMPILED_METADATA_VERSION=1.0\n"
|
| ).encode("utf-8")
|
|
|
|
|
| try:
|
| from ai_edge_litert import schema_py_generated as schema_fb
|
| import flatbuffers
|
|
|
| model = schema_fb.ModelT.InitFromObj(schema_fb.Model.GetRootAsModel(buf, 0))
|
|
|
| meta_buffer = schema_fb.BufferT()
|
| meta_buffer.data = list(metadata_content)
|
| model.buffers.append(meta_buffer)
|
| meta_buffer_idx = len(model.buffers) - 1
|
|
|
| meta_entry = schema_fb.MetadataT()
|
| meta_entry.name = f"mediatek_neuropilot_aot_{soc_model.lower()}"
|
| meta_entry.buffer = meta_buffer_idx
|
| if model.metadata is None:
|
| model.metadata = []
|
| model.metadata.append(meta_entry)
|
|
|
| builder = flatbuffers.Builder(1024 * 1024 * 20)
|
| builder.Finish(model.Pack(builder), b"TFL3")
|
| out_bytes = builder.Output()
|
| except Exception as e:
|
| print(f"[-] Erreur injection FlatBuffer MediaTek: {e}")
|
| return False
|
|
|
| with open(dst_path, "wb") as f:
|
| f.write(out_bytes)
|
|
|
| size_mb = len(out_bytes) / (1024 * 1024)
|
| print(f" [OK] MediaTek [{soc_model}] {dst_path.name} ({size_mb:.2f} MB)")
|
| return True
|
|
|
| def compile_mediatek_targets(targets=None):
|
| if targets is None:
|
| targets = list(MEDIATEK_GENERATIONS.keys())
|
|
|
| print("=" * 80)
|
| print(" Compiling MediaTek NeuroPilot AOT Targets (Dimensity)")
|
| print("=" * 80)
|
|
|
| for soc_model in targets:
|
| if soc_model not in MEDIATEK_GENERATIONS:
|
| print(f"[!] Target inconnue: {soc_model}")
|
| continue
|
|
|
| info = MEDIATEK_GENERATIONS[soc_model]
|
| is_w4 = (info["precision"] == "W4A8")
|
| print(f"\n[*] Target: MediaTek {soc_model} ({info['name']} - {info['apu']}) -> Priorite: {info['precision']}")
|
|
|
|
|
| v5_src = MODELS_ROOT / "MobileNetV5" / "tflite" / ("best_mobilenetv5_rocks_w4a8.tflite" if is_w4 else "best_mobilenetv5_rocks_int8.tflite")
|
| v5_out = MODELS_ROOT / "MobileNetV5" / "MediaTek" / soc_model / ("best_mobilenetv5_rocks_w4a8.tflite" if is_w4 else "best_mobilenetv5_rocks_int8.tflite")
|
| inject_mediatek_metadata(v5_src, v5_out, soc_model, info, "W4A8_FULL_INTEGER" if is_w4 else "PER_TENSOR_INT8_SIGNED")
|
|
|
| if is_w4:
|
| v5_int8_src = MODELS_ROOT / "MobileNetV5" / "tflite" / "best_mobilenetv5_rocks_int8.tflite"
|
| v5_int8_out = MODELS_ROOT / "MobileNetV5" / "MediaTek" / soc_model / "best_mobilenetv5_rocks_int8.tflite"
|
| inject_mediatek_metadata(v5_int8_src, v5_int8_out, soc_model, info, "PER_TENSOR_INT8_SIGNED")
|
|
|
|
|
| v4_src = MODELS_ROOT / "MobileNetV4" / "tflite" / ("rock_model_w4a8_MobileNetV4.tflite" if is_w4 else "rock_model_w8a8_MobileNetV4.tflite")
|
| v4_out = MODELS_ROOT / "MobileNetV4" / "MediaTek" / soc_model / ("rock_model_w4a8_MobileNetV4.tflite" if is_w4 else "rock_model_w8a8_MobileNetV4.tflite")
|
| inject_mediatek_metadata(v4_src, v4_out, soc_model, info, "W4A8_FULL_INTEGER" if is_w4 else "PER_TENSOR_INT8_SIGNED")
|
|
|
|
|
| scale_src = MODELS_ROOT / "ScaleDetection_MobileNetV4" / "tflite" / ("v4_small_scale_int4.tflite" if is_w4 else "v4_small_scale_int8.tflite")
|
| scale_out = MODELS_ROOT / "ScaleDetection_MobileNetV4" / "MediaTek" / soc_model / ("v4_small_scale_int4.tflite" if is_w4 else "v4_small_scale_int8.tflite")
|
| inject_mediatek_metadata(scale_src, scale_out, soc_model, info, "W4A8_FULL_INTEGER" if is_w4 else "PER_TENSOR_INT8_SIGNED")
|
|
|
|
|
| sesr_src = MODELS_ROOT / "SESR" / "tflite" / "sesr_m5.tflite"
|
| sesr_out = MODELS_ROOT / "SESR" / "MediaTek" / soc_model / "sesr_m5.tflite"
|
| inject_mediatek_metadata(sesr_src, sesr_out, soc_model, info, "PER_TENSOR_INT8_SIGNED")
|
|
|
| print("\n[OK] Compilation MediaTek terminee avec succes.")
|
|
|
| if __name__ == "__main__":
|
| parser = argparse.ArgumentParser(description="MediaTek NeuroPilot AOT Compilation")
|
| parser.add_argument("--soc", choices=list(MEDIATEK_GENERATIONS.keys()), default=None,
|
| help="Cible SoC specifique (ex: MT6989). Par defaut: Toutes les cibles.")
|
| args = parser.parse_args()
|
|
|
| targets = [args.soc] if args.soc else list(MEDIATEK_GENERATIONS.keys())
|
| compile_mediatek_targets(targets)
|
|
|
