Skip to content

Commit

Permalink
Dev refactor onediff comfy nodes (#796)
Browse files Browse the repository at this point in the history 
## Online quantization for comfyui

### Install

1. [OneDiff Installation
Guide](https://github.com/siliconflow/onediff/blob/main/README_ENTERPRISE.md#install-onediff-enterprise)
2. [OneDiffx Installation
Guide](https://github.com/siliconflow/onediff/tree/main/onediff_diffusers_extensions#install-and-setup)

## Usage:
| Option | Range | Default | Description |
| -------------------------------------- | ------ | ------- |
----------------------------------------------------------------------------
|
| quantized_conv_percentage | [0, 100] | 100 | Example value
representing 100% quantization for linear layers｜
| quantized_linear_percentage | [0, 100] | 100 | Example value
representing 100% quantization for convolutional layers |
| conv_compute_density_threshold | [0, ∞) | 100 | Computational density
threshold for quantizing convolutional modules to 100. |
| linear_compute_density_threshold | [0, ∞) | 300 | Computational
density threshold for quantizing linear modules to 300. |

Notes:

1. Specify the directory for saving graphsusing export
COMFYUI_ONEDIFF_SAVE_GRAPH_DIR="/path/to/save/graphs".
2. The log *.pt file is cached. Quantization result information can be
found in `cache_dir`/quantization_stats.json.

## Performance Comparison

###
[SDXL](https://huggingface.co/stabilityai/stable-diffusion-xl-base-1.0)
Updated on Mon 08 Apr 2024 


![quant_sdxl](https://github.com/siliconflow/onediff/assets/109639975/b8f8da75-944b-4553-aea3-69c19886af37)

| Accelerator | Baseline (non-optimized) | OneDiff(optimized) | OneDiff
Quant(optimized) |
| ----------------------- | ------------------------ |
------------------ | ------------------------ |
| NVIDIA GeForce RTX 3090 | 8.03 s | 4.44 s ( ~44.7%) | 3.34 s ( ~58.4%)
|

- torch `python -c "import torch; print(torch.__version__)"`: {version:
2.2.1+cu121}
- oneflow `python -m oneflow --doctor`: {git_commit: 710818c, version:
0.9.1.dev20240406+cu121, enterprise: True}
- Start comfyui command: `python main.py --gpu-only`

### SD1.5

![sd 1.5
](https://github.com/siliconflow/onediff/assets/109639975/49a8ab1b-e2be-4719-a962-33b813f5e83f)

### SVD

![svd_quant](https://github.com/siliconflow/onediff/assets/109639975/93ebe3d5-8413-4a7e-8b93-fd016f61abe9)
  • Loading branch information
@ccssu
ccssu authored Apr 11, 2024
1 parent 85d5666 commit e4d56a0
Show file tree
Hide file tree
Showing 66 changed files with 1,926 additions and 1,100 deletions.
79 changes: 12 additions & 67 deletions onediff_comfy_nodes/__init__.py
View file
Original file line number Diff line number Diff line change
@@ -1,84 +1,29 @@
"""OneDiff ComfyUI Speedup Module"""
from ._config import _USE_UNET_INT8
from ._nodes import (
ModelSpeedup,
ModelGraphLoader,
ModelGraphSaver,
VaeSpeedup,
VaeGraphLoader,
VaeGraphSaver,
SVDSpeedup,
ModuleDeepCacheSpeedup,
OneDiffCheckpointLoaderSimple,
OneDiffControlNetLoader,
OneDiffDeepCacheCheckpointLoaderSimple,
BatchSizePatcher,
)
from ._compare_node import CompareModel, ShowImageDiff

from ._config import *
from ._nodes import (ModelSpeedup, OneDiffApplyModelBooster,
OneDiffCheckpointLoaderSimple, OneDiffControlNetLoader,
VaeSpeedup)

NODE_CLASS_MAPPINGS = {
"ModelSpeedup": ModelSpeedup,
"CompareModel": CompareModel,
"ShowImageDiff": ShowImageDiff,
"ModelGraphLoader": ModelGraphLoader,
"ModelGraphSaver": ModelGraphSaver,
"VaeSpeedup": VaeSpeedup,
"VaeGraphSaver": VaeGraphSaver,
"VaeGraphLoader": VaeGraphLoader,
"SVDSpeedup": SVDSpeedup,
"ModuleDeepCacheSpeedup": ModuleDeepCacheSpeedup,
"OneDiffModelBooster": OneDiffApplyModelBooster,
"OneDiffCheckpointLoaderSimple": OneDiffCheckpointLoaderSimple,
"OneDiffControlNetLoader": OneDiffControlNetLoader,
"OneDiffDeepCacheCheckpointLoaderSimple": OneDiffDeepCacheCheckpointLoaderSimple,
"BatchSizePatcher": BatchSizePatcher,

}

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",
"SVDSpeedup": "SVD Speedup",
"ModuleDeepCacheSpeedup": "Model DeepCache Speedup",
"OneDiffModelBooster": "Apply Model Booster - OneDff",
"OneDiffCheckpointLoaderSimple": "Load Checkpoint - OneDiff",
"OneDiffControlNetLoader": "Load ControlNet Model - OneDiff",
"OneDiffDeepCacheCheckpointLoaderSimple": "Load Checkpoint - OneDiff DeepCache",
"BatchSizePatcher": "Batch Size Patcher",
}

from .extras_nodes import (nodes_compare, nodes_oneflow_booster, nodes_torch_compile_booster)

if _USE_UNET_INT8:
from ._nodes import UNETLoaderInt8, Quant8Model
from ._nodes import (
QuantKSampler,
OneDiffQuantCheckpointLoaderSimple,
OneDiffQuantCheckpointLoaderSimpleAdvanced,
ImageOnlyOneDiffQuantCheckpointLoaderAdvanced,
)

NODE_CLASS_MAPPINGS.update(
{
"UNETLoaderInt8": UNETLoaderInt8,
"Quant8Model": Quant8Model,
"OneDiffQuantCheckpointLoaderSimple": OneDiffQuantCheckpointLoaderSimple,
"OneDiffQuantCheckpointLoaderSimpleAdvanced": OneDiffQuantCheckpointLoaderSimpleAdvanced,
"ImageOnlyOneDiffQuantCheckpointLoaderAdvanced": ImageOnlyOneDiffQuantCheckpointLoaderAdvanced,
"QuantKSampler": QuantKSampler,
}
)
extras = [nodes_compare, nodes_oneflow_booster, nodes_torch_compile_booster]
for node in extras:
NODE_CLASS_MAPPINGS.update(node.NODE_CLASS_MAPPINGS)
NODE_DISPLAY_NAME_MAPPINGS.update(node.NODE_DISPLAY_NAME_MAPPINGS)

NODE_DISPLAY_NAME_MAPPINGS.update(
{
"UNETLoaderInt8": "UNET Loader Int8",
"Quant8Model": "Model Quantization(int8)",
"OneDiffQuantCheckpointLoaderSimple": "Load Checkpoint - OneDiff Quant",
"OneDiffQuantCheckpointLoaderSimpleAdvanced": "Load Checkpoint - OneDiff Quant Advanced",
"ImageOnlyOneDiffQuantCheckpointLoaderAdvanced": "Load Checkpoint - OneDiff Quant Advanced (img2vid)",
"QuantKSampler": "Quant K Sampler",
}
)
45 changes: 3 additions & 42 deletions onediff_comfy_nodes/_config.py
View file
Original file line number Diff line number Diff line change
@@ -1,44 +1,5 @@
"""Config file for comfyui-speedup, which will automatically"""
import os
import sys
from pathlib import Path
import folder_paths

ONEDIFF_QUANTIZED_OPTIMIZED_MODELS = "onediff_quant"
_USE_UNET_INT8 = True
COMFYUI_ROOT = Path(os.path.abspath(__file__)).parents[2]
COMFYUI_SPEEDUP_ROOT = Path(os.path.abspath(__file__)).parents[0]
INFER_COMPILER_REGISTRY = Path(COMFYUI_SPEEDUP_ROOT) / "infer_compiler_registry"
os.environ["COMFYUI_ROOT"] = str(COMFYUI_ROOT)
custom_nodes_path = os.path.join(COMFYUI_ROOT, "custom_nodes")
sys.path.insert(0, str(COMFYUI_ROOT))
sys.path.insert(0, str(COMFYUI_SPEEDUP_ROOT))
sys.path.insert(0, str(INFER_COMPILER_REGISTRY))
if custom_nodes_path not in sys.path:
sys.path.append(custom_nodes_path)

import register_comfy # load plugins

from onediff.infer_compiler.utils import is_community_version

if is_community_version():
_USE_UNET_INT8 = False

if _USE_UNET_INT8:
import register_onediff_quant # load plugins
from folder_paths import (folder_names_and_paths, models_dir,
supported_pt_extensions)

unet_int8_model_dir = Path(models_dir) / "unet_int8"
unet_int8_model_dir.mkdir(parents=True, exist_ok=True)
folder_names_and_paths["unet_int8"] = (
[str(unet_int8_model_dir)],
supported_pt_extensions,
)

opt_models_dir = Path(models_dir) / ONEDIFF_QUANTIZED_OPTIMIZED_MODELS
opt_models_dir.mkdir(parents=True, exist_ok=True)

folder_names_and_paths[ONEDIFF_QUANTIZED_OPTIMIZED_MODELS] = (
[str(opt_models_dir)],
supported_pt_extensions,
)
# https://github.com/comfyanonymous/ComfyUI/blob/master/folder_paths.py#L9
os.environ["COMFYUI_ROOT"] = folder_paths.base_path
Loading

0 comments on commit e4d56a0

Please sign in to comment.