aszc
8d28964831
* 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.
391 lines
12 KiB
Python
391 lines
12 KiB
Python
import os
|
|
|
|
from coremltools import ComputeUnit
|
|
|
|
import folder_paths
|
|
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 nodes import KSampler, LoraLoader, KSamplerAdvanced
|
|
|
|
from coreml_suite.models import (
|
|
add_sdxl_model_options,
|
|
is_sdxl,
|
|
get_model_patcher,
|
|
get_latent_image,
|
|
)
|
|
|
|
|
|
def _discover(fn_name, fallback):
|
|
"""Populate a converter dropdown from coreml_diffusion's discovery API.
|
|
|
|
Fails soft: if the package is missing, too old to expose ``fn_name``, or
|
|
errors, the node still registers with the fallback list instead of vanishing
|
|
from the menu. Evaluated on every INPUT_TYPES call, so installing a newer
|
|
coreml_diffusion surfaces new conversion types with no Suite change.
|
|
"""
|
|
try:
|
|
import coreml_diffusion
|
|
|
|
return getattr(coreml_diffusion, fn_name)()
|
|
except Exception as exc: # missing/old package, import error, etc.
|
|
logger.warning(
|
|
f"coreml_diffusion.{fn_name} unavailable ({exc}); "
|
|
f"using fallback {fallback}"
|
|
)
|
|
return fallback
|
|
|
|
|
|
class CoreMLSampler(COREML_NODE, KSampler):
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
old_required = KSampler.INPUT_TYPES()["required"].copy()
|
|
old_required.pop("model")
|
|
old_required.pop("negative")
|
|
old_required.pop("latent_image")
|
|
new_required = {"coreml_model": ("COREML_UNET",)}
|
|
return {
|
|
"required": new_required | old_required,
|
|
"optional": {"negative": ("CONDITIONING",), "latent_image": ("LATENT",)},
|
|
}
|
|
|
|
def sample(
|
|
self,
|
|
coreml_model,
|
|
seed,
|
|
steps,
|
|
cfg,
|
|
sampler_name,
|
|
scheduler,
|
|
positive,
|
|
negative=None,
|
|
latent_image=None,
|
|
denoise=1.0,
|
|
):
|
|
model_patcher = get_model_patcher(coreml_model)
|
|
latent_image = get_latent_image(coreml_model, latent_image)
|
|
|
|
if is_lcm(coreml_model):
|
|
negative = [[None, {}]]
|
|
positive[0][1]["control_apply_to_uncond"] = False
|
|
model_patcher = add_lcm_model_options(model_patcher, cfg, latent_image)
|
|
model_patcher = lcm_patch(model_patcher)
|
|
else:
|
|
assert (
|
|
negative is not None
|
|
), "Negative conditioning is optional only for LCM models."
|
|
|
|
if is_sdxl(coreml_model):
|
|
model_patcher = add_sdxl_model_options(model_patcher, positive, negative)
|
|
|
|
return super().sample(
|
|
model_patcher,
|
|
seed,
|
|
steps,
|
|
cfg,
|
|
sampler_name,
|
|
scheduler,
|
|
positive,
|
|
negative,
|
|
latent_image,
|
|
denoise,
|
|
)
|
|
|
|
|
|
class CoreMLSamplerAdvanced(COREML_NODE, KSamplerAdvanced):
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
old_required = KSamplerAdvanced.INPUT_TYPES()["required"].copy()
|
|
old_required.pop("model")
|
|
old_required.pop("negative")
|
|
old_required.pop("latent_image")
|
|
new_required = {"coreml_model": ("COREML_UNET",)}
|
|
return {
|
|
"required": new_required | old_required,
|
|
"optional": {"negative": ("CONDITIONING",), "latent_image": ("LATENT",)},
|
|
}
|
|
|
|
def sample(
|
|
self,
|
|
coreml_model,
|
|
add_noise,
|
|
noise_seed,
|
|
steps,
|
|
cfg,
|
|
sampler_name,
|
|
scheduler,
|
|
positive,
|
|
start_at_step,
|
|
end_at_step,
|
|
return_with_leftover_noise,
|
|
negative=None,
|
|
latent_image=None,
|
|
denoise=1.0,
|
|
):
|
|
model_patcher = get_model_patcher(coreml_model)
|
|
latent_image = get_latent_image(coreml_model, latent_image)
|
|
|
|
if is_lcm(coreml_model):
|
|
negative = [[None, {}]]
|
|
positive[0][1]["control_apply_to_uncond"] = False
|
|
model_patcher = add_lcm_model_options(model_patcher, cfg, latent_image)
|
|
model_patcher = lcm_patch(model_patcher)
|
|
else:
|
|
assert (
|
|
negative is not None
|
|
), "Negative conditioning is optional only for LCM models."
|
|
|
|
if is_sdxl(coreml_model):
|
|
model_patcher = add_sdxl_model_options(model_patcher, positive, negative)
|
|
|
|
return super().sample(
|
|
model_patcher,
|
|
add_noise,
|
|
noise_seed,
|
|
steps,
|
|
cfg,
|
|
sampler_name,
|
|
scheduler,
|
|
positive,
|
|
negative,
|
|
latent_image,
|
|
start_at_step,
|
|
end_at_step,
|
|
return_with_leftover_noise,
|
|
denoise,
|
|
)
|
|
|
|
|
|
class CoreMLLoader(COREML_NODE):
|
|
PACKAGE_DIRNAME = ""
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return {
|
|
"required": {
|
|
"coreml_name": (list(s.coreml_filenames().keys()),),
|
|
"compute_unit": (
|
|
[
|
|
ComputeUnit.CPU_AND_NE.name,
|
|
ComputeUnit.CPU_AND_GPU.name,
|
|
ComputeUnit.ALL.name,
|
|
ComputeUnit.CPU_ONLY.name,
|
|
],
|
|
),
|
|
}
|
|
}
|
|
|
|
FUNCTION = "load"
|
|
|
|
@classmethod
|
|
def coreml_filenames(cls):
|
|
extensions = (".mlpackage",)
|
|
all_paths = folder_paths.get_filename_list_(cls.PACKAGE_DIRNAME)[1]
|
|
coreml_paths = folder_paths.filter_files_extensions(all_paths, extensions)
|
|
|
|
return {os.path.split(p)[-1]: p for p in coreml_paths}
|
|
|
|
def load(self, coreml_name, compute_unit):
|
|
logger.info(f"Loading {coreml_name} to {compute_unit}")
|
|
|
|
coreml_path = self.coreml_filenames()[coreml_name]
|
|
|
|
return (CoreMLModel(coreml_path, compute_unit),)
|
|
|
|
|
|
class CoreMLLoaderUNet(CoreMLLoader):
|
|
PACKAGE_DIRNAME = "unet"
|
|
RETURN_TYPES = ("COREML_UNET",)
|
|
RETURN_NAMES = ("coreml_model",)
|
|
|
|
|
|
class CoreMLModelAdapter(COREML_NODE):
|
|
"""
|
|
Adapter Node to use CoreML models as Comfy models. This is an experimental
|
|
feature and may not work as expected.
|
|
"""
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return {
|
|
"required": {
|
|
"coreml_model": ("COREML_UNET",),
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES = ("MODEL",)
|
|
|
|
FUNCTION = "wrap"
|
|
CATEGORY = "Core ML Suite"
|
|
|
|
def wrap(self, coreml_model):
|
|
model_patcher = get_model_patcher(coreml_model)
|
|
return (model_patcher,)
|
|
|
|
|
|
class CoreMLConverter(COREML_NODE):
|
|
"""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"),),
|
|
"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}),
|
|
"attention_implementation": (
|
|
_discover(
|
|
"list_attention_impls",
|
|
["SPLIT_EINSUM", "SPLIT_EINSUM_V2", "ORIGINAL"],
|
|
),
|
|
),
|
|
"compute_unit": (
|
|
[
|
|
ComputeUnit.CPU_AND_NE.name,
|
|
ComputeUnit.CPU_AND_GPU.name,
|
|
ComputeUnit.ALL.name,
|
|
ComputeUnit.CPU_ONLY.name,
|
|
],
|
|
),
|
|
"controlnet_support": ("BOOLEAN", {"default": False}),
|
|
},
|
|
"optional": {
|
|
# k-means weight palettization. Kept optional so workflows
|
|
# that omit it still validate — ComfyUI rejects a prompt that
|
|
# omits any `required` input. When omitted it defaults to
|
|
# "none", identical to unquantized behavior and filename, so
|
|
# existing cached .mlpackages still resolve.
|
|
"quantize_nbits": (
|
|
_discover("list_quant_modes", ["none", "8", "6", "4"]),
|
|
{"default": "none"},
|
|
),
|
|
"lora_params": ("LORA_PARAMS",),
|
|
},
|
|
}
|
|
|
|
RETURN_TYPES = ("COREML_UNET",)
|
|
RETURN_NAMES = ("coreml_model",)
|
|
FUNCTION = "convert"
|
|
|
|
def convert(
|
|
self,
|
|
ckpt_name,
|
|
height,
|
|
width,
|
|
batch_size,
|
|
attention_implementation,
|
|
compute_unit,
|
|
controlnet_support,
|
|
quantize_nbits="none",
|
|
lora_params=None,
|
|
):
|
|
"""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.
|
|
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 "Load Core ML UNet" node.
|
|
"""
|
|
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])
|
|
lora_weights = [(self.lora_path(lora[0]), lora[1]) for lora in lora_params]
|
|
|
|
h = height
|
|
w = width
|
|
sample_size = (h // 8, w // 8)
|
|
import coreml_diffusion
|
|
|
|
out_name = coreml_diffusion.compose_out_name(
|
|
ckpt_name=ckpt_name,
|
|
batch_size=batch_size,
|
|
width=w,
|
|
height=h,
|
|
controlnet_support=controlnet_support,
|
|
attention_implementation=attention_implementation,
|
|
lora_names=coreml_diffusion.lora_names_from_params(lora_params),
|
|
quantize_nbits=quantize_nbits,
|
|
)
|
|
|
|
logger.info(f"Converting {ckpt_name} to {out_name}")
|
|
logger.info(f"Batch size: {batch_size}")
|
|
logger.info(f"Width: {w}, Height: {h}")
|
|
logger.info(f"ControlNet support: {controlnet_support}")
|
|
logger.info(f"Attention implementation: {attention_implementation}")
|
|
|
|
if lora_params:
|
|
logger.info("LoRAs used:")
|
|
for lora_param in lora_params:
|
|
logger.info(f" {lora_param[0]} - strength: {lora_param[1]}")
|
|
|
|
# Resolve the ComfyUI models/unet path here (a node concern); the package
|
|
# takes the output path as an injected argument.
|
|
unet_path = folder_paths.get_folder_paths("unet")[0]
|
|
unet_out_path = os.path.join(unet_path, f"{out_name}_unet.mlpackage")
|
|
ckpt_path = folder_paths.get_full_path("checkpoints", ckpt_name)
|
|
|
|
config_filename = ckpt_name.split(".")[0] + ".yaml"
|
|
config_path = folder_paths.get_full_path("configs", config_filename)
|
|
if config_path:
|
|
logger.info(f"Using config file {config_path}")
|
|
|
|
coreml_diffusion.convert(
|
|
ckpt_path,
|
|
None, # model_version auto-detected from the checkpoint
|
|
unet_out_path,
|
|
sample_size=sample_size,
|
|
batch_size=batch_size,
|
|
controlnet_support=controlnet_support,
|
|
lora_weights=lora_weights,
|
|
attn_impl=attention_implementation,
|
|
config_path=config_path,
|
|
quantize_nbits=quantize_nbits,
|
|
)
|
|
return (CoreMLModel(unet_out_path, compute_unit),)
|
|
|
|
@staticmethod
|
|
def lora_path(lora_name):
|
|
return folder_paths.get_full_path("loras", lora_name)
|
|
|
|
|
|
class COREML_LOAD_LORA(COREML_NODE, LoraLoader):
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
required = LoraLoader.INPUT_TYPES()["required"].copy()
|
|
required.pop("model")
|
|
return {
|
|
"required": required,
|
|
"optional": {"lora_params": ("LORA_PARAMS",)},
|
|
}
|
|
|
|
RETURN_TYPES = ("CLIP", "LORA_PARAMS")
|
|
RETURN_NAMES = ("CLIP", "lora_params")
|
|
|
|
def load_lora(
|
|
self, clip, lora_name, strength_model, strength_clip, lora_params=None
|
|
):
|
|
_, lora_clip = super().load_lora(
|
|
None, clip, lora_name, strength_model, strength_clip
|
|
)
|
|
|
|
lora_params = lora_params or {}
|
|
lora_params[lora_name] = (strength_model, strength_clip)
|
|
|
|
return lora_clip, lora_params
|