Feat/controlnet nodes

invokeai/app/invocations/baseinvocation.py

@@ -96,6 +96,7 @@ class UIConfig(TypedDict, total=False):
             "image",
             "latents",
             "model",
+            "control",
         ],
     ]
     tags: List[str]

invokeai/app/invocations/collections.py

@@ -22,6 +22,14 @@ class IntCollectionOutput(BaseInvocationOutput):
     # Outputs
     collection: list[int] = Field(default=[], description="The int collection")
 
+class FloatCollectionOutput(BaseInvocationOutput):
+    """A collection of floats"""
+
+    type: Literal["float_collection"] = "float_collection"
+
+    # Outputs
+    collection: list[float] = Field(default=[], description="The float collection")
+
 
 class RangeInvocation(BaseInvocation):
     """Creates a range of numbers from start to stop with step"""

invokeai/app/invocations/generate.py

@@ -4,7 +4,9 @@
 from typing import Literal, Optional, Union, get_args
 
 import numpy as np
+from diffusers import ControlNetModel
 from torch import Tensor
+import torch
 
 from pydantic import BaseModel, Field
 
@@ -58,6 +60,9 @@ class TextToImageInvocation(BaseInvocation, SDImageInvocation):
     cfg_scale: float = Field(default=7.5, ge=1, description="The Classifier-Free Guidance, higher values may result in a result closer to the prompt", )
     scheduler: SAMPLER_NAME_VALUES = Field(default="euler", description="The scheduler to use" )
     model:       str = Field(default="", description="The model to use (currently ignored)")
+    progress_images: bool = Field(default=False, description="Whether or not to produce progress images during generation",  )
+    control_model: Optional[str] = Field(default=None, description="The control model to use")
+    control_image: Optional[ImageField] = Field(default=None, description="The processed control image")
     # fmt: on
 
     # TODO: pass this an emitter method or something? or a session for dispatching?
@@ -78,17 +83,35 @@ def invoke(self, context: InvocationContext) -> ImageOutput:
         # Handle invalid model parameter
         model = choose_model(context.services.model_manager, self.model)
 
+        # loading controlnet image (currently requires pre-processed image)
+        control_image = (
+            None if self.control_image is None
+            else context.services.images.get(
+                self.control_image.image_type, self.control_image.image_name
+            )
+        )
+        # 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,
+                                                            torch_dtype=torch.float16).to("cuda")
+
         # Get the source node id (we are invoking the prepared node)
         graph_execution_state = context.services.graph_execution_manager.get(
             context.graph_execution_state_id
         )
         source_node_id = graph_execution_state.prepared_source_mapping[self.id]
 
-        outputs = Txt2Img(model).generate(
+        txt2img = Txt2Img(model, control_model=control_model)
+        outputs = txt2img.generate(
             prompt=self.prompt,
             step_callback=partial(self.dispatch_progress, context, source_node_id),
+            control_image=control_image,
             **self.dict(
-                exclude={"prompt"}
+                exclude={"prompt", "control_image" }
             ),  # Shorthand for passing all of the parameters above manually
         )
         # Outputs is an infinite iterator that will return a new InvokeAIGeneratorOutput object

invokeai/app/invocations/latent.py

@@ -1,8 +1,11 @@
 # Copyright (c) 2023 Kyle Schouviller (https://github.com/kyle0654)
 
 import random
-from typing import Literal, Optional, Union
 import einops
+from typing import Literal, Optional, Union, List
+
+from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion_controlnet import MultiControlNetModel
+
 from pydantic import BaseModel, Field, validator
 import torch
 
@@ -11,14 +14,18 @@
 from invokeai.app.util.misc import SEED_MAX, get_random_seed
 
 from invokeai.app.util.step_callback import stable_diffusion_step_callback
+from .controlnet_image_processors import ControlField
 
 from ...backend.model_management.model_manager import ModelManager
 from ...backend.util.devices import choose_torch_device, torch_dtype
 from ...backend.stable_diffusion.diffusion.shared_invokeai_diffusion import PostprocessingSettings
 from ...backend.image_util.seamless import configure_model_padding
 from ...backend.prompting.conditioning import get_uc_and_c_and_ec
+
 from ...backend.stable_diffusion.diffusers_pipeline import ConditioningData, StableDiffusionGeneratorPipeline, image_resized_to_grid_as_tensor
 from ...backend.stable_diffusion.schedulers import SCHEDULER_MAP
+from ...backend.stable_diffusion.diffusers_pipeline import ControlNetData
+
 from .baseinvocation import BaseInvocation, BaseInvocationOutput, InvocationContext, InvocationConfig
 import numpy as np
 from ..services.image_file_storage import ImageType
@@ -28,7 +35,7 @@
 from ...backend.stable_diffusion import PipelineIntermediateState
 from diffusers.schedulers import SchedulerMixin as Scheduler
 import diffusers
-from diffusers import DiffusionPipeline
+from diffusers import DiffusionPipeline, ControlNetModel
 
 
 class LatentsField(BaseModel):
@@ -84,13 +91,13 @@ def build_noise_output(latents_name: str, latents: torch.Tensor):
 
 def get_scheduler(scheduler_name:str, model: StableDiffusionGeneratorPipeline)->Scheduler:
     scheduler_class, scheduler_extra_config = SCHEDULER_MAP.get(scheduler_name, SCHEDULER_MAP['ddim'])
-    
+
     scheduler_config = model.scheduler.config
     if "_backup" in scheduler_config:
         scheduler_config = scheduler_config["_backup"]
     scheduler_config = {**scheduler_config, **scheduler_extra_config, "_backup": scheduler_config}
     scheduler = scheduler_class.from_config(scheduler_config)
-    
+
     # hack copied over from generate.py
     if not hasattr(scheduler, 'uses_inpainting_model'):
         scheduler.uses_inpainting_model = lambda: False
@@ -169,6 +176,7 @@ class TextToLatentsInvocation(BaseInvocation):
     cfg_scale: float = Field(default=7.5, gt=0, description="The Classifier-Free Guidance, higher values may result in a result closer to the prompt", )
     scheduler: SAMPLER_NAME_VALUES = Field(default="euler", description="The scheduler to use" )
     model:       str = Field(default="", description="The model to use (currently ignored)")
+    control: Union[ControlField, list[ControlField]] = Field(default=None, description="The control to use")
     # seamless:   bool = Field(default=False, description="Whether or not to generate an image that can tile without seams", )
     # seamless_axes: str = Field(default="", description="The axes to tile the image on, 'x' and/or 'y'")
     # fmt: on
@@ -179,7 +187,8 @@ class Config(InvocationConfig):
             "ui": {
                 "tags": ["latents", "image"],
                 "type_hints": {
-                  "model": "model"
+                  "model": "model",
+                  "control": "control",
                 }
             },
         }
@@ -238,6 +247,81 @@ def get_conditioning_data(self, context: InvocationContext, model: StableDiffusi
         ).add_scheduler_args_if_applicable(model.scheduler, eta=0.0)#ddim_eta)
         return conditioning_data
 
+    def prep_control_data(self,
+                          context: InvocationContext,
+                          model: StableDiffusionGeneratorPipeline, # really only need model for dtype and device
+                          control_input: List[ControlField],
+                          latents_shape: List[int],
+                          do_classifier_free_guidance: bool = True,
+                          ) -> List[ControlNetData]:
+        # assuming fixed dimensional scaling of 8:1 for image:latents
+        control_height_resize = latents_shape[2] * 8
+        control_width_resize = latents_shape[3] * 8
+        if control_input is None:
+            # print("control input is None")
+            control_list = None
+        elif isinstance(control_input, list) and len(control_input) == 0:
+            # print("control input is empty list")
+            control_list = None
+        elif isinstance(control_input, ControlField):
+            # print("control input is ControlField")
+            control_list = [control_input]
+        elif isinstance(control_input, list) and len(control_input) > 0 and isinstance(control_input[0], ControlField):
+            # print("control input is list[ControlField]")
+            control_list = control_input
+        else:
+            # print("input control is unrecognized:", type(self.control))
+            control_list = None
+        if (control_list is None):
+            control_data = None
+            # from above handling, any control that is not None should now be of type list[ControlField]
+        else:
+            # FIXME: add checks to skip entry if model or image is None
+            #        and if weight is None, populate with default 1.0?
+            control_data = []
+            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)
+                control_image_field = control_info.image
+                input_image = context.services.images.get_pil_image(control_image_field.image_type,
+                                                                    control_image_field.image_name)
+                # self.image.image_type, self.image.image_name
+                # FIXME: still need to test with different widths, heights, devices, dtypes
+                #        and add in batch_size, num_images_per_prompt?
+                #        and do real check for classifier_free_guidance?
+                # prepare_control_image should return torch.Tensor of shape(batch_size, 3, height, width)
+                control_image = model.prepare_control_image(
+                    image=input_image,
+                    do_classifier_free_guidance=do_classifier_free_guidance,
+                    width=control_width_resize,
+                    height=control_height_resize,
+                    # batch_size=batch_size * num_images_per_prompt,
+                    # num_images_per_prompt=num_images_per_prompt,
+                    device=control_model.device,
+                    dtype=control_model.dtype,
+                )
+                control_item = ControlNetData(model=control_model,
+                                              image_tensor=control_image,
+                                              weight=control_info.control_weight,
+                                              begin_step_percent=control_info.begin_step_percent,
+                                              end_step_percent=control_info.end_step_percent)
+                control_data.append(control_item)
+                # MultiControlNetModel has been refactored out, just need list[ControlNetData]
+        return control_data
 
     def invoke(self, context: InvocationContext) -> LatentsOutput:
         noise = context.services.latents.get(self.noise.latents_name)
@@ -252,14 +336,19 @@ def step_callback(state: PipelineIntermediateState):
         model = self.get_model(context.services.model_manager)
         conditioning_data = self.get_conditioning_data(context, model)
 
-        # TODO: Verify the noise is the right size
+        print("type of control input: ", type(self.control))
+        control_data = self.prep_control_data(model=model, context=context, control_input=self.control,
+                                              latents_shape=noise.shape,
+                                              do_classifier_free_guidance=(self.cfg_scale >= 1.0))
 
+        # TODO: Verify the noise is the right size
         result_latents, result_attention_map_saver = model.latents_from_embeddings(
             latents=torch.zeros_like(noise, dtype=torch_dtype(model.device)),
             noise=noise,
             num_inference_steps=self.steps,
             conditioning_data=conditioning_data,
-            callback=step_callback
+            control_data=control_data,  # list[ControlNetData]
+            callback=step_callback,
         )
 
         # https://discuss.huggingface.co/t/memory-usage-by-later-pipeline-stages/23699
@@ -285,7 +374,8 @@ class Config(InvocationConfig):
             "ui": {
                 "tags": ["latents"],
                 "type_hints": {
-                    "model": "model"
+                    "model": "model",
+                    "control": "control",
                 }
             },
         }
@@ -304,6 +394,11 @@ def step_callback(state: PipelineIntermediateState):
         model = self.get_model(context.services.model_manager)
         conditioning_data = self.get_conditioning_data(context, model)
 
+        print("type of control input: ", type(self.control))
+        control_data = self.prep_control_data(model=model, context=context, control_input=self.control,
+                                              latents_shape=noise.shape,
+                                              do_classifier_free_guidance=(self.cfg_scale >= 1.0))
+
         # TODO: Verify the noise is the right size
 
         initial_latents = latent if self.strength < 1.0 else torch.zeros_like(
@@ -318,6 +413,7 @@ def step_callback(state: PipelineIntermediateState):
             noise=noise,
             num_inference_steps=self.steps,
             conditioning_data=conditioning_data,
+            control_data=control_data,  # list[ControlNetData]
             callback=step_callback
         )
 
@@ -362,8 +458,14 @@ def invoke(self, context: InvocationContext) -> ImageOutput:
             np_image = model.decode_latents(latents)
             image = model.numpy_to_pil(np_image)[0]
 
+            # what happened to metadata?
+            # metadata = context.services.metadata.build_metadata(
+            #     session_id=context.graph_execution_state_id, node=self
+
             torch.cuda.empty_cache()
 
+            # new (post Image service refactor) way of using services to save image
+            #     and gnenerate unique image_name
             image_dto = context.services.images.create(
                 image=image,
                 image_type=ImageType.RESULT,
@@ -414,6 +516,7 @@ def invoke(self, context: InvocationContext) -> LatentsOutput:
         torch.cuda.empty_cache()
 
         name = f"{context.graph_execution_state_id}__{self.id}"
+        # context.services.latents.set(name, resized_latents)
         context.services.latents.save(name, resized_latents)
         return build_latents_output(latents_name=name, latents=resized_latents)
 
@@ -444,6 +547,7 @@ def invoke(self, context: InvocationContext) -> LatentsOutput:
         torch.cuda.empty_cache()
 
         name = f"{context.graph_execution_state_id}__{self.id}"
+        # context.services.latents.set(name, resized_latents)
         context.services.latents.save(name, resized_latents)
         return build_latents_output(latents_name=name, latents=resized_latents)
 
@@ -468,6 +572,9 @@ class Config(InvocationConfig):
 
     @torch.no_grad()
     def invoke(self, context: InvocationContext) -> LatentsOutput:
+        # image = context.services.images.get(
+        #     self.image.image_type, self.image.image_name
+        # )
         image = context.services.images.get_pil_image(
             self.image.image_type, self.image.image_name
         )
@@ -488,6 +595,6 @@ def invoke(self, context: InvocationContext) -> LatentsOutput:
         )
 
         name = f"{context.graph_execution_state_id}__{self.id}"
+        # context.services.latents.set(name, latents)
         context.services.latents.save(name, latents)
         return build_latents_output(latents_name=name, latents=latents)
-

invokeai/app/invocations/math.py

@@ -34,6 +34,15 @@ class IntOutput(BaseInvocationOutput):
     # fmt: on
 
 
+class FloatOutput(BaseInvocationOutput):
+    """A float output"""
+
+    # fmt: off
+    type: Literal["float_output"] = "float_output"
+    param: float = Field(default=None, description="The output float")
+    # fmt: on
+
+
 class AddInvocation(BaseInvocation, MathInvocationConfig):
     """Adds two numbers"""
 

invokeai/app/invocations/params.py

@@ -3,7 +3,7 @@
 from typing import Literal
 from pydantic import Field
 from .baseinvocation import BaseInvocation, BaseInvocationOutput, InvocationContext
-from .math import IntOutput
+from .math import IntOutput, FloatOutput
 
 # Pass-through parameter nodes - used by subgraphs
 
@@ -16,3 +16,13 @@ class ParamIntInvocation(BaseInvocation):
 
     def invoke(self, context: InvocationContext) -> IntOutput:
         return IntOutput(a=self.a)
+
+class ParamFloatInvocation(BaseInvocation):
+    """A float parameter"""
+    #fmt: off
+    type: Literal["param_float"] = "param_float"
+    param: float = Field(default=0.0, description="The float value")
+    #fmt: on
+
+    def invoke(self, context: InvocationContext) -> FloatOutput:
+        return FloatOutput(param=self.param)