Merge pull request #357 from TobyRoseman/dev

Fix Cartoonish Output for non-XL model

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):

python_coreml_stable_diffusion/pipeline.py

@@ -604,8 +604,7 @@ def get_coreml_pipe(pytorch_pipe,
             "tokenizer": pytorch_pipe.tokenizer,
             'tokenizer_2': pytorch_pipe.tokenizer_2,
             "scheduler": pytorch_pipe.scheduler if scheduler_override is None else scheduler_override,
-            "force_zeros_for_empty_prompt": force_zeros_for_empty_prompt,
-            'xl': True
+            'xl': True,
         }
 
         model_packages_to_load = ["text_encoder", "text_encoder_2", "unet", "vae_decoder"]
@@ -618,6 +617,8 @@ def get_coreml_pipe(pytorch_pipe,
         }
         model_packages_to_load = ["text_encoder", "unet", "vae_decoder"]
 
+    coreml_pipe_kwargs["force_zeros_for_empty_prompt"] = force_zeros_for_empty_prompt
+
     if getattr(pytorch_pipe, "safety_checker", None) is not None:
         model_packages_to_load.append("safety_checker")
     else:
@@ -713,7 +714,7 @@ def main(args):
 
     # Get Force Zeros Config if it exists
     force_zeros_for_empty_prompt: bool = False
-    if 'force_zeros_for_empty_prompt' in pytorch_pipe.config:
+    if 'xl' in args.model_version and 'force_zeros_for_empty_prompt' in pytorch_pipe.config:
         force_zeros_for_empty_prompt = pytorch_pipe.config['force_zeros_for_empty_prompt']
 
     coreml_pipe = get_coreml_pipe(

requirements.txt

@@ -1,8 +1,8 @@
-coremltools>=7.0
-diffusers[torch]
-diffusionkit
+coremltools>=8.0
+diffusers[torch]==0.30.2
+diffusionkit==0.4.0
 torch
-transformers==4.29.2
+transformers==4.44.2
 scipy
 scikit-learn
 pytest

setup.py

@@ -14,19 +14,19 @@
     long_description_content_type='text/markdown',
     author='Apple Inc.',
     install_requires=[
-        "coremltools>=7.0b2",
-        "diffusers[torch]",
+        "coremltools>=8.0",
+        "diffusers[torch]==0.30.2",
         "torch",
-        "transformers>=4.30.0",
-        "huggingface-hub",
+        "transformers==4.44.2",
+        "huggingface-hub==0.24.6",
         "scipy",
         "numpy<1.24",
         "pytest",
         "scikit-learn",
         "invisible-watermark",
         "safetensors",
         "matplotlib",
-        "diffusionkit",
+        "diffusionkit==0.4.0",
     ],
     packages=find_packages(),
     classifiers=[