Switching to ControlField for output from controlnet nodes.

This commit is contained in:
user1 2023-05-04 14:21:11 -07:00 committed by Kent Keirsey
parent 3a645c4e80
commit dfdf8e2ead
2 changed files with 19 additions and 8 deletions

View File

@ -51,8 +51,6 @@ class ControlOutput(BaseInvocationOutput):
# fmt: off
type: Literal["control_output"] = "control_output"
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
@ -100,6 +98,7 @@ class ImageProcessorInvocation(BaseInvocation, PILInvocationConfig):
def invoke(self, context: InvocationContext) -> ImageOutput:
raw_image = context.services.images.get(
self.image.image_type, self.image.image_name
)
# image type should be PIL.PngImagePlugin.PngImageFile ?
@ -117,6 +116,7 @@ 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_name=image_name,
image_type=image_type,

View File

@ -342,21 +342,32 @@ class TextToLatentsInvocation(BaseInvocation):
multi_control = MultiControlNetModel(control_models)
model.control_model = multi_control
# loading controlnet model
if (self.control_model is None or self.control_model==''):
control_model = None
print("type of control input: ", type(self.control))
if (self.control is None):
control_model_name = None
control_image_field = None
control_weight = None
else:
control_model_name = self.control.control_model
control_image_field = self.control.image
control_weight = self.control.control_weight
# # loading controlnet model
# if (self.control_model is None or self.control_model==''):
# control_model = None
# else:
# FIXME: change this to dropdown menu?
# FIXME: generalize so don't have to hardcode torch_dtype and device
control_model = ControlNetModel.from_pretrained(self.control_model,
control_model = ControlNetModel.from_pretrained(control_model_name,
torch_dtype=torch.float16).to("cuda")
model.control_model = control_model
# loading controlnet image (currently requires pre-processed image)
control_image = (
None if self.control_image is None
None if control_image_field is None
else context.services.images.get(
self.control_image.image_type, self.control_image.image_name
control_image_field.image_type, control_image_field.image_name
)
)