diff --git a/__init__.py b/__init__.py index 0be3898f..45dd511a 100644 --- a/__init__.py +++ b/__init__.py @@ -74,6 +74,13 @@ "Nodes `NunchakuPulidApply`,`NunchakuPulidLoader`, " "`NunchakuPuLIDLoaderV2` and `NunchakuFluxPuLIDApplyV2` import failed:" ) +try: + from .nodes.models.ipadapter import NunchakuFluxIPAdapterApply, NunchakuIPAdapterLoader + + NODE_CLASS_MAPPINGS["NunchakuFluxIPAdapterApply"] = NunchakuFluxIPAdapterApply + NODE_CLASS_MAPPINGS["NunchakuIPAdapterLoader"] = NunchakuIPAdapterLoader +except ImportError: + logger.exception("Nodes `NunchakuFluxIPAdapterApply` and `NunchakuIPAdapterLoader` import failed:") try: from .nodes.tools.merge_safetensors import NunchakuModelMerger diff --git a/docs/source/api/nodes.models.ipadapter.rst b/docs/source/api/nodes.models.ipadapter.rst new file mode 100644 index 00000000..32bd49d2 --- /dev/null +++ b/docs/source/api/nodes.models.ipadapter.rst @@ -0,0 +1,6 @@ +nodes.models.ipadapter +====================== + +.. automodule:: comfyui_nunchaku.nodes.models.ipadapter + :members: + :show-inheritance: diff --git a/docs/source/api/nodes.models.rst b/docs/source/api/nodes.models.rst index 0ed35a4a..146e9847 100644 --- a/docs/source/api/nodes.models.rst +++ b/docs/source/api/nodes.models.rst @@ -7,3 +7,5 @@ nodes.models nodes.models.flux nodes.models.pulid nodes.models.text_encoder + nodes.models.ipadapter + nodes.models.utils diff --git a/docs/source/api/nodes.models.utils.rst b/docs/source/api/nodes.models.utils.rst new file mode 100644 index 00000000..2bf454b0 --- /dev/null +++ b/docs/source/api/nodes.models.utils.rst @@ -0,0 +1,6 @@ +nodes.models.utils +================== + +.. automodule:: comfyui_nunchaku.nodes.models.utils + :members: + :show-inheritance: diff --git a/docs/source/nodes/ipadapter.rst b/docs/source/nodes/ipadapter.rst new file mode 100644 index 00000000..cd1d7fea --- /dev/null +++ b/docs/source/nodes/ipadapter.rst @@ -0,0 +1,55 @@ +IP-Adapter Nodes +================ + +.. _nunchaku-flux-ip-adapter-loader: + +Nunchaku IP-Adapter Loader +-------------------------- + +.. image:: https://huggingface.co/datasets/nunchaku-tech/cdn/resolve/main/ComfyUI-nunchaku/nodes/NunchakuIPAdapterLoader.png + +A node for loading IP-Adapter weights to Nunchaku FLUX models within ComfyUI. + +**Inputs:** + +- **model**: The Nunchaku FLUX model to inject IP-Adapter to. Make sure the model is loaded by :ref:`nunchaku-flux-dit-loader`. + +**Outputs:** + +- **model**: IP-Adapter injected Nunchaku FLUX model. +- **ipadapter_pipeline**: The loaded IP-Adapter pipeline, ready for use with IP-Adapter Apply nodes. + +.. warning:: + This node will automatically download the IP-Adapter and associated CLIP models from Hugging Face. + Custom model paths are not supported for now. + +.. seealso:: + API reference: :class:`~comfyui_nunchaku.nodes.models.ipadapter.NunchakuIPAdapterLoader`. + + Example workflow: :ref:`nunchaku-flux.1-ip-adapter-json`. + +.. _nunchaku-flux-ip-adapter-apply: + +Nunchaku IP-Adapter Apply +------------------------- + +.. image:: https://huggingface.co/datasets/nunchaku-tech/cdn/resolve/main/ComfyUI-nunchaku/nodes/NunchakuFluxIPAdapterApply.png + :alt: NunchakuFluxIPAdapterApply + +A node for applying IP-Adapter to a Nunchaku model using a given image and weight. + +**Inputs:** + +- **model**: The Nunchaku FLUX model to apply IP-Adapter to. Make sure the model is loaded by :ref:`nunchaku-flux-ip-adapter-loader`. +- **ipadapter_pipeline**: The IP-Adapter pipeline to apply. +- **image**: The image to apply IP-Adapter to. +- **weight**: The weight of the IP-Adapter. + +**Outputs:** + +- **model**: The Nunchaku FLUX model with IP-Adapter applied. + +.. seealso:: + API reference: :class:`~comfyui_nunchaku.nodes.models.ipadapter.NunchakuFluxIPAdapterApply`. + + Example workflow: :ref:`nunchaku-flux.1-ip-adapter-json`. diff --git a/docs/source/nodes/pulid.rst b/docs/source/nodes/pulid.rst index ec375bc2..70481dab 100644 --- a/docs/source/nodes/pulid.rst +++ b/docs/source/nodes/pulid.rst @@ -52,7 +52,6 @@ A node for loading the PuLID pipeline required for identity-preserving image gen **Outputs:** - **model**: PuLID injected Nunchaku FLUX model. - - **pulid_pipeline**: The loaded PuLID pipeline, ready for use with PuLID Apply nodes. diff --git a/docs/source/nodes/toc.rst b/docs/source/nodes/toc.rst index 7f482982..c995e7ba 100644 --- a/docs/source/nodes/toc.rst +++ b/docs/source/nodes/toc.rst @@ -8,5 +8,6 @@ Node Reference lora.rst text_encoder.rst pulid.rst + ipadapter.rst tools.rst preprocessors.rst diff --git a/docs/source/workflows/ipadapter.rst b/docs/source/workflows/ipadapter.rst new file mode 100644 index 00000000..669d7ac0 --- /dev/null +++ b/docs/source/workflows/ipadapter.rst @@ -0,0 +1,27 @@ +FLUX IP-Adapter +=============== + +.. _nunchaku-flux.1-ip-adapter-json: + +nunchaku-flux.1-ip-adapter.json +------------------------------- + +.. image:: https://huggingface.co/datasets/nunchaku-tech/cdn/resolve/main/ComfyUI-nunchaku/workflows/nunchaku-flux.1-ip-adapter.png + :alt: nunchaku-flux.1-ip-adapter.json + :target: https://github.com/nunchaku-tech/ComfyUI-nunchaku/blob/main/example_workflows/nunchaku-flux.1-ip-adapter.json + +Image-based prompting workflow using `IP-Adapter `__ and the Nunchaku FLUX.1-dev model. + +**Links:** + +- Workflow: :download:`nunchaku-flux.1-ip-adapter.json ` +- Example input image: :download:`monalisa.jpg ` + +.. warning:: + This workflow is experimental and currently requires a large amount of VRAM. + It will automatically download the IP-Adapter and its associated CLIP models + from `Hugging Face `__ to the default cache directory. + At this time, specifying custom model paths is not supported. + +.. seealso:: + See nodes :ref:`nunchaku-flux-ip-adapter-loader` and :ref:`nunchaku-flux-ip-adapter-apply`. diff --git a/docs/source/workflows/toc.rst b/docs/source/workflows/toc.rst index 60b51b58..f5e083bb 100644 --- a/docs/source/workflows/toc.rst +++ b/docs/source/workflows/toc.rst @@ -12,4 +12,5 @@ Example Workflows canny.rst depth.rst fill.rst + ipadapter.rst tools.rst diff --git a/example_workflows/nunchaku-flux.1-ip-adapter.json b/example_workflows/nunchaku-flux.1-ip-adapter.json new file mode 100644 index 00000000..19b560c9 --- /dev/null +++ b/example_workflows/nunchaku-flux.1-ip-adapter.json @@ -0,0 +1,1014 @@ +{ + "id": "f253212e-0ec7-40c5-9671-bafc52d66023", + "revision": 0, + "last_node_id": 51, + "last_link_id": 138, + "nodes": [ + { + "id": 26, + "type": "FluxGuidance", + "pos": [ + 533.9339599609375, + 118.7322998046875 + ], + "size": [ + 317.4000244140625, + 58 + ], + "flags": {}, + "order": 11, + "mode": 0, + "inputs": [ + { + "name": "conditioning", + "type": "CONDITIONING", + "link": 41 + } + ], + "outputs": [ + { + "name": "CONDITIONING", + "type": "CONDITIONING", + "slot_index": 0, + "links": [ + 42 + ] + } + ], + "properties": { + "cnr_id": "comfy-core", + "ver": "0.3.24", + "Node name for S&R": "FluxGuidance" + }, + "widgets_values": [ + 3.5 + ], + "color": "#233", + "bgcolor": "#355" + }, + { + "id": 16, + "type": "KSamplerSelect", + "pos": [ + 813.7789306640625, + 847.5765991210938 + ], + "size": [ + 315, + 58 + ], + "flags": {}, + "order": 0, + "mode": 0, + "inputs": [], + "outputs": [ + { + "name": "SAMPLER", + "type": "SAMPLER", + "links": [ + 19 + ] + } + ], + "properties": { + "cnr_id": "comfy-core", + "ver": "0.3.24", + "Node name for S&R": "KSamplerSelect" + }, + "widgets_values": [ + "euler" + ] + }, + { + "id": 8, + "type": "VAEDecode", + "pos": [ + 938.151611328125, + 299.3258056640625 + ], + "size": [ + 210, + 46 + ], + "flags": {}, + "order": 17, + "mode": 0, + "inputs": [ + { + "name": "samples", + "type": "LATENT", + "link": 24 + }, + { + "name": "vae", + "type": "VAE", + "link": 12 + } + ], + "outputs": [ + { + "name": "IMAGE", + "type": "IMAGE", + "slot_index": 0, + "links": [ + 9 + ] + } + ], + "properties": { + "cnr_id": "comfy-core", + "ver": "0.3.24", + "Node name for S&R": "VAEDecode" + }, + "widgets_values": [] + }, + { + "id": 22, + "type": "BasicGuider", + "pos": [ + 691.55615234375, + 303.7747497558594 + ], + "size": [ + 222.3482666015625, + 46 + ], + "flags": {}, + "order": 14, + "mode": 0, + "inputs": [ + { + "name": "model", + "type": "MODEL", + "link": 54 + }, + { + "name": "conditioning", + "type": "CONDITIONING", + "link": 42 + } + ], + "outputs": [ + { + "name": "GUIDER", + "type": "GUIDER", + "slot_index": 0, + "links": [ + 30 + ] + } + ], + "properties": { + "cnr_id": "comfy-core", + "ver": "0.3.24", + "Node name for S&R": "BasicGuider" + }, + "widgets_values": [] + }, + { + "id": 13, + "type": "SamplerCustomAdvanced", + "pos": [ + 867.9028930664062, + 126.43718719482422 + ], + "size": [ + 272.3617858886719, + 124.53733825683594 + ], + "flags": {}, + "order": 16, + "mode": 0, + "inputs": [ + { + "name": "noise", + "type": "NOISE", + "link": 37 + }, + { + "name": "guider", + "type": "GUIDER", + "link": 30 + }, + { + "name": "sampler", + "type": "SAMPLER", + "link": 19 + }, + { + "name": "sigmas", + "type": "SIGMAS", + "link": 20 + }, + { + "name": "latent_image", + "type": "LATENT", + "link": 116 + } + ], + "outputs": [ + { + "name": "output", + "type": "LATENT", + "slot_index": 0, + "links": [ + 24 + ] + }, + { + "name": "denoised_output", + "type": "LATENT", + "links": null + } + ], + "properties": { + "cnr_id": "comfy-core", + "ver": "0.3.24", + "Node name for S&R": "SamplerCustomAdvanced" + }, + "widgets_values": [] + }, + { + "id": 34, + "type": "PrimitiveNode", + "pos": [ + 696.5196533203125, + 397.441650390625 + ], + "size": [ + 210, + 82 + ], + "flags": {}, + "order": 1, + "mode": 0, + "inputs": [], + "outputs": [ + { + "name": "INT", + "type": "INT", + "widget": { + "name": "width" + }, + "slot_index": 0, + "links": [ + 112, + 115 + ] + } + ], + "title": "width", + "properties": { + "Run widget replace on values": false + }, + "widgets_values": [ + 1024, + "fixed" + ], + "color": "#323", + "bgcolor": "#535" + }, + { + "id": 35, + "type": "PrimitiveNode", + "pos": [ + 937.396484375, + 392.3027038574219 + ], + "size": [ + 210, + 86.4900131225586 + ], + "flags": {}, + "order": 2, + "mode": 0, + "inputs": [], + "outputs": [ + { + "name": "INT", + "type": "INT", + "widget": { + "name": "height" + }, + "slot_index": 0, + "links": [ + 113, + 114 + ] + } + ], + "title": "height", + "properties": { + "Run widget replace on values": false + }, + "widgets_values": [ + 1024, + "fixed" + ], + "color": "#323", + "bgcolor": "#535" + }, + { + "id": 30, + "type": "ModelSamplingFlux", + "pos": [ + 879.5872802734375, + 1107.096923828125 + ], + "size": [ + 210, + 130 + ], + "flags": {}, + "order": 13, + "mode": 0, + "inputs": [ + { + "name": "model", + "type": "MODEL", + "link": 136 + }, + { + "name": "width", + "type": "INT", + "widget": { + "name": "width" + }, + "link": 115 + }, + { + "name": "height", + "type": "INT", + "widget": { + "name": "height" + }, + "link": 114 + } + ], + "outputs": [ + { + "name": "MODEL", + "type": "MODEL", + "slot_index": 0, + "links": [ + 54, + 55 + ] + } + ], + "properties": { + "cnr_id": "comfy-core", + "ver": "0.3.24", + "Node name for S&R": "ModelSamplingFlux" + }, + "widgets_values": [ + 1.15, + 0.5, + 1024, + 1024 + ] + }, + { + "id": 17, + "type": "BasicScheduler", + "pos": [ + 818.5535888671875, + 948.8093872070312 + ], + "size": [ + 315, + 106 + ], + "flags": {}, + "order": 15, + "mode": 0, + "inputs": [ + { + "name": "model", + "type": "MODEL", + "link": 55 + } + ], + "outputs": [ + { + "name": "SIGMAS", + "type": "SIGMAS", + "links": [ + 20 + ] + } + ], + "properties": { + "cnr_id": "comfy-core", + "ver": "0.3.24", + "Node name for S&R": "BasicScheduler" + }, + "widgets_values": [ + "simple", + 50, + 1 + ] + }, + { + "id": 48, + "type": "NunchakuIPAdapterLoader", + "pos": [ + 204.7442169189453, + 766.4160766601562 + ], + "size": [ + 221.23690795898438, + 46 + ], + "flags": {}, + "order": 10, + "mode": 0, + "inputs": [ + { + "name": "model", + "type": "MODEL", + "link": 133 + } + ], + "outputs": [ + { + "name": "MODEL", + "type": "MODEL", + "links": [ + 134 + ] + }, + { + "name": "IPADAPTER_PIPELINE", + "type": "IPADAPTER_PIPELINE", + "links": [ + 135 + ] + } + ], + "properties": { + "aux_id": "nunchaku-tech/ComfyUI-nunchaku", + "ver": "a87c20120f155ffaf39a4b7f12094499690e70e2", + "Node name for S&R": "NunchakuIPAdapterLoader", + "cnr_id": "ComfyUI-nunchaku" + }, + "widgets_values": [] + }, + { + "id": 50, + "type": "LoadImage", + "pos": [ + 496.5161437988281, + 871.9661865234375 + ], + "size": [ + 270, + 314 + ], + "flags": {}, + "order": 3, + "mode": 0, + "inputs": [], + "outputs": [ + { + "name": "IMAGE", + "type": "IMAGE", + "links": [ + 137 + ] + }, + { + "name": "MASK", + "type": "MASK", + "links": null + } + ], + "properties": { + "cnr_id": "comfy-core", + "ver": "0.3.40", + "Node name for S&R": "LoadImage" + }, + "widgets_values": [ + "monalisa.jpg", + "image" + ] + }, + { + "id": 6, + "type": "CLIPTextEncode", + "pos": [ + 261.212158203125, + 225.61419677734375 + ], + "size": [ + 422.84503173828125, + 164.31304931640625 + ], + "flags": {}, + "order": 9, + "mode": 0, + "inputs": [ + { + "name": "clip", + "type": "CLIP", + "link": 138 + } + ], + "outputs": [ + { + "name": "CONDITIONING", + "type": "CONDITIONING", + "slot_index": 0, + "links": [ + 41 + ] + } + ], + "title": "CLIP Text Encode (Positive Prompt)", + "properties": { + "cnr_id": "comfy-core", + "ver": "0.3.24", + "Node name for S&R": "CLIPTextEncode" + }, + "widgets_values": [ + "holding an sign saying 'SVDQuant is fast!'" + ], + "color": "#232", + "bgcolor": "#353" + }, + { + "id": 51, + "type": "NunchakuTextEncoderLoaderV2", + "pos": [ + -87.10687255859375, + 237.3723602294922 + ], + "size": [ + 317.0367126464844, + 130 + ], + "flags": {}, + "order": 4, + "mode": 0, + "inputs": [], + "outputs": [ + { + "name": "CLIP", + "type": "CLIP", + "links": [ + 138 + ] + } + ], + "properties": { + "aux_id": "nunchaku-tech/ComfyUI-nunchaku", + "ver": "4b53bcf64d12c34b338cfde8abd3b52a37fe71a0", + "Node name for S&R": "NunchakuTextEncoderLoaderV2" + }, + "widgets_values": [ + "flux.1", + "t5xxl_fp16.safetensors", + "clip_l.safetensors", + 512 + ] + }, + { + "id": 49, + "type": "NunchakuFluxIPAdapterApply", + "pos": [ + 476.3758239746094, + 728.8355102539062 + ], + "size": [ + 312.3404235839844, + 98 + ], + "flags": {}, + "order": 12, + "mode": 0, + "inputs": [ + { + "name": "model", + "type": "MODEL", + "link": 134 + }, + { + "name": "ipadapter_pipeline", + "type": "IPADAPTER_PIPELINE", + "link": 135 + }, + { + "name": "image", + "type": "IMAGE", + "link": 137 + } + ], + "outputs": [ + { + "name": "MODEL", + "type": "MODEL", + "links": [ + 136 + ] + } + ], + "properties": { + "aux_id": "nunchaku-tech/ComfyUI-nunchaku", + "ver": "a87c20120f155ffaf39a4b7f12094499690e70e2", + "Node name for S&R": "NunchakuFluxIPAdapterApply", + "cnr_id": "ComfyUI-nunchaku" + }, + "widgets_values": [ + 1.1500000000000001 + ] + }, + { + "id": 25, + "type": "RandomNoise", + "pos": [ + 819.1885986328125, + 709.9674072265625 + ], + "size": [ + 315, + 82 + ], + "flags": {}, + "order": 5, + "mode": 0, + "inputs": [], + "outputs": [ + { + "name": "NOISE", + "type": "NOISE", + "links": [ + 37 + ] + } + ], + "properties": { + "cnr_id": "comfy-core", + "ver": "0.3.24", + "Node name for S&R": "RandomNoise" + }, + "widgets_values": [ + 952572216486857, + "fixed" + ], + "color": "#2a363b", + "bgcolor": "#3f5159" + }, + { + "id": 9, + "type": "SaveImage", + "pos": [ + 1170.8551025390625, + 128.45352172851562 + ], + "size": [ + 985.3012084960938, + 1060.3828125 + ], + "flags": {}, + "order": 18, + "mode": 0, + "inputs": [ + { + "name": "images", + "type": "IMAGE", + "link": 9 + } + ], + "outputs": [], + "properties": { + "cnr_id": "comfy-core", + "ver": "0.3.24" + }, + "widgets_values": [ + "ComfyUI" + ] + }, + { + "id": 27, + "type": "EmptySD3LatentImage", + "pos": [ + 824.7072143554688, + 515.9119873046875 + ], + "size": [ + 315, + 126 + ], + "flags": {}, + "order": 8, + "mode": 0, + "inputs": [ + { + "name": "width", + "type": "INT", + "widget": { + "name": "width" + }, + "link": 112 + }, + { + "name": "height", + "type": "INT", + "widget": { + "name": "height" + }, + "link": 113 + } + ], + "outputs": [ + { + "name": "LATENT", + "type": "LATENT", + "slot_index": 0, + "links": [ + 116 + ] + } + ], + "properties": { + "cnr_id": "comfy-core", + "ver": "0.3.24", + "Node name for S&R": "EmptySD3LatentImage" + }, + "widgets_values": [ + 1024, + 1024, + 1 + ] + }, + { + "id": 45, + "type": "NunchakuFluxDiTLoader", + "pos": [ + 137.85166931152344, + 462.1759338378906 + ], + "size": [ + 315, + 202 + ], + "flags": {}, + "order": 6, + "mode": 0, + "inputs": [], + "outputs": [ + { + "name": "MODEL", + "type": "MODEL", + "links": [ + 133 + ] + } + ], + "properties": { + "aux_id": "mit-han-lab/ComfyUI-nunchaku", + "ver": "dba78a94a199a1888c43b3c343d0ceec4843afc5", + "Node name for S&R": "NunchakuFluxDiTLoader", + "cnr_id": "ComfyUI-nunchaku" + }, + "widgets_values": [ + "svdq-fp4_r32-flux.1-dev.safetensors", + 0.10000000000000002, + "nunchaku-fp16", + "auto", + 0, + "bfloat16", + "enabled" + ] + }, + { + "id": 10, + "type": "VAELoader", + "pos": [ + 489.78961181640625, + 528.6732177734375 + ], + "size": [ + 311.81634521484375, + 60.429901123046875 + ], + "flags": {}, + "order": 7, + "mode": 0, + "inputs": [], + "outputs": [ + { + "name": "VAE", + "type": "VAE", + "slot_index": 0, + "links": [ + 12 + ] + } + ], + "properties": { + "cnr_id": "comfy-core", + "ver": "0.3.24", + "Node name for S&R": "VAELoader" + }, + "widgets_values": [ + "ae.safetensors" + ] + } + ], + "links": [ + [ + 9, + 8, + 0, + 9, + 0, + "IMAGE" + ], + [ + 12, + 10, + 0, + 8, + 1, + "VAE" + ], + [ + 19, + 16, + 0, + 13, + 2, + "SAMPLER" + ], + [ + 20, + 17, + 0, + 13, + 3, + "SIGMAS" + ], + [ + 24, + 13, + 0, + 8, + 0, + "LATENT" + ], + [ + 30, + 22, + 0, + 13, + 1, + "GUIDER" + ], + [ + 37, + 25, + 0, + 13, + 0, + "NOISE" + ], + [ + 41, + 6, + 0, + 26, + 0, + "CONDITIONING" + ], + [ + 42, + 26, + 0, + 22, + 1, + "CONDITIONING" + ], + [ + 54, + 30, + 0, + 22, + 0, + "MODEL" + ], + [ + 55, + 30, + 0, + 17, + 0, + "MODEL" + ], + [ + 112, + 34, + 0, + 27, + 0, + "INT" + ], + [ + 113, + 35, + 0, + 27, + 1, + "INT" + ], + [ + 114, + 35, + 0, + 30, + 2, + "INT" + ], + [ + 115, + 34, + 0, + 30, + 1, + "INT" + ], + [ + 116, + 27, + 0, + 13, + 4, + "LATENT" + ], + [ + 133, + 45, + 0, + 48, + 0, + "MODEL" + ], + [ + 134, + 48, + 0, + 49, + 0, + "MODEL" + ], + [ + 135, + 48, + 1, + 49, + 1, + "IPADAPTER_PIPELINE" + ], + [ + 136, + 49, + 0, + 30, + 0, + "MODEL" + ], + [ + 137, + 50, + 0, + 49, + 2, + "IMAGE" + ], + [ + 138, + 51, + 0, + 6, + 0, + "CLIP" + ] + ], + "groups": [], + "config": {}, + "extra": { + "ds": { + "scale": 0.8954302432552416, + "offset": [ + 137.3477047670991, + -39.91926725644665 + ] + }, + "frontendVersion": "1.23.4", + "groupNodes": {}, + "node_versions": { + "comfy-core": "0.3.24" + } + }, + "version": 0.4 +} diff --git a/nodes/models/ipadapter.py b/nodes/models/ipadapter.py new file mode 100644 index 00000000..51c95eb6 --- /dev/null +++ b/nodes/models/ipadapter.py @@ -0,0 +1,221 @@ +""" +This module provides nodes load and apply IP-Adapter models +to FLUX pipelines, enabling image-based conditioning for generative models. +""" + +import logging +import os +from typing import Any, List, Optional + +import torch +from diffusers import FluxPipeline +from torchvision import transforms + +from nunchaku.models.ip_adapter.diffusers_adapters import apply_IPA_on_pipe +from nunchaku.models.ip_adapter.utils import undo_all_mods_on_transformer + +from .utils import set_extra_config_model_path + +log_level = os.getenv("LOG_LEVEL", "INFO").upper() + +logging.basicConfig(level=getattr(logging, log_level, logging.INFO), format="%(asctime)s - %(levelname)s - %(message)s") +logger = logging.getLogger(__name__) + + +class IPAFluxPipelineWrapper(FluxPipeline): + """ + FluxPipeline wrapper with IP-Adapter support. + """ + + @torch.no_grad() + def get_image_embeds( + self, + num_images_per_prompt: int = 1, + ip_adapter_image: Optional[Any] = None, + ip_adapter_image_embeds: Optional[List[torch.Tensor]] = None, + negative_ip_adapter_image: Optional[Any] = None, + negative_ip_adapter_image_embeds: Optional[List[torch.Tensor]] = None, + ) -> (Optional[torch.Tensor], Optional[torch.Tensor]): + """ + Compute image and negative image embeddings for IP-Adapter. + + Parameters + ---------- + num_images_per_prompt : int, optional + Number of images per prompt (default is 1). + ip_adapter_image : Any, optional + Input image for positive conditioning. + ip_adapter_image_embeds : list of torch.Tensor, optional + Precomputed positive image embeddings. + negative_ip_adapter_image : Any, optional + Input image for negative conditioning. + negative_ip_adapter_image_embeds : list of torch.Tensor, optional + Precomputed negative image embeddings. + + Returns + ------- + image_embeds : torch.Tensor or None + Positive image embeddings. + negative_image_embeds : torch.Tensor or None + Negative image embeddings. + """ + batch_size = 1 + + device = self.transformer.device + + image_embeds = None + if ip_adapter_image is not None or ip_adapter_image_embeds is not None: + image_embeds = self.prepare_ip_adapter_image_embeds( + ip_adapter_image=ip_adapter_image, + ip_adapter_image_embeds=ip_adapter_image_embeds, + device=device, + num_images_per_prompt=batch_size * num_images_per_prompt, + ) + image_embeds = self.transformer.encoder_hid_proj(image_embeds) + + negative_image_embeds = None + if negative_ip_adapter_image is not None or negative_ip_adapter_image_embeds is not None: + negative_image_embeds = self.prepare_ip_adapter_image_embeds( + ip_adapter_image=negative_ip_adapter_image, + ip_adapter_image_embeds=negative_ip_adapter_image_embeds, + device=device, + num_images_per_prompt=batch_size * num_images_per_prompt, + ) + negative_image_embeds = self.transformer.encoder_hid_proj(negative_image_embeds) + + return image_embeds, negative_image_embeds + + +set_extra_config_model_path("ipadapter", "ipadapter") +set_extra_config_model_path("clip", "clip") + + +class NunchakuIPAdapterLoader: + """ + Node for loading Nunchaku IP-Adapter pipelines. + + .. warning:: + This node will automatically download the IP-Adapter and associated CLIP models from Hugging Face. + Custom model paths are not supported for now. + """ + + @classmethod + def INPUT_TYPES(s): + """ + Defines the input types and tooltips for the node. + + Returns + ------- + dict + A dictionary specifying the required inputs and their descriptions for the node interface. + """ + return { + "required": { + "model": ("MODEL", {"tooltip": "The nunchaku model."}), + } + } + + RETURN_TYPES = ("MODEL", "IPADAPTER_PIPELINE") + FUNCTION = "load" + CATEGORY = "Nunchaku" + TITLE = "Nunchaku IP-Adapter Loader" + + def load(self, model): + """ + Load the IP-Adapter pipeline and attach it to the given model. + + Parameters + ---------- + model : object + The Nunchaku model to which the IP-Adapter will be attached. + It should be loaded with :class:`~comfyui_nunchaku.nodes.models.flux.NunchakuFluxDiTLoader`. + + Returns + ------- + tuple + The original model and the loaded IP-Adapter pipeline. + """ + device = model.model.diffusion_model.model.device + pipeline = IPAFluxPipelineWrapper.from_pretrained( + "black-forest-labs/FLUX.1-dev", transformer=model.model.diffusion_model.model, torch_dtype=torch.bfloat16 + ).to(device) + + pipeline.load_ip_adapter( + pretrained_model_name_or_path_or_dict="XLabs-AI/flux-ip-adapter-v2", + weight_name="ip_adapter.safetensors", + image_encoder_pretrained_model_name_or_path="openai/clip-vit-large-patch14", + ) + return (model, pipeline) + + +class NunchakuFluxIPAdapterApply: + """ + Node for applying IP-Adapter to a Nunchaku model using a given image and weight. + """ + + @classmethod + def INPUT_TYPES(s): + """ + Defines the input types and tooltips for the node. + + Returns + ------- + dict + A dictionary specifying the required inputs and their descriptions for the node interface. + """ + return { + "required": { + "model": ("MODEL",), + "ipadapter_pipeline": ("IPADAPTER_PIPELINE",), + "image": ("IMAGE",), + "weight": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 5.0, "step": 0.05}), + }, + } + + RETURN_TYPES = ("MODEL",) + FUNCTION = "apply_ipa" + CATEGORY = "Nunchaku" + TITLE = "Nunchaku FLUX IP-Adapter Apply" + + def apply_ipa( + self, + model, + ipadapter_pipeline: IPAFluxPipelineWrapper, + image, + weight: float, + ): + """ + Apply the IP-Adapter to the given model using the provided image and weight. + + Parameters + ---------- + model : object + The Nunchaku model to modify. + ipadapter_pipeline : IPAFluxPipelineWrapper + The IP-Adapter pipeline. + image : torch.Tensor + The input image tensor. + weight : float + The scale/weight for the IP-Adapter. + + Returns + ------- + tuple + The modified model. + """ + to_pil_transformer = transforms.ToPILImage() + image_tensor_chw = image[0].permute(2, 0, 1) + pil_image = to_pil_transformer(image_tensor_chw) + + image_embeds, _ = ipadapter_pipeline.get_image_embeds( + ip_adapter_image=pil_image, + ) + + undo_all_mods_on_transformer(ipadapter_pipeline.transformer) + apply_IPA_on_pipe(ipadapter_pipeline, ip_adapter_scale=weight, repo_id="XLabs-AI/flux-ip-adapter-v2") + + ipadapter_pipeline.transformer.transformer_blocks[0].set_ip_hidden_states(image_embeds=image_embeds) + + model.model.diffusion_model.model = ipadapter_pipeline.transformer + + return (model,) diff --git a/nodes/models/pulid.py b/nodes/models/pulid.py index a58c9673..2e27d24a 100644 --- a/nodes/models/pulid.py +++ b/nodes/models/pulid.py @@ -22,6 +22,7 @@ from nunchaku.pipeline.pipeline_flux_pulid import PuLIDPipeline from ...wrappers.flux import ComfyFluxWrapper +from .utils import set_extra_config_model_path # Get log level from environment variable (default to INFO) log_level = os.getenv("LOG_LEVEL", "INFO").upper() @@ -31,29 +32,6 @@ logger = logging.getLogger(__name__) -def set_extra_config_model_path(extra_config_models_dir_key, models_dir_name: str): - """ - Register an extra model directory (`pulid`, `insightface`, `facexlib`) with ComfyUI's folder_paths. - - Parameters - ---------- - extra_config_models_dir_key : str - The key to register the model directory under. - models_dir_name : str - The name of the subdirectory to use for models. - """ - models_dir_default = os.path.join(folder_paths.models_dir, models_dir_name) - if extra_config_models_dir_key not in folder_paths.folder_names_and_paths: - folder_paths.folder_names_and_paths[extra_config_models_dir_key] = ( - [os.path.join(folder_paths.models_dir, models_dir_name)], - folder_paths.supported_pt_extensions, - ) - else: - if not os.path.exists(models_dir_default): - os.makedirs(models_dir_default, exist_ok=True) - folder_paths.add_model_folder_path(extra_config_models_dir_key, models_dir_default, is_default=True) - - set_extra_config_model_path("pulid", "pulid") set_extra_config_model_path("insightface", "insightface") set_extra_config_model_path("facexlib", "facexlib") diff --git a/nodes/models/utils.py b/nodes/models/utils.py new file mode 100644 index 00000000..cb1861f1 --- /dev/null +++ b/nodes/models/utils.py @@ -0,0 +1,26 @@ +import os + +import folder_paths + + +def set_extra_config_model_path(extra_config_models_dir_key, models_dir_name: str): + """ + Register an extra model directory (e.g., ``pulid``, ``insightface``, ``facexlib``) with ComfyUI's folder_paths. + + Parameters + ---------- + extra_config_models_dir_key : str + The key to register the model directory under. + models_dir_name : str + The name of the subdirectory to use for models. + """ + models_dir_default = os.path.join(folder_paths.models_dir, models_dir_name) + if extra_config_models_dir_key not in folder_paths.folder_names_and_paths: + folder_paths.folder_names_and_paths[extra_config_models_dir_key] = ( + [os.path.join(folder_paths.models_dir, models_dir_name)], + folder_paths.supported_pt_extensions, + ) + else: + if not os.path.exists(models_dir_default): + os.makedirs(models_dir_default, exist_ok=True) + folder_paths.add_model_folder_path(extra_config_models_dir_key, models_dir_default, is_default=True) diff --git a/requirements.txt b/requirements.txt index 1aa029a8..5d79ee12 100644 --- a/requirements.txt +++ b/requirements.txt @@ -1,9 +1,10 @@ -diffusers>=0.33.1 +diffusers>=0.34 accelerate sentencepiece protobuf huggingface_hub tomli +peft>=0.15 # PuLID dependencies insightface diff --git a/test_data/images.yaml b/test_data/images.yaml index 9c6791a1..6ed728b2 100644 --- a/test_data/images.yaml +++ b/test_data/images.yaml @@ -7,6 +7,7 @@ images: - strawberry-mask.png - mushroom_depth.webp - removal.png + - monalisa.jpg - base_url: "https://github.com/ToTheBeginning/PuLID/blob/main/example_inputs/{filename}?raw=true" files: - lecun.jpg diff --git a/tests/scripts/nunchaku-flux1-ip-adapter.py b/tests/scripts/nunchaku-flux1-ip-adapter.py new file mode 100644 index 00000000..34b1c2fb --- /dev/null +++ b/tests/scripts/nunchaku-flux1-ip-adapter.py @@ -0,0 +1,231 @@ +import os +import sys +from typing import Any, Mapping, Sequence, Union + +import torch + +from nunchaku.utils import get_precision + + +def get_value_at_index(obj: Union[Sequence, Mapping], index: int) -> Any: + """Returns the value at the given index of a sequence or mapping. + + If the object is a sequence (like list or string), returns the value at the given index. + If the object is a mapping (like a dictionary), returns the value at the index-th key. + + Some return a dictionary, in these cases, we look for the "results" key + + Args: + obj (Union[Sequence, Mapping]): The object to retrieve the value from. + index (int): The index of the value to retrieve. + + Returns: + Any: The value at the given index. + + Raises: + IndexError: If the index is out of bounds for the object and the object is not a mapping. + """ + try: + return obj[index] + except KeyError: + return obj["result"][index] + + +def find_path(name: str, path: str = None) -> str: + """ + Recursively looks at parent folders starting from the given path until it finds the given name. + Returns the path as a Path object if found, or None otherwise. + """ + # If no path is given, use the current working directory + if path is None: + path = os.getcwd() + + # Check if the current directory contains the name + if name in os.listdir(path): + path_name = os.path.join(path, name) + print(f"{name} found: {path_name}") + return path_name + + # Get the parent directory + parent_directory = os.path.dirname(path) + + # If the parent directory is the same as the current directory, we've reached the root and stop the search + if parent_directory == path: + return None + + # Recursively call the function with the parent directory + return find_path(name, parent_directory) + + +def add_comfyui_directory_to_sys_path() -> None: + """ + Add 'ComfyUI' to the sys.path + """ + comfyui_path = find_path("ComfyUI") + if comfyui_path is not None and os.path.isdir(comfyui_path): + sys.path.append(comfyui_path) + print(f"'{comfyui_path}' added to sys.path") + + +def add_extra_model_paths() -> None: + """ + Parse the optional extra_model_paths.yaml file and add the parsed paths to the sys.path. + """ + try: + from main import load_extra_path_config + except ImportError: + print("Could not import load_extra_path_config from main.py. Looking in utils.extra_config instead.") + from utils.extra_config import load_extra_path_config + + extra_model_paths = find_path("extra_model_paths.yaml") + + if extra_model_paths is not None: + load_extra_path_config(extra_model_paths) + else: + print("Could not find the extra_model_paths config file.") + + +add_comfyui_directory_to_sys_path() +add_extra_model_paths() + + +def import_custom_nodes() -> None: + """Find all custom nodes in the custom_nodes folder and add those node objects to NODE_CLASS_MAPPINGS + + This function sets up a new asyncio event loop, initializes the PromptServer, + creates a PromptQueue, and initializes the custom nodes. + """ + import asyncio + + import execution + import server + + from nodes import init_extra_nodes + + # Creating a new event loop and setting it as the default loop + loop = asyncio.new_event_loop() + asyncio.set_event_loop(loop) + + # Creating an instance of PromptServer with the loop + server_instance = server.PromptServer(loop) + execution.PromptQueue(server_instance) + + # Initializing custom nodes + init_extra_nodes() + + +from nodes import NODE_CLASS_MAPPINGS + + +def main(precision: str): + import_custom_nodes() + with torch.inference_mode(): + nunchakutextencoderloaderv2 = NODE_CLASS_MAPPINGS["NunchakuTextEncoderLoaderV2"]() + nunchakutextencoderloaderv2_51 = nunchakutextencoderloaderv2.load_text_encoder( + model_type="flux.1", + text_encoder1="t5xxl_fp8_e4m3fn_scaled.safetensors", + text_encoder2="clip_l.safetensors", + t5_min_length=512, + ) + + cliptextencode = NODE_CLASS_MAPPINGS["CLIPTextEncode"]() + cliptextencode_6 = cliptextencode.encode( + text="holding an sign saying 'SVDQuant is fast!'", + clip=get_value_at_index(nunchakutextencoderloaderv2_51, 0), + ) + + vaeloader = NODE_CLASS_MAPPINGS["VAELoader"]() + vaeloader_10 = vaeloader.load_vae(vae_name="ae.safetensors") + + ksamplerselect = NODE_CLASS_MAPPINGS["KSamplerSelect"]() + ksamplerselect_16 = ksamplerselect.get_sampler(sampler_name="euler") + + randomnoise = NODE_CLASS_MAPPINGS["RandomNoise"]() + randomnoise_25 = randomnoise.get_noise(noise_seed=952572216486857) + + emptysd3latentimage = NODE_CLASS_MAPPINGS["EmptySD3LatentImage"]() + emptysd3latentimage_27 = emptysd3latentimage.generate(width=1024, height=1024, batch_size=1) + + nunchakufluxditloader = NODE_CLASS_MAPPINGS["NunchakuFluxDiTLoader"]() + nunchakufluxditloader_45 = nunchakufluxditloader.load_model( + model_path=f"svdq-{precision}_r32-flux.1-dev.safetensors", + cache_threshold=0.10000000000000002, + attention="nunchaku-fp16", + cpu_offload="enable", + device_id=0, + data_type="bfloat16", + i2f_mode="enabled", + ) + + loadimage = NODE_CLASS_MAPPINGS["LoadImage"]() + loadimage_50 = loadimage.load_image(image="monalisa.jpg") + + nunchakuipadapterloader = NODE_CLASS_MAPPINGS["NunchakuIPAdapterLoader"]() + nunchakufluxipadapterapply = NODE_CLASS_MAPPINGS["NunchakuFluxIPAdapterApply"]() + modelsamplingflux = NODE_CLASS_MAPPINGS["ModelSamplingFlux"]() + fluxguidance = NODE_CLASS_MAPPINGS["FluxGuidance"]() + basicguider = NODE_CLASS_MAPPINGS["BasicGuider"]() + basicscheduler = NODE_CLASS_MAPPINGS["BasicScheduler"]() + samplercustomadvanced = NODE_CLASS_MAPPINGS["SamplerCustomAdvanced"]() + vaedecode = NODE_CLASS_MAPPINGS["VAEDecode"]() + saveimage = NODE_CLASS_MAPPINGS["SaveImage"]() + + for q in range(1): + nunchakuipadapterloader_48 = nunchakuipadapterloader.load( + model=get_value_at_index(nunchakufluxditloader_45, 0) + ) + + nunchakufluxipadapterapply_49 = nunchakufluxipadapterapply.apply_ipa( + weight=1.1500000000000001, + model=get_value_at_index(nunchakuipadapterloader_48, 0), + ipadapter_pipeline=get_value_at_index(nunchakuipadapterloader_48, 1), + image=get_value_at_index(loadimage_50, 0), + ) + + modelsamplingflux_30 = modelsamplingflux.patch( + max_shift=1.15, + base_shift=0.5, + width=1024, + height=1024, + model=get_value_at_index(nunchakufluxipadapterapply_49, 0), + ) + + fluxguidance_26 = fluxguidance.append(guidance=3.5, conditioning=get_value_at_index(cliptextencode_6, 0)) + + basicguider_22 = basicguider.get_guider( + model=get_value_at_index(modelsamplingflux_30, 0), + conditioning=get_value_at_index(fluxguidance_26, 0), + ) + + basicscheduler_17 = basicscheduler.get_sigmas( + scheduler="simple", + steps=50, + denoise=1, + model=get_value_at_index(modelsamplingflux_30, 0), + ) + + samplercustomadvanced_13 = samplercustomadvanced.sample( + noise=get_value_at_index(randomnoise_25, 0), + guider=get_value_at_index(basicguider_22, 0), + sampler=get_value_at_index(ksamplerselect_16, 0), + sigmas=get_value_at_index(basicscheduler_17, 0), + latent_image=get_value_at_index(emptysd3latentimage_27, 0), + ) + + vaedecode_8 = vaedecode.decode( + samples=get_value_at_index(samplercustomadvanced_13, 0), + vae=get_value_at_index(vaeloader_10, 0), + ) + + saveimage_9 = saveimage.save_images(filename_prefix="ComfyUI", images=get_value_at_index(vaedecode_8, 0)) + + filename = saveimage_9["ui"]["images"][0]["filename"] + path = os.path.join("output", filename) + with open("image_path.txt", "w") as f: + f.write(path) + print(path) + return path + + +if __name__ == "__main__": + main(get_precision()) diff --git a/tests/test_workflows.py b/tests/test_workflows.py index db2ee2fb..712c4038 100644 --- a/tests/test_workflows.py +++ b/tests/test_workflows.py @@ -35,6 +35,7 @@ ("nunchaku-flux1-dev-pulid.py", 0.9, 0.194, 15.8), ("nunchaku-flux1-kontext-dev.py", 0.9, 0.1, 18.3), ("nunchaku-flux1-kontext-dev-turbo_lora.py", 0.87, 0.13, 18.8), + ("nunchaku-flux1-ip-adapter.py", 0.6, 0.25, 19), ], ) @pytest.mark.flaky(reruns=2, reruns_delay=0) @@ -43,7 +44,7 @@ def test_workflows(script_name: str, expected_clip_iqa: float, expected_lpips: f torch.cuda.empty_cache() script_path = os.path.join(script_dir, script_name) - result = subprocess.run(["python", script_path], capture_output=True, text=True) + result = subprocess.run(["python", script_path]) print(f"Running {script_path} -> Return code: {result.returncode}") assert result.returncode == 0, f"{script_path} failed with code {result.returncode}" diff --git a/utils.py b/utils.py index 1c2eab98..83737452 100644 --- a/utils.py +++ b/utils.py @@ -82,4 +82,4 @@ def get_plugin_version() -> str: return project_version -supported_versions = ["v0.3.1"] +supported_versions = ["v0.3.1", "v0.3.2"]