From 2d0fb855b12dbfcbf20d6c4c29db8161ed23c91c Mon Sep 17 00:00:00 2001 From: Musa Mahmood Date: Fri, 16 Oct 2026 19:57:16 -0400 Subject: [PATCH] [0.1.29] macOS phase 5: SigLIP 2 labels on the GPU through native Core ML --- README.md | 5 + check_targets.jai | 2 +- first.jai | 16 +++- modules/Core_ML/module.jai | 176 ++++++++++++++++++++++++++++++++++ modules/Labels/models.config | 3 + modules/Labels/module.jai | 73 ++++++++++++-- modules/onnxruntime/cuda.jai | 7 +- src/app/labels_ui.jai | 4 +- src/app/ocr_ui.jai | 2 +- src/test/labels_tests.jai | 29 +++++- tools/export_siglip_coreml.py | 115 ++++++++++++++++++++++ 11 files changed, 419 insertions(+), 13 deletions(-) create mode 100644 modules/Core_ML/module.jai create mode 100644 tools/export_siglip_coreml.py diff --git a/README.md b/README.md index 63b2252..37eeb28 100644 --- a/README.md +++ b/README.md @@ -28,6 +28,11 @@ see [Screen capture](#screen-capture). place (`C:\Program Files\NVIDIA\CUDNN`) is found without changing PATH. The OCR window says what, if anything, is missing. + **On macOS**, labelling runs on the GPU through Core ML, about ten times as fast as the CPU, from + a Core ML copy of SigLIP 2's image model made once with `tools/export_siglip_coreml.py` (its first + lines say how to set up the Python it needs). Without it, labels are made on the CPU. OCR runs on + the CPU: Core ML through onnxruntime was measured slower there. + 3. **Build and run:** `jai first.jai - release`, then `bin\mcap.exe`. The build copies onnxruntime and the models beside the program. diff --git a/check_targets.jai b/check_targets.jai index 20efd7a..c70023d 100644 --- a/check_targets.jai +++ b/check_targets.jai @@ -49,7 +49,7 @@ APP_NAME :: "mcap"; APP_FOLDER_NAME :: "mcap"; GIT_BRANCH :: "check"; GIT_HASH :: "-"; -VERSION :: "0.1.28"; +VERSION :: "0.1.29"; BUILD_DATE :: "-"; DEBUG :: true; COMPILE_TESTS :: true; diff --git a/first.jai b/first.jai index 1f3c5ee..b43f186 100644 --- a/first.jai +++ b/first.jai @@ -8,7 +8,7 @@ APP_NAME :: "mcap"; APP_FOLDER_NAME :: "mcap"; APP_BUNDLE_IDENTIFIER :: "com.example.mcap"; -VERSION :: "0.1.28"; +VERSION :: "0.1.29"; // Where everything goes: bin/, or bin-headless/ for builds made with no_crash_dialog, so automated // builds never write over the ones you run by hand. build_directory := "bin"; @@ -259,6 +259,7 @@ finish_main_build :: (build: Pending_Build) { clone_if_changed(tprint("modules/%/models.config", it.module), tprint("%/models.config", folder)); models := tprint("modules/%/models", it.module); if file_exists(models) for file_list(models) clone_if_changed(it, tprint("%/%", folder, path_filename(it))); + if file_exists(models) clone_packages(models, folder); } // An .iconset of every size, made with sips, then iconutil: sips -s format icns fails on @@ -311,6 +312,18 @@ finish_main_build :: (build: Pending_Build) { return "ad hoc"; } + // Core ML packages (.mlpackage, folders) in from, cloned whole into to: made afresh each build, which + // as clones costs no time or space. + clone_packages :: (from: string, to: string) { + visit_files(from, false, to, (info: *File_Visit_Info, to: string) { + if !ends_with(info.short_name, ".mlpackage") return; + target := tprint("%/%", to, info.short_name); + run_command("rm", "-rf", target); + result := run_command("cp", "-c", "-R", info.full_name, target); + if result.exit_code != 0 log_error("Could not copy % to %.", info.full_name, target); + }, visit_files = false, visit_directories = true); + } + // As copy_if_changed, but an APFS clone (cp -c), which shares the original's blocks. clone_if_changed :: (from: string, to: string) { from_time, from_size, from_ok := file_modtime_and_size(from); @@ -383,6 +396,7 @@ copy_assets :: () { copy_if_changed("modules/Labels/models.config", tprint("%/label-models/models.config", build_directory)); if file_exists("modules/Labels/models") { for file_list("modules/Labels/models") copy_if_changed(it, tprint("%/label-models/%", build_directory, path_filename(it))); + #if OS == .MACOS clone_packages("modules/Labels/models", tprint("%/label-models", build_directory)); } else { log("Labels: the models have not been fetched; run jai fetch_models.jai in modules/Labels for labels."); } diff --git a/modules/Core_ML/module.jai b/modules/Core_ML/module.jai new file mode 100644 index 0000000..3127ba7 --- /dev/null +++ b/modules/Core_ML/module.jai @@ -0,0 +1,176 @@ +// Core_ML: a Core ML model (an .mlpackage made by coremltools) run on one float tensor, on macOS. +// The models that run faster here than through onnxruntime on the CPU (tools/export_siglip_coreml.py +// says how much) use it; the rest stay with modules/onnxruntime. +// +// model, ok, error := coreml_load(path, "pixel_values", "pooler_output"); +// ok, error = coreml_run(model, input, s64.[1, 3, 384, 384], output); // output: the floats it gives +// coreml_free(model); +// +// Loading compiles the package (Core ML's compileModelAtURL, a fraction of a second) into a temporary +// .mlmodelc, removed by coreml_free. The model runs on the CPU and GPU, never the Neural Engine +// unless asked: with float16 weights it can be badly wrong there (SigLIP 2: cosine 0.72 to the CPU). +// A model runs one call at a time (a mutex), from any thread. + +Coreml_Units :: enum s64 { // MLComputeUnits + CPU_ONLY :: 0; + CPU_AND_GPU :: 1; + ALL :: 2; + CPU_AND_ENGINE :: 3; +} + +Coreml_Model :: struct { + model: *NSObject; // MLModel, retained + compiled: *NSObject; // the .mlmodelc's NSURL, retained; removed on free + input_name, output_name: string; + mutex: Mutex; +} + +coreml_load :: (path: string, input_name: string, output_name: string, units := Coreml_Units.CPU_AND_GPU) -> *Coreml_Model, ok: bool, error: string { + init_objective_c(); // before anything Objective-C, the pool included; once per process does it all + pool := objc_new(NSAutoreleasePool); + defer NSAutoreleasePool.drain(pool); + + if !file_exists(path) return null, false, tprint("% is not there.", path); + package := NSURL.fileURLWithPath(NSString.getTempString(path)); + + error: *NSError; + func_compile :: #type (*void, Selector, *NSURL, **NSError) -> *NSURL #c_call; + compiled := (cast(func_compile) objc_msgSend)(objc_getClass("MLModel"), selector("compileModelAtURL:error:"), package, *error); + if !compiled return null, false, tprint("Core ML could not compile %: %", path, error_text(error)); + + configuration := objc_init(objc_alloc(objc_getClass("MLModelConfiguration"))); + defer release(configuration); + func_units :: #type (*void, Selector, s64) -> void #c_call; + (cast(func_units) objc_msgSend)(configuration, selector("setComputeUnits:"), cast(s64) units); + + func_load :: #type (*void, Selector, *NSURL, *NSObject, **NSError) -> *NSObject #c_call; + model := (cast(func_load) objc_msgSend)(objc_getClass("MLModel"), selector("modelWithContentsOfURL:configuration:error:"), compiled, configuration, *error); + if !model { + remove_compiled(compiled); + return null, false, tprint("Core ML could not load %: %", path, error_text(error)); + } + + result := New(Coreml_Model,, context.default_allocator); + result.model = xx retain(model); + result.compiled = xx retain(compiled); + result.input_name = copy_string(input_name,, context.default_allocator); + result.output_name = copy_string(output_name,, context.default_allocator); + init(*result.mutex); + return result, true, ""; +} + +coreml_free :: (model: *Coreml_Model) { + if !model return; + pool := objc_new(NSAutoreleasePool); + defer NSAutoreleasePool.drain(pool); + release(model.model); + remove_compiled(cast(*NSURL) model.compiled); + release(model.compiled); + destroy(*model.mutex); + free(model.input_name,, context.default_allocator); + free(model.output_name,, context.default_allocator); + free(model,, context.default_allocator); +} + +// Runs the model on input, a float32 tensor of shape, and copies its output (float32 or float16) into +// output as float32. ok is false, with why, when it fails or gives a different number of values. +coreml_run :: (model: *Coreml_Model, input: [] float, shape: [] s64, output: [] float) -> ok: bool, error: string { + lock(*model.mutex); + defer unlock(*model.mutex); + pool := objc_new(NSAutoreleasePool); + defer NSAutoreleasePool.drain(pool); + + // The input, over the caller's floats: packed, so each stride is the product of the sizes after it. + numbers :: (values: [] s64) -> *NSObject { + objects := NewArray(values.count, *NSObject,, temp); + func_number :: #type (*void, Selector, s64) -> *NSObject #c_call; + for values objects[it_index] = (cast(func_number) objc_msgSend)(objc_getClass("NSNumber"), selector("numberWithLongLong:"), it); + func_array :: #type (*void, Selector, **NSObject, u64) -> *NSObject #c_call; + return (cast(func_array) objc_msgSend)(objc_getClass("NSArray"), selector("arrayWithObjects:count:"), objects.data, xx objects.count); + } + strides := NewArray(shape.count, s64,, temp); + step: s64 = 1; + for < shape { strides[it_index] = step; step *= it; } + if step != input.count return false, tprint("the input has % values, not % for its shape %", input.count, step, shape); + + error: *NSError; + func_array_init :: #type (*void, Selector, *void, *NSObject, s64, *NSObject, *void, **NSError) -> *NSObject #c_call; + array := (cast(func_array_init) objc_msgSend)(objc_alloc(objc_getClass("MLMultiArray")), selector("initWithDataPointer:shape:dataType:strides:deallocator:error:"), + input.data, numbers(shape), MULTIARRAY_FLOAT32, numbers(strides), null, *error); + if !array return false, tprint("Core ML could not take the input: %", error_text(error)); + autorelease(array); + + func_value :: #type (*void, Selector, *NSObject) -> *NSObject #c_call; + value := (cast(func_value) objc_msgSend)(objc_getClass("MLFeatureValue"), selector("featureValueWithMultiArray:"), array); + key := NSString.getTempString(model.input_name); + func_dictionary :: #type (*void, Selector, *NSObject, *NSObject) -> *NSObject #c_call; + dictionary := (cast(func_dictionary) objc_msgSend)(objc_getClass("NSDictionary"), selector("dictionaryWithObject:forKey:"), value, key); + func_provider :: #type (*void, Selector, *NSObject, **NSError) -> *NSObject #c_call; + features := (cast(func_provider) objc_msgSend)(objc_alloc(objc_getClass("MLDictionaryFeatureProvider")), selector("initWithDictionary:error:"), dictionary, *error); + if !features return false, tprint("Core ML could not take the input: %", error_text(error)); + autorelease(features); + + func_predict :: #type (*void, Selector, *NSObject, **NSError) -> *NSObject #c_call; + result := (cast(func_predict) objc_msgSend)(model.model, selector("predictionFromFeatures:error:"), features, *error); + if !result return false, tprint("Core ML could not run the model: %", error_text(error)); + + out_value := (cast(func_value) objc_msgSend)(result, selector("featureValueForName:"), NSString.getTempString(model.output_name)); + out_array := ifx out_value then objc_msgSend_typed(out_value, selector("multiArrayValue")) else null; + if !out_array return false, tprint("the model gives no %", model.output_name); + count := cast(s64) objc_msgSend_typed(out_array, selector("count")); + if count != output.count return false, tprint("the model gives % values, not %", count, output.count); + func_type :: #type (*void, Selector) -> s64 #c_call; + type := (cast(func_type) objc_msgSend)(out_array, selector("dataType")); + // The output is packed (one row of the pooled vector here); dataPointer is read before the pool drains. + data := objc_msgSend_typed(out_array, selector("dataPointer")); + if type == MULTIARRAY_FLOAT32 { + memcpy(output.data, data, count * size_of(float)); + } else if type == MULTIARRAY_FLOAT16 { + half := cast(*u16) data; + for 0..count - 1 output[it] = float16_to_float32(half[it]); + } else { + return false, tprint("the model's output is of type 0x%, not float32 or float16", formatInt(type, base = 16)); + } + return true, ""; +} + +#scope_file + +MULTIARRAY_FLOAT32 :: 0x10000 | 32; // MLMultiArrayDataTypeFloat32 +MULTIARRAY_FLOAT16 :: 0x10000 | 16; // MLMultiArrayDataTypeFloat16 + +error_text :: (error: *NSError) -> string { + if !error return "no reason given"; + return copy_string(to_string(NSError.localizedDescription(error)),, temp); +} + +remove_compiled :: (url: *NSURL) { + func :: #type (*void, Selector, *NSURL, **NSError) -> BOOL #c_call; + (cast(func) objc_msgSend)(NSFileManager.defaultManager(), selector("removeItemAtURL:error:"), url, null); +} + +float16_to_float32 :: (h: u16) -> float { + sign := (cast(u32) h & 0x8000) << 16; + exponent := (cast(u32) h >> 10) & 0x1F; + mantissa := cast(u32) h & 0x3FF; + bits: u32; + if exponent == 0 { + if mantissa == 0 bits = sign; + else { // subnormal: normalised + e: u32 = 113; + while !(mantissa & 0x400) { mantissa <<= 1; e -= 1; } + bits = sign | (e << 23) | ((mantissa & 0x3FF) << 13); + } + } else if exponent == 31 bits = sign | 0x7F80_0000 | (mantissa << 13); + else bits = sign | ((exponent + 112) << 23) | (mantissa << 13); + return (cast(*float) *bits).*; +} + +CoreML_Framework :: #library,system,link_always "CoreML"; + +#import "Basic"; +#import "File"; +#import "File_Utilities"; +#import "Thread"; +#import "Objective_C"; +#import "Objective_C/AppKit"; diff --git a/modules/Labels/models.config b/modules/Labels/models.config index 2809b75..f49645b 100644 --- a/modules/Labels/models.config +++ b/modules/Labels/models.config @@ -10,6 +10,8 @@ # logit_scale, logit_bias a phrase's score for an image is sigmoid(e^logit_scale * cosine + logit_bias), # from the trained model (they are not in the ONNX encoders) # image_batch images run through the vision model together +# vision_coreml on macOS, the vision model as a Core ML package, quoted, run on the GPU +# (tools/export_siglip_coreml.py makes it); without it, onnxruntime runs it [model siglip2-base-patch16-384] title: "SigLIP 2 base, 384 px" @@ -21,3 +23,4 @@ max_tokens: 64 logit_scale: 4.726024 logit_bias: -16.772243 image_batch: 8 +vision_coreml: "siglip2-base-patch16-384-vision.mlpackage" diff --git a/modules/Labels/module.jai b/modules/Labels/module.jai index a7eb2c6..bb5d9ed 100644 --- a/modules/Labels/module.jai +++ b/modules/Labels/module.jai @@ -8,9 +8,11 @@ // // The models are listed in models_folder/models.config (fetch_models.jai downloads them) and run // with onnxruntime (modules/onnxruntime): the vision model on an NVIDIA GPU through CUDA when it -// can, else on the CPU; the text model, small inputs and only a few at a time, on the CPU. +// can, else on the CPU; the text model, small inputs and only a few at a time, on the CPU. On macOS +// the vision model runs on the GPU through Core ML (modules/Core_ML) when its vision_coreml package +// is there (tools/export_siglip_coreml.py makes it): about ten times the CPU's speed. -Labels_Provider :: enum { CPU; CUDA; } +Labels_Provider :: enum { CPU; CUDA; COREML; } // Tightly packed or not: stride is the bytes from one row to the next. Labels_Image :: struct { @@ -22,6 +24,7 @@ Labels_Model_Info :: struct { name: string; title: string; vision, text, tokenizer: string; // file names in the models folder + vision_coreml: string; // the vision model as a Core ML package, for macOS; "" for none image_size: s32 = 384; max_tokens: s32 = 64; logit_scale: float = 0; @@ -36,6 +39,7 @@ Labels_Engine :: struct { env: *OrtEnv; memory_info: *OrtMemoryInfo; vision, text: Ort_Session; + #if OS == .MACOS coreml: *Coreml_Model; // the vision model, when it runs through Core ML tokenizer: *Gemma_Tokenizer; dimensions: s32; // of every vector } @@ -75,6 +79,7 @@ labels_list_models :: (models_folder: string) -> [] Labels_Model_Info, error: st case "vision"; model.vision = copy_string(value); case "text"; model.text = copy_string(value); case "tokenizer"; model.tokenizer = copy_string(value); + case "vision_coreml"; model.vision_coreml = copy_string(value); case "image_size"; if number_ok model.image_size = cast(s32) number; case "max_tokens"; if number_ok model.max_tokens = cast(s32) number; case "logit_scale"; if number_ok model.logit_scale = number; @@ -113,11 +118,23 @@ labels_create :: (models_folder: string, model_name: string, prefer_gpu := true) cores := cast(s32) max(get_number_of_processors(.ALL_LOGICAL), 1); vision_path := tprint("%/%", models_folder, model.vision); use_gpu := false; - if prefer_gpu { + #if OS == .MACOS if prefer_gpu && model.vision_coreml && !file_exists(tprint("%/%", models_folder, model.vision_coreml)) { + engine.gpu_explanation = sprint("Labelling runs on the CPU: % is not there; tools/export_siglip_coreml.py makes it.", model.vision_coreml); + } else if prefer_gpu && model.vision_coreml { + // Core ML on the GPU; on any failure, onnxruntime on the CPU below, with the reason. + coreml_path := tprint("%/%", models_folder, model.vision_coreml); + engine.coreml, ok, error = coreml_load(coreml_path, "pixel_values", "pooler_output"); + if ok engine.provider = .COREML; // its first run is below, once the vectors' size is known + else engine.gpu_explanation = sprint("Labelling runs on the CPU: Core ML could not load % on the GPU:\n%", model.vision_coreml, error); + } + #if OS == .MACOS if prefer_gpu && !model.vision_coreml { + engine.gpu_explanation = sprint("Labelling runs on the CPU: there is no Core ML model for it; tools/export_siglip_coreml.py makes one."); + } + if prefer_gpu && OS != .MACOS { available, explanation := ort_cuda_check("Labelling"); use_gpu = available; if !available engine.gpu_explanation = copy_string(explanation); - } else { + } else if !prefer_gpu { engine.gpu_explanation = copy_string("Labelling runs on the CPU, as it was asked to."); } if use_gpu { @@ -133,7 +150,7 @@ labels_create :: (models_folder: string, model_name: string, prefer_gpu := true) use_gpu = false; } } - if !use_gpu { + if !use_gpu && engine.provider != .COREML { engine.vision, ok, error = ort_create_session(engine.env, vision_path, false, cores, output = "pooler_output", missing_hint = HINT); if !ok { labels_destroy(engine); return null, false, error; } engine.provider = .CPU; @@ -144,6 +161,19 @@ labels_create :: (models_folder: string, model_name: string, prefer_gpu := true) // The vectors' size, from a run of the text model on an empty text. probe := labels_embed_texts(engine, .[""],, temp); if !engine.dimensions { labels_destroy(engine); return null, false, "The text model gave no vector."; } + + #if OS == .MACOS if engine.provider == .COREML { + // Core ML's first run; when it fails, onnxruntime on the CPU after all. + ok, error = labels_warm_up(engine); + if !ok { + coreml_free(engine.coreml); + engine.coreml = null; + engine.provider = .CPU; + engine.gpu_explanation = sprint("Labelling runs on the CPU: Core ML could not run % on the GPU:\n%", model.vision_coreml, error); + engine.vision, ok, error = ort_create_session(engine.env, vision_path, false, cores, output = "pooler_output", missing_hint = HINT); + if !ok { labels_destroy(engine); return null, false, error; } + } + } return engine, true, ""; } @@ -152,6 +182,7 @@ labels_destroy :: (engine: *Labels_Engine) { if engine.provider == .CUDA ort_gpu_lock(); ort_release_session(*engine.vision); if engine.provider == .CUDA ort_gpu_unlock(); + #if OS == .MACOS coreml_free(engine.coreml); ort_release_session(*engine.text); tokenizer_free(engine.tokenizer,, context.default_allocator); if engine.memory_info ort.api.ReleaseMemoryInfo(engine.memory_info); @@ -161,7 +192,16 @@ labels_destroy :: (engine: *Labels_Engine) { } labels_gpu_status :: (engine: *Labels_Engine) -> in_use: bool, explanation: string { - return engine.provider == .CUDA, engine.gpu_explanation; + return engine.provider == .CUDA || engine.provider == .COREML, engine.gpu_explanation; +} + +// What runs the vision model, for people and the files it labels: "CUDA", "Core ML" or "CPU". +labels_provider_name :: (engine: *Labels_Engine) -> string { + if engine.provider == { + case .CUDA; return "CUDA"; + case .COREML; return "Core ML"; + } + return "CPU"; } // One unit vector per image, engine.dimensions long, one after another; from allocator. An image @@ -176,6 +216,18 @@ labels_embed_images :: (engine: *Labels_Engine, images: [] Labels_Image, allocat input := NewArray(plane * 3 * batch, float,, context.default_allocator); defer array_free(input,, context.default_allocator); + #if OS == .MACOS if engine.coreml { + // One image a run: the Core ML model takes [1, 3, size, size]. + out := NewArray(dims, float,, temp); + for images { + labels_prepare_image(it, size, input.data); + ok, run_error := coreml_run(engine.coreml, array_view(input, 0, plane * 3), s64.[1, 3, size, size], out); + if !ok { error = tprint("%", run_error); continue; } + labels_normalize_into(out.data, dims, result.data + it_index * dims); + } + return result, error; + } + first := 0; while first < images.count { count := min(batch, images.count - first); @@ -279,6 +331,13 @@ labels_warm_up :: (engine: *Labels_Engine) -> ok: bool, error: string { size := engine.model.image_size; input := NewArray(3 * size * size, float, initialized = true,, context.default_allocator); defer array_free(input,, context.default_allocator); + #if OS == .MACOS if engine.coreml { + out := NewArray(engine.dimensions, float,, context.default_allocator); + defer array_free(out,, context.default_allocator); + ok, error := coreml_run(engine.coreml, input, s64.[1, 3, size, size], out); + if !ok return false, tprint("The first run on the GPU failed: %", error); + return true, ""; + } out, ok, error := ort_run_float(*engine.vision, engine.memory_info, input, s64.[1, 3, size, size]); if !ok return false, tprint("The first run on the GPU failed: %", error); ort_release_output(*out); @@ -374,8 +433,10 @@ labels_strip_comment :: (line: string) -> string { #import "Basic"; #import "String"; #import "File"; +#import "File_Utilities"; #import "Math"; #import "System"; #import "Hash_Table"; #import "Float16"; #import "onnxruntime"; +#if OS == .MACOS #import "Core_ML"; diff --git a/modules/onnxruntime/cuda.jai b/modules/onnxruntime/cuda.jai index b9ca657..14ddd87 100644 --- a/modules/onnxruntime/cuda.jai +++ b/modules/onnxruntime/cuda.jai @@ -49,8 +49,13 @@ Cuda_Part :: enum { DRIVER; TOOLKIT; CUDNN; } // Whether the CUDA provider should be tried, and when not, an explanation for people: what is // missing, and the steps that fix it, starting " runs on the CPU". Temp. ort_cuda_check :: (what: string) -> available: bool, explanation: string { + // CoreML, onnxruntime's provider for Apple's GPU and Neural Engine, was measured on an M5 Pro + // (onnxruntime 1.30.0) and is slower than the CPU for these models: OCR's take a new image size + // nearly every run, which CoreML compiles for each time, and SigLIP 2 converts only with every + // dimension fixed (its ML Program conversion fails on the patch embedding), then runs at about + // half the CPU's speed. So on macOS the models run on the CPU. #if OS == .MACOS { - return false, tprint("% runs on the CPU: onnxruntime has no CUDA on macOS.", what); + return false, tprint("% runs on the CPU: onnxruntime has no CUDA on macOS, and its CoreML provider was measured slower than the CPU on Apple silicon.", what); } provider_path := tprint("%/%", path_strip_filename(ort.library_path), CUDA_PROVIDER_LIBRARY); diff --git a/src/app/labels_ui.jai b/src/app/labels_ui.jai index 8b89b0f..9cd629e 100644 --- a/src/app/labels_ui.jai +++ b/src/app/labels_ui.jai @@ -244,7 +244,7 @@ labels_ui_task :: (task: Ocr_Task) { engine, ok, error := labels_create(labels_models_folder(), LABELS_MODEL); free(labels_ui.engine_error); labels_ui.engine_error = ifx ok then "" else copy_string(error); - if ok log("[Labels] % on %.%", LABELS_MODEL, engine.provider, ifx engine.gpu_explanation then tprint(" %", engine.gpu_explanation) else ""); + if ok log("[Labels] % on %.%", LABELS_MODEL, labels_provider_name(engine), ifx engine.gpu_explanation then tprint(" %", engine.gpu_explanation) else ""); else log_error("[Labels] Could not start %: %", LABELS_MODEL, error); labels_ui.engine = engine; labels_ui.engine_tried = true; @@ -273,7 +273,7 @@ labels_window_contents :: () { using ImGui; TextWrapped("Run `jai fetch_models.jai` in modules/Labels, then rebuild."); } else if labels_ui.engine { in_use, explanation := labels_gpu_status(labels_ui.engine); - Text("%: images on %, text on the CPU", labels_ui.engine.model.title, ifx in_use then "the GPU (CUDA)" else "the CPU"); + Text("%: images on %, text on the CPU", labels_ui.engine.model.title, ifx in_use then tprint("the GPU (%)", labels_provider_name(labels_ui.engine)) else "the CPU"); if explanation text_with_links(explanation); } else if !labels_ui.engine_tried { TextColored(app_theme.palette.dim, "Not loaded: it loads when a capture is read or labelled, or a like: search runs, and unloads after % minutes unused.", MODELS_IDLE_UNLOAD_SECONDS / 60); diff --git a/src/app/ocr_ui.jai b/src/app/ocr_ui.jai index 7f1d4d8..fdc6091 100644 --- a/src/app/ocr_ui.jai +++ b/src/app/ocr_ui.jai @@ -446,7 +446,7 @@ ocr_read_and_store :: (items: [] Scan_Item, set: *Label_Set) -> [] string { vectors, label_error := labels_embed_images(labels, label_images,, temp); if label_error for label_of if !errors[it] errors[it] = sprint("labelling failed: %", label_error); dims := labels.dimensions; - label_engine := tprint("onnxruntime %, %", ort.version, ifx labels.provider == .CUDA then "CUDA" else "CPU"); + label_engine := ifx labels.provider == .COREML then "Core ML, GPU" else tprint("onnxruntime %, %", ort.version, labels_provider_name(labels)); for n: 0..items.count - 1 { if errors[n] continue; vector: [] float; diff --git a/src/test/labels_tests.jai b/src/test/labels_tests.jai index 07a3f69..ee5dcc4 100644 --- a/src/test/labels_tests.jai +++ b/src/test/labels_tests.jai @@ -83,7 +83,34 @@ labels_tests :: () { expect(created, "the label engine starts: %", error); if !created return; defer labels_destroy(engine); - emit(" images on %; % dimensions\n", ifx engine.provider == .CUDA then "CUDA" else "the CPU", engine.dimensions); + emit(" images on %; % dimensions\n", labels_provider_name(engine), engine.dimensions); + + #if OS == .MACOS if engine.provider == .COREML { + // Core ML's vectors against onnxruntime's on the CPU, for an image of bands and a gradient: the + // float16 model on the GPU is within cosine 0.99999 of the CPU's on real images. + group("core ml"); + cpu, cpu_created, cpu_error := labels_create(folder, LABELS_MODEL, prefer_gpu = false); + expect(cpu_created && cpu.provider == .CPU, "a CPU engine to compare with starts: %", cpu_error); + if cpu_created { + defer labels_destroy(cpu); + W :: 320; H :: 200; + pixels := NewArray(W * H * 4, u8,, temp); + for y: 0..H - 1 for x: 0..W - 1 { + p := pixels.data + (y * W + x) * 4; + p[0] = cast,trunc(u8) (x * 255 / W); + p[1] = cast(u8) (ifx (y / 25) % 2 then 220 else 40); + p[2] = cast,trunc(u8) (y * 255 / H); + p[3] = 255; + } + image := Labels_Image.{pixels.data, W, H, W * 4}; + on_gpu, gpu_error := labels_embed_images(engine, .[image],, temp); + on_cpu, cpu_run_error := labels_embed_images(cpu, .[image],, temp); + expect(!gpu_error && !cpu_run_error, "both make the image's vector (% / %)", gpu_error, cpu_run_error); + cosine: float; + for 0..engine.dimensions - 1 cosine += on_gpu[it] * on_cpu[it]; + expect(cosine > 0.999, "Core ML's vector is the CPU's: cosine %", cosine); + } + } group("tokenizer"); for KNOWN { diff --git a/tools/export_siglip_coreml.py b/tools/export_siglip_coreml.py new file mode 100644 index 0000000..c0bf75f --- /dev/null +++ b/tools/export_siglip_coreml.py @@ -0,0 +1,115 @@ +# Converts SigLIP 2's vision encoder to a Core ML model for modules/Labels on macOS, then checks it +# against the ONNX model mcap runs elsewhere, on a folder of images. Run once, on a Mac; mcap only uses +# the .mlpackage it writes (into modules/Labels/models, where the build copies it from). +# +# python3 -m venv ~/mcap-coreml +# ~/mcap-coreml/bin/pip install torch==2.7.0 transformers coremltools==9.0 onnxruntime pillow sentencepiece protobuf +# ~/mcap-coreml/bin/python tools/export_siglip_coreml.py --check FOLDER_OF_IMAGES +# +# torch 2.7.0 is the newest coremltools 9.0 was tested with. The model is the vision tower of +# google/siglip2-base-patch16-384, from pixel_values [1, 3, 384, 384] (RGB resized to 384 x 384, +# scaled to -1..1) to pooler_output [1, 768], as the ONNX one; an ML Program in float16, for the GPU. +# Measured on an M5 Pro: 6.7 ms an image on the GPU, against 67 ms for onnxruntime on the CPU, with +# vectors within cosine 0.99994 of the CPU's. The Neural Engine gets float16 badly wrong for this +# model (cosine 0.72), so mcap runs it on the CPU and GPU only. + +import argparse, glob, os, sys, time +import numpy as np +import torch +from PIL import Image +import coremltools as ct +from transformers import AutoModel +from transformers.utils import logging + +logging.set_verbosity_error() + +CHECKPOINT = "google/siglip2-base-patch16-384" +SIZE = 384 +HERE = os.path.dirname(os.path.abspath(__file__)) +MODELS = os.path.join(HERE, "..", "modules", "Labels", "models") +OUTPUT = os.path.join(MODELS, "siglip2-base-patch16-384-vision.mlpackage") +ONNX = os.path.join(MODELS, "siglip2-base-patch16-384-vision.onnx") + + +# coremltools 9.0 casts a folded one-element tensor with int(array), which NumPy 2 refuses ("only +# 0-dimensional arrays can be converted"); the same cast, with the element taken out first. +from coremltools.converters.mil.frontend.torch import ops as ct_ops +_original_cast = ct_ops._cast +def _cast_one_element(context, node, dtype, dtype_name): + x = ct_ops._get_inputs(context, node, expected=1)[0] + if x.can_be_folded_to_const() and np.ndim(x.val) > 0 and np.size(x.val) == 1: + context.add(ct_ops.mb.const(val=dtype(np.asarray(x.val).item()), name=node.name), node.name) + return + _original_cast(context, node, dtype, dtype_name) +ct_ops._cast = _cast_one_element + + +class Pooled(torch.nn.Module): + """The vision tower, from pixel values to the pooled vector, as the ONNX graph has it.""" + + def __init__(self, model): + super().__init__() + self.model = model + + def forward(self, pixel_values): + return self.model(pixel_values=pixel_values).pooler_output + + +def load_images(folder, limit): + paths = sorted(p for p in glob.glob(os.path.join(folder, "**", "*"), recursive=True) + if p.lower().endswith((".png", ".jpg", ".jpeg", ".webp")))[:limit] + images = [] + for p in paths: + image = Image.open(p).convert("RGB").resize((SIZE, SIZE), Image.BICUBIC) + a = np.asarray(image, dtype=np.float32) / 255.0 + images.append(((a - 0.5) / 0.5).transpose(2, 0, 1)) + return np.stack(images).astype(np.float32) + + +def worst_cosine(a, b): + a = a / np.linalg.norm(a, axis=-1, keepdims=True) + b = b / np.linalg.norm(b, axis=-1, keepdims=True) + return float(np.min(np.sum(a * b, axis=-1))) + + +def main(): + parser = argparse.ArgumentParser() + parser.add_argument("--check", metavar="FOLDER", help="images to check the Core ML model against the ONNX one on") + parser.add_argument("--limit", type=int, default=32) + args = parser.parse_args() + + model = Pooled(AutoModel.from_pretrained(CHECKPOINT, torch_dtype=torch.float32).vision_model.eval()).eval() + example = torch.zeros(1, 3, SIZE, SIZE) + with torch.no_grad(): + traced = torch.jit.trace(model, example) + start = time.perf_counter() + ml = ct.convert(traced, inputs=[ct.TensorType(name="pixel_values", shape=(1, 3, SIZE, SIZE))], + outputs=[ct.TensorType(name="pooler_output")], + convert_to="mlprogram", compute_precision=ct.precision.FLOAT16, + minimum_deployment_target=ct.target.macOS14) + ml.short_description = f"{CHECKPOINT} vision encoder (pooler_output), float16, for mcap" + ml.save(OUTPUT) + print(f"Wrote {OUTPUT} in {time.perf_counter() - start:.1f} s.") + + if not args.check: + return + import onnxruntime as ort + images = load_images(args.check, args.limit) + if not len(images): + sys.exit(f"No images in {args.check}.") + session = ort.InferenceSession(ONNX, providers=["CPUExecutionProvider"]) + name = session.get_inputs()[0].name + reference = np.concatenate([session.run(["pooler_output"], {name: images[i:i + 1]})[0] for i in range(len(images))]) + coreml = ct.models.MLModel(OUTPUT, compute_units=ct.ComputeUnit.CPU_AND_GPU) + coreml.predict({"pixel_values": images[:1]}) + start = time.perf_counter() + got = np.concatenate([coreml.predict({"pixel_values": images[i:i + 1]})["pooler_output"] for i in range(len(images))]) + per = (time.perf_counter() - start) / len(images) + cosine = worst_cosine(got.astype(np.float32), reference) + print(f"{len(images)} images: {per * 1000:.1f} ms an image on the GPU, worst cosine to the ONNX model on the CPU {cosine:.5f}.") + if cosine < 0.999: + sys.exit("The Core ML model is too far from the ONNX one.") + + +if __name__ == "__main__": + main()