LCM Converter works

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"]

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

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

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"),)

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)