mirror of
https://github.com/invoke-ai/InvokeAI
synced 2024-08-30 20:32:17 +00:00
Added more preprocessor nodes for:
MidasDepth
ZoeDepth
MLSD
NormalBae
Pidi
LineartAnime
ContentShuffle
Removed pil_output options, ControlNet preprocessors should always output as PIL. Removed diagnostics and other general cleanup.
This commit is contained in:
user1
Kent Keirsey
parent
5acbbeecaa
commit
0e027ec3ef
@ -1,10 +1,9 @@
|
||||
# InvokeAI nodes for ControlNet image preprocessors
|
||||
# initial implementation by Gregg Helt, 2023
|
||||
# heavily leverages controlnet_aux package: https://github.com/patrickvonplaten/controlnet_aux
|
||||
|
||||
import numpy as np
|
||||
|
||||
from typing import Literal, Optional, Union, List
|
||||
from PIL import Image, ImageFilter, ImageOps
|
||||
from pydantic import BaseModel, Field
|
||||
|
||||
from ..models.image import ImageField, ImageType
|
||||
@ -26,83 +25,23 @@ from controlnet_aux import (
|
||||
OpenposeDetector,
|
||||
PidiNetDetector,
|
||||
ContentShuffleDetector,
|
||||
ZoeDetector,
|
||||
)
|
||||
ZoeDetector)
|
||||
|
||||
from .image import ImageOutput, build_image_output, PILInvocationConfig
|
||||
|
||||
CONTROLNET_DEFAULT_MODELS = [
|
||||
###########################################
|
||||
# lllyasviel sd v1.5, ControlNet v1.0 models
|
||||
##############################################
|
||||
"lllyasviel/sd-controlnet-canny",
|
||||
"lllyasviel/sd-controlnet-depth",
|
||||
"lllyasviel/sd-controlnet-hed",
|
||||
"lllyasviel/sd-controlnet-seg",
|
||||
"lllyasviel/sd-controlnet-openpose",
|
||||
"lllyasviel/sd-controlnet-scribble",
|
||||
"lllyasviel/sd-controlnet-normal",
|
||||
"lllyasviel/sd-controlnet-mlsd",
|
||||
|
||||
#############################################
|
||||
# lllyasviel sd v1.5, ControlNet v1.1 models
|
||||
#############################################
|
||||
"lllyasviel/control_v11p_sd15_canny",
|
||||
"lllyasviel/control_v11p_sd15_openpose",
|
||||
"lllyasviel/control_v11p_sd15_seg",
|
||||
# "lllyasviel/control_v11p_sd15_depth", # broken
|
||||
"lllyasviel/control_v11f1p_sd15_depth",
|
||||
"lllyasviel/control_v11p_sd15_normalbae",
|
||||
"lllyasviel/control_v11p_sd15_scribble",
|
||||
"lllyasviel/control_v11p_sd15_mlsd",
|
||||
"lllyasviel/control_v11p_sd15_softedge",
|
||||
"lllyasviel/control_v11p_sd15s2_lineart_anime",
|
||||
"lllyasviel/control_v11p_sd15_lineart",
|
||||
"lllyasviel/control_v11p_sd15_inpaint",
|
||||
# "lllyasviel/control_v11u_sd15_tile",
|
||||
# problem (temporary?) with huffingface "lllyasviel/control_v11u_sd15_tile",
|
||||
# so for now replace "lllyasviel/control_v11f1e_sd15_tile",
|
||||
"lllyasviel/control_v11e_sd15_shuffle",
|
||||
"lllyasviel/control_v11e_sd15_ip2p",
|
||||
"lllyasviel/control_v11f1e_sd15_tile",
|
||||
|
||||
#################################################
|
||||
# thibaud sd v2.1 models (ControlNet v1.0? or v1.1?
|
||||
##################################################
|
||||
"thibaud/controlnet-sd21-openpose-diffusers",
|
||||
"thibaud/controlnet-sd21-canny-diffusers",
|
||||
"thibaud/controlnet-sd21-depth-diffusers",
|
||||
"thibaud/controlnet-sd21-scribble-diffusers",
|
||||
"thibaud/controlnet-sd21-hed-diffusers",
|
||||
"thibaud/controlnet-sd21-zoedepth-diffusers",
|
||||
"thibaud/controlnet-sd21-color-diffusers",
|
||||
"thibaud/controlnet-sd21-openposev2-diffusers",
|
||||
"thibaud/controlnet-sd21-lineart-diffusers",
|
||||
"thibaud/controlnet-sd21-normalbae-diffusers",
|
||||
"thibaud/controlnet-sd21-ade20k-diffusers",
|
||||
|
||||
##############################################
|
||||
# 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
|
||||
]
|
||||
|
||||
CONTROLNET_NAME_VALUES = Literal[tuple(CONTROLNET_DEFAULT_MODELS)]
|
||||
|
||||
class ControlField(BaseModel):
|
||||
image: ImageField = Field(default=None, description="processed image")
|
||||
control_model: Optional[str] = Field(default=None, description="control model used")
|
||||
control_weight: Optional[float] = Field(default=1, description="weight given to controlnet")
|
||||
begin_step_percent: float = Field(default=0, ge=0, le=1,
|
||||
description="% of total steps at which controlnet is first applied")
|
||||
end_step_percent: float = Field(default=1, ge=0, le=1,
|
||||
description="% of total steps at which controlnet is last applied")
|
||||
# width: Optional[int] = Field(default=None, description="The width of the image in pixels")
|
||||
# height: Optional[int] = Field(default=None, description="The height of the image in pixels")
|
||||
# mode: Optional[str] = Field(default=None, description="The mode of the image")
|
||||
control_model: Optional[str] = Field(default=None, description="The control model used")
|
||||
control_weight: Optional[float] = Field(default=None, description="The control weight used")
|
||||
|
||||
class Config:
|
||||
schema_extra = {
|
||||
"required": ["image", "control_model", "control_weight", "begin_step_percent", "end_step_percent"]
|
||||
"required": ["image", "control_model", "control_weight"]
|
||||
# "required": ["type", "image", "width", "height", "mode"]
|
||||
}
|
||||
|
||||
|
||||
@ -110,50 +49,29 @@ class ControlOutput(BaseInvocationOutput):
|
||||
"""node output for ControlNet info"""
|
||||
# fmt: off
|
||||
type: Literal["control_output"] = "control_output"
|
||||
control: ControlField = Field(default=None, description="The control info dict")
|
||||
control: Optional[ControlField] = Field(default=None, description="The control info dict")
|
||||
raw_processed_image: ImageField = Field(default=None,
|
||||
description="outputs just the image info (also included in control output)")
|
||||
# fmt: on
|
||||
|
||||
|
||||
class PreprocessedControlInvocation(BaseInvocation, PILInvocationConfig):
|
||||
"""Base class for invocations that preprocess images for ControlNet"""
|
||||
|
||||
class ControlNetInvocation(BaseInvocation):
|
||||
"""Collects ControlNet info to pass to other nodes"""
|
||||
# fmt: off
|
||||
type: Literal["controlnet"] = "controlnet"
|
||||
# Inputs
|
||||
image: ImageField = Field(default=None, description="image to process")
|
||||
control_model: CONTROLNET_NAME_VALUES = Field(default="lllyasviel/sd-controlnet-canny",
|
||||
description="control model used")
|
||||
control_weight: float = Field(default=1.0, ge=0, le=1, description="weight given to controlnet")
|
||||
# TODO: add support in backend core for begin_step_percent, end_step_percent, guess_mode
|
||||
begin_step_percent: float = Field(default=0, ge=0, le=1,
|
||||
description="% of total steps at which controlnet is first applied")
|
||||
end_step_percent: float = Field(default=1, ge=0, le=1,
|
||||
description="% of total steps at which controlnet is last applied")
|
||||
# fmt: on
|
||||
|
||||
|
||||
def invoke(self, context: InvocationContext) -> ControlOutput:
|
||||
|
||||
return ControlOutput(
|
||||
control=ControlField(
|
||||
image=self.image,
|
||||
control_model=self.control_model,
|
||||
control_weight=self.control_weight,
|
||||
begin_step_percent=self.begin_step_percent,
|
||||
end_step_percent=self.end_step_percent,
|
||||
),
|
||||
)
|
||||
|
||||
# TODO: move image processors to separate file (image_analysis.py
|
||||
class ImageProcessorInvocation(BaseInvocation, PILInvocationConfig):
|
||||
# This super class handles invoke() call, which in turn calls run_processor(image)
|
||||
# subclasses override run_processor() instead of implementing their own invoke()
|
||||
class PreprocessedControlNetInvocation(BaseInvocation, PILInvocationConfig):
|
||||
"""Base class for invocations that preprocess images for ControlNet"""
|
||||
|
||||
# fmt: off
|
||||
type: Literal["image_processor"] = "image_processor"
|
||||
type: Literal["preprocessed_control"] = "preprocessed_control"
|
||||
# Inputs
|
||||
image: ImageField = Field(default=None, description="image to process")
|
||||
control_model: str = Field(default=None, description="control model to use")
|
||||
control_weight: float = Field(default=0.5, ge=0, le=1, description="control weight")
|
||||
# TODO: support additional ControlNet parameters (mostly just passthroughs to other nodes with ControlField inputs)
|
||||
# begin_step_percent: float = Field(default=0, ge=0, le=1,
|
||||
# description="% of total steps at which controlnet is first applied")
|
||||
# end_step_percent: float = Field(default=1, ge=0, le=1,
|
||||
# description="% of total steps at which controlnet is last applied")
|
||||
# guess_mode: bool = Field(default=False, description="use guess mode (controlnet ignores prompt)")
|
||||
# fmt: on
|
||||
|
||||
|
||||
@ -161,12 +79,12 @@ class ImageProcessorInvocation(BaseInvocation, PILInvocationConfig):
|
||||
# superclass just passes through image without processing
|
||||
return image
|
||||
|
||||
def invoke(self, context: InvocationContext) -> ImageOutput:
|
||||
raw_image = context.services.images.get(
|
||||
def invoke(self, context: InvocationContext) -> ControlOutput:
|
||||
image = context.services.images.get(
|
||||
self.image.image_type, self.image.image_name
|
||||
)
|
||||
# image type should be PIL.PngImagePlugin.PngImageFile ?
|
||||
processed_image = self.run_processor(raw_image)
|
||||
processed_image = self.run_processor(image)
|
||||
# currently can't see processed image in node UI without a showImage node,
|
||||
# so for now setting image_type to RESULT instead of INTERMEDIATE so will get saved in gallery
|
||||
# image_type = ImageType.INTERMEDIATE
|
||||
@ -180,22 +98,24 @@ class ImageProcessorInvocation(BaseInvocation, PILInvocationConfig):
|
||||
context.services.images.save(image_type, image_name, processed_image, metadata)
|
||||
|
||||
"""Builds an ImageOutput and its ImageField"""
|
||||
processed_image_field = ImageField(
|
||||
image_field = ImageField(
|
||||
image_name=image_name,
|
||||
image_type=image_type,
|
||||
)
|
||||
return ImageOutput(
|
||||
image=processed_image_field,
|
||||
width=processed_image.width,
|
||||
height=processed_image.height,
|
||||
mode=processed_image.mode,
|
||||
return ControlOutput(
|
||||
control=ControlField(
|
||||
image=image_field,
|
||||
control_model=self.control_model,
|
||||
control_weight=self.control_weight,
|
||||
),
|
||||
raw_processed_image=image_field,
|
||||
)
|
||||
|
||||
|
||||
class CannyImageProcessorInvocation(ImageProcessorInvocation, PILInvocationConfig):
|
||||
class CannyControlInvocation(PreprocessedControlNetInvocation, PILInvocationConfig):
|
||||
"""Canny edge detection for ControlNet"""
|
||||
# fmt: off
|
||||
type: Literal["canny_image_processor"] = "canny_image_processor"
|
||||
type: Literal["cannycontrol"] = "cannycontrol"
|
||||
# Input
|
||||
low_threshold: float = Field(default=100, ge=0, description="low threshold of Canny pixel gradient")
|
||||
high_threshold: float = Field(default=200, ge=0, description="high threshold of Canny pixel gradient")
|
||||
@ -207,15 +127,14 @@ class CannyImageProcessorInvocation(ImageProcessorInvocation, PILInvocationConfi
|
||||
return processed_image
|
||||
|
||||
|
||||
class HedImageprocessorInvocation(ImageProcessorInvocation, PILInvocationConfig):
|
||||
class HedControlNetInvocation(PreprocessedControlNetInvocation, PILInvocationConfig):
|
||||
"""Applies HED edge detection to image"""
|
||||
# fmt: off
|
||||
type: Literal["hed_image_processor"] = "hed_image_processor"
|
||||
type: Literal["hed_control"] = "hed_control"
|
||||
# Inputs
|
||||
detect_resolution: int = Field(default=512, ge=0, description="pixel resolution for edge detection")
|
||||
image_resolution: int = Field(default=512, ge=0, description="pixel resolution for output image")
|
||||
# safe not supported in controlnet_aux v0.0.3
|
||||
# safe: bool = Field(default=False, description="whether to use safe mode")
|
||||
safe: bool = Field(default=False, description="whether to use safe mode")
|
||||
scribble: bool = Field(default=False, description="whether to use scribble mode")
|
||||
# fmt: on
|
||||
|
||||
@ -224,17 +143,16 @@ class HedImageprocessorInvocation(ImageProcessorInvocation, PILInvocationConfig)
|
||||
processed_image = hed_processor(image,
|
||||
detect_resolution=self.detect_resolution,
|
||||
image_resolution=self.image_resolution,
|
||||
# safe not supported in controlnet_aux v0.0.3
|
||||
# safe=self.safe,
|
||||
safe=self.safe,
|
||||
scribble=self.scribble,
|
||||
)
|
||||
return processed_image
|
||||
|
||||
|
||||
class LineartImageProcessorInvocation(ImageProcessorInvocation, PILInvocationConfig):
|
||||
class LineartControlInvocation(PreprocessedControlNetInvocation, PILInvocationConfig):
|
||||
"""Applies line art processing to image"""
|
||||
# fmt: off
|
||||
type: Literal["lineart_image_processor"] = "lineart_image_processor"
|
||||
type: Literal["lineart_control"] = "lineart_control"
|
||||
# Inputs
|
||||
detect_resolution: int = Field(default=512, ge=0, description="pixel resolution for edge detection")
|
||||
image_resolution: int = Field(default=512, ge=0, description="pixel resolution for output image")
|
||||
@ -250,10 +168,10 @@ class LineartImageProcessorInvocation(ImageProcessorInvocation, PILInvocationCon
|
||||
return processed_image
|
||||
|
||||
|
||||
class LineartAnimeImageProcessorInvocation(ImageProcessorInvocation, PILInvocationConfig):
|
||||
class LineartAnimeControlInvocation(PreprocessedControlNetInvocation, PILInvocationConfig):
|
||||
"""Applies line art anime processing to image"""
|
||||
# fmt: off
|
||||
type: Literal["lineart_anime_image_processor"] = "lineart_anime_image_processor"
|
||||
type: Literal["lineart_anime_control"] = "lineart_anime_control"
|
||||
# Inputs
|
||||
detect_resolution: int = Field(default=512, ge=0, description="pixel resolution for edge detection")
|
||||
image_resolution: int = Field(default=512, ge=0, description="pixel resolution for output image")
|
||||
@ -268,10 +186,10 @@ class LineartAnimeImageProcessorInvocation(ImageProcessorInvocation, PILInvocati
|
||||
return processed_image
|
||||
|
||||
|
||||
class OpenposeImageProcessorInvocation(ImageProcessorInvocation, PILInvocationConfig):
|
||||
class OpenposeControlInvocation(PreprocessedControlNetInvocation, PILInvocationConfig):
|
||||
"""Applies Openpose processing to image"""
|
||||
# fmt: off
|
||||
type: Literal["openpose_image_processor"] = "openpose_image_processor"
|
||||
type: Literal["openpose_control"] = "openpose_control"
|
||||
# Inputs
|
||||
hand_and_face: bool = Field(default=False, description="whether to use hands and face mode")
|
||||
detect_resolution: int = Field(default=512, ge=0, description="pixel resolution for edge detection")
|
||||
@ -288,15 +206,14 @@ class OpenposeImageProcessorInvocation(ImageProcessorInvocation, PILInvocationCo
|
||||
return processed_image
|
||||
|
||||
|
||||
class MidasDepthImageProcessorInvocation(ImageProcessorInvocation, PILInvocationConfig):
|
||||
class MidasDepthControlInvocation(PreprocessedControlNetInvocation, PILInvocationConfig):
|
||||
"""Applies Midas depth processing to image"""
|
||||
# fmt: off
|
||||
type: Literal["midas_depth_image_processor"] = "midas_depth_image_processor"
|
||||
type: Literal["midas_control"] = "midas_control"
|
||||
# Inputs
|
||||
a_mult: float = Field(default=2.0, ge=0, description="Midas parameter a = amult * PI")
|
||||
bg_th: float = Field(default=0.1, ge=0, description="Midas parameter bg_th")
|
||||
# depth_and_normal not supported in controlnet_aux v0.0.3
|
||||
# depth_and_normal: bool = Field(default=False, description="whether to use depth and normal mode")
|
||||
depth_and_normal: bool = Field(default=False, description="whether to use depth and normal mode")
|
||||
# fmt: on
|
||||
|
||||
def run_processor(self, image):
|
||||
@ -304,16 +221,14 @@ class MidasDepthImageProcessorInvocation(ImageProcessorInvocation, PILInvocation
|
||||
processed_image = midas_processor(image,
|
||||
a=np.pi * self.a_mult,
|
||||
bg_th=self.bg_th,
|
||||
# dept_and_normal not supported in controlnet_aux v0.0.3
|
||||
# depth_and_normal=self.depth_and_normal,
|
||||
)
|
||||
depth_and_normal=self.depth_and_normal)
|
||||
return processed_image
|
||||
|
||||
|
||||
class NormalbaeImageProcessorInvocation(ImageProcessorInvocation, PILInvocationConfig):
|
||||
class NormalbaeControlNetInvocation(PreprocessedControlNetInvocation, PILInvocationConfig):
|
||||
"""Applies NormalBae processing to image"""
|
||||
# fmt: off
|
||||
type: Literal["normalbae_image_processor"] = "normalbae_image_processor"
|
||||
type: Literal["normalbae_control"] = "normalbae_control"
|
||||
# Inputs
|
||||
detect_resolution: int = Field(default=512, ge=0, description="pixel resolution for edge detection")
|
||||
image_resolution: int = Field(default=512, ge=0, description="pixel resolution for output image")
|
||||
@ -327,10 +242,10 @@ class NormalbaeImageProcessorInvocation(ImageProcessorInvocation, PILInvocationC
|
||||
return processed_image
|
||||
|
||||
|
||||
class MlsdImageProcessorInvocation(ImageProcessorInvocation, PILInvocationConfig):
|
||||
class MLSDControlNetInvocation(PreprocessedControlNetInvocation, PILInvocationConfig):
|
||||
"""Applies MLSD processing to image"""
|
||||
# fmt: off
|
||||
type: Literal["mlsd_image_processor"] = "mlsd_image_processor"
|
||||
type: Literal["mlsd_control"] = "mlsd_control"
|
||||
# Inputs
|
||||
detect_resolution: int = Field(default=512, ge=0, description="pixel resolution for edge detection")
|
||||
image_resolution: int = Field(default=512, ge=0, description="pixel resolution for output image")
|
||||
@ -348,10 +263,10 @@ class MlsdImageProcessorInvocation(ImageProcessorInvocation, PILInvocationConfig
|
||||
return processed_image
|
||||
|
||||
|
||||
class PidiImageProcessorInvocation(ImageProcessorInvocation, PILInvocationConfig):
|
||||
class PidiControlNetInvocation(PreprocessedControlNetInvocation, PILInvocationConfig):
|
||||
"""Applies PIDI processing to image"""
|
||||
# fmt: off
|
||||
type: Literal["pidi_image_processor"] = "pidi_image_processor"
|
||||
type: Literal["pidi_control"] = "pidi_control"
|
||||
# Inputs
|
||||
detect_resolution: int = Field(default=512, ge=0, description="pixel resolution for edge detection")
|
||||
image_resolution: int = Field(default=512, ge=0, description="pixel resolution for output image")
|
||||
@ -369,16 +284,16 @@ class PidiImageProcessorInvocation(ImageProcessorInvocation, PILInvocationConfig
|
||||
return processed_image
|
||||
|
||||
|
||||
class ContentShuffleImageProcessorInvocation(ImageProcessorInvocation, PILInvocationConfig):
|
||||
class ContentShuffleControlInvocation(PreprocessedControlNetInvocation, PILInvocationConfig):
|
||||
"""Applies content shuffle processing to image"""
|
||||
# fmt: off
|
||||
type: Literal["content_shuffle_image_processor"] = "content_shuffle_image_processor"
|
||||
type: Literal["content_shuffle_control"] = "content_shuffle_control"
|
||||
# Inputs
|
||||
detect_resolution: int = Field(default=512, ge=0, description="pixel resolution for edge detection")
|
||||
image_resolution: int = Field(default=512, ge=0, description="pixel resolution for output image")
|
||||
h: Union[int | None] = Field(default=512, ge=0, description="content shuffle h parameter")
|
||||
w: Union[int | None] = Field(default=512, ge=0, description="content shuffle w parameter")
|
||||
f: Union[int | None] = Field(default=256, ge=0, description="cont")
|
||||
h: Union[int | None] = Field(default=None, ge=0, description="content shuffle h parameter")
|
||||
w: Union[int | None] = Field(default=None, ge=0, description="content shuffle w parameter")
|
||||
f: Union[int | None] = Field(default=None, ge=0, description="cont")
|
||||
# fmt: on
|
||||
|
||||
def run_processor(self, image):
|
||||
@ -393,10 +308,10 @@ class ContentShuffleImageProcessorInvocation(ImageProcessorInvocation, PILInvoca
|
||||
return processed_image
|
||||
|
||||
|
||||
class ZoeDepthImageProcessorInvocation(ImageProcessorInvocation, PILInvocationConfig):
|
||||
class ZoeDepthControlInvocation(PreprocessedControlNetInvocation, PILInvocationConfig):
|
||||
"""Applies Zoe depth processing to image"""
|
||||
# fmt: off
|
||||
type: Literal["zoe_depth_image_processor"] = "zoe_depth_image_processor"
|
||||
type: Literal["zoe_depth_control"] = "zoe_depth_control"
|
||||
# fmt: on
|
||||
|
||||
def run_processor(self, image):
|
||||
|
||||
Loading…
Reference in New Issue
Block a user