test_quantize_custom_model

.github/workflows/examples.yml

@@ -209,6 +209,8 @@ jobs:
         if: matrix.test-suite == 'diffusers_examples'
         run: |
           docker exec -w /src/onediff/onediff_diffusers_extensions ${{ env.CONTAINER_NAME }} python3 -m pip install -e .
+      - if: matrix.test-suite == 'diffusers_examples' && startsWith(matrix.image, 'onediff-pro')
+        run: docker exec -w /src/onediff ${{ env.CONTAINER_NAME }} python3 tests/test_quantize_custom_model.py
       - if: matrix.test-suite == 'diffusers_examples' && startsWith(matrix.image, 'onediff-pro')
         run: docker exec -w /src/onediff/onediff_diffusers_extensions ${{ env.CONTAINER_NAME }} python3 examples/text_to_image_sdxl_enterprise.py --model /share_nfs/hf_models/stable-diffusion-xl-base-1.0-int8 --width 512 --height 512 --saved_image output_enterprise_sdxl.png
       - if: matrix.test-suite == 'diffusers_examples' && startsWith(matrix.image, 'onediff-pro')

onediff_comfy_nodes/__init__.py

@@ -15,7 +15,7 @@
     BatchSizePatcher,
 )
 from ._compare_node import CompareModel, ShowImageDiff
-
+from ._nodes import OneDiffModelOptimizer, OneDiffDeepcacheOptimizer, OneDiffOnlineQuantizationOptimizer
 
 NODE_CLASS_MAPPINGS = {
     "ModelSpeedup": ModelSpeedup,
@@ -32,6 +32,9 @@
     "OneDiffControlNetLoader": OneDiffControlNetLoader,
     "OneDiffDeepCacheCheckpointLoaderSimple": OneDiffDeepCacheCheckpointLoaderSimple,
     "BatchSizePatcher": BatchSizePatcher,
+    "OneDiffModelOptimizer": OneDiffModelOptimizer,
+    "OneDiffDeepcacheOptimizer": OneDiffDeepcacheOptimizer,
+    "OneDiffOnlineQuantizationOptimizer": OneDiffOnlineQuantizationOptimizer,
 }
 
 NODE_DISPLAY_NAME_MAPPINGS = {
@@ -49,6 +52,9 @@
     "OneDiffControlNetLoader": "Load ControlNet Model - OneDiff",
     "OneDiffDeepCacheCheckpointLoaderSimple": "Load Checkpoint - OneDiff DeepCache",
     "BatchSizePatcher": "Batch Size Patcher",
+    "OneDiffModelOptimizer": "Model Optimizer - OneDiff",
+    "OneDiffDeepcacheOptimizer": "DeepCache Optimizer - OneDiff",
+    "OneDiffOnlineQuantizationOptimizer": "Online Quantization Optimizer - OneDiff",
 }
 
 

onediff_comfy_nodes/_nodes.py

@@ -385,7 +385,7 @@ def deep_cache_convert(
 from nodes import CheckpointLoaderSimple, ControlNetLoader
 from comfy.controlnet import ControlLora, ControlNet
 from .modules.onediff_controlnet import OneDiffControlLora
-
+from .modules.optimizer_strategy import OptimizerStrategy, DeepcacheOptimizerExecutor, OnelineQuantizationOptimizerExecutor
 
 class OneDiffControlNetLoader(ControlNetLoader):
     CATEGORY = "OneDiff/Loaders"
@@ -429,24 +429,32 @@ def INPUT_TYPES(s):
             "required": {
                 "ckpt_name": (folder_paths.get_filename_list("checkpoints"),),
                 "vae_speedup": (["disable", "enable"],),
+            },
+            "optional": {
+                    "model_optimizer": ("MODEL_OPTIMIZER",),
             }
         }
 
     CATEGORY = "OneDiff/Loaders"
     FUNCTION = "onediff_load_checkpoint"
 
     def onediff_load_checkpoint(
-        self, ckpt_name, vae_speedup, output_vae=True, output_clip=True
+        self, ckpt_name, vae_speedup="disable", output_vae=True, output_clip=True, model_optimizer: callable=None,
     ):
         # CheckpointLoaderSimple.load_checkpoint
         modelpatcher, clip, vae = self.load_checkpoint(
             ckpt_name, output_vae, output_clip
         )
+
         unet_graph_file = generate_graph_path(ckpt_name, modelpatcher.model)
+        if model_optimizer is not None:
+            modelpatcher = model_optimizer(modelpatcher)
+
         modelpatcher.model.diffusion_model = compoile_unet(
-            modelpatcher.model.diffusion_model, unet_graph_file
-        )
+                modelpatcher.model.diffusion_model, unet_graph_file)
+
         modelpatcher.model._register_state_dict_hook(state_dict_hook)
+
         if vae_speedup == "enable":
             file_path = generate_graph_path(ckpt_name, vae.first_stage_model)
             vae.first_stage_model = oneflow_compile(
@@ -463,6 +471,181 @@ def onediff_load_checkpoint(
         return modelpatcher, clip, vae
 
 
+class OneDiffDeepcacheOptimizer:
+    @classmethod
+    def INPUT_TYPES(s):
+        return {
+            "required": {
+                "cache_interval": (
+                    "INT",
+                    {
+                        "default": 3,
+                        "min": 1,
+                        "max": 1000,
+                        "step": 1,
+                        "display": "number",
+                    },
+                ),
+                "cache_layer_id": (
+                    "INT",
+                    {"default": 0, "min": 0, "max": 12, "step": 1, "display": "number"},
+                ),
+                "cache_block_id": (
+                    "INT",
+                    {"default": 1, "min": 0, "max": 12, "step": 1, "display": "number"},
+                ),
+                "start_step": (
+                    "INT",
+                    {
+                        "default": 0,
+                        "min": 0,
+                        "max": 1000,
+                        "step": 1,
+                        "display": "number",
+                    },
+                ),
+                "end_step": (
+                    "INT",
+                    {"default": 1000, "min": 0, "max": 1000, "step": 0.1},
+                ),
+            },
+        }
+
+    CATEGORY = "OneDiff/Optimizer"
+    RETURN_TYPES = ("DeepPCacheOptimizer",)
+    FUNCTION = "apply"
+
+    def apply(
+        self,
+        cache_interval=3,
+        cache_layer_id=0,
+        cache_block_id=1,
+        start_step=0,
+        end_step=1000,
+    ):
+        return (
+            DeepcacheOptimizerExecutor(
+                cache_interval=cache_interval,
+                cache_layer_id=cache_layer_id,
+                cache_block_id=cache_block_id,
+                start_step=start_step,
+                end_step=end_step,
+            ),
+        )
+
+
+class OneDiffOnlineQuantizationOptimizer:
+    @classmethod
+    def INPUT_TYPES(s):
+        return {
+            "required": {
+                "quantized_conv_percentage": (
+                    "INT",
+                    {
+                        "default": 70,
+                        "min": 0,  # Minimum value
+                        "max": 100,  # Maximum value
+                        "step": 1,  # Slider's step
+                        "display": "slider",  # Cosmetic only: display as "number" or "slider"
+                    },
+                ),
+                "quantized_linear_percentage": (
+                    "INT",
+                    {
+                        "default": 80,
+                        "min": 0,  # Minimum value
+                        "max": 100,  # Maximum value
+                        "step": 1,  # Slider's step
+                        "display": "slider",  # Cosmetic only: display as "number" or "slider"
+                    },
+                ),
+                "conv_compute_density_threshold":(
+                   "INT",
+                    {
+                        "default": 100,
+                        "min": 0,  # Minimum value
+                        "max": 2000,  # Maximum value
+                        "step": 10,  # Slider's step
+                        "display": "number",  # Cosmetic only: display as "number" or "slider"
+                    },
+                ),
+                "linear_compute_density_threshold":(
+                    "INT",
+                    {
+                        "default": 300,
+                        "min": 0,  # Minimum value
+                        "max": 2000,  # Maximum value
+                        "step": 10,  # Slider's step
+                        "display": "number",  # Cosmetic only: display as "number" or "slider"
+                    },
+                )
+            },
+        }
+
+    CATEGORY = "OneDiff/Optimizer"
+    RETURN_TYPES = ("QuantizationOptimizer",)
+    FUNCTION = "apply"
+
+    def apply(self, quantized_conv_percentage=0, quantized_linear_percentage=0, conv_compute_density_threshold=0, linear_compute_density_threshold=0):
+        return (
+            OnelineQuantizationOptimizerExecutor(
+                conv_percentage=quantized_conv_percentage,
+                linear_percentage=quantized_linear_percentage,
+                conv_compute_density_threshold = conv_compute_density_threshold,
+                linear_compute_density_threshold = linear_compute_density_threshold,
+            ),
+        )
+
+
+class OneDiffModelOptimizer:
+    """Main class responsible for optimizing models."""
+
+    @classmethod
+    def INPUT_TYPES(s):
+        return {
+            "required": {},
+            "optional": {
+                "quantization_optimizer": ("QuantizationOptimizer",),
+                "deeppcache_optimizer": ("DeepPCacheOptimizer",),
+            },
+        }
+
+    CATEGORY = "OneDiff/Optimization"
+    RETURN_TYPES = ("MODEL_OPTIMIZER",)
+    FUNCTION = "optimize_model"
+
+    def optimize_model(self, quantization_optimizer:OptimizerStrategy =None, deeppcache_optimizer:OptimizerStrategy=None):
+        """Apply the optimization technique to the model."""
+
+        def apply_optimizer(model, ckpt_name=""):
+            def set_compiled_options(module: DeployableModule, graph_file="unet"):
+                assert isinstance(module, DeployableModule)
+                compile_options = {
+                    "graph_file": graph_file,
+                    "graph_file_device": model_management.get_torch_device(),
+                }
+                module._deployable_module_options.update(compile_options)
+
+            if deeppcache_optimizer is not None:
+                model = deeppcache_optimizer.apply(model)
+                graph_file = generate_graph_path(ckpt_name, model.fast_deep_cache_unet._torch_module)
+                set_compiled_options(model.fast_deep_cache_unet, graph_file)
+                graph_file = generate_graph_path(ckpt_name, model.deep_cache_unet._torch_module)
+                set_compiled_options(model.deep_cache_unet, graph_file)
+
+            if quantization_optimizer is not None:
+                model = quantization_optimizer.apply(model)
+                diff_model: DeployableModule = model.model.diffusion_model
+                graph_file = generate_graph_path(ckpt_name, model.model)
+                set_compiled_options(diff_model, graph_file)
+                quant_config = diff_model._deployable_module_quant_config
+                quant_config.cache_dir = os.path.dirname(graph_file)
+
+            return model
+
+        return (apply_optimizer,)
+
+
 class OneDiffDeepCacheCheckpointLoaderSimple(CheckpointLoaderSimple):
     @classmethod
     def INPUT_TYPES(s):
@@ -754,7 +937,7 @@ def set_cache_filename(self, model, latent_image):
 
 
 
-    
+
 if _USE_UNET_INT8:
     from .utils.quant_ksampler_tools import (
         KSampleQuantumBase,

onediff_comfy_nodes/modules/optimizer_strategy/__init__.py

@@ -0,0 +1,3 @@
+from .deepcache_optimizer import DeepcacheOptimizerExecutor
+from .quantization_optimizer import OnelineQuantizationOptimizerExecutor
+from .optimizer_strategy import OptimizerStrategy

onediff_comfy_nodes/modules/optimizer_strategy/deepcache_optimizer.py

@@ -0,0 +1,40 @@
+from dataclasses import dataclass
+from functools import singledispatchmethod
+
+from comfy.model_patcher import ModelPatcher
+
+from ...utils.deep_cache_speedup import deep_cache_speedup
+from .optimizer_strategy import OptimizerStrategy
+
+
+@dataclass
+class DeepcacheOptimizerExecutor(OptimizerStrategy):
+    cache_interval: int = 3
+    cache_layer_id: int = 0
+    cache_block_id: int = 1
+    start_step: int = 0
+    end_step: int = 1000
+
+    @singledispatchmethod
+    def apply(self, model):
+        print(
+            "DeepcacheOptimizerExecutor.apply: not implemented for model type:",
+            type(model),
+        )
+        return model
+
+    @apply.register(ModelPatcher)
+    def _(self, model):
+        return deep_cache_speedup(
+            model=model,
+            use_graph=True,
+            cache_interval=self.cache_interval,
+            cache_layer_id=self.cache_layer_id,
+            cache_block_id=self.cache_block_id,
+            start_step=self.start_step,
+            end_step=self.end_step,
+            use_oneflow_deepcache_speedup_modelpatcher=False,
+        )[0]
+
+
+

onediff_comfy_nodes/modules/optimizer_strategy/optimizer_strategy.py

@@ -0,0 +1,13 @@
+from abc import ABC, abstractmethod
+
+from comfy.model_patcher import ModelPatcher
+
+
+class OptimizerStrategy(ABC):
+    """Interface for optimization strategies."""
+
+    @abstractmethod
+    def apply(self, model: ModelPatcher):
+        """Apply the optimization strategy to the model."""
+        pass
+