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