From eb6e3c337cc9cb1cf43c40a7ae95baf8f8ccce51 Mon Sep 17 00:00:00 2001
From: Disty0 <disty@disty.xyz>
Date: Wed, 19 Jun 2024 22:42:00 +0300
Subject: [PATCH] NNCF ControlNet support

---
 modules/control/units/controlnet.py | 13 +++++++++++--
 modules/shared.py                   |  2 +-
 2 files changed, 12 insertions(+), 3 deletions(-)

diff --git a/modules/control/units/controlnet.py b/modules/control/units/controlnet.py
index e911bc6ff..b57005f3c 100644
--- a/modules/control/units/controlnet.py
+++ b/modules/control/units/controlnet.py
@@ -172,10 +172,19 @@ class ControlNet():
                 self.load_safetensors(model_path)
             else:
                 self.model = ControlNetModel.from_pretrained(model_path, **self.load_config)
-            if self.device is not None:
-                self.model.to(self.device)
             if self.dtype is not None:
                 self.model.to(self.dtype)
+            if "ControlNet" in opts.nncf_compress_weights:
+                try:
+                    log.debug(f'Control {what} model NNCF Compress: id="{model_id}"')
+                    from installer import install
+                    install('nncf==2.7.0', quiet=True)
+                    from modules.sd_models_compile import nncf_compress_model
+                    self.model = nncf_compress_model(self.model)
+                except Exception as e:
+                    log.error(f'Control {what} model NNCF Compression failed: id="{model_id}" error={e}')
+            if self.device is not None:
+                self.model.to(self.device)
             t1 = time.time()
             self.model_id = model_id
             log.debug(f'Control {what} model loaded: id="{model_id}" path="{model_path}" time={t1-t0:.2f}')
diff --git a/modules/shared.py b/modules/shared.py
index fcdf9bbb8..de2cebf64 100644
--- a/modules/shared.py
+++ b/modules/shared.py
@@ -449,7 +449,7 @@ options_templates.update(options_section(('cuda', "Compute Settings"), {
     "deep_cache_interval": OptionInfo(3, "DeepCache cache interval", gr.Slider, {"minimum": 1, "maximum": 10, "step": 1}),
 
     "nncf_sep": OptionInfo("<h2>Model Compress</h2>", "", gr.HTML),
-    "nncf_compress_weights": OptionInfo([], "Compress Model weights with NNCF", gr.CheckboxGroup, {"choices": ["Model", "VAE", "Text Encoder"], "visible": native}),
+    "nncf_compress_weights": OptionInfo([], "Compress Model weights with NNCF", gr.CheckboxGroup, {"choices": ["Model", "VAE", "Text Encoder", "ControlNet"], "visible": native}),
 
     "ipex_sep": OptionInfo("<h2>IPEX</h2>", "", gr.HTML, {"visible": devices.backend == "ipex"}),
     "ipex_optimize": OptionInfo([], "IPEX Optimize for Intel GPUs", gr.CheckboxGroup, {"choices": ["Model", "VAE", "Text Encoder", "Upscaler"], "visible": devices.backend == "ipex"}),