kamp/IOPaint
1
0
Fork 0
Code Issues Pull Requests Releases 1 Activity

change controlnet depth preprocessor

Browse Source
This commit is contained in:
Qing 2024-01-10 21:25:51 +08:00
parent dfcb9017f1
commit a869982d13
2 changed files with 5 additions and 6 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

3
iopaint/model/controlnet.py
View File

@ -93,8 +93,9 @@ class ControlNet(DiffusionInpaintModel):
model_info.path,
controlnet=controlnet,
load_safety_checker=not disable_nsfw_checker,
torch_dtype=torch_dtype,
**model_kwargs,
).to(torch_dtype)
)
else:
self.model = handle_from_pretrained_exceptions(
PipeClass.from_pretrained,

8
iopaint/model/helper/controlnet_preprocess.py
View File

@ -23,11 +23,9 @@ def make_openpose_control_image(image: np.ndarray) -> Image:
def make_depth_control_image(image: np.ndarray) -> Image:
from transformers import pipeline
depth_estimator = pipeline("depth-estimation")
depth_image = depth_estimator(PIL.Image.fromarray(image))["depth"]
depth_image = np.array(depth_image)
from controlnet_aux import MidasDetector
midas = MidasDetector.from_pretrained("lllyasviel/Annotators")
depth_image = midas(image)
depth_image = depth_image[:, :, None]
depth_image = np.concatenate([depth_image, depth_image, depth_image], axis=2)
control_image = PIL.Image.fromarray(depth_image)