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

feat: modernize toolchain, refactor core, add tiered CI and opt-in quantization

Browse Source 
Modernizes ComfyUI-CoreMLSuite onto Python 3.12 / torch 2.7 / coremltools 9
with a characterization-test safety net. The default conversion path is
unchanged; existing saved workflows produce identical output.

- Toolchain bump (Python 3.12, torch 2.7, coremltools 9, numpy <2) with the
  blocking upstream pins overridden.
- Framework-free logic moved into coreml_suite/core/ (no comfy/coremltools
  imports); old module paths re-export from there.
- Opt-in quantize_nbits dropdown (none|8|6|4) for k-means weight
  palettization; default none is byte-for-byte identical to before.
- Tiered CI: Tier 0 (Linux unit), Tier 1 (macOS-ARM smoke), Tier 2
  (self-hosted Apple Silicon golden-image check on the ANE).
This commit is contained in:
aszc
2026-05-25 19:11:49 +02:00
committed by GitHub
parent 7678a07ed5
commit 02b6e8ece3
37 changed files with 3819 additions and 416 deletions
Download Patch File Download Diff File Expand all files Collapse all files
coreml_suite/controlnet.py
+13 -61
View File
@@ -1,62 +1,14 @@
from itertools import chain
from math import ceil
"""Compatibility shim — re-exports from coreml_suite.core.controlnet."""
from coreml_suite.core.controlnet import (
chunk_control,
expand_inputs,
extract_residual_kwargs,
no_control,
)
import numpy as np
import torch
from coreml_suite.latents import chunk_batch
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(expected_inputs, control):
if "additional_residual_0" not in expected_inputs.keys():
return {}
if control is None:
return no_control(expected_inputs)
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(expected_inputs):
shapes_dict = {
k: v["shape"] for k, v in expected_inputs.items() if k.startswith("additional")
}
residual_kwargs = {
k: torch.zeros(*shape).cpu().numpy().astype(dtype=np.float16)
for k, shape in shapes_dict.items()
}
return residual_kwargs
def chunk_control(cn, target_size):
if cn is None:
return [None] * target_size
num_chunks = ceil(cn["output"][0].shape[0] / target_size)
out = [{"output": [], "middle": []} for _ in range(num_chunks)]
for k, v in cn.items():
for i, x in enumerate(v):
chunks = chunk_batch(x, (target_size, *x.shape[1:]))
for j, chunk in enumerate(chunks):
out[j][k].append(chunk)
return out
__all__ = [
"chunk_control",
"expand_inputs",
"extract_residual_kwargs",
"no_control",
]
coreml_suite/converter.py
+21
View File
@@ -258,6 +258,7 @@ def convert_unet(
batch_size: int = 1,
sample_size: tuple[int, int] = (64, 64),
controlnet_support: bool = False,
quantize_nbits: str = "none",
):
coreml_unet = get_unet(model_version, ref_pipe)
ref_unet = ref_pipe.unet
@@ -305,6 +306,24 @@ def convert_unet(
del traced_unet
gc.collect()
if quantize_nbits != "none":
# Opt-in k-means weight palettization. The default path
# (quantize_nbits="none") leaves the traced UNet untouched.
from coremltools.optimize.coreml import (
OpPalettizerConfig,
OptimizationConfig,
palettize_weights,
)
nbits = int(quantize_nbits)
logger.info(f"Palettizing UNet weights to {nbits}-bit (kmeans)..")
t0 = time.time()
cfg = OptimizationConfig(
global_config=OpPalettizerConfig(mode="kmeans", nbits=nbits)
)
coreml_unet = palettize_weights(coreml_unet, config=cfg)
logger.info(f"Palettization took {time.time() - t0:.1f}s")
coreml_unet.save(unet_out_path)
logger.info(f"Saved unet into {unet_out_path}")
@@ -319,6 +338,7 @@ def convert(
lora_weights: list[tuple[Union[str, os.PathLike], float]] = None,
attn_impl: str = AttentionImplementations.SPLIT_EINSUM.name,
config_path: str = None,
quantize_nbits: str = "none",
):
if os.path.exists(unet_out_path):
logger.info(f"Found existing model at {unet_out_path}! Skipping..")
@@ -344,6 +364,7 @@ def convert(
batch_size,
sample_size,
controlnet_support,
quantize_nbits=quantize_nbits,
)
coreml_suite/core/__init__.py
+10
View File
@@ -0,0 +1,10 @@
"""Framework-free pure-logic core of ComfyUI-CoreMLSuite.
Modules under this package must NOT import `comfy`, `coremltools`,
`python_coreml_stable_diffusion`, `folder_paths`, `nodes`, or any other
ComfyUI / Apple runtime. Only `numpy` and `torch` are allowed.
The thin adapters in `coreml_suite.{latents,controlnet,models}` keep the
old public import paths working so `coreml_suite/nodes.py` and downstream
ComfyUI workflows are unchanged.
"""
coreml_suite/core/controlnet.py
+67
View File
@@ -0,0 +1,67 @@
"""Pure helpers around the ControlNet residual inputs of the Core ML UNet.
Re-exported by coreml_suite.controlnet. Characterization tests cover
shapes, dtype (fp16), and zero-fill fallback.
"""
from itertools import chain
from math import ceil
import numpy as np
import torch
from coreml_suite.core.latents import chunk_batch
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(expected_inputs, control):
if "additional_residual_0" not in expected_inputs.keys():
return {}
if control is None:
return no_control(expected_inputs)
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(expected_inputs):
shapes_dict = {
k: v["shape"] for k, v in expected_inputs.items() if k.startswith("additional")
}
residual_kwargs = {
k: torch.zeros(*shape).cpu().numpy().astype(dtype=np.float16)
for k, shape in shapes_dict.items()
}
return residual_kwargs
def chunk_control(cn, target_size):
if cn is None:
return [None] * target_size
num_chunks = ceil(cn["output"][0].shape[0] / target_size)
out = [{"output": [], "middle": []} for _ in range(num_chunks)]
for k, v in cn.items():
for i, x in enumerate(v):
chunks = chunk_batch(x, (target_size, *x.shape[1:]))
for j, chunk in enumerate(chunks):
out[j][k].append(chunk)
return out
coreml_suite/core/inputs.py
+113
View File
@@ -0,0 +1,113 @@
"""Pure transform from torch sampler inputs to Core ML UNet kwargs.
Characterization tests cover SD1.5 / SDXL base / SDXL refiner / LCM
variants and the chunked-batch fan-out.
"""
import numpy as np
import torch
from coreml_suite.core.controlnet import extract_residual_kwargs, chunk_control
from coreml_suite.core.latents import chunk_batch
class CoreMLInputs:
def __init__(self, x, t, context, control, **kwargs):
self.x = x
self.t = t
self.context = context
self.control = control
self.time_ids = kwargs.get("time_ids")
self.text_embeds = kwargs.get("text_embeds")
self.ts_cond = kwargs.get("timestep_cond")
def coreml_kwargs(self, expected_inputs):
sample = self.x.cpu().numpy().astype(np.float16)
context = self.context.cpu().numpy().astype(np.float16)
context = context.transpose(0, 2, 1)[:, :, None, :]
t = self.t.cpu().numpy().astype(np.float16)
model_input_kwargs = {
"sample": sample,
"encoder_hidden_states": context,
"timestep": t,
}
residual_kwargs = extract_residual_kwargs(expected_inputs, self.control)
model_input_kwargs |= residual_kwargs
# LCM
if self.ts_cond is not None:
model_input_kwargs["timestep_cond"] = (
self.ts_cond.cpu().numpy().astype(np.float16)
)
# SDXL
if "text_embeds" in expected_inputs:
model_input_kwargs["text_embeds"] = (
self.text_embeds.cpu().numpy().astype(np.float16)
)
if "time_ids" in expected_inputs:
model_input_kwargs["time_ids"] = (
self.time_ids.cpu().numpy().astype(np.float16)
)
return model_input_kwargs
def chunks(self, expected_inputs):
sample_shape = expected_inputs["sample"]["shape"]
timestep_shape = expected_inputs["timestep"]["shape"]
hidden_shape = expected_inputs["encoder_hidden_states"]["shape"]
context_shape = (hidden_shape[0], hidden_shape[3], hidden_shape[1])
chunked_x = chunk_batch(self.x, sample_shape)
ts = list(torch.full((len(chunked_x), timestep_shape[0]), self.t[0]))
chunked_context = chunk_batch(self.context, context_shape)
chunked_control = [None] * len(chunked_x)
if self.control is not None:
chunked_control = chunk_control(self.control, sample_shape[0])
chunked_ts_cond = [None] * len(chunked_x)
if self.ts_cond is not None:
ts_cond_shape = expected_inputs["timestep_cond"]["shape"]
chunked_ts_cond = chunk_batch(self.ts_cond, ts_cond_shape)
chunked_time_ids = [None] * len(chunked_x)
if expected_inputs.get("time_ids") is not None:
time_ids_shape = expected_inputs["time_ids"]["shape"]
if self.time_ids is None:
self.time_ids = torch.zeros(len(chunked_x), *time_ids_shape[1:]).to(
self.x.device
)
chunked_time_ids = chunk_batch(self.time_ids, time_ids_shape)
chunked_text_embeds = [None] * len(chunked_x)
if expected_inputs.get("text_embeds") is not None:
text_embeds_shape = expected_inputs["text_embeds"]["shape"]
if self.text_embeds is None:
self.text_embeds = torch.zeros(
len(chunked_x), *text_embeds_shape[1:]
).to(self.x.device)
chunked_text_embeds = chunk_batch(self.text_embeds, text_embeds_shape)
return [
CoreMLInputs(
x,
t,
context,
control,
timestep_cond=ts_cond,
time_ids=time_ids,
text_embeds=text_embeds,
)
for x, t, context, control, ts_cond, time_ids, text_embeds in zip(
chunked_x,
ts,
chunked_context,
chunked_control,
chunked_ts_cond,
chunked_time_ids,
chunked_text_embeds,
)
]
coreml_suite/core/latents.py
+42
View File
@@ -0,0 +1,42 @@
"""Pure batch-chunking helpers for Core ML's fixed-shape UNet inputs.
Re-exported by coreml_suite.latents. Characterization tests cover the
contract (padding-zero regions, truncation in merge_chunks,
identity-passthrough when shape already matches).
"""
import torch
def chunk_batch(input_tensor, target_shape):
if input_tensor.shape == target_shape:
return [input_tensor]
batch_size = input_tensor.shape[0]
target_batch_size = target_shape[0]
num_chunks = batch_size // target_batch_size
if num_chunks == 0:
padding = torch.zeros(target_batch_size - batch_size, *target_shape[1:]).to(
input_tensor.device
)
return [torch.cat((input_tensor, padding), dim=0)]
mod = batch_size % target_batch_size
if mod != 0:
chunks = list(torch.chunk(input_tensor[:-mod], num_chunks))
padding = torch.zeros(target_batch_size - mod, *target_shape[1:]).to(
input_tensor.device
)
padded = torch.cat((input_tensor[-mod:], padding), dim=0)
chunks.append(padded)
return chunks
chunks = list(torch.chunk(input_tensor, num_chunks))
return chunks
def merge_chunks(chunks, orig_shape):
merged = torch.cat(chunks, dim=0)
if merged.shape == orig_shape:
return merged
return merged[: orig_shape[0]]
coreml_suite/core/naming.py
+68
View File
@@ -0,0 +1,68 @@
"""Pure out_name composition for the Core ML UNet artifact.
Extracted from CoreMLConverter.convert so the filename contract
can be tested + reused without instantiating the node. The string is the
cache key: every workflow that references a converted .mlpackage depends
on it staying byte-for-byte identical.
"""
from typing import Iterable, Tuple
ATTN_SUFFIX = {
"SPLIT_EINSUM": "se",
"SPLIT_EINSUM_V2": "se2",
"ORIGINAL": "orig",
}
# Palettization bits. "none" = no quantization (default; keeps the
# unquantized filename intact so existing workflows still resolve their
# cached .mlpackage). Numeric values append a `_q<bits>` suffix.
QUANT_NBITS_VALUES = ("none", "8", "6", "4")
def compose_out_name(
*,
ckpt_name: str,
batch_size: int,
width: int,
height: int,
controlnet_support: bool,
attention_implementation: str,
lora_names: Iterable[str] = (),
quantize_nbits: str = "none",
) -> str:
"""Build the .mlpackage stem from convert() parameters.
Locked behaviour (characterization tests):
- first '.' in ckpt_name wins (`a.b.c.safetensors` -> `a`)
- spaces collapse to underscores
- LoRA names are taken stem-only, sorted, joined with '_' and
prefixed with '_' when present (caller is expected to pass a
sorted list; we sort defensively)
- controlnet adds `_cn`
- attn suffix is `_se` | `_se2` | `_orig`
Quantization:
- quantize_nbits "none" (default) appends nothing — existing
unquantized .mlpackages keep the old filename
- "4" / "6" / "8" appends `_q<bits>` after the attn suffix
"""
if quantize_nbits not in QUANT_NBITS_VALUES:
raise ValueError(
f"quantize_nbits={quantize_nbits!r} not in {QUANT_NBITS_VALUES}"
)
stem = ckpt_name.split(".")[0]
sorted_names = sorted(lora_names)
lora_str = "_" + "_".join(name.split(".")[0] for name in sorted_names) if sorted_names else ""
cn_suffix = "_cn" if controlnet_support else ""
attn_suffix = "_" + ATTN_SUFFIX[attention_implementation]
quant_suffix = f"_q{quantize_nbits}" if quantize_nbits != "none" else ""
out_name = (
f"{stem}{lora_str}_{batch_size}x{width}x{height}"
f"{cn_suffix}{attn_suffix}{quant_suffix}"
)
return out_name.replace(" ", "_")
def lora_names_from_params(lora_params: Iterable[Tuple[str, float]]) -> list[str]:
"""Mirror the sort applied inside CoreMLConverter.convert."""
return [name for name, _ in sorted(lora_params, key=lambda pair: pair[0])]
coreml_suite/core/sdxl.py
+91
View File
@@ -0,0 +1,91 @@
"""Pure SDXL detection + time_ids/text_embeds assembly.
The framework-coupled adapter `add_sdxl_model_options` lives in models.py
and delegates the math here. Characterization tests cover base (len 6) vs
refiner (len 5) and the closure free-vars produced by
`sdxl_model_function_wrapper`.
"""
import torch
def is_sdxl(coreml_model):
return (
"time_ids" in coreml_model.expected_inputs
and "text_embeds" in coreml_model.expected_inputs
)
def is_sdxl_base(coreml_model):
return (
is_sdxl(coreml_model)
and coreml_model.expected_inputs["time_ids"]["shape"][1] == 6
)
def is_sdxl_refiner(coreml_model):
return (
is_sdxl(coreml_model)
and coreml_model.expected_inputs["time_ids"]["shape"][1] == 5
)
def build_sdxl_time_ids(pos_dict, neg_dict, *, is_base: bool, is_refiner: bool):
"""Compose the (2, N) time_ids tensor for the SDXL Core ML UNet.
- base: N=6 -> [h, w, crop_h, crop_w, target_h, target_w]
- refiner: N=5 -> [h, w, crop_h, crop_w, aesthetic_score]
- neither: N=4 -> [h, w, crop_h, crop_w] (edge case kept for parity)
"""
pos_time_ids = [
pos_dict.get("height", 768),
pos_dict.get("width", 768),
pos_dict.get("crop_h", 0),
pos_dict.get("crop_w", 0),
]
neg_time_ids = [
neg_dict.get("height", 768),
neg_dict.get("width", 768),
neg_dict.get("crop_h", 0),
neg_dict.get("crop_w", 0),
]
if is_base:
pos_time_ids += [
pos_dict.get("target_height", 768),
pos_dict.get("target_width", 768),
]
neg_time_ids += [
neg_dict.get("target_height", 768),
neg_dict.get("target_width", 768),
]
if is_refiner:
pos_time_ids += [pos_dict.get("aesthetic_score", 6)]
neg_time_ids += [neg_dict.get("aesthetic_score", 2.5)]
return torch.tensor([pos_time_ids, neg_time_ids])
def build_sdxl_text_embeds(pos_pooled, neg_pooled):
"""Concat pos then neg along the batch dim. Locked contract."""
return torch.cat((pos_pooled, neg_pooled))
def sdxl_model_function_wrapper(time_ids, text_embeds, refiner=False):
def wrapper(model_function, params):
x = params["input"]
t = params["timestep"]
c = params["c"]
context = c.get("c_crossattn")
if context is None:
return torch.zeros_like(x)
if refiner and context is not None:
# converted refiner accepts only g clip
c["c_crossattn"] = context[:, :, 768:]
return model_function(x, t, **c, time_ids=time_ids, text_embeds=text_embeds)
return wrapper
coreml_suite/latents.py
+3 -35
View File
@@ -1,36 +1,4 @@
import torch
"""Compatibility shim — re-exports from coreml_suite.core.latents."""
from coreml_suite.core.latents import chunk_batch, merge_chunks
def chunk_batch(input_tensor, target_shape):
if input_tensor.shape == target_shape:
return [input_tensor]
batch_size = input_tensor.shape[0]
target_batch_size = target_shape[0]
num_chunks = batch_size // target_batch_size
if num_chunks == 0:
padding = torch.zeros(target_batch_size - batch_size, *target_shape[1:]).to(
input_tensor.device
)
return [torch.cat((input_tensor, padding), dim=0)]
mod = batch_size % target_batch_size
if mod != 0:
chunks = list(torch.chunk(input_tensor[:-mod], num_chunks))
padding = torch.zeros(target_batch_size - mod, *target_shape[1:]).to(
input_tensor.device
)
padded = torch.cat((input_tensor[-mod:], padding), dim=0)
chunks.append(padded)
return chunks
chunks = list(torch.chunk(input_tensor, num_chunks))
return chunks
def merge_chunks(chunks, orig_shape):
merged = torch.cat(chunks, dim=0)
if merged.shape == orig_shape:
return merged
return merged[: orig_shape[0]]
__all__ = ["chunk_batch", "merge_chunks"]
coreml_suite/models.py
+40 -188
View File
@@ -1,15 +1,44 @@
import numpy as np
"""Framework-coupled glue between Core ML UNets and ComfyUI's sampler stack.
Pure math (CoreMLInputs, SDXL detection, time_ids/text_embeds assembly,
sdxl_model_function_wrapper) lives in coreml_suite.core.*.
This module is what touches comfy.*: model_base, ModelPatcher, the
diffusion_model wrapper, and the maintainer-facing add_sdxl_model_options
adapter.
"""
import torch
from comfy import model_base
from comfy.model_management import get_torch_device
from comfy.model_patcher import ModelPatcher
from coreml_suite.config import get_model_config, ModelVersion
from coreml_suite.controlnet import extract_residual_kwargs, chunk_control
from coreml_suite.latents import chunk_batch, merge_chunks
from coreml_suite.core.inputs import CoreMLInputs
from coreml_suite.core.latents import merge_chunks
from coreml_suite.core.sdxl import (
build_sdxl_text_embeds,
build_sdxl_time_ids,
is_sdxl,
is_sdxl_base,
is_sdxl_refiner,
sdxl_model_function_wrapper,
)
from coreml_suite.lcm.utils import is_lcm
from coreml_suite.logger import logger
__all__ = [
"CoreMLInputs",
"CoreMLModelWrapper",
"CoreMLModelWrapperLCM",
"add_sdxl_model_options",
"get_latent_image",
"get_model_patcher",
"is_sdxl",
"is_sdxl_base",
"is_sdxl_refiner",
"sdxl_model_function_wrapper",
]
class CoreMLModelWrapper:
def __init__(self, coreml_model):
@@ -68,204 +97,27 @@ class CoreMLModelWrapperLCM(CoreMLModelWrapper):
self.config = None
class CoreMLInputs:
def __init__(self, x, t, context, control, **kwargs):
self.x = x
self.t = t
self.context = context
self.control = control
self.time_ids = kwargs.get("time_ids")
self.text_embeds = kwargs.get("text_embeds")
self.ts_cond = kwargs.get("timestep_cond")
def coreml_kwargs(self, expected_inputs):
sample = self.x.cpu().numpy().astype(np.float16)
context = self.context.cpu().numpy().astype(np.float16)
context = context.transpose(0, 2, 1)[:, :, None, :]
t = self.t.cpu().numpy().astype(np.float16)
model_input_kwargs = {
"sample": sample,
"encoder_hidden_states": context,
"timestep": t,
}
residual_kwargs = extract_residual_kwargs(expected_inputs, self.control)
model_input_kwargs |= residual_kwargs
# LCM
if self.ts_cond is not None:
model_input_kwargs["timestep_cond"] = (
self.ts_cond.cpu().numpy().astype(np.float16)
)
# SDXL
if "text_embeds" in expected_inputs:
model_input_kwargs["text_embeds"] = (
self.text_embeds.cpu().numpy().astype(np.float16)
)
if "time_ids" in expected_inputs:
model_input_kwargs["time_ids"] = (
self.time_ids.cpu().numpy().astype(np.float16)
)
return model_input_kwargs
def chunks(self, expected_inputs):
sample_shape = expected_inputs["sample"]["shape"]
timestep_shape = expected_inputs["timestep"]["shape"]
hidden_shape = expected_inputs["encoder_hidden_states"]["shape"]
context_shape = (hidden_shape[0], hidden_shape[3], hidden_shape[1])
chunked_x = chunk_batch(self.x, sample_shape)
ts = list(torch.full((len(chunked_x), timestep_shape[0]), self.t[0]))
chunked_context = chunk_batch(self.context, context_shape)
chunked_control = [None] * len(chunked_x)
if self.control is not None:
chunked_control = chunk_control(self.control, sample_shape[0])
chunked_ts_cond = [None] * len(chunked_x)
if self.ts_cond is not None:
ts_cond_shape = expected_inputs["timestep_cond"]["shape"]
chunked_ts_cond = chunk_batch(self.ts_cond, ts_cond_shape)
chunked_time_ids = [None] * len(chunked_x)
if expected_inputs.get("time_ids") is not None:
time_ids_shape = expected_inputs["time_ids"]["shape"]
if self.time_ids is None:
self.time_ids = torch.zeros(len(chunked_x), *time_ids_shape[1:]).to(
self.x.device
)
chunked_time_ids = chunk_batch(self.time_ids, time_ids_shape)
chunked_text_embeds = [None] * len(chunked_x)
if expected_inputs.get("text_embeds") is not None:
text_embeds_shape = expected_inputs["text_embeds"]["shape"]
if self.text_embeds is None:
self.text_embeds = torch.zeros(
len(chunked_x), *text_embeds_shape[1:]
).to(self.x.device)
chunked_text_embeds = chunk_batch(self.text_embeds, text_embeds_shape)
return [
CoreMLInputs(
x,
t,
context,
control,
timestep_cond=ts_cond,
time_ids=time_ids,
text_embeds=text_embeds,
)
for x, t, context, control, ts_cond, time_ids, text_embeds in zip(
chunked_x,
ts,
chunked_context,
chunked_control,
chunked_ts_cond,
chunked_time_ids,
chunked_text_embeds,
)
]
def is_sdxl(coreml_model):
return (
"time_ids" in coreml_model.expected_inputs
and "text_embeds" in coreml_model.expected_inputs
)
def is_sdxl_base(coreml_model):
return (
is_sdxl(coreml_model)
and coreml_model.expected_inputs["time_ids"]["shape"][1] == 6
)
def is_sdxl_refiner(coreml_model):
return (
is_sdxl(coreml_model)
and coreml_model.expected_inputs["time_ids"]["shape"][1] == 5
)
def sdxl_model_function_wrapper(time_ids, text_embeds, refiner=False):
def wrapper(model_function, params):
x = params["input"]
t = params["timestep"]
c = params["c"]
context = c.get("c_crossattn")
if context is None:
return torch.zeros_like(x)
if refiner and context is not None:
# converted refiner accepts only g clip
c["c_crossattn"] = context[:, :, 768:]
return model_function(x, t, **c, time_ids=time_ids, text_embeds=text_embeds)
return wrapper
def add_sdxl_model_options(model_patcher, positive, negative):
mp = model_patcher.clone()
pos_dict = positive[0][1]
neg_dict = negative[0][1]
pos_pooled = pos_dict["pooled_output"]
neg_pooled = neg_dict["pooled_output"]
pos_time_ids = [
pos_dict.get("height", 768),
pos_dict.get("width", 768),
pos_dict.get("crop_h", 0),
pos_dict.get("crop_w", 0),
]
neg_time_ids = [
neg_dict.get("height", 768),
neg_dict.get("width", 768),
neg_dict.get("crop_h", 0),
neg_dict.get("crop_w", 0),
]
if model_patcher.model.diffusion_model.is_sdxl_base:
pos_time_ids += [
pos_dict.get("target_height", 768),
pos_dict.get("target_width", 768),
]
neg_time_ids += [
neg_dict.get("target_height", 768),
neg_dict.get("target_width", 768),
]
is_base = model_patcher.model.diffusion_model.is_sdxl_base
is_refiner = model_patcher.model.diffusion_model.is_sdxl_refiner
if is_refiner:
pos_time_ids += [
pos_dict.get("aesthetic_score", 6),
]
neg_time_ids += [
neg_dict.get("aesthetic_score", 2.5),
]
time_ids = build_sdxl_time_ids(
pos_dict, neg_dict, is_base=is_base, is_refiner=is_refiner
)
text_embeds = build_sdxl_text_embeds(
pos_dict["pooled_output"], neg_dict["pooled_output"]
)
time_ids = torch.tensor([pos_time_ids, neg_time_ids])
text_embeds = torch.cat((pos_pooled, neg_pooled))
model_options = {
mp.model_options |= {
"model_function_wrapper": sdxl_model_function_wrapper(
time_ids, text_embeds, is_refiner
),
}
mp.model_options |= model_options
return mp
coreml_suite/nodes.py
+22 -16
View File
@@ -8,6 +8,11 @@ import folder_paths
from coreml_suite import COREML_NODE
from coreml_suite import converter
from coreml_suite.config import ModelVersion
from coreml_suite.core.naming import (
QUANT_NBITS_VALUES,
compose_out_name,
lora_names_from_params,
)
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
@@ -244,6 +249,12 @@ class CoreMLConverter(COREML_NODE):
"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": (list(QUANT_NBITS_VALUES), {"default": "none"}),
"lora_params": ("LORA_PARAMS",),
},
}
@@ -262,6 +273,7 @@ class CoreMLConverter(COREML_NODE):
attention_implementation,
compute_unit,
controlnet_support,
quantize_nbits="none",
lora_params=None,
):
"""Converts a LCM model to Core ML.
@@ -288,24 +300,17 @@ class CoreMLConverter(COREML_NODE):
h = height
w = width
sample_size = (h // 8, w // 8)
batch_size = batch_size
cn_support_str = "_cn" if controlnet_support else ""
lora_str = (
"_" + "_".join(lora_param[0].split(".")[0] for lora_param in lora_params)
if lora_params
else ""
out_name = 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=lora_names_from_params(lora_params),
quantize_nbits=quantize_nbits,
)
attn_str = (
"_"
+ {"SPLIT_EINSUM": "se", "SPLIT_EINSUM_V2": "se2", "ORIGINAL": "orig"}[
attention_implementation
]
)
out_name = f"{ckpt_name.split('.')[0]}{lora_str}_{batch_size}x{w}x{h}{cn_support_str}{attn_str}"
out_name = out_name.replace(" ", "_")
logger.info(f"Converting {ckpt_name} to {out_name}")
logger.info(f"Batch size: {batch_size}")
logger.info(f"Width: {w}, Height: {h}")
@@ -335,6 +340,7 @@ class CoreMLConverter(COREML_NODE):
lora_weights=lora_weights,
attn_impl=attention_implementation,
config_path=config_path,
quantize_nbits=quantize_nbits,
)
unet_target_path = converter.compile_model(
out_path=unet_out_path, out_name=out_name, submodule_name="unet"