feat: restore cn hijack

controlnet_ext/__init__.py

@@ -1,3 +1,7 @@
 from .controlnet_ext import ControlNetExt, controlnet_exists, get_cn_inpaint_models
 
-__all__ = ["ControlNetExt", "controlnet_exists", "get_cn_inpaint_models"]
+__all__ = [
+    "ControlNetExt",
+    "controlnet_exists",
+    "get_cn_inpaint_models",
+]

controlnet_ext/restore.py

@@ -0,0 +1,33 @@
+from __future__ import annotations
+
+from modules import img2img, processing
+
+
+def cn_restore_unet_hook(p, cn_latest_network):
+    if cn_latest_network is not None:
+        unet = p.sd_model.model.diffusion_model
+        cn_latest_network.restore(unet)
+
+
+class CNHijackRestore:
+    def __init__(self):
+        self.process = hasattr(processing, "__controlnet_original_process_images_inner")
+        self.img2img = hasattr(img2img, "__controlnet_original_process_batch")
+
+    def __enter__(self):
+        if self.process:
+            self.orig_process = processing.process_images_inner
+            processing.process_images_inner = getattr(
+                processing, "__controlnet_original_process_images_inner"
+            )
+        if self.img2img:
+            self.orig_img2img = img2img.process_batch
+            img2img.process_batch = getattr(
+                img2img, "__controlnet_original_process_batch"
+            )
+
+    def __exit__(self, *args, **kwargs):
+        if self.process:
+            processing.process_images_inner = self.orig_process
+        if self.img2img:
+            img2img.process_batch = self.orig_img2img