diff --git a/README.md b/README.md index 7851e6a..c1e8cea 100644 --- a/README.md +++ b/README.md @@ -4,7 +4,7 @@ Welcome! I've developed a set of custom nodes for ComfyUI that allows you to use Core ML models in your ComfyUI workflows. -These models are designed to leverage the Apple Neural Engine (ANE) on Apple Silicon (M1/M2) machines, +These models are designed to leverage the Apple Neural Engine (ANE) on Apple Silicon (M1/M2) machines, thereby enhancing your workflows and improving performance. If you're not sure how to obtain these models, you can download them @@ -24,11 +24,11 @@ To start using custom nodes in your ComfyUI, follow these simple steps: 2. Install the dependencies: You'll need to use a package manager like pip to do this. That's it! You're now ready to start enhancing your ComfyUI workflows with Core ML models. + - Check [Installation](#installation) for more details on installation. - Check [How to use](#how-to-use) for more details on how to use the custom nodes. - Check [Example Workflows](#example-workflows) for some example workflows. - ## Glossary - **Core ML**: A machine learning framework developed by Apple. It's used to run machine learning models on Apple @@ -38,24 +38,27 @@ That's it! You're now ready to start enhancing your ComfyUI workflows with Core - **mlpackage**: A Core ML model packaged in a directory. This is the default format for Core ML models. - **ANE**: Apple Neural Engine. A hardware accelerator for machine learning tasks on Apple devices. - **Compute Unit**: A Core ML option that allows you to specify the hardware on which the model should run. - - **CPU_AND_ANE**: A Core ML compute unit option that allows the model to run on both the CPU and ANE. This is the - default option. - - **CPU_AND_GPU**: A Core ML compute unit option that allows the model to run on both the CPU and GPU. - - **CPU_ONLY**: A Core ML compute unit option that allows the model to run on the CPU only. - - **ALL**: A Core ML compute unit option that allows the model to run on all available hardware. -- **CLIP**: Contrastive Language-Image Pre-training. A model that learns visual concepts from natural language + - **CPU_AND_ANE**: A Core ML compute unit option that allows the model to run on both the CPU and ANE. This is the + default option. + - **CPU_AND_GPU**: A Core ML compute unit option that allows the model to run on both the CPU and GPU. + - **CPU_ONLY**: A Core ML compute unit option that allows the model to run on the CPU only. + - **ALL**: A Core ML compute unit option that allows the model to run on all available hardware. +- **CLIP**: Contrastive Language-Image Pre-training. A model that learns visual concepts from natural language supervision. It's used as a text encoder in Stable Diffusion. -- **VAE**: Variational Autoencoder. A model that learns a latent representation of images. It's used as a prior in +- **VAE**: Variational Autoencoder. A model that learns a latent representation of images. It's used as a prior in Stable Diffusion. - **Checkpoint**: A file that contains the weights of a model. It's used to load models in Stable Diffusion. > [!NOTE] > Note on Compute Units: -> For the model to run on the ANE, the model must be converted with the `--attention-implementation SPLIT_EINSUM` option. +> For the model to run on the ANE, the model must be converted with the `--attention-implementation SPLIT_EINSUM` +> option. > Models converted with `--attention-implementation ORIGINAL` will run on GPU instead of ANE. ## Features + These custom nodes come with a host of features, including: + - Loading Core ML Unet models - Support for ControlNet - Support for ANE (Apple Neural Engine) @@ -63,8 +66,8 @@ These custom nodes come with a host of features, including: - Support for `mlmodelc` and `mlpackage` files > [!NOTE] -> Please note that using Core ML models can take a bit longer to load initially. -> For the best experience, I recommend using the compiled models +> Please note that using Core ML models can take a bit longer to load initially. +> For the best experience, I recommend using the compiled models > (.mlmodelc files) instead of the .mlpackage files. > [!NOTE] @@ -90,73 +93,112 @@ The installation process is simple! ## How to use Once you've installed the custom nodes, you can start using them in your ComfyUI workflows. -To do this, you need to add the nodes to your workflow. You can do this by double-clicking on the workflow canvas and -selecting the nodes from the list of available nodes. You can also use the search bar to find the nodes. -The list of available nodes is given below. You can also find the nodes in the `CoreML Suite` category. +To do this, you need to add the nodes to your workflow. You can do this by right-clicking on the workflow canvas and +selecting the nodes from the list of available nodes (the nodes are in the `Core ML Suite` category). +You can also double-click the canvas and use the search bar to find the nodes. The list of available nodes is given +below. ### Available Nodes -#### CoreML UNet Loader (`CoreMLUnetLoader`) +#### Core ML UNet Loader (`CoreMLUnetLoader`) -![CoreMLUnetLoader](https://github.com/aszc-dev/ComfyUI-CoreMLSuite/assets/24932801/2bd10f73-4103-4860-894c-b6a6e56c6546) +![CoreMLUnetLoader](./assets/unet_loader.png?raw=true) This node allows you to load a Core ML UNet model and use it in your ComfyUI workflow. Place the converted .mlpackage or .mlmodelc file in ComfyUI's `models/unet` directory and use the node to load the model. The output of the -node is a `MODEL` object similar to standard ComfyUI models. +node is a `coreml_model` object that can be used with the Core ML Sampler. + +- **Inputs**: + - **model_name**: The name of the model to load. This should be the name of the .mlpackage or .mlmodelc file. + - **compute_unit**: The hardware on which the model should run. This can be one of the following: + - `CPU_AND_ANE`: The model will run on both the CPU and ANE. This is the default option. It works best with + models + converted with `--attention-implementation SPLIT_EINSUM` or `--attention-implementation SPLIT_EINSUM_V2`. + - `CPU_AND_GPU`: The model will run on both the CPU and GPU. It works best with models converted with + `--attention-implementation ORIGINAL`. + - `CPU_ONLY`: The model will run on the CPU only. + - `ALL`: The model will run on all available hardware. +- **Outputs**: + - **coreml_model**: A Core ML model that can be used with the Core ML Sampler. + > [!NOTE] > Some models are designed to support ControlNet. If you're using such a model, > make sure to provide a ControlNet input; otherwise, the model will use random noise as ControlNet input. +#### Core ML Sampler (`CoreMLSampler`) + +![CoreMLSampler](./assets/sampler.png?raw=true) + +This node allows you to generate images using a Core ML model. The node takes a Core ML model as input and outputs a +latent image similar to the latent image output by the KSampler. This means that you can use the +resulting latent as you normally would in your workflow. + +- **Inputs**: + - **coreml_model**: The Core ML model to use for sampling. This should be the output of the Core ML UNet Loader. + - **latent_image** [optional]: The latent image to use for sampling. If provided, should be of the same size as the + input of the Core ML model. If not provided, the node will create a latent suitable for the Core ML model used. + Useful in img2img workflows. + - ... _(the rest of the inputs are the same as the KSampler)_ +- **Outputs**: + - **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. + ### Example Workflows > [!NOTE] > The models used are just an example. Feel free to experiment with different models and see what works best for you. - #### Basic txt2img with Core ML UNet loader -This is a basic txt2img workflow that uses the Core ML UNet loader to load a Core ML UNet model. The CLIP and VAE models -are loaded using the standard ComfyUI nodes. In the first example, the text encoder (CLIP) and VAE models are loaded -separately. In the second example, the text encoder and VAE models are loaded from the checkpoint file. Note that you can use any CLIP or VAE model -as long as it's compatible with Stable Diffusion v1.5. + +This is a basic txt2img workflow that uses the Core ML UNet loader to load a model. The CLIP and VAE models +are loaded using the standard ComfyUI nodes. In the first example, the text encoder (CLIP) and VAE models are loaded +separately. In the second example, the text encoder and VAE models are loaded from the checkpoint file. Note that you +can use any CLIP or VAE model as long as it's compatible with Stable Diffusion v1.5. 1. **Loading text encoder (CLIP) and VAE models separately** - - This workflow uses CLIP and VAE models available - [here](https://huggingface.co/runwayml/stable-diffusion-v1-5/blob/main/text_encoder/model.safetensors) and - [here](https://huggingface.co/runwayml/stable-diffusion-v1-5/blob/main/vae/diffusion_pytorch_model.safetensors). - Once downloaded, place the models in the`models/clip` and `models/vae` directories respectively. - - The Core ML UNet model is available - [here](https://huggingface.co/coreml-community/coreml-stable-diffusion-v1-5_cn/blob/main/split_einsum/stable-diffusion-_v1-5_split-einsum_cn.zip). - Once downloaded, place the model in the `models/unet` directory. -![coreml-unet+clip+vae](https://github.com/aszc-dev/ComfyUI-CoreMLSuite/assets/24932801/ff7b8d75-37ea-4da9-a258-829edd6eb1b7) - - + - This workflow uses CLIP and VAE models available + [here](https://huggingface.co/runwayml/stable-diffusion-v1-5/blob/main/text_encoder/model.safetensors) and + [here](https://huggingface.co/runwayml/stable-diffusion-v1-5/blob/main/vae/diffusion_pytorch_model.safetensors). + Once downloaded, place the models in the`models/clip` and `models/vae` directories respectively. + - The Core ML UNet model is available + [here](https://huggingface.co/coreml-community/coreml-stable-diffusion-v1-5_cn/blob/main/split_einsum/stable-diffusion-_v1-5_split-einsum_cn.zip). + Once downloaded, place the model in the `models/unet` directory. + ![coreml-unet+clip+vae](./assets/unet+sampler+clip+vae.png?raw=true) 2. **Loading text encoder (CLIP) and VAE models from checkpoint file** - - This workflow loads the CLIP and VAE models from the checkpoint file available - [here](https://huggingface.co/runwayml/stable-diffusion-v1-5/blob/main/v1-5-pruned-emaonly.safetensors). - Once downloaded, place the model in the`models/checkpoints` directory. - - The Core ML UNet model is available - [here](https://huggingface.co/coreml-community/coreml-stable-diffusion-v1-5_cn/blob/main/split_einsum/stable-diffusion-_v1-5_split-einsum_cn.zip). - Once downloaded, place the model in the `models/unet` directory. -![coreml-unet+checkpoint](https://github.com/aszc-dev/ComfyUI-CoreMLSuite/assets/24932801/0c8b4f65-9bde-4b0d-936b-5bb27023d2ce) + - This workflow loads the CLIP and VAE models from the checkpoint file available + [here](https://huggingface.co/runwayml/stable-diffusion-v1-5/blob/main/v1-5-pruned-emaonly.safetensors). + Once downloaded, place the model in the`models/checkpoints` directory. + - The Core ML UNet model is available + [here](https://huggingface.co/coreml-community/coreml-stable-diffusion-v1-5_cn/blob/main/split_einsum/stable-diffusion-_v1-5_split-einsum_cn.zip). + Once downloaded, place the model in the `models/unet` directory. + ![coreml-unet+checkpoint](./assets/unet+sampler+checkpoint.png?raw=true) +#### ControlNet with Core ML UNet loader -#### ControlNet with Core ML UNet loader - -(Coming soon) +This workflow uses the Core ML UNet loader to load a Core ML UNet model that supports ControlNet. The ControlNet is +being loaded using the standard ComfyUI nodes. Please refer to +the [basic txt2img workflow](#basic-txt2img-with-core-ml-unet-loader) for more details on how to load the CLIP and VAE +models. +The ControlNet model used in this workflow is available +[here](https://huggingface.co/lllyasviel/ControlNet-v1-1/blob/main/control_v11p_sd15_lineart.pth). +Once downloaded, place the model in the `models/controlnet` directory. +![coreml-unet+controlnet](./assets/unet+sampler+controlnet.png?raw=true) ## Limitations -- Core ML models are fixed in terms of their inputs and outputs. -This means you'll need to use latent images of the same size as the input of the model (512x512 is the default for SD1.5). -However, you can convert the model to a different input size using tools available -in the [apple/ml-stable-diffusion](https://github.com/apple/ml-stable-diffusion) repository. +- Core ML models are fixed in terms of their inputs and outputs. + This means you'll need to use latent images of the same size as the input of the model (512x512 is the default for + SD1.5). + However, you can convert the model to a different input size using tools available + in the [apple/ml-stable-diffusion](https://github.com/apple/ml-stable-diffusion) repository. - For now, only Stable Diffusion v1.5 is supported. - LoRA is not supported yet. -[^1]: Unless [EnumeratedShapes](https://apple.github.io/coremltools/docs-guides/source/flexible-inputs.html#select-from-predetermined-shapes) +[^1]: +Unless [EnumeratedShapes](https://apple.github.io/coremltools/docs-guides/source/flexible-inputs.html#select-from-predetermined-shapes) is used during conversion. Needs more testing. ## Support -I'm here to help! If you have any questions or suggestions, don't hesitate to open an issue and I'll do my best +I'm here to help! If you have any questions or suggestions, don't hesitate to open an issue and I'll do my best to assist you. diff --git a/assets/sampler.png b/assets/sampler.png new file mode 100644 index 0000000..9f00939 Binary files /dev/null and b/assets/sampler.png differ diff --git a/assets/snake.png b/assets/snake.png new file mode 100644 index 0000000..c74839e Binary files /dev/null and b/assets/snake.png differ diff --git a/assets/unet+sampler+checkpoint.png b/assets/unet+sampler+checkpoint.png new file mode 100644 index 0000000..5ea0672 Binary files /dev/null and b/assets/unet+sampler+checkpoint.png differ diff --git a/assets/unet+sampler+clip+vae.png b/assets/unet+sampler+clip+vae.png new file mode 100644 index 0000000..ebc9138 Binary files /dev/null and b/assets/unet+sampler+clip+vae.png differ diff --git a/assets/unet+sampler+controlnet.png b/assets/unet+sampler+controlnet.png new file mode 100644 index 0000000..132c6c3 Binary files /dev/null and b/assets/unet+sampler+controlnet.png differ diff --git a/assets/unet_loader.png b/assets/unet_loader.png new file mode 100644 index 0000000..46d931a Binary files /dev/null and b/assets/unet_loader.png differ diff --git a/coreml_suite/loaders.py b/coreml_suite/loaders.py index 8e81832..8d201c6 100644 --- a/coreml_suite/loaders.py +++ b/coreml_suite/loaders.py @@ -28,7 +28,7 @@ class CoreMLLoader: } FUNCTION = "load" - CATEGORY = "CoreML Suite" + CATEGORY = "Core ML Suite" @classmethod def coreml_filenames(cls): diff --git a/coreml_suite/samplers.py b/coreml_suite/samplers.py index 84cc1f6..89f44db 100644 --- a/coreml_suite/samplers.py +++ b/coreml_suite/samplers.py @@ -37,7 +37,7 @@ class CoreMLSampler(KSampler): "optional": {"latent_image": ("LATENT",)}, } - CATEGORY = "CoreML Suite" + CATEGORY = "Core ML Suite" def sample( self,