From 771e83bf33e7f0bfde56b04a7e045da1e8e6486f Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Adrian=20Szczepa=C5=84ski?= Date: Wed, 18 Oct 2023 15:39:09 +0200 Subject: [PATCH] Core ML Unet with torch ControlNet works --- loaders.py | 46 ++++++++++++++++++++++----------------------- utils.py | 55 ++++++++++++++++++++++++++++++++++++++++++++++++++++++ 2 files changed, 78 insertions(+), 23 deletions(-) create mode 100644 utils.py diff --git a/loaders.py b/loaders.py index 383008f..13f62fe 100644 --- a/loaders.py +++ b/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) diff --git a/utils.py b/utils.py new file mode 100644 index 0000000..c78e48b --- /dev/null +++ b/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