Added Mediapipe image processor for use as ControlNet preprocessor.

Also hacked in ability to specify HF subfolder when loading ControlNet models from string.
This commit is contained in:
user1 2023-05-23 16:21:13 -07:00 committed by Kent Keirsey
parent be79d088c0
commit 8960ceb98b
2 changed files with 31 additions and 4 deletions

View File

@ -84,9 +84,11 @@ CONTROLNET_DEFAULT_MODELS = [
##############################################
# ControlNetMediaPipeface, ControlNet v1.1
##############################################
"CrucibleAI/ControlNetMediaPipeFace",# SD 2.1?
# diffusion_sd15 needs to be passed to from_pretrained() as subfolder arg
# ["CrucibleAI/ControlNetMediaPipeFace", "diffusion_sd15"], # SD 1.5
# diffusion_sd15 needs to be passed to from_pretrained() as subfolder arg
# hacked t2l to split to model & subfolder if format is "model,subfolder"
"CrucibleAI/ControlNetMediaPipeFace,diffusion_sd15", # SD 1.5
"CrucibleAI/ControlNetMediaPipeFace", # SD 2.1?
]
CONTROLNET_NAME_VALUES = Literal[tuple(CONTROLNET_DEFAULT_MODELS)]
@ -403,3 +405,17 @@ class ContentShuffleImageProcessorInvocation(ImageProcessorInvocation, PILInvoca
# processed_image = zoe_depth_processor(image)
# return processed_image
class MediapipeFaceProcessorInvocation(ImageProcessorInvocation, PILInvocationConfig):
"""Applies mediapipe face processing to image"""
# fmt: off
type: Literal["mediapipe_face_processor"] = "mediapipe_face_processor"
# Inputs
max_faces: int = Field(default=1, ge=1, description="maximum number of faces to detect")
min_confidence: float = Field(default=0.5, ge=0, le=1, description="minimum confidence for face detection")
# fmt: on
def run_processor(self, image):
mediapipe_face_processor = MediapipeFaceDetector()
processed_image = mediapipe_face_processor(image)
return processed_image

View File

@ -277,6 +277,17 @@ class TextToLatentsInvocation(BaseInvocation):
control_models = []
for control_info in control_list:
# handle control models
if ("," in control_info.control_model):
control_model_split = control_info.control_model.split(",")
control_name = control_model_split[0]
control_subfolder = control_model_split[1]
print("Using HF model subfolders")
print(" control_name: ", control_name)
print(" control_subfolder: ", control_subfolder)
control_model = ControlNetModel.from_pretrained(control_name,
subfolder=control_subfolder,
torch_dtype=model.unet.dtype).to(model.device)
else:
control_model = ControlNetModel.from_pretrained(control_info.control_model,
torch_dtype=model.unet.dtype).to(model.device)
control_models.append(control_model)