diff --git a/.gitignore b/.gitignore index 5733a59..d60dca5 100644 --- a/.gitignore +++ b/.gitignore @@ -3,3 +3,4 @@ __pycache__/ models/ .venv/ test_results/ +.claude/ diff --git a/__init__.py b/__init__.py index b38e219..198c202 100644 --- a/__init__.py +++ b/__init__.py @@ -11,9 +11,6 @@ CoreMLConverter, COREML_LOAD_LORA, ) -from coreml_suite.lcm import ( - COREML_CONVERT_LCM, -) NODE_CLASS_MAPPINGS = { "CoreMLUNetLoader": CoreMLLoaderUNet, @@ -22,7 +19,6 @@ "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 @@ "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", } diff --git a/coreml_suite/lcm/__init__.py b/coreml_suite/lcm/__init__.py index 4432285..527f673 100644 --- a/coreml_suite/lcm/__init__.py +++ b/coreml_suite/lcm/__init__.py @@ -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``. +""" diff --git a/coreml_suite/lcm/converter.py b/coreml_suite/lcm/converter.py deleted file mode 100644 index 3214375..0000000 --- a/coreml_suite/lcm/converter.py +++ /dev/null @@ -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) diff --git a/coreml_suite/lcm/nodes.py b/coreml_suite/lcm/nodes.py deleted file mode 100644 index e7d03a7..0000000 --- a/coreml_suite/lcm/nodes.py +++ /dev/null @@ -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),) diff --git a/coreml_suite/lcm/unet.py b/coreml_suite/lcm/unet.py deleted file mode 100644 index 115a3d7..0000000 --- a/coreml_suite/lcm/unet.py +++ /dev/null @@ -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,) diff --git a/coreml_suite/nodes.py b/coreml_suite/nodes.py index 75009d7..96d7e20 100644 --- a/coreml_suite/nodes.py +++ b/coreml_suite/nodes.py @@ -7,7 +7,6 @@ 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 @@ def wrap(self, coreml_model): 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 @@ def INPUT_TYPES(cls): def convert( self, ckpt_name, - model_version, height, width, batch_size, @@ -284,9 +286,11 @@ def convert( 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 @@ def convert( 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 @@ def convert( 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 @@ def convert( 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, diff --git a/pyproject.toml b/pyproject.toml index fcd1bb7..0f63486 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -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" diff --git a/requirements.txt b/requirements.txt index 642e1c5..ccfaa9c 100644 --- a/requirements.txt +++ b/requirements.txt @@ -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 diff --git a/tests/integration/workflows/e2e-1.5-basic-conversion.json b/tests/integration/workflows/e2e-1.5-basic-conversion.json index 97336c4..1de6aed 100644 --- a/tests/integration/workflows/e2e-1.5-basic-conversion.json +++ b/tests/integration/workflows/e2e-1.5-basic-conversion.json @@ -107,7 +107,6 @@ "10": { "inputs": { "ckpt_name": "dreamshaper_8.safetensors", - "model_version": "SD15", "height": 512, "width": 512, "batch_size": 1, diff --git a/uv.lock b/uv.lock index 5696883..8489a67 100644 --- a/uv.lock +++ b/uv.lock @@ -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]]