Core ML Unet with torch ControlNet works

loaders.py

@@ -1,16 +1,16 @@
-import os
-
 import numpy as np
 import torch
 from coremltools import ComputeUnit
 from python_coreml_stable_diffusion.coreml_model import CoreMLModel
 
 import folder_paths
-from comfy import supported_models_base
+from comfy import supported_models_base, model_management
 from comfy.latent_formats import SD15
 from comfy.model_base import BaseModel
 from comfy.model_patcher import ModelPatcher
 from .logger import logger
+from .utils import expand_inputs, extract_residual_kwargs
+
 
 class CoreMLLoader:
     PACKAGE_DIRNAME = ""
@@ -19,7 +19,8 @@ class CoreMLLoader:
     def INPUT_TYPES(s):
         return {
             "required": {
-                "coreml_name": (list(s.coreml_filenames().keys()),)
+                "coreml_name": (list(s.coreml_filenames().keys()),),
+                "compute_unit": (list(ComputeUnit.__members__.keys()),)
             }
         }
 
@@ -36,9 +37,7 @@ class CoreMLLoader:
             if p.endswith((".mlpackage", ".mlmodelc"))
         }
 
-    def load(self, coreml_name):
-        compute_unit = ComputeUnit.ALL.name
-
+    def load(self, coreml_name, compute_unit):
         logger.info(f"Loading {coreml_name}")
 
         coreml_path = self.coreml_filenames()[coreml_name]
@@ -59,7 +58,7 @@ class CoreMLLoaderCkpt(CoreMLLoader):
     PACKAGE_DIRNAME = "checkpoints"
     RETURN_TYPES = ("MODEL", "CLIP", "VAE")
 
-    def load(self, coreml_name):
+    def load(self, coreml_name, compute_unit):
         # TODO: Implement this
         pass
 
@@ -68,7 +67,7 @@ class CoreMLLoaderTextEncoder(CoreMLLoader):
     PACKAGE_DIRNAME = "clip"
     RETURN_TYPES = ("CLIP",)
 
-    def load(self, coreml_name):
+    def load(self, coreml_name, compute_unit):
         # TODO: Implement this
         pass
 
@@ -77,17 +76,19 @@ class CoreMLLoaderUNet(CoreMLLoader):
     PACKAGE_DIRNAME = "unet"
     RETURN_TYPES = ("MODEL",)
 
-    def load(self, coreml_name):
-        coreml_model = super().load(coreml_name)[0]
+    def load(self, coreml_name, compute_unit):
+        coreml_model = super().load(coreml_name, compute_unit)[0]
 
-        return (ModelPatcher(coreml_model, None, None),)
+        return (
+            ModelPatcher(coreml_model, model_management.get_torch_device(),
+                         None),)
 
 
 class CoreMLLoaderVAE(CoreMLLoader):
     PACKAGE_DIRNAME = "vae"
     RETURN_TYPES = ("VAE",)
 
-    def load(self, coreml_name):
+    def load(self, coreml_name, compute_unit):
         # TODO: Implement this
         pass
 
@@ -101,26 +102,25 @@ class CoreMLModelWrapper(BaseModel):
 
     def apply_model(self, x, t, c_concat=None, c_crossattn=None, c_adm=None,
                     control=None, transformer_options={}):
-        # Some models use timestep, some timesteps
-        # Normally we could use positional arguments, but CoreMLModel
-        # requires kwargs
-        timesteps_key = [name for name in
-                         self.diffusion_model.expected_inputs.keys()
-                         if name.startswith("timestep")][0]
-
         sample = x.cpu().numpy().astype(np.float16)
+
         context = c_crossattn.cpu().numpy().astype(np.float16)
         context = context.transpose(0, 2, 1)[:, :, None, :]
+
         t = t.cpu().numpy().astype(np.float16)
 
         model_input_kwargs = {
             "sample": sample,
             "encoder_hidden_states": context,
-            timesteps_key: t,
+            "timestep": t,
         }
-        np_out = self.diffusion_model(**model_input_kwargs)["noise_pred"]
+        residual_kwargs = extract_residual_kwargs(self.diffusion_model,
+                                                  control)
+        model_input_kwargs |= residual_kwargs
+        model_input_kwargs = expand_inputs(model_input_kwargs)
 
-        return torch.from_numpy(np_out).float()
+        np_out = self.diffusion_model(**model_input_kwargs)["noise_pred"]
+        return torch.from_numpy(np_out).to(x.device)
 
     def get_dtype(self):
         # Hardcoding torch-compatible dtype (used for memory allocation)

utils.py

@@ -0,0 +1,55 @@
+from itertools import chain
+
+import numpy as np
+import torch
+
+from .logger import logger
+
+
+def expand_inputs(inputs):
+    expanded = inputs.copy()
+    for k, v in inputs.items():
+        if isinstance(v, np.ndarray):
+            expanded[k] = np.concatenate([v] * 2) if v.shape[0] == 1 else v
+        elif isinstance(v, torch.Tensor):
+            expanded[k] = torch.cat([v] * 2) if v.shape[0] == 1 else v
+        elif isinstance(v, list):
+            expanded[k] = v * 2 if len(v) == 1 else v
+        elif isinstance(v, dict):
+            expand_inputs(v)
+    return expanded
+
+
+def extract_residual_kwargs(model, control):
+    if ("additional_residual_0" not in model.expected_inputs.keys()):
+        return {}
+    if control is None:
+        return no_control(model)
+
+    residual_kwargs = {
+        "additional_residual_{}".format(i): r.cpu().numpy().astype(
+            np.float16)
+        for i, r in
+        enumerate(chain(control["output"], control["middle"]))
+    }
+    return residual_kwargs
+
+
+def no_control(model):
+    # Dirty hack to get the expected input shape when doing partial ControlNet
+    # 0.18215 is the latent scale factor (IDK, it kinda works)
+    # TODO: Find a better way to do this or tweak the values
+
+    logger.warning("No ControlNet input, despite the model supports it. "
+                   "Using random noise as ControlNet residuals. "
+                   "For better results, please use a ControlNet or a model "
+                   "that does not support ControlNet.")
+    residuals_names = [name for name in model.expected_inputs.keys()
+                       if name.startswith("additional_residual")]
+    residual_kwargs = {
+        "additional_residual_{}".format(i): 0.18215 * torch.randn(
+            *model.expected_inputs["additional_residual_{}".format(i)][
+                "shape"]).cpu().numpy().astype(dtype=np.float16)
+        for i in range(len(residuals_names))
+    }
+    return residual_kwargs