Dev refactor onediff comfy nodes

onediff_comfy_nodes/__init__.py

@@ -1,84 +1,30 @@
 """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, OneDiffApplyModelOptimizer,
+                     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,
+    "OneDiffModelOptimizer": OneDiffApplyModelOptimizer, 
     "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",
+    "OneDiffModelOptimizer": "Apply Model Optimizer - 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_optimizer,
+                           nodes_torch_compile_optimizer)
 
-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_optimizer, nodes_torch_compile_optimizer]
+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",
-        }
-    )

onediff_comfy_nodes/_config.py

@@ -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