Mirrored_Repos/InvokeAI
1
0
Fork 0
mirror of https://github.com/invoke-ai/InvokeAI synced 2024-08-30 20:32:17 +00:00
Code Issues Packages Projects Releases Wiki Activity

Refactored controlnet nodes: split out controlnet stuff into separate node, stripped controlnet stuff form image processing/analysis nodes.

Browse Source
This commit is contained in:
user1 2023-05-05 21:41:07 -07:00 committed by Kent Keirsey
parent db27263bc2
commit 18e6a2b410
1 changed files with 56 additions and 40 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

94
invokeai/app/invocations/controlnet_image_processors.py
View File

@ -50,18 +50,15 @@ class ControlOutput(BaseInvocationOutput):
# fmt: off # fmt: off
type: Literal["control_output"] = "control_output" type: Literal["control_output"] = "control_output"
control: Optional[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, # raw_processed_image: ImageField = Field(default=None,
description="outputs just the image info (also included in control output)") # description="outputs just the image info (also included in control output)")
# fmt: on # fmt: on
# This super class handles invoke() call, which in turn calls run_processor(image) class ControlNetInvocation(BaseInvocation):
# subclasses override run_processor() instead of implementing their own invoke() """Collects ControlNet info to pass to other nodes"""
class PreprocessedControlNetInvocation(BaseInvocation, PILInvocationConfig):
"""Base class for invocations that preprocess images for ControlNet"""
# fmt: off # fmt: off
type: Literal["preprocessed_control"] = "preprocessed_control" type: Literal["controlnet"] = "controlnet"
# Inputs # Inputs
image: ImageField = Field(default=None, description="image to process") image: ImageField = Field(default=None, description="image to process")
control_model: str = Field(default=None, description="control model to use") control_model: str = Field(default=None, description="control model to use")
@ -75,11 +72,32 @@ class PreprocessedControlNetInvocation(BaseInvocation, PILInvocationConfig):
# fmt: on # 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,
),
)
# TODO: move image processors to separate file (image_analysis.py
class ImageProcessorInvocation(BaseInvocation, PILInvocationConfig):
"""Base class for invocations that preprocess images for ControlNet"""
# fmt: off
type: Literal["image_processor"] = "image_processor"
# Inputs
image: ImageField = Field(default=None, description="image to process")
# fmt: on
def run_processor(self, image): def run_processor(self, image):
# superclass just passes through image without processing # superclass just passes through image without processing
return image return image
def invoke(self, context: InvocationContext) -> ControlOutput: def invoke(self, context: InvocationContext) -> ImageOutput:
raw_image = context.services.images.get( raw_image = context.services.images.get(
self.image.image_type, self.image.image_name self.image.image_type, self.image.image_name
) )
@ -98,24 +116,22 @@ class PreprocessedControlNetInvocation(BaseInvocation, PILInvocationConfig):
context.services.images.save(image_type, image_name, processed_image, metadata) context.services.images.save(image_type, image_name, processed_image, metadata)
"""Builds an ImageOutput and its ImageField""" """Builds an ImageOutput and its ImageField"""
image_field = ImageField( processed_image_field = ImageField(
image_name=image_name, image_name=image_name,
image_type=image_type, image_type=image_type,
) )
return ControlOutput( return ImageOutput(
control=ControlField( image=processed_image_field,
image=image_field, width=processed_image.width,
control_model=self.control_model, height=processed_image.height,
control_weight=self.control_weight, mode=processed_image.mode,
),
raw_processed_image=image_field,
) )
class CannyControlInvocation(PreprocessedControlNetInvocation, PILInvocationConfig): class CannyImageProcessorInvocation(ImageProcessorInvocation, PILInvocationConfig):
"""Canny edge detection for ControlNet""" """Canny edge detection for ControlNet"""
# fmt: off # fmt: off
type: Literal["cannycontrol"] = "cannycontrol" type: Literal["canny_image_processor"] = "canny_image_processor"
# Input # Input
low_threshold: float = Field(default=100, ge=0, description="low threshold of Canny pixel gradient") 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") high_threshold: float = Field(default=200, ge=0, description="high threshold of Canny pixel gradient")
@ -127,10 +143,10 @@ class CannyControlInvocation(PreprocessedControlNetInvocation, PILInvocationConf
return processed_image return processed_image
class HedControlNetInvocation(PreprocessedControlNetInvocation, PILInvocationConfig): class HedImageprocessorInvocation(ImageProcessorInvocation, PILInvocationConfig):
"""Applies HED edge detection to image""" """Applies HED edge detection to image"""
# fmt: off # fmt: off
type: Literal["hed_control"] = "hed_control" type: Literal["hed_image_processor"] = "hed_image_processor"
# Inputs # Inputs
detect_resolution: int = Field(default=512, ge=0, description="pixel resolution for edge detection") 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") image_resolution: int = Field(default=512, ge=0, description="pixel resolution for output image")
@ -149,10 +165,10 @@ class HedControlNetInvocation(PreprocessedControlNetInvocation, PILInvocationCon
return processed_image return processed_image
class LineartControlInvocation(PreprocessedControlNetInvocation, PILInvocationConfig): class LineartImageProcessorInvocation(ImageProcessorInvocation, PILInvocationConfig):
"""Applies line art processing to image""" """Applies line art processing to image"""
# fmt: off # fmt: off
type: Literal["lineart_control"] = "lineart_control" type: Literal["lineart_image_processor"] = "lineart_image_processor"
# Inputs # Inputs
detect_resolution: int = Field(default=512, ge=0, description="pixel resolution for edge detection") 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") image_resolution: int = Field(default=512, ge=0, description="pixel resolution for output image")
@ -168,10 +184,10 @@ class LineartControlInvocation(PreprocessedControlNetInvocation, PILInvocationCo
return processed_image return processed_image
class LineartAnimeControlInvocation(PreprocessedControlNetInvocation, PILInvocationConfig): class LineartAnimeImageProcessorInvocation(ImageProcessorInvocation, PILInvocationConfig):
"""Applies line art anime processing to image""" """Applies line art anime processing to image"""
# fmt: off # fmt: off
type: Literal["lineart_anime_control"] = "lineart_anime_control" type: Literal["lineart_anime_image_processor"] = "lineart_anime_image_processor"
# Inputs # Inputs
detect_resolution: int = Field(default=512, ge=0, description="pixel resolution for edge detection") 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") image_resolution: int = Field(default=512, ge=0, description="pixel resolution for output image")
@ -186,10 +202,10 @@ class LineartAnimeControlInvocation(PreprocessedControlNetInvocation, PILInvocat
return processed_image return processed_image
class OpenposeControlInvocation(PreprocessedControlNetInvocation, PILInvocationConfig): class OpenposeImageProcessorInvocation(ImageProcessorInvocation, PILInvocationConfig):
"""Applies Openpose processing to image""" """Applies Openpose processing to image"""
# fmt: off # fmt: off
type: Literal["openpose_control"] = "openpose_control" type: Literal["openpose_image_processor"] = "openpose_image_processor"
# Inputs # Inputs
hand_and_face: bool = Field(default=False, description="whether to use hands and face mode") 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") detect_resolution: int = Field(default=512, ge=0, description="pixel resolution for edge detection")
@ -206,10 +222,10 @@ class OpenposeControlInvocation(PreprocessedControlNetInvocation, PILInvocationC
return processed_image return processed_image
class MidasDepthControlInvocation(PreprocessedControlNetInvocation, PILInvocationConfig): class MidasDepthImageProcessorInvocation(ImageProcessorInvocation, PILInvocationConfig):
"""Applies Midas depth processing to image""" """Applies Midas depth processing to image"""
# fmt: off # fmt: off
type: Literal["midas_control"] = "midas_control" type: Literal["midas_depth_image_processor"] = "midas_depth_image_processor"
# Inputs # Inputs
a_mult: float = Field(default=2.0, ge=0, description="Midas parameter a = amult * PI") 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") bg_th: float = Field(default=0.1, ge=0, description="Midas parameter bg_th")
@ -225,10 +241,10 @@ class MidasDepthControlInvocation(PreprocessedControlNetInvocation, PILInvocatio
return processed_image return processed_image
class NormalbaeControlNetInvocation(PreprocessedControlNetInvocation, PILInvocationConfig): class NormalbaeImageProcessorInvocation(ImageProcessorInvocation, PILInvocationConfig):
"""Applies NormalBae processing to image""" """Applies NormalBae processing to image"""
# fmt: off # fmt: off
type: Literal["normalbae_control"] = "normalbae_control" type: Literal["normalbae_image_processor"] = "normalbae_image_processor"
# Inputs # Inputs
detect_resolution: int = Field(default=512, ge=0, description="pixel resolution for edge detection") 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") image_resolution: int = Field(default=512, ge=0, description="pixel resolution for output image")
@ -242,10 +258,10 @@ class NormalbaeControlNetInvocation(PreprocessedControlNetInvocation, PILInvocat
return processed_image return processed_image
class MLSDControlNetInvocation(PreprocessedControlNetInvocation, PILInvocationConfig): class MlsdImageProcessorInvocation(ImageProcessorInvocation, PILInvocationConfig):
"""Applies MLSD processing to image""" """Applies MLSD processing to image"""
# fmt: off # fmt: off
type: Literal["mlsd_control"] = "mlsd_control" type: Literal["mlsd_image_processor"] = "mlsd_image_processor"
# Inputs # Inputs
detect_resolution: int = Field(default=512, ge=0, description="pixel resolution for edge detection") 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") image_resolution: int = Field(default=512, ge=0, description="pixel resolution for output image")
@ -263,10 +279,10 @@ class MLSDControlNetInvocation(PreprocessedControlNetInvocation, PILInvocationCo
return processed_image return processed_image
class PidiControlNetInvocation(PreprocessedControlNetInvocation, PILInvocationConfig): class PidiImageProcessorInvocation(ImageProcessorInvocation, PILInvocationConfig):
"""Applies PIDI processing to image""" """Applies PIDI processing to image"""
# fmt: off # fmt: off
type: Literal["pidi_control"] = "pidi_control" type: Literal["pidi_image_processor"] = "pidi_image_processor"
# Inputs # Inputs
detect_resolution: int = Field(default=512, ge=0, description="pixel resolution for edge detection") 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") image_resolution: int = Field(default=512, ge=0, description="pixel resolution for output image")
@ -284,10 +300,10 @@ class PidiControlNetInvocation(PreprocessedControlNetInvocation, PILInvocationCo
return processed_image return processed_image
class ContentShuffleControlInvocation(PreprocessedControlNetInvocation, PILInvocationConfig): class ContentShuffleImageProcessorInvocation(ImageProcessorInvocation, PILInvocationConfig):
"""Applies content shuffle processing to image""" """Applies content shuffle processing to image"""
# fmt: off # fmt: off
type: Literal["content_shuffle_control"] = "content_shuffle_control" type: Literal["content_shuffle_image_processor"] = "content_shuffle_image_processor"
# Inputs # Inputs
detect_resolution: int = Field(default=512, ge=0, description="pixel resolution for edge detection") 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") image_resolution: int = Field(default=512, ge=0, description="pixel resolution for output image")
@ -308,10 +324,10 @@ class ContentShuffleControlInvocation(PreprocessedControlNetInvocation, PILInvoc
return processed_image return processed_image
class ZoeDepthControlInvocation(PreprocessedControlNetInvocation, PILInvocationConfig): class ZoeDepthImageProcessorInvocation(ImageProcessorInvocation, PILInvocationConfig):
"""Applies Zoe depth processing to image""" """Applies Zoe depth processing to image"""
# fmt: off # fmt: off
type: Literal["zoe_depth_control"] = "zoe_depth_control" type: Literal["zoe_depth_image_processor"] = "zoe_depth_image_processor"
# fmt: on # fmt: on
def run_processor(self, image): def run_processor(self, image):