nodes/aszc-dev-ComfyUI-CoreMLSuite
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

feat(convert): consolidate into one auto-detecting converter node (#67)

Browse Source 
* feat(lcm): convert any full-distill LCM checkpoint

- COREML_CONVERT_LCM gains a ckpt_name input: any checkpoint from the
  checkpoints folder, with the canonical SimianLuo single file as the
  default auto-download entry, so workflows saved before this input
  existed keep the old behavior
- conversion routes through the unified
  coreml_diffusion.convert(model_version=LCM) path; the bespoke
  trace/convert pipeline in lcm/converter.py and the dead
  UNet2DConditionModelLCM wrapper are removed
- output naming via compose_out_name; existing cached LCM .mlpackages
  reconvert once due to the new _se attention suffix in the name
- LCM-LoRA merged checkpoints (plain SD1.5 architecture, no guidance
  embedding) are rejected by the package with a pointer to the standard
  converter node + LCM scheduler
- requires coreml-diffusion>=0.1.4 (generic LCM conversion fix)

Verified against a local ComfyUI checkout: the default entry resolves
the Hugging Face single file and cache-hits a previously converted
.mlpackage exposing timestep_cond; an LCM-LoRA merge raises the
explanatory ValueError.

* feat(convert): consolidate conversion into one auto-detecting node

The standard CoreMLConverter now auto-detects the model version from the
checkpoint (coreml-diffusion>=0.1.5, convert(model_version=None)), so:

- the model_version dropdown is gone — one node converts SD15 / SDXL / SDXL
  refiner / full-distill LCM, the version inferred from the UNet architecture
- the dedicated "Core ML LCM Converter" node, its single-model autodownload,
  and coreml_suite/lcm/nodes.py are removed; the converter UX was previously
  inconsistent (LCM only reachable through a separate autodownload-only node,
  while the standard converter did not list LCM at all)

lcm/utils.py (sampler-side timestep_cond patching) is unchanged — runtime LCM
support still keys off the converted UNet exposing timestep_cond.

diffusers is dropped from the dependencies (no longer imported directly after
the LCM converter removal). The e2e workflow fixture drops its now-invalid
model_version input.

* fix(deps): require coreml-diffusion>=0.1.5, keep requires-python <3.13

The auto-detect consolidation needs convert(model_version=None) from
coreml-diffusion 0.1.5. requires-python stays pinned to <3.13 to match the
library (coremltools-driven); relaxing it past the library's own cap makes the
dependency unresolvable for the 3.13+ range.
This commit is contained in:
aszc
2026-06-13 14:11:20 +02:00
committed by GitHub
parent a199e749bc
commit 8d28964831
11 changed files with 35 additions and 343 deletions
Download Patch File Download Diff File Expand all files Collapse all files
.gitignore
+1
View File
@@ -3,3 +3,4 @@ __pycache__/
models/
.venv/
test_results/
.claude/
__init__.py
-5
View File
@@ -11,9 +11,6 @@ from coreml_suite.nodes import (
CoreMLConverter,
COREML_LOAD_LORA,
)
from coreml_suite.lcm import (
COREML_CONVERT_LCM,
)
NODE_CLASS_MAPPINGS = {
"CoreMLUNetLoader": CoreMLLoaderUNet,
@@ -22,7 +19,6 @@ NODE_CLASS_MAPPINGS = {
"CoreMLModelAdapter": CoreMLModelAdapter,
"Core ML LoRA Loader": COREML_LOAD_LORA,
"Core ML Converter": CoreMLConverter,
"Core ML LCM Converter": COREML_CONVERT_LCM,
}
NODE_DISPLAY_NAME_MAPPINGS = {
"CoreMLUNetLoader": "Load Core ML UNet",
@@ -31,5 +27,4 @@ NODE_DISPLAY_NAME_MAPPINGS = {
"CoreMLModelAdapter": "Core ML Adapter (Experimental)",
"Core ML LoRA Loader": "Load LoRA to use with Core ML",
"Core ML Converter": "Convert Checkpoint to Core ML",
"Core ML LCM Converter": "Convert LCM to Core ML",
}
coreml_suite/lcm/__init__.py
+7 -2
View File
@@ -1,3 +1,8 @@
from .nodes import COREML_CONVERT_LCM
"""LCM runtime support (sampler-side).
__all__ = ["COREML_CONVERT_LCM"]
The dedicated LCM converter node was removed once the standard ``CoreMLConverter``
gained model-version auto-detection (full-distill LCM is detected from the
checkpoint). What remains here is runtime sampling support — ``utils`` patches the
model sampling and supplies the guidance embedding when a converted UNet exposes
``timestep_cond``.
"""
coreml_suite/lcm/converter.py
-144
View File
@@ -1,144 +0,0 @@
"""LCM-specific conversion orchestration (comfy-side).
E2 deduped the generic helpers (input building, Core ML export, residual-shape
calc) into ``coreml_diffusion.convert`` — this file now imports them instead of
carrying near-identical copies. What stays here is the genuinely LCM-specific
path: the hardcoded ``SimianLuo/LCM_Dreamshaper_v7`` download and the scheduler
that supplies the trace timestep. Consolidating that into the unified
``coreml_diffusion.convert(model_version=LCM, ...)`` path is a behavior change
deferred to E-LCM (it needs its own golden anchor).
``get_scheduler`` keeps using ``comfy.model_management`` because it runs on the
comfy side; the conversion package itself stays comfy-free.
"""
import gc
import logging
import os
import torch
from diffusers import UNet2DConditionModel, LCMScheduler
from diffusers.loaders import LoraLoaderMixin
from coreml_diffusion.conversion.attention import apply_attention_implementation
from coreml_diffusion.conversion.unet import CoreMLUNetWrapper
from coreml_diffusion.convert import (
add_cnet_support,
convert_to_coreml,
get_coreml_inputs,
get_encoder_hidden_states_shape,
get_inputs_spec,
get_sample_input,
lcm_inputs,
)
from coreml_diffusion import ModelVersion
logging.basicConfig()
logger = logging.getLogger(__name__)
logger.setLevel(logging.DEBUG)
MODEL_VERSION = "SimianLuo/LCM_Dreamshaper_v7"
MODEL_NAME = MODEL_VERSION.split("/")[-1] + "_4k"
def get_unets():
ref_unet = UNet2DConditionModel.from_pretrained(
MODEL_VERSION,
subfolder="unet",
device_map=None,
low_cpu_mem_usage=False,
)
cml_unet = CoreMLUNetWrapper(
apply_attention_implementation(ref_unet.eval(), "SPLIT_EINSUM"),
ModelVersion.LCM,
)
return cml_unet, ref_unet
def get_scheduler():
from comfy.model_management import get_torch_device
scheduler = LCMScheduler.from_pretrained(MODEL_VERSION, subfolder="scheduler")
scheduler.set_timesteps(50, get_torch_device(), 50)
return scheduler
def get_out_path(submodule_name, model_name):
from folder_paths import get_folder_paths
fname = f"{model_name}_{submodule_name}.mlpackage"
unet_path = get_folder_paths(submodule_name)[0]
out_path = os.path.join(unet_path, fname)
return out_path
def convert(
out_path: str,
batch_size: int = 1,
sample_size: tuple[int, int] = (64, 64),
controlnet_support: bool = False,
lora_paths: list[str] = None,
):
lora_paths = lora_paths or []
coreml_unet, ref_unet = get_unets()
for lora_path in lora_paths:
lora_sd, network_alphas = LoraLoaderMixin.lora_state_dict(lora_path)
LoraLoaderMixin.load_lora_into_unet(lora_sd, network_alphas, ref_unet)
ref_unet.fuse_lora()
sample_shape = (
batch_size, # B
ref_unet.config.in_channels, # C
sample_size[0], # H
sample_size[1], # W
)
encoder_hidden_states_shape = get_encoder_hidden_states_shape(ref_unet, batch_size)
scheduler = get_scheduler()
sample_inputs = get_sample_input(
batch_size, encoder_hidden_states_shape, sample_shape, scheduler=scheduler
)
sample_inputs |= lcm_inputs(sample_inputs)
if controlnet_support:
sample_inputs |= add_cnet_support(sample_shape, ref_unet)
sample_inputs_spec = get_inputs_spec(sample_inputs)
logger.info(f"Sample UNet inputs spec: {sample_inputs_spec}")
logger.info("JIT tracing..")
traced_unet = torch.jit.trace(
coreml_unet, example_inputs=list(sample_inputs.values())
)
logger.info("Done.")
coreml_sample_inputs = get_coreml_inputs(sample_inputs)
coreml_unet = convert_to_coreml(
"unet", traced_unet, coreml_sample_inputs, ["noise_pred"], out_path
)
del traced_unet
gc.collect()
coreml_unet.save(out_path)
logger.info(f"Saved unet into {out_path}")
if __name__ == "__main__":
h = 512
w = 512
sample_size = (h // 8, w // 8)
batch_size = 4
cn_support_str = "_cn" if True else ""
out_name = f"{MODEL_NAME}_{batch_size}x{w}x{h}{cn_support_str}"
out_path = get_out_path("unet", f"{out_name}")
if not os.path.exists(out_path):
convert(out_path=out_path, sample_size=sample_size, batch_size=batch_size)
coreml_suite/lcm/nodes.py
-70
View File
@@ -1,70 +0,0 @@
import os
from coremltools import ComputeUnit
from coreml_suite import COREML_NODE
from coreml_suite.coreml_model import CoreMLModel
class COREML_CONVERT_LCM(COREML_NODE):
"""Converts a LCM model to Core ML."""
@classmethod
def INPUT_TYPES(cls):
return {
"required": {
"height": ("INT", {"default": 512, "min": 512, "max": 768, "step": 8}),
"width": ("INT", {"default": 512, "min": 512, "max": 768, "step": 8}),
"batch_size": ("INT", {"default": 1, "min": 1, "max": 64}),
"compute_unit": (
[
ComputeUnit.CPU_AND_NE.name,
ComputeUnit.CPU_AND_GPU.name,
ComputeUnit.ALL.name,
ComputeUnit.CPU_ONLY.name,
],
),
"controlnet_support": ("BOOLEAN", {"default": False}),
}
}
RETURN_TYPES = ("COREML_UNET",)
RETURN_NAMES = ("coreml_model",)
FUNCTION = "convert"
def convert(self, height, width, batch_size, compute_unit, controlnet_support):
"""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:
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.
"""
from coreml_suite.lcm import converter as lcm_converter
h = height
w = width
sample_size = (h // 8, w // 8)
batch_size = batch_size
cn_support_str = "_cn" if controlnet_support else ""
out_name = f"{lcm_converter.MODEL_NAME}_{batch_size}x{w}x{h}{cn_support_str}"
out_path = lcm_converter.get_out_path("unet", f"{out_name}")
if not os.path.exists(out_path):
lcm_converter.convert(
out_path=out_path,
sample_size=sample_size,
batch_size=batch_size,
controlnet_support=controlnet_support,
)
return (CoreMLModel(out_path, compute_unit),)
coreml_suite/lcm/unet.py
-98
View File
@@ -1,98 +0,0 @@
from diffusers import UNet2DConditionModel
from diffusers.models.embeddings import TimestepEmbedding
class UNet2DConditionModelLCM(UNet2DConditionModel):
def __init__(
self,
time_cond_proj_dim=None,
**kwargs,
):
super().__init__(**kwargs)
timestep_input_dim = self.config.block_out_channels[0]
time_embed_dim = self.config.block_out_channels[0] * 4
time_embedding = TimestepEmbedding(
timestep_input_dim, time_embed_dim, cond_proj_dim=time_cond_proj_dim
)
self.time_embedding = time_embedding
def forward(
self,
sample,
timestep,
encoder_hidden_states,
timestep_cond,
*additional_residuals,
):
# 0. Project (or look-up) time embeddings
t_emb = self.time_proj(timestep)
emb = self.time_embedding(t_emb, timestep_cond)
# 1. center input if necessary
if self.config.center_input_sample:
sample = 2 * sample - 1.0
# 2. pre-process
sample = self.conv_in(sample)
# 3. down
down_block_res_samples = (sample,)
for downsample_block in self.down_blocks:
if (
hasattr(downsample_block, "attentions")
and downsample_block.attentions is not None
):
sample, res_samples = downsample_block(
hidden_states=sample,
temb=emb,
encoder_hidden_states=encoder_hidden_states,
)
else:
sample, res_samples = downsample_block(hidden_states=sample, temb=emb)
down_block_res_samples += res_samples
if additional_residuals:
new_down_block_res_samples = ()
for i, down_block_res_sample in enumerate(down_block_res_samples):
down_block_res_sample = down_block_res_sample + additional_residuals[i]
new_down_block_res_samples += (down_block_res_sample,)
down_block_res_samples = new_down_block_res_samples
# 4. mid
sample = self.mid_block(
sample, emb, encoder_hidden_states=encoder_hidden_states
)
if additional_residuals:
sample = sample + additional_residuals[-1]
# 5. up
for upsample_block in self.up_blocks:
res_samples = down_block_res_samples[-len(upsample_block.resnets) :]
down_block_res_samples = down_block_res_samples[
: -len(upsample_block.resnets)
]
if (
hasattr(upsample_block, "attentions")
and upsample_block.attentions is not None
):
sample = upsample_block(
hidden_states=sample,
temb=emb,
res_hidden_states_tuple=res_samples,
encoder_hidden_states=encoder_hidden_states,
)
else:
sample = upsample_block(
hidden_states=sample, temb=emb, res_hidden_states_tuple=res_samples
)
# 6. post-process
sample = self.conv_norm_out(sample)
sample = self.conv_act(sample)
sample = self.conv_out(sample)
return (sample,)
coreml_suite/nodes.py
+12 -10
View File
@@ -7,7 +7,6 @@ from coreml_suite import COREML_NODE
from coreml_suite.coreml_model import CoreMLModel
from coreml_suite.lcm.utils import add_lcm_model_options, lcm_patch, is_lcm
from coreml_suite.logger import logger
from coreml_diffusion import ModelVersion
from nodes import KSampler, LoraLoader, KSamplerAdvanced
from coreml_suite.models import (
@@ -226,14 +225,18 @@ class CoreMLModelAdapter(COREML_NODE):
class CoreMLConverter(COREML_NODE):
"""Converts a LCM model to Core ML."""
"""Converts a Stable Diffusion checkpoint (UNet) to Core ML.
The model version (SD15 / SDXL / SDXL refiner / LCM) is auto-detected from
the checkpoint's architecture, so there is no version dropdown — one node
converts every supported family, including full-distill LCM.
"""
@classmethod
def INPUT_TYPES(cls):
return {
"required": {
"ckpt_name": (folder_paths.get_filename_list("checkpoints"),),
"model_version": (_discover("list_model_versions", ["SD15", "SDXL"]),),
"height": ("INT", {"default": 512, "min": 8, "step": 8}),
"width": ("INT", {"default": 512, "min": 8, "step": 8}),
"batch_size": ("INT", {"default": 1, "min": 1, "max": 64}),
@@ -274,7 +277,6 @@ class CoreMLConverter(COREML_NODE):
def convert(
self,
ckpt_name,
model_version,
height,
width,
batch_size,
@@ -284,9 +286,11 @@ class CoreMLConverter(COREML_NODE):
quantize_nbits="none",
lora_params=None,
):
"""Converts a LCM model to Core ML.
"""Converts a checkpoint's UNet to Core ML.
Args:
ckpt_name (str): Checkpoint to convert; its model version is
auto-detected from the weights.
height (int): Height of the target image.
width (int): Width of the target image.
batch_size (int): Batch size.
@@ -296,10 +300,8 @@ class CoreMLConverter(COREML_NODE):
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.
can be loaded with the "Load Core ML UNet" node.
"""
model_version = ModelVersion[model_version]
lora_params = lora_params or {}
lora_params = [(k, v[0]) for k, v in lora_params.items()]
lora_params = sorted(lora_params, key=lambda lora: lora[0])
@@ -328,7 +330,7 @@ class CoreMLConverter(COREML_NODE):
logger.info(f"Attention implementation: {attention_implementation}")
if lora_params:
logger.info(f"LoRAs used:")
logger.info("LoRAs used:")
for lora_param in lora_params:
logger.info(f" {lora_param[0]} - strength: {lora_param[1]}")
@@ -345,7 +347,7 @@ class CoreMLConverter(COREML_NODE):
coreml_diffusion.convert(
ckpt_path,
model_version,
None, # model_version auto-detected from the checkpoint
unet_out_path,
sample_size=sample_size,
batch_size=batch_size,
pyproject.toml
+9 -5
View File
@@ -1,25 +1,29 @@
[build-system]
requires = ["hatchling"]
build-backend = "hatchling.build"
[project]
name = "comfyui-coremlsuite"
description = "This extension contains a set of custom nodes for ComfyUI that allow you to use Core ML models in your ComfyUI workflows."
version = "2.1.2"
license = "MIT"
requires-python = ">=3.12,<3.13"
packages = [{ include = "coreml_suite" }]
dependencies = [
# torch is provided by the host (ComfyUI) and intentionally left unpinned
# here: a hard torch cap would downgrade the host's torch and break its
# torchvision/torchaudio ABI. coreml-diffusion pulls torch>=2.7 transitively.
"coreml-diffusion>=0.1.1,<0.2",
# >=0.1.5: model-version auto-detection (convert(model_version=None)).
"coreml-diffusion>=0.1.5,<0.2",
"coremltools>=9,<10",
"numpy>=2,<3",
# diffusers is still imported directly by the comfy-side LCM converter
# (coreml_suite/lcm/converter.py) until E-LCM folds it into the package.
"diffusers>=0.30",
]
[project.urls]
Repository = "https://github.com/aszc-dev/ComfyUI-CoreMLSuite"
[tool.hatch.build.targets.wheel]
packages = ["coreml_suite"]
[tool.comfy]
PublisherId = "aszc-dev"
DisplayName = "ComfyUI-CoreMLSuite"
requirements.txt
+1 -1
View File
@@ -1,4 +1,4 @@
coreml-diffusion>=0.1.1,<0.2
coreml-diffusion>=0.1.4,<0.2
coremltools>=9,<10
numpy>=2,<3
diffusers>=0.30
tests/integration/workflows/e2e-1.5-basic-conversion.json
-1
View File
@@ -107,7 +107,6 @@
"10": {
"inputs": {
"ckpt_name": "dreamshaper_8.safetensors",
"model_version": "SD15",
"height": 512,
"width": 512,
"batch_size": 1,
uv.lock
Generated
+5 -7
View File
@@ -186,11 +186,10 @@ wheels = [
[[package]]
name = "comfyui-coremlsuite"
version = "2.1.2"
source = { virtual = "." }
source = { editable = "." }
dependencies = [
{ name = "coreml-diffusion" },
{ name = "coremltools" },
{ name = "diffusers" },
{ name = "numpy" },
]
@@ -218,9 +217,8 @@ dev = [
[package.metadata]
requires-dist = [
{ name = "coreml-diffusion", specifier = ">=0.1.1,<0.2" },
{ name = "coreml-diffusion", specifier = ">=0.1.5,<0.2" },
{ name = "coremltools", specifier = ">=9,<10" },
{ name = "diffusers", specifier = ">=0.30" },
{ name = "numpy", specifier = ">=2,<3" },
]
@@ -257,7 +255,7 @@ wheels = [
[[package]]
name = "coreml-diffusion"
version = "0.1.1"
version = "0.1.5"
source = { registry = "https://pypi.org/simple" }
dependencies = [
{ name = "coremltools" },
@@ -268,9 +266,9 @@ dependencies = [
{ name = "torch" },
{ name = "transformers" },
]
sdist = { url = "https://files.pythonhosted.org/packages/b9/d9/fbdfd6b87668c33711483447d1bf4669522daa32fab472dfc3477c24b60d/coreml_diffusion-0.1.1.tar.gz", hash = "sha256:8e1d5aee727c35a38b7693d17cc9711f818c4996dfaa2ade8340faf448097b77", size = 499670, upload-time = "2026-05-27T03:54:09.093Z" }
sdist = { url = "https://files.pythonhosted.org/packages/46/34/816457497ad7f039e79ced38b76703c908b7fe12cc5f755e8b88a46b443c/coreml_diffusion-0.1.5.tar.gz", hash = "sha256:4b6ca2182e1ee18d4d50ee526949887c7524b69855f8372621473a7016d1fd35", size = 930907, upload-time = "2026-06-13T12:06:21.398Z" }
wheels = [
{ url = "https://files.pythonhosted.org/packages/e3/17/094f310fc8f4ba144ca872f786a985a45abbe947adf43d881075edda5dab/coreml_diffusion-0.1.1-py3-none-any.whl", hash = "sha256:4bfeac8004d71825c145d302154aa87e7cc3d6fec78c101b221a211178a08f1f", size = 23882, upload-time = "2026-05-27T03:54:08.053Z" },
{ url = "https://files.pythonhosted.org/packages/aa/0d/56ced91b9f3e6135f377340a6fce347ea371b7816073421096a61a2e9e27/coreml_diffusion-0.1.5-py3-none-any.whl", hash = "sha256:6aff81a69a8d79d40d49fe1ab95cd964360ea025f297731f5ee6fcc22f46cdb1", size = 37368, upload-time = "2026-06-13T12:06:20.333Z" },
]
[[package]]