Add Onediff ComfyUI nodes (#397)

diffusers 中主要修改： - .github/workflows/examples.yml - .gitignore - tests/comfy-docker-compose.yml --------- Co-authored-by: Yao Chi <[email protected]>
siliconflow · Dec 5, 2023 · 70bf429 · 70bf429
1 parent 659f194
commit 70bf429
Show file tree

Hide file tree

Showing 33 changed files with 1,720 additions and 24 deletions.
diff --git a/.github/workflows/examples.yml b/.github/workflows/examples.yml
@@ -11,7 +11,6 @@ env:
   REGION_ID: cn-beijing
   ACR_ORG: registry.cn-beijing.aliyuncs.com/oneflow
   COMFYUI_SRC_DIR: ComfyUI
-  COMFYUI_SPEEDUP_DIR: comfyui-speedup
   SDXL_BASE: /share_nfs/hf_models/sd_xl_base_1.0.safetensors
   UNET_INT8: /share_nfs/hf_models/unet_int8
   SELENIUM_CONTAINER_NAME: selenium-test
@@ -54,13 +53,6 @@ jobs:
         with:
           repository: comfyanonymous/ComfyUI
           path: ${{ env.COMFYUI_SRC_DIR }}
-      - name: Checkout ConfyUI Speedup Nodes
-        uses: actions/checkout@v2
-        if: matrix.test-suite == 'comfy'
-        with:
-          repository: ${{ secrets.ONEFLOW_PRIV_ORG }}/comfyui-speedup
-          token: ${{ secrets.ONEFLOW_PRIV_GH_TOKEN }}
-          path: ${{ env.COMFYUI_SPEEDUP_DIR }}
 
       - run: |
           docker rm -f ${{ env.CONTAINER_NAME }} || true
@@ -79,7 +71,6 @@ jobs:
           ACR_ORG: ${{ env.ACR_ORG }}
           MATRIX_IMAGE: ${{ matrix.image }}
           COMFYUI_SRC_DIR: ${{ env.COMFYUI_SRC_DIR }}
-          COMFYUI_SPEEDUP_DIR: ${{ env.COMFYUI_SPEEDUP_DIR }}
           SDXL_BASE: ${{ env.SDXL_BASE }}
           UNET_INT8: ${{ env.UNET_INT8 }}
           SELENIUM_IMAGE: ${{ env.SELENIUM_IMAGE }}
@@ -93,8 +84,6 @@ jobs:
           CONTAINER_NAME: ${{ env.CONTAINER_NAME }}
           ACR_ORG: ${{ env.ACR_ORG }}
           MATRIX_IMAGE: ${{ matrix.image }}
-          COMFYUI_SRC_DIR: ${{ env.COMFYUI_SRC_DIR }}
-          COMFYUI_SPEEDUP_DIR: ${{ env.COMFYUI_SPEEDUP_DIR }}
           SDXL_BASE: ${{ env.SDXL_BASE }}
           UNET_INT8: ${{ env.UNET_INT8 }}
 
@@ -143,7 +132,6 @@ jobs:
           ACR_ORG: ${{ env.ACR_ORG }}
           MATRIX_IMAGE: ${{ matrix.image }}
           COMFYUI_SRC_DIR: ${{ env.COMFYUI_SRC_DIR }}
-          COMFYUI_SPEEDUP_DIR: ${{ env.COMFYUI_SPEEDUP_DIR }}
           SDXL_BASE: ${{ env.SDXL_BASE }}
           UNET_INT8: ${{ env.UNET_INT8 }}
           SELENIUM_IMAGE: ${{ env.SELENIUM_IMAGE }}
@@ -157,7 +145,5 @@ jobs:
           CONTAINER_NAME: ${{ env.CONTAINER_NAME }}
           ACR_ORG: ${{ env.ACR_ORG }}
           MATRIX_IMAGE: ${{ matrix.image }}
-          COMFYUI_SRC_DIR: ${{ env.COMFYUI_SRC_DIR }}
-          COMFYUI_SPEEDUP_DIR: ${{ env.COMFYUI_SPEEDUP_DIR }}
           SDXL_BASE: ${{ env.SDXL_BASE }}
           UNET_INT8: ${{ env.UNET_INT8 }}
diff --git a/.gitignore b/.gitignore
@@ -174,5 +174,6 @@ log
 unet_graphs
 *.json
 
-ComfyUI/
-comfyui-speedup/
+# onediff_comfy_nodes
+*.pt
+*.graph
diff --git a/README.md b/README.md
@@ -36,7 +36,8 @@ Updated on Nov 6, 2023.
 
 ## Business inquiry on OneDiff Enterprise Edition
 
-If you need **LoRA**, [**ComfyUI**](https://github.com/Oneflow-Inc/onediff/wiki/Run-ComfyUI-with-OneDiff), **quant** support or any other more advanced features, please send an email to [email protected] . Tell us about your **use case, deployment scale and requirements**! 
+If you need **unrestricted multiple resolution**, **quantization** support or any other more advanced features, please send an email to [email protected] . Tell us about your **use case, deployment scale and requirements**! 
+
 |                      | OneDiff Community   | OneDiff Enterprise|
 | -------------------- | ------------------- | ----------- |
 | diffusers            | Yes                 | Yes         |
@@ -45,7 +46,7 @@ If you need **LoRA**, [**ComfyUI**](https://github.com/Oneflow-Inc/onediff/wiki/
 | LCM                  | Limited             | Yes         |
 | Multiple Resolutions | Limited             | Yes         |
 | Technical Support    | Community           | Yes         |
-| ComfyUI              |                     | Yes         |
+| ComfyUI              | Community           | Yes         |
 | Quantization         |                     | Yes         |
 | Source Code Access   |                     | Yes         |
 

diff --git a/onediff_comfy_nodes/.gitignore b/onediff_comfy_nodes/.gitignore
@@ -0,0 +1,6 @@
+__pycache__/
+*.py[cod]
+temp/
+*.png
+*.pt
+*.graph
diff --git a/onediff_comfy_nodes/README.md b/onediff_comfy_nodes/README.md
@@ -0,0 +1,124 @@
+# OneDiff ComfyUI Nodes
+
+- [Installation Guide](#installation-guide)
+- [Nodes Usage](#nodes-usage)
+  - [Model Acceleration](#model-acceleration)
+    - [Model Speedup](#model-speedup)
+    - [Model Graph Saver](#model-graph-saver)
+    - [Model Graph Loader](#model-graph-loader)
+  - [Quantization](#quantization)
+  - [VAE Acceleration](#vae-acceleration)
+  - [Image Distinction Scanner](#image-distinction-scanner)
+
+
+## Installation Guide
+
+1. Install and set up [ComfyUI](https://github.com/comfyanonymous/ComfyUI)
+
+
+2. Install PyTorch and OneFlow
+
+Install PyTorch:
+
+```bash
+pip install torch torchvision torchaudio
+```
+
+Install OneFlow Community(CUDA 11.x)
+
+```bash
+pip install --pre oneflow -f https://oneflow-pro.oss-cn-beijing.aliyuncs.com/branch/community/cu118
+```
+
+Install OneFlow Community(CUDA 12.x)
+
+```bash
+pip install --pre oneflow -f https://oneflow-pro.oss-cn-beijing.aliyuncs.com/branch/community/cu121
+```
+
+3. Intall onediff
+
+```bash
+https://github.com/Oneflow-Inc/onediff.git
+cd onediff && pip install -e .
+```
+
+5. Install onediff_comfy_nodes for ComfyUI
+
+```bash
+cd onediff
+cp -r onediff_comfy_nodes path/to/ComfyUI/custom_nodes/
+```
+
+6. (Optional) Advanced features
+
+If you need **unrestricted multiple resolution**, **quantization** support or any other more advanced features, please send an email to [email protected] . Tell us about your **use case, deployment scale and requirements**! 
+
+
+
+
+## Nodes Usage
+
+**Note** All the images in this section can be loaded directly into ComfyUI.
+
+### Model Acceleration
+
+#### Model Speedup
+
+The "Model Speedup" node takes a model as input and outputs an optimized model.
+
+If the `static_mode` is `enabled` (which is the default), it will take some time to compile before the first inference.
+
+If `static_model` is `disabled`, there is no need for additional compilation time before the first inference, but the inference speed will be slower compared to `enabled`, albeit slightly.
+
+![](workflows/model-speedup.png)
+
+#### Model Graph Saver
+
+The optimized model from the "Model Speedup" node can be saved to "graph" by the "Model Graph Saver" node, allowing it to be used in other scenarios without the need for recompilation.
+
+![](workflows/model-graph-saver.png)
+
+You can set different file name prefixes for different types of models.
+
+#### Model Graph Loader
+
+The "Model Graph Loader" node is used to load graph files from the disk, thus saving the time required for the initial compilation.
+
+![](workflows/model-graph-loader.png)
+
+### Quantization
+
+The "UNet Loader Int8" node is used to load quantized models. Quantized models need to be used in conjunction with the "Model Speedup" node.
+
+![](workflows/int8-speedup.png)
+
+The compilation result of the quantized model can also be saved as a graph and loaded when needed.
+
+
+![quantized model saver](workflows/int8-graph-saver.png)
+
+![quantized model loader](workflows/int8-graph-loader.png)
+
+
+### VAE Acceleration
+
+The VAE nodes used for accelerating, saving, and loading VAE graphs operate in a manner very similar to the usage of Model nodes.
+
+Omitting specific details here, the following workflow can be loaded and tested.
+
+**VAE Speedup and Graph Saver**
+
+![](workflows/vae-graph-saver.png)
+
+**VAE Speedup and Graph Loader**
+
+![](workflows/vae-graph-loader.png)
+
+### Image Distinction Scanner
+
+The "Image Distinction Scanner" node is used to compare the differences between two images and visualize the resulting variances.
+
+![](workflows/image-distinction-scanner.png)
+
+
diff --git a/onediff_comfy_nodes/__init__.py b/onediff_comfy_nodes/__init__.py
@@ -0,0 +1,40 @@
+"""OneFlow ComfyUI Speedup Module"""
+from ._config import _USE_UNET_INT8
+from ._nodes import (
+    ModelSpeedup,
+    ModelGraphLoader,
+    ModelGraphSaver,
+    VaeSpeedup,
+    VaeGraphLoader,
+    VaeGraphSaver,
+)
+from ._compare_node import CompareModel, ShowImageDiff
+
+
+NODE_CLASS_MAPPINGS = {
+    "ModelSpeedup": ModelSpeedup,
+    "CompareModel": CompareModel,
+    "ShowImageDiff": ShowImageDiff,
+    "ModelGraphLoader": ModelGraphLoader,
+    "ModelGraphSaver": ModelGraphSaver,
+    "VaeSpeedup": VaeSpeedup,
+    "VaeGraphSaver": VaeGraphSaver,
+    "VaeGraphLoader": VaeGraphLoader,
+}
+
+NODE_DISPLAY_NAME_MAPPINGS = {
+    "ModelSpeedup": "Model Speedup",
+    "CompareModel": "Model Weight Comparator",
+    "ShowImageDiff": "Image Distinction Scanner",
+    "ModelGraphLoader": "Model Graph Loader",
+    "ModelGraphSaver": "Model Graph Saver",
+    "VaeSpeedup": "VAE Speedup",
+    "VaeGraphLoader": "VAE Graph Loader",
+    "VaeGraphSaver": "VAE Graph Saver",
+}
+
+if _USE_UNET_INT8:
+    from ._nodes import UNETLoaderInt8
+
+    NODE_CLASS_MAPPINGS.update({"UNETLoaderInt8": UNETLoaderInt8})
+    NODE_DISPLAY_NAME_MAPPINGS.update({"UNETLoaderInt8": "UNET Loader Int8"})