Public/automatic
1
0
Fork 0
mirror of https://github.com/vladmandic/automatic synced 2026-09-19 17:24:32 +02:00
Code Issues Packages Projects Releases Wiki Activity

refactor all control processors to support unload and offload

Browse Source
This commit is contained in:
Vladimir Mandic
2024-01-25 11:45:47 -05:00
parent c19f95141a
commit 7aa6876759
27 changed files with 154 additions and 350 deletions
Download Patch File Download Diff File Expand all files Collapse all files
modules/control/proc/canny.py
-6
View File
@@ -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
modules/control/proc/depth_anything/__init__.py
+3
View File
@@ -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)
modules/control/proc/dpt.py
+20 -13
View File
@@ -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
modules/control/proc/dwpose/__init__.py
-11
View File
@@ -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
modules/control/proc/edge.py
-5
View File
@@ -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
modules/control/proc/glpn.py
+20 -16
View File
@@ -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
modules/control/proc/hed.py
+14 -31
View File
@@ -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
modules/control/proc/leres/__init__.py
+5 -17
View File
@@ -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
modules/control/proc/lineart.py
+5 -21
View File
@@ -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
modules/control/proc/lineart_anime.py
+5 -24
View File
@@ -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
modules/control/proc/marigold/__init__.py
+3
View File
@@ -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:
modules/control/proc/mediapipe_face.py
-20
View File
@@ -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
modules/control/proc/midas/__init__.py
+6 -14
View File
@@ -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:
modules/control/proc/mlsd/__init__.py
+6 -22
View File
@@ -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
modules/control/proc/normalbae/__init__.py
+5 -21
View File
@@ -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
modules/control/proc/openpose/__init__.py
+10 -20
View File
@@ -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
modules/control/proc/pidi.py
+14 -30
View File
@@ -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
modules/control/proc/pidi/LICENSE
-21
View File
@@ -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.
modules/control/proc/segment_anything/__init__.py
+8 -21
View File
@@ -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
modules/control/proc/segment_anything/automatic_mask_generator.py
+1 -9
View File
@@ -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
modules/control/proc/segment_anything/predictor.py
+1 -4
View File
@@ -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
modules/control/proc/zoe/__init__.py
+5 -13
View File
@@ -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