diff --git a/coreml_suite/lcm/__init__.py b/coreml_suite/lcm/__init__.py index 7ac4c9c..8ae122a 100644 --- a/coreml_suite/lcm/__init__.py +++ b/coreml_suite/lcm/__init__.py @@ -1,4 +1,4 @@ -from lcm_sampler import CoreMLSamplerLCM +from .lcm_sampler import CoreMLSamplerLCM from .nodes import CoreMLConverterLCM __all__ = ["CoreMLSamplerLCM", "CoreMLConverterLCM"] diff --git a/coreml_suite/lcm/lcm_converter.py b/coreml_suite/lcm/lcm_converter.py index 54c5446..6356586 100644 --- a/coreml_suite/lcm/lcm_converter.py +++ b/coreml_suite/lcm/lcm_converter.py @@ -25,7 +25,7 @@ MODEL_NAME = MODEL_VERSION.split("/")[-1] + "_4k" import python_coreml_stable_diffusion.unet as unet -unet.ATTENTION_IMPLEMENTATION_IN_EFFECT = unet.AttentionImplementations.SPLIT_EINSUM_V2 +unet.ATTENTION_IMPLEMENTATION_IN_EFFECT = unet.AttentionImplementations.SPLIT_EINSUM def get_unets(): diff --git a/coreml_suite/lcm/lcm_sampler.py b/coreml_suite/lcm/lcm_sampler.py index 12c9d73..1c9c8ca 100644 --- a/coreml_suite/lcm/lcm_sampler.py +++ b/coreml_suite/lcm/lcm_sampler.py @@ -6,6 +6,7 @@ import torch from comfy.model_management import get_torch_device from coreml_suite.lcm.lcm_pipeline import LatentConsistencyModelPipeline from coreml_suite.lcm.lcm_scheduler import LCMScheduler +from coreml_suite.models import get_model_config, CoreMLModelWrapperLCM class CoreMLSamplerLCM: @@ -19,7 +20,7 @@ class CoreMLSamplerLCM: def INPUT_TYPES(s): return { "required": { - "model": ("MODEL",), + "coreml_model": ("COREML_UNET",), "seed": ("INT", {"default": 0, "min": 0, "max": 0xFFFFFFFFFFFFFFFF}), "steps": ("INT", {"default": 4, "min": 1, "max": 10000}), "cfg": ( @@ -46,7 +47,7 @@ class CoreMLSamplerLCM: def sample( self, - model, + coreml_model, seed, steps, cfg, @@ -56,6 +57,10 @@ class CoreMLSamplerLCM: num_images, use_fp16, ): + + model_config = get_model_config() + wrapped_model = CoreMLModelWrapperLCM(model_config, coreml_model) + if self.pipe is None: self.pipe = LatentConsistencyModelPipeline.from_pretrained( pretrained_model_name_or_path="SimianLuo/LCM_Dreamshaper_v7", @@ -68,7 +73,7 @@ class CoreMLSamplerLCM: else: self.pipe.to(torch_device=get_torch_device(), torch_dtype=torch.float32) - coreml_unet = model.model + coreml_unet = wrapped_model coreml_unet.config = self.pipe.unet.config self.pipe.unet = coreml_unet diff --git a/coreml_suite/lcm/nodes.py b/coreml_suite/lcm/nodes.py index 980e1a0..9028668 100644 --- a/coreml_suite/lcm/nodes.py +++ b/coreml_suite/lcm/nodes.py @@ -1,15 +1,14 @@ import os +from coremltools import ComputeUnit +from python_coreml_stable_diffusion.coreml_model import CoreMLModel + from coreml_suite.lcm import lcm_converter class CoreMLConverterLCM: """Converts a LCM model to Core ML.""" - RETURN_TYPES = ("COMBO",) - RETURN_NAMES = ("model_name",) - FUNCTION = "convert" - @classmethod def INPUT_TYPES(cls): return { @@ -17,19 +16,30 @@ class CoreMLConverterLCM: "height": ("INT", {"default": 512, "min": 512, "max": 768, "step": 8}), "width": ("INT", {"default": 512, "min": 512, "max": 768, "step": 8}), "batch_size": ("INT", {"default": 4, "min": 1, "max": 64}), + "compute_unit": ([ + ComputeUnit.CPU_AND_NE.name, + ComputeUnit.CPU_AND_GPU.name, + ComputeUnit.ALL.name, + ComputeUnit.CPU_ONLY.name, + ],) } } - def convert(self, height, width, batch_size): + RETURN_TYPES = ("COREML_UNET",) + RETURN_NAMES = ("coreml_model",) + FUNCTION = "convert" + + def convert(self, height, width, batch_size, compute_unit): """Converts a LCM model to Core ML. Args: height (int): Height of the target image. width (int): Width of the target image. batch_size (int): Batch size. + compute_unit (str): Compute unit to use when loading the model. Returns: - MODEL: The converted Core ML model. + coreml_model: The converted Core ML model. The converted model is also saved to "models/unet" directory and can be loaded with the "LCMCoreMLLoaderUNet" node. @@ -49,4 +59,4 @@ class CoreMLConverterLCM: ) target_path = lcm_converter.compile_model(out_path=out_path, out_name=out_name) - return (target_path.split("/")[-1],) + return (CoreMLModel(target_path, compute_unit, "compiled"),) diff --git a/coreml_suite/models.py b/coreml_suite/models.py index a85fdc2..24f455c 100644 --- a/coreml_suite/models.py +++ b/coreml_suite/models.py @@ -51,7 +51,7 @@ class CoreMLModelWrapper(BaseModel): return merged_out def _apply_model(self, x, t, c_concat=None, c_crossattn=None, c_adm=None, - control=None, transformer_options={}): + control=None, transformer_options={}): model_input_kwargs = self.prepare_inputs(x, t, c_crossattn, control) np_out = self.diffusion_model(**model_input_kwargs)["noise_pred"] @@ -99,17 +99,18 @@ class CoreMLModelWrapper(BaseModel): def expected_inputs(self): return self.diffusion_model.expected_inputs + class CoreMLModelWrapperLCM(CoreMLModelWrapper): def __init__(self, model_config, coreml_model): super().__init__(model_config, coreml_model) self.config = None def _apply_model(self, x, t, c_concat=None, c_crossattn=None, c_adm=None, - control=None, transformer_options={}): + control=None, transformer_options={}): model_input_kwargs = self.prepare_inputs(x, t, c_crossattn, control) np_out = self.diffusion_model(**model_input_kwargs)["noise_pred"] return (torch.from_numpy(np_out).to(x.device),) - def __call__(self, latents, t, encoder_hidden_states, **kwargs): - return self.apply_model(latents, t, c_crossattn=encoder_hidden_states, **kwargs) + def __call__(self, latents, ts, encoder_hidden_states, **kwargs): + return self.apply_model(latents, ts, c_crossattn=encoder_hidden_states)