diff --git a/javascript/sdnext.css b/javascript/sdnext.css index cd2e92c3f..2a82e5add 100644 --- a/javascript/sdnext.css +++ b/javascript/sdnext.css @@ -194,7 +194,7 @@ table.settings-value-table td { padding: 0.4em; border: 1px solid #ccc; max-widt #txt2img_extra_networks, #img2img_extra_networks, #control_extra_networks { width: 0; } .extra-networks > div { margin: 0; border-bottom: none !important; gap: 0.3em 0; } .extra-networks .second-line { display: flex; width: -moz-available; width: -webkit-fill-available; gap: 0.3em; box-shadow: var(--input-shadow); } -.extra-networks .search { flex: 1; } +.extra-networks .search { flex: 1; height: 4em; } .extra-networks .description { flex: 3; } .extra-networks .tab-nav>button { margin-right: 0; height: 24px; padding: 2px 4px 2px 4px; } .extra-networks .buttons { position: absolute; right: 0; margin: -4px; background: var(--background-color); } @@ -222,7 +222,7 @@ table.settings-value-table td { padding: 0.4em; border: 1px solid #ccc; max-widt .extra-network-cards .card-list .tag { color: var(--primary-500); margin-left: 0.8em; } .extra-details-close { position: fixed; top: 0.2em; right: 0.2em; z-index: 99; background: var(--button-secondary-background-fill) !important; } #txt2img_description, #img2img_description, #control_description { max-height: 63px; overflow-y: auto !important; } -#txt2img_description>label>textarea, #img2img_description>label>textarea, #control_description>label>textarea { font-size: var(--text-sm) } +#txt2img_description>label>textarea, #img2img_description>label>textarea, #control_description>label>textarea { font-size: var(--text-xs); height: 6em; } #txt2img_extra_details>div, #img2img_extra_details>div { overflow-y: auto; min-height: 40vh; max-height: 80vh; align-self: flex-start; } #txt2img_extra_details, #img2img_extra_details { position: fixed; bottom: 50%; left: 50%; transform: translate(-50%, 50%); padding: 0.8em; border: var(--block-border-width) solid var(--highlight-color) !important; diff --git a/modules/control/proc/canny.py b/modules/control/proc/canny.py index e68673d88..1e4bb3176 100644 --- a/modules/control/proc/canny.py +++ b/modules/control/proc/canny.py @@ -17,19 +17,13 @@ class CannyDetector: output_type = output_type or "pil" else: output_type = output_type or "np" - input_image = HWC3(input_image) input_image = resize_image(input_image, detect_resolution) - detected_map = cv2.Canny(input_image, low_threshold, high_threshold) detected_map = HWC3(detected_map) - img = resize_image(input_image, image_resolution) H, W, _C = img.shape - detected_map = cv2.resize(detected_map, (W, H), interpolation=cv2.INTER_LINEAR) - if output_type == "pil": detected_map = Image.fromarray(detected_map) - return detected_map diff --git a/modules/control/proc/depth_anything/__init__.py b/modules/control/proc/depth_anything/__init__.py index 7af45405e..cf7ee92e2 100644 --- a/modules/control/proc/depth_anything/__init__.py +++ b/modules/control/proc/depth_anything/__init__.py @@ -4,6 +4,7 @@ import torch.nn.functional as F import numpy as np from PIL import Image from modules import devices, masking +from modules.shared import opts class DepthAnythingDetector: @@ -54,6 +55,8 @@ class DepthAnythingDetector: image = torch.from_numpy(image).unsqueeze(0).to(devices.device) with devices.inference_context(): depth = self.model(image) + if opts.control_move_processor: + self.model.to('cpu') depth = F.interpolate(depth[None], (h, w), mode="bilinear", align_corners=False)[0, 0] depth = (depth - depth.min()) / (depth.max() - depth.min()) * 255.0 depth = depth.cpu().numpy().astype(np.uint8) diff --git a/modules/control/proc/dpt.py b/modules/control/proc/dpt.py index c297fc80c..f726bf58f 100644 --- a/modules/control/proc/dpt.py +++ b/modules/control/proc/dpt.py @@ -3,24 +3,29 @@ import numpy as np import torch from transformers import AutoImageProcessor, DPTForDepthEstimation from modules import devices +from modules.shared import opts image_processor: AutoImageProcessor = None -dpt_model: DPTForDepthEstimation = None class DPTDetector: - def __call__(self, input_image=None): - global image_processor, dpt_model # pylint: disable=global-statement - from modules.control.processors import cache_dir - if image_processor is None: - image_processor = AutoImageProcessor.from_pretrained("Intel/dpt-large", cache_dir=cache_dir) - if dpt_model is None: - dpt_model = DPTForDepthEstimation.from_pretrained("Intel/dpt-large", cache_dir=cache_dir) + def __init__(self, model=None, processor=None): + self.model = model + self.processor = processor + def __call__(self, input_image=None): + from modules.control.processors import cache_dir + if self.processor is None: + self.processor = AutoImageProcessor.from_pretrained("Intel/dpt-large", cache_dir=cache_dir) + if self.model is None: + self.model = DPTForDepthEstimation.from_pretrained("Intel/dpt-large", cache_dir=cache_dir) + + self.model.to(devices.device) with devices.inference_context(): - inputs = image_processor(images=input_image, return_tensors="pt") - outputs = dpt_model(**inputs) + inputs = self.processor(images=input_image, return_tensors="pt") + inputs.to(devices.device) + outputs = self.model(**inputs) predicted_depth = outputs.predicted_depth prediction = torch.nn.functional.interpolate( predicted_depth.unsqueeze(1), @@ -30,6 +35,8 @@ class DPTDetector: ) output = prediction.squeeze().cpu().numpy() formatted = (output * 255 / np.max(output)).astype("uint8") - depth = Image.fromarray(formatted) - depth = depth.convert('RGB') - return depth + if opts.control_move_processor: + self.model.to('cpu') + depth = Image.fromarray(formatted) + depth = depth.convert('RGB') + return depth diff --git a/modules/control/proc/dwpose/__init__.py b/modules/control/proc/dwpose/__init__.py index a0c5c513b..d5009ade6 100644 --- a/modules/control/proc/dwpose/__init__.py +++ b/modules/control/proc/dwpose/__init__.py @@ -32,7 +32,6 @@ def draw_pose(pose, H, W): class DWposeDetector: def __init__(self, det_config=None, det_ckpt=None, pose_config=None, pose_ckpt=None, device="cpu"): from .wholebody import Wholebody - self.pose_estimation = Wholebody(det_config, det_ckpt, pose_config, pose_ckpt, device) def to(self, device): @@ -62,29 +61,19 @@ class DWposeDetector: score[i][j] = int(18*i+j) else: score[i][j] = -1 - un_visible = subset < min_confidence candidate[un_visible] = -1 - _foot = candidate[:,18:24] - faces = candidate[:,24:92] - hands = candidate[:,92:113] hands = np.vstack([hands, candidate[:,113:]]) - bodies = dict(candidate=body, subset=score) pose = dict(bodies=bodies, hands=hands, faces=faces) - detected_map = draw_pose(pose, H, W) detected_map = HWC3(detected_map) - img = resize_image(input_image, image_resolution) H, W, _C = img.shape - detected_map = cv2.resize(detected_map, (W, H), interpolation=cv2.INTER_LINEAR) - if output_type == "pil": detected_map = Image.fromarray(detected_map) - return detected_map diff --git a/modules/control/proc/edge.py b/modules/control/proc/edge.py index d068383c1..bd46ff491 100644 --- a/modules/control/proc/edge.py +++ b/modules/control/proc/edge.py @@ -34,13 +34,11 @@ class EdgeDetector: input_image = kwargs.pop("img") if input_image is None: raise ValueError("input_image must be defined.") - if not isinstance(input_image, np.ndarray): input_image = np.array(input_image, dtype=np.uint8) output_type = output_type or "pil" else: output_type = output_type or "np" - input_image = HWC3(input_image) input_image = resize_image(input_image, detect_resolution) img_gray = cv2.cvtColor(input_image, cv2.COLOR_BGR2GRAY) @@ -52,12 +50,9 @@ class EdgeDetector: edge_map = np.expand_dims(edge_map, axis=2) edge_map = cv2.cvtColor(edge_map, cv2.COLOR_GRAY2BGR).astype(np.uint8) edge_map = HWC3(edge_map) - img = resize_image(input_image, image_resolution) H, W, _C = img.shape edge_map = cv2.resize(edge_map, (W, H), interpolation=cv2.INTER_LINEAR) - if output_type == "pil": edge_map = Image.fromarray(edge_map) - return edge_map diff --git a/modules/control/proc/glpn.py b/modules/control/proc/glpn.py index 13e9ee3e2..f6c87096c 100644 --- a/modules/control/proc/glpn.py +++ b/modules/control/proc/glpn.py @@ -3,24 +3,26 @@ import numpy as np import torch from transformers import AutoImageProcessor, GLPNForDepthEstimation from modules import devices - - -image_processor: AutoImageProcessor = None -glpn_model: GLPNForDepthEstimation = None +from modules.shared import opts class GLPNDetector: - def __call__(self, input_image=None): - global image_processor, glpn_model # pylint: disable=global-statement - from modules.control.processors import cache_dir - if image_processor is None: - image_processor = AutoImageProcessor.from_pretrained("vinvino02/glpn-kitti", cache_dir=cache_dir) - if glpn_model is None: - glpn_model = GLPNForDepthEstimation.from_pretrained("vinvino02/glpn-kitti", cache_dir=cache_dir) + def __init__(self, model=None, processor=None): + self.model = model + self.processor = processor + def __call__(self, input_image=None): + from modules.control.processors import cache_dir + if self.processor is None: + self.processor = AutoImageProcessor.from_pretrained("vinvino02/glpn-kitti", cache_dir=cache_dir) + if self.model is None: + self.model = GLPNForDepthEstimation.from_pretrained("vinvino02/glpn-kitti", cache_dir=cache_dir) + + self.model.to(devices.device) with devices.inference_context(): - inputs = image_processor(images=input_image, return_tensors="pt") - outputs = glpn_model(**inputs) + inputs = self.processor(images=input_image, return_tensors="pt") + inputs.to(devices.device) + outputs = self.model(**inputs) predicted_depth = outputs.predicted_depth prediction = torch.nn.functional.interpolate( predicted_depth.unsqueeze(1), @@ -30,6 +32,8 @@ class GLPNDetector: ) output = prediction.squeeze().cpu().numpy() formatted = 255 - (output * 255 / np.max(output)).astype("uint8") - depth = Image.fromarray(formatted) - depth = depth.convert('RGB') - return depth + if opts.control_move_processor: + self.model.to('cpu') + depth = Image.fromarray(formatted) + depth = depth.convert('RGB') + return depth diff --git a/modules/control/proc/hed.py b/modules/control/proc/hed.py index 9504e627e..e0144c41f 100644 --- a/modules/control/proc/hed.py +++ b/modules/control/proc/hed.py @@ -6,15 +6,14 @@ # and in this way it works better for gradio's RGB protocol import os -import warnings - import cv2 import numpy as np import torch from einops import rearrange from huggingface_hub import hf_hub_download from PIL import Image - +from modules import devices +from modules.shared import opts from modules.control.util import HWC3, nms, resize_image, safe_step @@ -57,49 +56,37 @@ class ControlNetHED_Apache2(torch.nn.Module): # pylint: disable=abstract-method return projection1, projection2, projection3, projection4, projection5 class HEDdetector: - def __init__(self, netNetwork): - self.netNetwork = netNetwork + def __init__(self, model): + self.model = model @classmethod def from_pretrained(cls, pretrained_model_or_path, filename=None, cache_dir=None): filename = filename or "ControlNetHED.pth" - if os.path.isdir(pretrained_model_or_path): model_path = os.path.join(pretrained_model_or_path, filename) else: model_path = hf_hub_download(pretrained_model_or_path, filename, cache_dir=cache_dir) - - netNetwork = ControlNetHED_Apache2() - netNetwork.load_state_dict(torch.load(model_path, map_location='cpu')) - netNetwork.float().eval() - - return cls(netNetwork) + model = ControlNetHED_Apache2() + model.load_state_dict(torch.load(model_path, map_location='cpu')) + model.float().eval() + return cls(model) def to(self, device): - self.netNetwork.to(device) + self.model.to(device) return self def __call__(self, input_image, detect_resolution=512, image_resolution=512, safe=False, output_type="pil", scribble=False, **kwargs): - if "return_pil" in kwargs: - warnings.warn("return_pil is deprecated. Use output_type instead.", DeprecationWarning) - output_type = "pil" if kwargs["return_pil"] else "np" - if type(output_type) is bool: - warnings.warn("Passing `True` or `False` to `output_type` is deprecated and will raise an error in future versions") - if output_type: - output_type = "pil" - - device = next(iter(self.netNetwork.parameters())).device + self.model.to(devices.device) + device = next(iter(self.model.parameters())).device if not isinstance(input_image, np.ndarray): input_image = np.array(input_image, dtype=np.uint8) - input_image = HWC3(input_image) input_image = resize_image(input_image, detect_resolution) - assert input_image.ndim == 3 H, W, _C = input_image.shape image_hed = torch.from_numpy(input_image.copy()).float().to(device) image_hed = rearrange(image_hed, 'h w c -> 1 c h w') - edges = self.netNetwork(image_hed) + edges = self.model(image_hed) edges = [e.detach().cpu().numpy().astype(np.float32)[0, 0] for e in edges] edges = [cv2.resize(e, (W, H), interpolation=cv2.INTER_LINEAR) for e in edges] edges = np.stack(edges, axis=2) @@ -107,22 +94,18 @@ class HEDdetector: if safe: edge = safe_step(edge) edge = (edge * 255.0).clip(0, 255).astype(np.uint8) - detected_map = edge detected_map = HWC3(detected_map) - img = resize_image(input_image, image_resolution) H, W, _C = img.shape - detected_map = cv2.resize(detected_map, (W, H), interpolation=cv2.INTER_LINEAR) - if scribble: detected_map = nms(detected_map, 127, 3.0) detected_map = cv2.GaussianBlur(detected_map, (0, 0), 3.0) detected_map[detected_map > 4] = 255 detected_map[detected_map < 255] = 0 - + if opts.control_move_processor: + self.model.to('cpu') if output_type == "pil": detected_map = Image.fromarray(detected_map) - return detected_map diff --git a/modules/control/proc/leres/__init__.py b/modules/control/proc/leres/__init__.py index 3a62882bd..63040bc49 100644 --- a/modules/control/proc/leres/__init__.py +++ b/modules/control/proc/leres/__init__.py @@ -1,11 +1,11 @@ import os - import cv2 import numpy as np import torch from huggingface_hub import hf_hub_download from PIL import Image - +from modules import devices +from modules.shared import opts from modules.control.util import HWC3, resize_image from .leres.depthmap import estimateboost, estimateleres from .leres.multi_depth_model_woauxi import RelDepthModel @@ -49,60 +49,48 @@ class LeresDetector: return self def __call__(self, input_image, thr_a=0, thr_b=0, boost=False, detect_resolution=512, image_resolution=512, output_type="pil"): + self.model.to(devices.device) # device = next(iter(self.model.parameters())).device if not isinstance(input_image, np.ndarray): input_image = np.array(input_image, dtype=np.uint8) - input_image = HWC3(input_image) input_image = resize_image(input_image, detect_resolution) - assert input_image.ndim == 3 height, width, _dim = input_image.shape - if boost: depth = estimateboost(input_image, self.model, 0, self.pix2pixmodel, max(width, height)) else: depth = estimateleres(input_image, self.model, width, height) - numbytes=2 depth_min = depth.min() depth_max = depth.max() max_val = (2**(8*numbytes))-1 - # check output before normalizing and mapping to 16 bit if depth_max - depth_min > np.finfo("float").eps: out = max_val * (depth - depth_min) / (depth_max - depth_min) else: out = np.zeros(depth.shape) - # single channel, 16 bit image depth_image = out.astype("uint16") - # convert to uint8 depth_image = cv2.convertScaleAbs(depth_image, alpha=255.0/65535.0) - # remove near if thr_a != 0: thr_a = thr_a/100*255 depth_image = cv2.threshold(depth_image, thr_a, 255, cv2.THRESH_TOZERO)[1] - # invert image depth_image = cv2.bitwise_not(depth_image) - # remove bg if thr_b != 0: thr_b = thr_b/100*255 depth_image = cv2.threshold(depth_image, thr_b, 255, cv2.THRESH_TOZERO)[1] - detected_map = depth_image detected_map = HWC3(detected_map) - img = resize_image(input_image, image_resolution) H, W, _C = img.shape - detected_map = cv2.resize(detected_map, (W, H), interpolation=cv2.INTER_LINEAR) - + if opts.control_move_processor: + self.model.to('cpu') if output_type == "pil": detected_map = Image.fromarray(detected_map) - return detected_map diff --git a/modules/control/proc/lineart.py b/modules/control/proc/lineart.py index 7f7aef10a..afd55d9a7 100644 --- a/modules/control/proc/lineart.py +++ b/modules/control/proc/lineart.py @@ -1,6 +1,4 @@ import os -import warnings - import cv2 import numpy as np import torch @@ -8,9 +6,9 @@ import torch.nn as nn from einops import rearrange from huggingface_hub import hf_hub_download from PIL import Image - +from modules import devices +from modules.shared import opts from modules.control.util import HWC3, resize_image - norm_layer = nn.InstanceNorm2d @@ -124,21 +122,12 @@ class LineartDetector: return self def __call__(self, input_image, coarse=False, detect_resolution=512, image_resolution=512, output_type="pil", **kwargs): - if "return_pil" in kwargs: - warnings.warn("return_pil is deprecated. Use output_type instead.", DeprecationWarning) - output_type = "pil" if kwargs["return_pil"] else "np" - if type(output_type) is bool: - warnings.warn("Passing `True` or `False` to `output_type` is deprecated and will raise an error in future versions") - if output_type: - output_type = "pil" - + self.model.to(devices.device) device = next(iter(self.model.parameters())).device if not isinstance(input_image, np.ndarray): input_image = np.array(input_image, dtype=np.uint8) - input_image = HWC3(input_image) input_image = resize_image(input_image, detect_resolution) - model = self.model_coarse if coarse else self.model assert input_image.ndim == 3 image = input_image @@ -146,21 +135,16 @@ class LineartDetector: image = image / 255.0 image = rearrange(image, 'h w c -> 1 c h w') line = model(image)[0][0] - line = line.cpu().numpy() line = (line * 255.0).clip(0, 255).astype(np.uint8) - detected_map = line - detected_map = HWC3(detected_map) - img = resize_image(input_image, image_resolution) H, W, _C = img.shape - detected_map = cv2.resize(detected_map, (W, H), interpolation=cv2.INTER_LINEAR) detected_map = 255 - detected_map - + if opts.control_move_processor: + self.model.to('cpu') if output_type == "pil": detected_map = Image.fromarray(detected_map) - return detected_map diff --git a/modules/control/proc/lineart_anime.py b/modules/control/proc/lineart_anime.py index 9eb4fcc09..541fcad61 100644 --- a/modules/control/proc/lineart_anime.py +++ b/modules/control/proc/lineart_anime.py @@ -1,7 +1,5 @@ import functools import os -import warnings - import cv2 import numpy as np import torch @@ -9,7 +7,8 @@ import torch.nn as nn from einops import rearrange from huggingface_hub import hf_hub_download from PIL import Image - +from modules import devices +from modules.shared import opts from modules.control.util import HWC3, resize_image @@ -120,12 +119,10 @@ class LineartAnimeDetector: @classmethod def from_pretrained(cls, pretrained_model_or_path, filename=None, cache_dir=None): filename = filename or "netG.pth" - if os.path.isdir(pretrained_model_or_path): model_path = os.path.join(pretrained_model_or_path, filename) else: model_path = hf_hub_download(pretrained_model_or_path, filename, cache_dir=cache_dir) - norm_layer = functools.partial(nn.InstanceNorm2d, affine=False, track_running_stats=False) net = UnetGenerator(3, 1, 8, 64, norm_layer=norm_layer, use_dropout=False) ckpt = torch.load(model_path) @@ -135,7 +132,6 @@ class LineartAnimeDetector: del ckpt[key] net.load_state_dict(ckpt) net.eval() - return cls(net) def to(self, device): @@ -143,21 +139,12 @@ class LineartAnimeDetector: return self def __call__(self, input_image, detect_resolution=512, image_resolution=512, output_type="pil", **kwargs): - if "return_pil" in kwargs: - warnings.warn("return_pil is deprecated. Use output_type instead.", DeprecationWarning) - output_type = "pil" if kwargs["return_pil"] else "np" - if type(output_type) is bool: - warnings.warn("Passing `True` or `False` to `output_type` is deprecated and will raise an error in future versions") - if output_type: - output_type = "pil" - + self.model.to(devices.device) device = next(iter(self.model.parameters())).device if not isinstance(input_image, np.ndarray): input_image = np.array(input_image, dtype=np.uint8) - input_image = HWC3(input_image) input_image = resize_image(input_image, detect_resolution) - H, W, _C = input_image.shape Hn = 256 * int(np.ceil(float(H) / 256.0)) Wn = 256 * int(np.ceil(float(W) / 256.0)) @@ -165,24 +152,18 @@ class LineartAnimeDetector: image_feed = torch.from_numpy(img).float().to(device) image_feed = image_feed / 127.5 - 1.0 image_feed = rearrange(image_feed, 'h w c -> 1 c h w') - line = self.model(image_feed)[0, 0] * 127.5 + 127.5 line = line.cpu().numpy() - line = cv2.resize(line, (W, H), interpolation=cv2.INTER_CUBIC) line = line.clip(0, 255).astype(np.uint8) - detected_map = line - detected_map = HWC3(detected_map) - img = resize_image(input_image, image_resolution) H, W, _C = img.shape - detected_map = cv2.resize(detected_map, (W, H), interpolation=cv2.INTER_LINEAR) detected_map = 255 - detected_map - + if opts.control_move_processor: + self.model.to('cpu') if output_type == "pil": detected_map = Image.fromarray(detected_map) - return detected_map diff --git a/modules/control/proc/marigold/__init__.py b/modules/control/proc/marigold/__init__.py index ede957786..6ea346430 100644 --- a/modules/control/proc/marigold/__init__.py +++ b/modules/control/proc/marigold/__init__.py @@ -1,6 +1,7 @@ from PIL import Image from modules.control.util import HWC3, resize_image from modules import devices +from modules.shared import opts from .marigold_pipeline import MarigoldPipeline @@ -39,6 +40,8 @@ class MarigoldDetector: show_progress_bar=True, ) depth_map = res.depth_colored if color_map != 'None' else res.depth_np + if opts.control_move_processor: + self.model.to('cpu') if output_type == "pil": return Image.fromarray(depth_map) else: diff --git a/modules/control/proc/mediapipe_face.py b/modules/control/proc/mediapipe_face.py index 187f14765..7c26ff410 100644 --- a/modules/control/proc/mediapipe_face.py +++ b/modules/control/proc/mediapipe_face.py @@ -1,4 +1,3 @@ -import warnings from typing import Union import cv2 import numpy as np @@ -15,37 +14,18 @@ class MediapipeFaceDetector: detect_resolution: int = 512, image_resolution: int = 512, **kwargs): - from .mediapipe_face_util import generate_annotation - if "image" in kwargs: - warnings.warn("image is deprecated, please use `input_image=...` instead.", DeprecationWarning) - input_image = kwargs.pop("image") if input_image is None: raise ValueError("input_image must be defined.") - - if "return_pil" in kwargs: - warnings.warn("return_pil is deprecated. Use output_type instead.", DeprecationWarning) - output_type = "pil" if kwargs["return_pil"] else "np" - if type(output_type) is bool: - warnings.warn("Passing `True` or `False` to `output_type` is deprecated and will raise an error in future versions") - if output_type: - output_type = "pil" - if not isinstance(input_image, np.ndarray): input_image = np.array(input_image, dtype=np.uint8) - input_image = HWC3(input_image) input_image = resize_image(input_image, detect_resolution) - detected_map = generate_annotation(input_image, max_faces, min_confidence) detected_map = HWC3(detected_map) - img = resize_image(input_image, image_resolution) H, W, _C = img.shape - detected_map = cv2.resize(detected_map, (W, H), interpolation=cv2.INTER_LINEAR) - if output_type == "pil": detected_map = Image.fromarray(detected_map) - return detected_map diff --git a/modules/control/proc/midas/__init__.py b/modules/control/proc/midas/__init__.py index 13a5ad061..ba6d1b0e5 100644 --- a/modules/control/proc/midas/__init__.py +++ b/modules/control/proc/midas/__init__.py @@ -6,8 +6,9 @@ import torch from einops import rearrange from huggingface_hub import hf_hub_download from PIL import Image - from modules.control.util import HWC3, resize_image +from modules import devices +from modules.shared import opts from .api import MiDaSInference @@ -21,14 +22,11 @@ class MidasDetector: filename = filename or "annotator/ckpts/dpt_hybrid-midas-501f0c75.pt" else: filename = filename or "dpt_hybrid-midas-501f0c75.pt" - if os.path.isdir(pretrained_model_or_path): model_path = os.path.join(pretrained_model_or_path, filename) else: model_path = hf_hub_download(pretrained_model_or_path, filename, cache_dir=cache_dir) - model = MiDaSInference(model_type=model_type, model_path=model_path) - return cls(model) @@ -37,16 +35,15 @@ class MidasDetector: return self def __call__(self, input_image, a=np.pi * 2.0, bg_th=0.1, depth_and_normal=False, detect_resolution=512, image_resolution=512, output_type=None): + self.model.to(devices.device) device = next(iter(self.model.parameters())).device if not isinstance(input_image, np.ndarray): input_image = np.array(input_image, dtype=np.uint8) output_type = output_type or "pil" else: output_type = output_type or "np" - input_image = HWC3(input_image) input_image = resize_image(input_image, detect_resolution) - assert input_image.ndim == 3 image_depth = input_image image_depth = torch.from_numpy(image_depth).float() @@ -54,13 +51,11 @@ class MidasDetector: image_depth = image_depth / 127.5 - 1.0 image_depth = rearrange(image_depth, 'h w c -> 1 c h w') depth = self.model(image_depth)[0] - depth_pt = depth.clone() depth_pt -= torch.min(depth_pt) depth_pt /= torch.max(depth_pt) depth_pt = depth_pt.cpu().numpy() depth_image = (depth_pt * 255.0).clip(0, 255).astype(np.uint8) - if depth_and_normal: depth_np = depth.cpu().numpy() x = cv2.Sobel(depth_np, cv2.CV_32F, 1, 0, ksize=3) @@ -71,23 +66,20 @@ class MidasDetector: normal = np.stack([x, y, z], axis=2) normal /= np.sum(normal ** 2.0, axis=2, keepdims=True) ** 0.5 normal_image = (normal * 127.5 + 127.5).clip(0, 255).astype(np.uint8)[:, :, ::-1] - depth_image = HWC3(depth_image) if depth_and_normal: normal_image = HWC3(normal_image) - img = resize_image(input_image, image_resolution) - H, W, C = img.shape - + H, W, _C = img.shape depth_image = cv2.resize(depth_image, (W, H), interpolation=cv2.INTER_LINEAR) if depth_and_normal: normal_image = cv2.resize(normal_image, (W, H), interpolation=cv2.INTER_LINEAR) - if output_type == "pil": depth_image = Image.fromarray(depth_image) if depth_and_normal: normal_image = Image.fromarray(normal_image) - + if opts.control_move_processor: + self.model.to('cpu') if depth_and_normal: return depth_image, normal_image else: diff --git a/modules/control/proc/mlsd/__init__.py b/modules/control/proc/mlsd/__init__.py index 456e1050d..ea26c5b0d 100644 --- a/modules/control/proc/mlsd/__init__.py +++ b/modules/control/proc/mlsd/__init__.py @@ -1,12 +1,11 @@ import os -import warnings - import cv2 import numpy as np import torch from huggingface_hub import hf_hub_download from PIL import Image - +from modules import devices +from modules.shared import opts from modules.control.util import HWC3, resize_image from .models.mbv2_mlsd_large import MobileV2_MLSD_Large from .utils import pred_lines @@ -22,16 +21,13 @@ class MLSDdetector: filename = filename or "annotator/ckpts/mlsd_large_512_fp32.pth" else: filename = filename or "mlsd_large_512_fp32.pth" - if os.path.isdir(pretrained_model_or_path): model_path = os.path.join(pretrained_model_or_path, filename) else: model_path = hf_hub_download(pretrained_model_or_path, filename, cache_dir=cache_dir) - model = MobileV2_MLSD_Large() model.load_state_dict(torch.load(model_path), strict=True) model.eval() - return cls(model) def to(self, device): @@ -39,20 +35,11 @@ class MLSDdetector: return self def __call__(self, input_image, thr_v=0.1, thr_d=0.1, detect_resolution=512, image_resolution=512, output_type="pil", **kwargs): - if "return_pil" in kwargs: - warnings.warn("return_pil is deprecated. Use output_type instead.", DeprecationWarning) - output_type = "pil" if kwargs["return_pil"] else "np" - if type(output_type) is bool: - warnings.warn("Passing `True` or `False` to `output_type` is deprecated and will raise an error in future versions") - if output_type: - output_type = "pil" - + self.model.to(devices.device) if not isinstance(input_image, np.ndarray): input_image = np.array(input_image, dtype=np.uint8) - input_image = HWC3(input_image) input_image = resize_image(input_image, detect_resolution) - assert input_image.ndim == 3 img = input_image img_output = np.zeros_like(img) @@ -63,16 +50,13 @@ class MLSDdetector: cv2.line(img_output, (x_start, y_start), (x_end, y_end), [255, 255, 255], 1) except Exception: pass - detected_map = img_output[:, :, 0] detected_map = HWC3(detected_map) - img = resize_image(input_image, image_resolution) - H, W, C = img.shape - + H, W, _C = img.shape detected_map = cv2.resize(detected_map, (W, H), interpolation=cv2.INTER_LINEAR) - if output_type == "pil": detected_map = Image.fromarray(detected_map) - + if opts.control_move_processor: + self.model.to('cpu') return detected_map diff --git a/modules/control/proc/normalbae/__init__.py b/modules/control/proc/normalbae/__init__.py index 852189f68..ba10570c6 100644 --- a/modules/control/proc/normalbae/__init__.py +++ b/modules/control/proc/normalbae/__init__.py @@ -1,7 +1,5 @@ import os import types -import warnings - import cv2 import numpy as np import torch @@ -9,7 +7,8 @@ import torchvision.transforms as transforms from einops import rearrange from huggingface_hub import hf_hub_download from PIL import Image - +from modules import devices +from modules.shared import opts from modules.control.util import HWC3, resize_image from .nets.NNET import NNET @@ -25,7 +24,6 @@ def load_checkpoint(fpath, model): load_dict[k_] = v else: load_dict[k] = v - model.load_state_dict(load_dict) return model @@ -37,12 +35,10 @@ class NormalBaeDetector: @classmethod def from_pretrained(cls, pretrained_model_or_path, filename=None, cache_dir=None): filename = filename or "scannet.pt" - if os.path.isdir(pretrained_model_or_path): model_path = os.path.join(pretrained_model_or_path, filename) else: model_path = hf_hub_download(pretrained_model_or_path, filename, cache_dir=cache_dir) - args = types.SimpleNamespace() args.mode = 'client' args.architecture = 'BN' @@ -52,7 +48,6 @@ class NormalBaeDetector: model = NNET(args) model = load_checkpoint(model_path, model) model.eval() - return cls(model) def to(self, device): @@ -61,14 +56,7 @@ class NormalBaeDetector: def __call__(self, input_image, detect_resolution=512, image_resolution=512, output_type="pil", **kwargs): - if "return_pil" in kwargs: - warnings.warn("return_pil is deprecated. Use output_type instead.", DeprecationWarning) - output_type = "pil" if kwargs["return_pil"] else "np" - if type(output_type) is bool: - warnings.warn("Passing `True` or `False` to `output_type` is deprecated and will raise an error in future versions") - if output_type: - output_type = "pil" - + self.model.to(devices.device) device = next(iter(self.model.parameters())).device if not isinstance(input_image, np.ndarray): input_image = np.array(input_image, dtype=np.uint8) @@ -89,19 +77,15 @@ class NormalBaeDetector: # d = torch.maximum(d, torch.ones_like(d) * 1e-5) # normal /= d normal = ((normal + 1) * 0.5).clip(0, 1) - normal = rearrange(normal[0], 'c h w -> h w c').cpu().numpy() normal_image = (normal * 255.0).clip(0, 255).astype(np.uint8) - detected_map = normal_image detected_map = HWC3(detected_map) - img = resize_image(input_image, image_resolution) H, W, _C = img.shape - detected_map = cv2.resize(detected_map, (W, H), interpolation=cv2.INTER_LINEAR) - if output_type == "pil": detected_map = Image.fromarray(detected_map) - + if opts.control_move_processor: + self.model.to('cpu') return detected_map diff --git a/modules/control/proc/openpose/__init__.py b/modules/control/proc/openpose/__init__.py index 398b7ec40..80649e213 100644 --- a/modules/control/proc/openpose/__init__.py +++ b/modules/control/proc/openpose/__init__.py @@ -6,27 +6,23 @@ # 5th Edited by ControlNet (Improved JSON serialization/deserialization, and lots of bug fixs) # This preprocessor is licensed by CMU for non-commercial use only. - import os - os.environ["KMP_DUPLICATE_LIB_OK"] = "TRUE" - -import json import warnings -from typing import Callable, List, NamedTuple, Tuple, Union - +from typing import List, NamedTuple, Tuple, Union import cv2 import numpy as np -import torch from huggingface_hub import hf_hub_download from PIL import Image - +from modules import devices +from modules.shared import opts from modules.control.util import HWC3, resize_image from . import util from .body import Body, BodyResult, Keypoint from .face import Face from .hand import Hand + HandResult = List[Keypoint] FaceResult = List[Keypoint] @@ -169,7 +165,7 @@ class OpenposeDetector: List[PoseResult]: A list of PoseResult objects containing the detected poses. """ oriImg = oriImg[:, :, ::-1].copy() - H, W, C = oriImg.shape + H, W, _C = oriImg.shape candidate, subset = self.body_estimation(oriImg) bodies = self.body_estimation.format_body_result(candidate, subset) @@ -196,11 +192,11 @@ class OpenposeDetector: return results def __call__(self, input_image, detect_resolution=512, image_resolution=512, include_body=True, include_hand=False, include_face=False, hand_and_face=None, output_type="pil", **kwargs): + self.to(devices.device) if hand_and_face is not None: warnings.warn("hand_and_face is deprecated. Use include_hand and include_face instead.", DeprecationWarning) include_hand = hand_and_face include_face = hand_and_face - if "return_pil" in kwargs: warnings.warn("return_pil is deprecated. Use output_type instead.", DeprecationWarning) output_type = "pil" if kwargs["return_pil"] else "np" @@ -208,26 +204,20 @@ class OpenposeDetector: warnings.warn("Passing `True` or `False` to `output_type` is deprecated and will raise an error in future versions") if output_type: output_type = "pil" - if not isinstance(input_image, np.ndarray): input_image = np.array(input_image, dtype=np.uint8) - input_image = HWC3(input_image) input_image = resize_image(input_image, detect_resolution) - H, W, C = input_image.shape - + H, W, _C = input_image.shape poses = self.detect_poses(input_image, include_hand, include_face) canvas = draw_poses(poses, H, W, draw_body=include_body, draw_hand=include_hand, draw_face=include_face) - detected_map = canvas detected_map = HWC3(detected_map) - img = resize_image(input_image, image_resolution) - H, W, C = img.shape - + H, W, _C = img.shape detected_map = cv2.resize(detected_map, (W, H), interpolation=cv2.INTER_LINEAR) - + if opts.control_move_processor: + self.to('cpu') if output_type == "pil": detected_map = Image.fromarray(detected_map) - return detected_map diff --git a/modules/control/proc/pidi.py b/modules/control/proc/pidi.py index 4c661b923..078525f2b 100644 --- a/modules/control/proc/pidi.py +++ b/modules/control/proc/pidi.py @@ -1,53 +1,41 @@ import os -import warnings - import cv2 import numpy as np import torch from einops import rearrange from huggingface_hub import hf_hub_download from PIL import Image - +from modules import devices +from modules.shared import opts from modules.control.util import HWC3, nms, resize_image, safe_step from .pidi_model import pidinet class PidiNetDetector: - def __init__(self, netNetwork): - self.netNetwork = netNetwork + def __init__(self, model): + self.model = model @classmethod def from_pretrained(cls, pretrained_model_or_path, filename=None, cache_dir=None): filename = filename or "table5_pidinet.pth" - if os.path.isdir(pretrained_model_or_path): model_path = os.path.join(pretrained_model_or_path, filename) else: model_path = hf_hub_download(pretrained_model_or_path, filename, cache_dir=cache_dir) - - netNetwork = pidinet() - netNetwork.load_state_dict({k.replace('module.', ''): v for k, v in torch.load(model_path)['state_dict'].items()}) - netNetwork.eval() - - return cls(netNetwork) + model = pidinet() + model.load_state_dict({k.replace('module.', ''): v for k, v in torch.load(model_path)['state_dict'].items()}) + model.eval() + return cls(model) def to(self, device): - self.netNetwork.to(device) + self.model.to(device) return self def __call__(self, input_image, detect_resolution=512, image_resolution=512, safe=False, output_type="pil", scribble=False, apply_filter=False, **kwargs): - if "return_pil" in kwargs: - warnings.warn("return_pil is deprecated. Use output_type instead.", DeprecationWarning) - output_type = "pil" if kwargs["return_pil"] else "np" - if type(output_type) is bool: - warnings.warn("Passing `True` or `False` to `output_type` is deprecated and will raise an error in future versions") - if output_type: - output_type = "pil" - - device = next(iter(self.netNetwork.parameters())).device + self.model.to(devices.device) + device = next(iter(self.model.parameters())).device if not isinstance(input_image, np.ndarray): input_image = np.array(input_image, dtype=np.uint8) - input_image = HWC3(input_image) input_image = resize_image(input_image, detect_resolution) assert input_image.ndim == 3 @@ -55,29 +43,25 @@ class PidiNetDetector: image_pidi = torch.from_numpy(input_image).float().to(device) image_pidi = image_pidi / 255.0 image_pidi = rearrange(image_pidi, 'h w c -> 1 c h w') - edge = self.netNetwork(image_pidi)[-1] + edge = self.model(image_pidi)[-1] edge = edge.cpu().numpy() if apply_filter: edge = edge > 0.5 if safe: edge = safe_step(edge) edge = (edge * 255.0).clip(0, 255).astype(np.uint8) - detected_map = edge[0, 0] detected_map = HWC3(detected_map) - img = resize_image(input_image, image_resolution) H, W, _C = img.shape - detected_map = cv2.resize(detected_map, (W, H), interpolation=cv2.INTER_LINEAR) - if scribble: detected_map = nms(detected_map, 127, 3.0) detected_map = cv2.GaussianBlur(detected_map, (0, 0), 3.0) detected_map[detected_map > 4] = 255 detected_map[detected_map < 255] = 0 - + if opts.control_move_processor: + self.model.to('cpu') if output_type == "pil": detected_map = Image.fromarray(detected_map) - return detected_map diff --git a/modules/control/proc/pidi/LICENSE b/modules/control/proc/pidi/LICENSE deleted file mode 100644 index 913b6cf92..000000000 --- a/modules/control/proc/pidi/LICENSE +++ /dev/null @@ -1,21 +0,0 @@ -It is just for research purpose, and commercial use should be contacted with authors first. - -Copyright (c) 2021 Zhuo Su - -Permission is hereby granted, free of charge, to any person obtaining a copy -of this software and associated documentation files (the "Software"), to deal -in the Software without restriction, including without limitation the rights -to use, copy, modify, merge, publish, distribute, sublicense, and/or sell -copies of the Software, and to permit persons to whom the Software is -furnished to do so, subject to the following conditions: - -The above copyright notice and this permission notice shall be included in all -copies or substantial portions of the Software. - -THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR -IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, -FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE -AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER -LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, -OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE -SOFTWARE. \ No newline at end of file diff --git a/modules/control/proc/segment_anything/__init__.py b/modules/control/proc/segment_anything/__init__.py index bcd7195c7..d698b46fc 100644 --- a/modules/control/proc/segment_anything/__init__.py +++ b/modules/control/proc/segment_anything/__init__.py @@ -4,16 +4,15 @@ # This source code is licensed under the license found in the # LICENSE file in the root directory of this source tree. -import os import warnings from typing import Union - import cv2 import numpy as np import torch from huggingface_hub import hf_hub_download from PIL import Image - +from modules import devices +from modules.shared import opts from modules.control.util import HWC3, resize_image from .automatic_mask_generator import SamAutomaticMaskGenerator from .build_sam import sam_model_registry @@ -21,7 +20,7 @@ from .build_sam import sam_model_registry class SamDetector: def __init__(self, mask_generator: SamAutomaticMaskGenerator = None): - self.mask_generator = mask_generator + self.model = mask_generator @classmethod def from_pretrained(cls, model_path, filename, model_type, cache_dir=None): @@ -30,14 +29,9 @@ class SamDetector: download weights from https://github.com/facebookresearch/segment-anything """ model_path = hf_hub_download(model_path, filename, cache_dir=cache_dir) - sam = sam_model_registry[model_type](checkpoint=model_path) - - if torch.cuda.is_available(): - sam.to("cuda") - + sam.to(devices.device) mask_generator = SamAutomaticMaskGenerator(sam) - return cls(mask_generator) @@ -55,37 +49,30 @@ class SamDetector: for i in range(3): img[:,:,i] = gen.integers(255, dtype=np.uint8) final_img.paste(Image.fromarray(img, mode="RGB"), (0, 0), Image.fromarray(np.uint8(m*255))) - return np.array(final_img, dtype=np.uint8) def __call__(self, input_image: Union[np.ndarray, Image.Image]=None, detect_resolution=512, image_resolution=512, output_type="pil", **kwargs) -> Image.Image: if "image" in kwargs: warnings.warn("image is deprecated, please use `input_image=...` instead.", DeprecationWarning) input_image = kwargs.pop("image") - if input_image is None: raise ValueError("input_image must be defined.") - if not isinstance(input_image, np.ndarray): input_image = np.array(input_image, dtype=np.uint8) - input_image = HWC3(input_image) input_image = resize_image(input_image, detect_resolution) - # Generate Masks - masks = self.mask_generator.generate(input_image) + self.model.predictor.model.to(devices.device) + masks = self.model.generate(input_image) + if opts.control_move_processor: + self.model.predictor.model.to('cpu') # Create map image_map = self.show_anns(masks) - detected_map = image_map detected_map = HWC3(detected_map) - img = resize_image(input_image, image_resolution) H, W, _C = img.shape - detected_map = cv2.resize(detected_map, (W, H), interpolation=cv2.INTER_LINEAR) - if output_type == "pil": detected_map = Image.fromarray(detected_map) - return detected_map diff --git a/modules/control/proc/segment_anything/automatic_mask_generator.py b/modules/control/proc/segment_anything/automatic_mask_generator.py index a5029053e..ba0287862 100644 --- a/modules/control/proc/segment_anything/automatic_mask_generator.py +++ b/modules/control/proc/segment_anything/automatic_mask_generator.py @@ -4,12 +4,10 @@ # This source code is licensed under the license found in the # LICENSE file in the root directory of this source tree. +from typing import Any, Dict, List, Optional, Tuple import numpy as np import torch from torchvision.ops.boxes import batched_nms, box_area # type: ignore - -from typing import Any, Dict, List, Optional, Tuple - from .modeling import Sam from .predictor import SamPredictor from .utils.amg import ( @@ -114,12 +112,6 @@ class SamAutomaticMaskGenerator: "uncompressed_rle", "coco_rle", ], f"Unknown output_mode {output_mode}." - if output_mode == "coco_rle": - from pycocotools import mask as mask_utils # type: ignore - - if min_mask_region_area > 0: - import cv2 # type: ignore - self.predictor = SamPredictor(model) self.points_per_batch = points_per_batch self.pred_iou_thresh = pred_iou_thresh diff --git a/modules/control/proc/segment_anything/predictor.py b/modules/control/proc/segment_anything/predictor.py index 742a34ef1..cafb8ea18 100644 --- a/modules/control/proc/segment_anything/predictor.py +++ b/modules/control/proc/segment_anything/predictor.py @@ -4,13 +4,10 @@ # This source code is licensed under the license found in the # LICENSE file in the root directory of this source tree. +from typing import Optional, Tuple import numpy as np import torch - from .modeling import Sam - -from typing import Optional, Tuple - from .utils.transforms import ResizeLongestSide diff --git a/modules/control/proc/zoe/__init__.py b/modules/control/proc/zoe/__init__.py index e140496a4..bb18201e0 100644 --- a/modules/control/proc/zoe/__init__.py +++ b/modules/control/proc/zoe/__init__.py @@ -7,7 +7,8 @@ from einops import rearrange from huggingface_hub import hf_hub_download from PIL import Image import safetensors - +from modules import devices +from modules.shared import opts from modules.control.util import HWC3, resize_image from .zoedepth.models.zoedepth.zoedepth_v1 import ZoeDepth from .zoedepth.models.zoedepth_nk.zoedepth_nk_v1 import ZoeDepthNK @@ -25,7 +26,6 @@ class ZoeDetector: model_path = os.path.join(pretrained_model_or_path, filename) else: model_path = hf_hub_download(pretrained_model_or_path, filename, cache_dir=cache_dir) - if model_type == "zoedepth": model_cls = ZoeDepth elif model_type == "zoedepth_nk": @@ -53,45 +53,37 @@ class ZoeDetector: return self def __call__(self, input_image, detect_resolution=512, image_resolution=512, output_type=None, gamma_corrected=False): + self.model.to(devices.device) device = next(iter(self.model.parameters())).device if not isinstance(input_image, np.ndarray): input_image = np.array(input_image, dtype=np.uint8) output_type = output_type or "pil" else: output_type = output_type or "np" - input_image = HWC3(input_image) input_image = resize_image(input_image, detect_resolution) - assert input_image.ndim == 3 image_depth = input_image image_depth = torch.from_numpy(image_depth).float().to(device) image_depth = image_depth / 255.0 image_depth = rearrange(image_depth, 'h w c -> 1 c h w') depth = self.model.infer(image_depth) - + if opts.control_move_processor: + self.model.to('cpu') depth = depth[0, 0].cpu().numpy() - vmin = np.percentile(depth, 2) vmax = np.percentile(depth, 85) - depth -= vmin depth /= vmax - vmin depth = 1.0 - depth - if gamma_corrected: depth = np.power(depth, 2.2) depth_image = (depth * 255.0).clip(0, 255).astype(np.uint8) - detected_map = depth_image detected_map = HWC3(detected_map) - img = resize_image(input_image, image_resolution) H, W, _C = img.shape - detected_map = cv2.resize(detected_map, (W, H), interpolation=cv2.INTER_LINEAR) - if output_type == "pil": detected_map = Image.fromarray(detected_map) - return detected_map diff --git a/modules/control/processors.py b/modules/control/processors.py index f300bcd2c..99e9e9e0b 100644 --- a/modules/control/processors.py +++ b/modules/control/processors.py @@ -131,12 +131,13 @@ class Processor(): if processor_id is not None: self.load() - def reset(self): + def reset(self, processor_id: str = None): if self.model is not None: log.debug(f'Control Processor unloaded: id="{self.processor_id}"') self.model = None - self.processor_id = None + self.processor_id = processor_id self.override = None + devices.torch_gc() self.load_config = { 'cache_dir': cache_dir } def config(self, processor_id = None): @@ -204,20 +205,22 @@ class Processor(): return f'Processor load filed: {processor_id}' def __call__(self, image_input: Image, mode: str = 'RGB'): + if self.processor_id is None or self.processor_id == 'None': + return image_input if self.override is not None: image_input = self.override image_process = image_input if image_input is None: log.error('Control Processor: no input') return image_process - if self.model is None: - # log.error('Control Processor: model not loaded') - return image_process if config[self.processor_id].get('dirty', False): processor_id = self.processor_id config[processor_id].pop('dirty') self.reset() self.load(processor_id) + if self.model is None: + # log.error('Control Processor: model not loaded') + return image_process try: t0 = time.time() kwargs = config.get(self.processor_id, {}).get('params', None) diff --git a/modules/control/run.py b/modules/control/run.py index ac4181989..dfcf72a1f 100644 --- a/modules/control/run.py +++ b/modules/control/run.py @@ -4,7 +4,6 @@ import math from typing import List, Union import cv2 import numpy as np -import diffusers from PIL import Image from modules.control import util from modules.control import unit @@ -384,8 +383,12 @@ def control_run(units: List[unit.Unit], inputs, inits, mask, unit_type: str, is_ masked_image = masking.run_mask(input_image=input_image, input_mask=mask, return_type='Masked') if mask is not None else input_image for i, process in enumerate(active_process): # list[image] image_mode = 'L' if unit_type == 'adapter' and len(active_model) > i and ('Canny' in active_model[i].model_id or 'Sketch' in active_model[i].model_id) else 'RGB' # t2iadapter canny and sketch work in grayscale only - debug(f'Control: process={[process.processor_id for p in active_process]} i={i} image={p.image}') - p.image.append(process(masked_image, image_mode)) + debug(f'Control: process="{process.processor_id}" i={i} image={p.image}') + processed_image = process(masked_image, image_mode) + p.image.append(processed_image) + if shared.opts.control_unload_processor: + processors.config[process.processor_id]['dirty'] = True # to force reload + process.model = None if p.image is not None and len(p.image) > 0: p.init_images = p.image diff --git a/modules/control/test.py b/modules/control/test.py index 8345997a0..fcc0eb31d 100644 --- a/modules/control/test.py +++ b/modules/control/test.py @@ -21,7 +21,8 @@ def test_processors(image): processor_id = f'{processor_id} error' else: output = processor(image) - processor.reset() + if shared.opts.control_unload_processor: + processor.reset() if output.size != image.size: output = output.resize(image.size, Image.Resampling.LANCZOS) if output.mode != image.mode: diff --git a/modules/shared.py b/modules/shared.py index 0141a52da..1e8e67049 100644 --- a/modules/shared.py +++ b/modules/shared.py @@ -629,6 +629,11 @@ options_templates.update(options_section(('postprocessing', "Postprocessing"), { "upscaler_tile_overlap": OptionInfo(8, "Upscaler tile overlap", gr.Slider, {"minimum": 0, "maximum": 64, "step": 1}), })) +options_templates.update(options_section(('control', "Control"), { + "control_move_processor": OptionInfo(False, "Processor move to CPU after use"), + "control_unload_processor": OptionInfo(False, "Processor unload after use"), +})) + options_templates.update(options_section(('training', "Training"), { "unload_models_when_training": OptionInfo(False, "Move VAE and CLIP to RAM when training"), "pin_memory": OptionInfo(True, "Pin training dataset to memory"),