On macOS 15+, set fast prediction optimization hint for the unet.

Author: TobyRoseman

python_coreml_stable_diffusion/coreml_model.py

@@ -16,13 +16,28 @@
 
 import os
 import time
+import subprocess
+import sys
+
+
+def _macos_version():
+    """
+    Returns macOS version as a tuple of integers. On non-Macs, returns an empty tuple.
+    """
+    if sys.platform == "darwin":
+        try:
+            ver_str = subprocess.run(["sw_vers", "-productVersion"], stdout=subprocess.PIPE).stdout.decode('utf-8').strip('\n')
+            return tuple([int(v) for v in ver_str.split(".")])
+        except:
+            raise Exception("Unable to determine the macOS version")
+    return ()
 
 
 class CoreMLModel:
     """ Wrapper for running CoreML models using coremltools
     """
 
-    def __init__(self, model_path, compute_unit, sources='packages'):
+    def __init__(self, model_path, compute_unit, sources='packages', optimization_hints=None):
 
         logger.info(f"Loading {model_path}")
 
@@ -31,7 +46,10 @@ def __init__(self, model_path, compute_unit, sources='packages'):
             assert os.path.exists(model_path) and model_path.endswith(".mlpackage")
 
             self.model = ct.models.MLModel(
-                model_path, compute_units=ct.ComputeUnit[compute_unit])
+                model_path,
+                compute_units=ct.ComputeUnit[compute_unit],
+                optimization_hints=optimization_hints,
+            )
             DTYPE_MAP = {
                 65552: np.float16,
                 65568: np.float32,
@@ -47,7 +65,11 @@ def __init__(self, model_path, compute_unit, sources='packages'):
         elif sources == 'compiled':
             assert os.path.exists(model_path) and model_path.endswith(".mlmodelc")
 
-            self.model = ct.models.CompiledMLModel(model_path, ct.ComputeUnit[compute_unit])
+            self.model = ct.models.CompiledMLModel(
+                model_path,
+                compute_units=ct.ComputeUnit[compute_unit],
+                optimization_hints=optimization_hints,
+            )
 
             # Grab expected inputs from metadata.json
             with open(os.path.join(model_path, 'metadata.json'), 'r') as f:
@@ -170,7 +192,15 @@ def _load_mlpackage(submodule_name,
             raise FileNotFoundError(
                 f"{submodule_name} CoreML model doesn't exist at {mlpackage_path}")
 
-    return CoreMLModel(mlpackage_path, compute_unit, sources=sources)
+    # On macOS 15+, set fast prediction optimization hint for the unet.
+    optimization_hints = None
+    if submodule_name == "unet" and _macos_version() >= (15, 0):
+        optimization_hints = {"specializationStrategy": ct.SpecializationStrategy.FastPrediction}
+
+    return CoreMLModel(mlpackage_path,
+                       compute_unit,
+                       sources=sources,
+                       optimization_hints=optimization_hints)
 
 
 def _load_mlpackage_controlnet(mlpackages_dir, model_version, compute_unit):