diff --git a/README.md b/README.md index 39da520..b530bd9 100644 --- a/README.md +++ b/README.md @@ -143,6 +143,23 @@ resulting latent as you normally would in your workflow. - **LATENT**: The latent image output by the Core ML model. This can be decoded using a VAE Decoder or used as input to the next node in your workflow. +#### Core ML Adapter (Experimental) (`CoreMLModelAdapter`) + +![CoreMLModelAdapter](./assets/adapter.png?raw=true) + +This node allows you to use a Core ML as a standard ComfyUI model. This is an experimental node and may not work with +all models and nodes. Please use with caution and pay attention to the expected inputs of the model. + +- **Input**: + - **coreml_model**: The Core ML model to use as a ComfyUI model. +- **Output**: + - **MODEL**: The Core ML model wrapped in a ComfyUI model. + +> [!NOTE] +> While this approach allows you to use Core ML models with many ComfyUI nodes (both standard and custom), the +> expected inputs of the model will not be checked, which may cause errors. Please make sure to use a model compatible +> with the expected parameters. + ### Example Workflows > [!NOTE] diff --git a/__init__.py b/__init__.py index b259a08..8a0d5c9 100644 --- a/__init__.py +++ b/__init__.py @@ -3,13 +3,15 @@ import sys sys.path.append(os.path.dirname(__file__)) -from coreml_suite.nodes import CoreMLLoaderUNet, CoreMLSampler +from coreml_suite.nodes import CoreMLLoaderUNet, CoreMLSampler, CoreMLModelAdapter NODE_CLASS_MAPPINGS = { "CoreMLUNetLoader": CoreMLLoaderUNet, "CoreMLSampler": CoreMLSampler, + "CoreMLModelAdapter": CoreMLModelAdapter, } NODE_DISPLAY_NAME_MAPPINGS = { "CoreMLUNetLoader": "Load Core ML UNet", "CoreMLSampler": "Core ML Sampler", + "CoreMLModelAdapter": "Core ML Adapter (Experimental)", } diff --git a/assets/adapter.png b/assets/adapter.png new file mode 100644 index 0000000..df21d92 Binary files /dev/null and b/assets/adapter.png differ diff --git a/coreml_suite/nodes.py b/coreml_suite/nodes.py index 7246a6f..d1e58bb 100644 --- a/coreml_suite/nodes.py +++ b/coreml_suite/nodes.py @@ -28,24 +28,24 @@ class CoreMLSampler(KSampler): CATEGORY = "Core ML Suite" def sample( - self, - coreml_model, - seed, - steps, - cfg, - sampler_name, - scheduler, - positive, - negative, - latent_image=None, - denoise=1.0, + self, + coreml_model, + seed, + steps, + cfg, + sampler_name, + scheduler, + positive, + negative, + latent_image=None, + denoise=1.0, ): sample_shape = coreml_model.expected_inputs["sample"]["shape"] latent_image = reshape_latent_image(latent_image, sample_shape) latent_image["samples"] = latent_image["samples"][0:1] model_config = get_model_config() - wrapped_model = CoreMLModelWrapper(model_config, coreml_model) + wrapped_model = CoreMLModelAdapter(model_config, coreml_model) model = ModelPatcher(wrapped_model, get_torch_device(), None) return super().sample( @@ -99,9 +99,6 @@ class CoreMLLoader: sources = "compiled" if coreml_name.endswith(".mlmodelc") else "packages" - return self._load(coreml_path, compute_unit, sources) - - def _load(self, coreml_path, compute_unit, sources): return (CoreMLModel(coreml_path, compute_unit, sources),) @@ -109,3 +106,29 @@ class CoreMLLoaderUNet(CoreMLLoader): PACKAGE_DIRNAME = "unet" RETURN_TYPES = ("COREML_UNET",) RETURN_NAMES = ("coreml_model",) + + +class CoreMLModelAdapter: + """ + 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_config = get_model_config() + wrapped_model = CoreMLModelWrapper(model_config, coreml_model) + patched_model = ModelPatcher(wrapped_model, get_torch_device(), None) + return (patched_model,)