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

changes to base class for controlnet nodes

Browse Source
This commit is contained in:
user1 2023-05-04 17:06:49 -07:00 committed by Kent Keirsey
parent f2f4c37f19
commit dc12fa6cd6
1 changed files with 62 additions and 47 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

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

@ -9,12 +9,24 @@ from .baseinvocation import (
InvocationConfig, InvocationConfig,
) )
from controlnet_aux import CannyDetector from controlnet_aux import (
CannyDetector,
HEDdetector,
LineartDetector,
LineartAnimeDetector,
MidasDetector,
MLSDdetector,
NormalBaeDetector,
OpenposeDetector,
PidiNetDetector,
ContentShuffleDetector,
# StyleShuffleDetector,
ZoeDetector)
from .image import ImageOutput, build_image_output, PILInvocationConfig from .image import ImageOutput, build_image_output, PILInvocationConfig
class ControlField(BaseModel): class ControlField(BaseModel):
image: ImageField = Field(default=None, description="processed image") image: ImageField = Field(default=None, description="processed image")
# width: Optional[int] = Field(default=None, description="The width of the image in pixels") # 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") # height: Optional[int] = Field(default=None, description="The height of the image in pixels")
@ -38,66 +50,69 @@ class ControlOutput(BaseInvocationOutput):
# image: ImageField = Field(default=None, description="outputs just them image info (which is also included in control output)") # image: ImageField = Field(default=None, description="outputs just them image info (which is also included in control output)")
# fmt: on # fmt: on
class PreprocessedControlInvocation(BaseInvocation, PILInvocationConfig): class PreprocessedControlInvocation(BaseInvocation, PILInvocationConfig):
"""Base class for invocations that preprocess images for ControlNet""" """Base class for invocations that preprocess images for ControlNet"""
# fmt: off # fmt: off
type: Literal["preprocessed_control"] = "preprocessed_control" type: Literal["preprocessed_control"] = "preprocessed_control"
# 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")
control_weight: float = Field(default=0.5, ge=0, le=1, description="control weight") control_weight: float = Field(default=0.5, ge=0, le=1, description="control weight")
# 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
# This super class handles invoke() call, which in turn calls run_processor(image) # begin_step_percent: float = Field(default=0, ge=0, le=1,
# subclasses override run_processor instead of implementing their own invoke() # description="% of total steps at which controlnet is first applied")
def run_processor(self, image): # end_step_percent: float = Field(default=1, ge=0, le=1,
# super class pass through of image # description="% of total steps at which controlnet is last applied")
return image # guess_mode: bool = Field(default=False, description="use guess mode (controlnet ignores prompt)")
# fmt: on
def invoke(self, context: InvocationContext) -> ControlOutput: # This super class handles invoke() call, which in turn calls run_processor(image)
image = context.services.images.get( # subclasses override run_processor instead of implementing their own invoke()
self.image.image_type, self.image.image_name def run_processor(self, image):
) # superclass just passes through image without processing
# image type should be PIL.PngImagePlugin.PngImageFile ? return image
processed_image = self.run_processor(image)
image_type = ImageType.INTERMEDIATE
image_name = context.services.images.create_name(
context.graph_execution_state_id, self.id
)
metadata = context.services.metadata.build_metadata(
session_id=context.graph_execution_state_id, node=self
)
context.services.images.save(image_type, image_name, processed_image, metadata)
"""Builds an ImageOutput and its ImageField""" def invoke(self, context: InvocationContext) -> ControlOutput:
image_field = ImageField( image = context.services.images.get(
image_name=image_name, self.image.image_type, self.image.image_name
image_type=image_type, )
) # image type should be PIL.PngImagePlugin.PngImageFile ?
return ControlOutput( processed_image = self.run_processor(image)
control=ControlField( image_type = ImageType.INTERMEDIATE
image=image_field, image_name = context.services.images.create_name(
control_model=self.control_model, context.graph_execution_state_id, self.id
control_weight=self.control_weight, )
) metadata = context.services.metadata.build_metadata(
) session_id=context.graph_execution_state_id, node=self
)
context.services.images.save(image_type, image_name, processed_image, metadata)
"""Builds an ImageOutput and its ImageField"""
image_field = ImageField(
image_name=image_name,
image_type=image_type,
)
return ControlOutput(
control=ControlField(
image=image_field,
control_model=self.control_model,
control_weight=self.control_weight,
)
)
class CannyControlInvocation(PreprocessedControlInvocation, PILInvocationConfig): class CannyControlInvocation(PreprocessedControlInvocation, PILInvocationConfig):
"""Canny edge detection for ControlNet""" """Canny edge detection for ControlNet"""
# fmt: off # fmt: off
type: Literal["cannycontrol"] = "cannycontrol" type: Literal["canny_control"] = "canny_control"
# Inputs # Inputs
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")
# fmt: on # fmt: on
def run_processor(self, image): def run_processor(self, image):