mirror of
https://github.com/Bing-su/adetailer.git
synced 2026-01-26 19:29:54 +00:00
50 lines
1.6 KiB
Python
50 lines
1.6 KiB
Python
from __future__ import annotations
|
|
|
|
from contextlib import contextmanager
|
|
|
|
from modules import img2img, processing, shared
|
|
|
|
|
|
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
|
|
|
|
|
|
@contextmanager
|
|
def cn_allow_script_control():
|
|
orig = False
|
|
if "control_net_allow_script_control" in shared.opts.data:
|
|
try:
|
|
orig = shared.opts.data["control_net_allow_script_control"]
|
|
shared.opts.data["control_net_allow_script_control"] = True
|
|
yield
|
|
finally:
|
|
shared.opts.data["control_net_allow_script_control"] = orig
|
|
else:
|
|
yield
|