Stable diffusion controlnet support

optimum/exporters/openvino/model_configs.py

@@ -604,6 +604,90 @@ def outputs(self) -> Dict[str, Dict[int, str]]:
  }
 
 
+@register_in_tasks_manager("unet-controlnet", *["semantic-segmentation"], library_name="diffusers")
+class UNetControlNetOpenVINOConfig(UNetOnnxConfig):
+ @property
+ def inputs(self) -> Dict[str, Dict[int, str]]:
+ common_inputs = {
+ "sample": {0: "batch_size", 2: "height", 3: "width"},
+ "timestep": {0: "steps"},
+ "encoder_hidden_states": {0: "batch_size", 1: "sequence_length"},
+ "mid_block_additional_residual": {0: "batch_size", 2: "height", 3: "width"},
+ }
+ for a in range(1, 25, 2):
+ if a == 23:
+ common_inputs["down_block_additional_residual"] = {0: "batch_size", 2: "height", 3: "width"}
+ break
+ else:
+ common_inputs[f"down_block_additional_residual.{a}"] = {0: "batch_size", 2: "height", 3: "width"}
+ # TODO : add text_image, image and image_embeds
+ if getattr(self._normalized_config, "addition_embed_type", None) == "text_time":
+ common_inputs["text_embeds"] = {0: "batch_size"}
+ common_inputs["time_ids"] = {0: "batch_size"}
+
+ if getattr(self._normalized_config, "time_cond_proj_dim", None) is not None:
+ common_inputs["timestep_cond"] = {0: "batch_size"}
+ return common_inputs
+
+ @property
+ def outputs(self) -> Dict[str, Dict[int, str]]:
+ return {
+ "out_sample": {0: "batch_size", 2: "height", 3: "width"},
+ }
+
+ def generate_dummy_inputs(self, framework: str = "pt", **kwargs):
+ dummy_inputs_generators = self._create_dummy_input_generator_classes(**kwargs)
+
+ dummy_inputs = {}
+ for input_name in self.inputs:
+ for dummy_input_gen in dummy_inputs_generators:
+ if dummy_input_gen.supports_input(input_name):
+ dummy_inputs[input_name] = dummy_input_gen.generate(
+ input_name, framework=framework, int_dtype=self.int_dtype, float_dtype=self.float_dtype
+ )
+ break
+
+ import torch
+
+ dummy_inputs["down_block_additional_residual_1"] = torch.randn(2, 320, 64, 64)
+ dummy_inputs["down_block_additional_residual_3"] = torch.randn(2, 320, 64, 64)
+ dummy_inputs["down_block_additional_residual_5"] = torch.randn(2, 320, 64, 64)
+ dummy_inputs["down_block_additional_residual_7"] = torch.randn(2, 320, 32, 32)
+ dummy_inputs["down_block_additional_residual_9"] = torch.randn(2, 640, 32, 32)
+ dummy_inputs["down_block_additional_residual_11"] = torch.randn(2, 640, 32, 32)
+ dummy_inputs["down_block_additional_residual_13"] = torch.randn(2, 640, 16, 16)
+ dummy_inputs["down_block_additional_residual_15"] = torch.randn(2, 1280, 16, 16)
+ dummy_inputs["down_block_additional_residual_17"] = torch.randn(2, 1280, 16, 16)
+ dummy_inputs["down_block_additional_residual_19"] = torch.randn(2, 1280, 8, 8)
+ dummy_inputs["down_block_additional_residual_21"] = torch.randn(2, 1280, 8, 8)
+ dummy_inputs["down_block_additional_residual"] = torch.randn(2, 1280, 8, 8)
+ dummy_inputs["mid_block_additional_residual"] = torch.randn(2, 1280, 8, 8)
+ dummy_inputs["encoder_hidden_states"] = dummy_inputs["encoder_hidden_states"][0]
+
+ if getattr(self._normalized_config, "addition_embed_type", None) == "text_time":
+ dummy_inputs["added_cond_kwargs"] = {
+ "text_embeds": dummy_inputs.pop("text_embeds"),
+ "time_ids": dummy_inputs.pop("time_ids"),
+ }
+
+ return dummy_inputs
+
+ def rename_ambiguous_inputs(self, inputs) -> Dict[str, Dict[int, str]]:
+ """
+ Updates the input names of the model to export.
+ Override the function when the model input names are ambiguous or too generic.
+
+ Returns:
+ `Dict[str, Dict[int, str]]`: Updated inputs.
+ """
+ new_inputs = {}
+ for name, v in inputs.items():
+ if name.startswith("down_block_additional_residual"):
+ new_inputs[name.replace(".", "_")] = v
+ else:
+ new_inputs[name] = v
+ return new_inputs
+
 @register_in_tasks_manager("vae-encoder", *["semantic-segmentation"], library_name="diffusers")
 class VaeEncoderOpenVINOConfig(VaeEncoderOnnxConfig):
  @property

optimum/exporters/openvino/model_patcher.py

@@ -382,9 +382,9 @@ def _llama_gemma_update_causal_mask_legacy(self, attention_mask, input_tensor, c
  offset = 0
  mask_shape = attention_mask.shape
  mask_slice = (attention_mask.eq(0.0)).to(dtype=dtype) * min_dtype
- causal_mask[
- : mask_shape[0], : mask_shape[1], offset : mask_shape[2] + offset, : mask_shape[3]
- ] = mask_slice
+ causal_mask[: mask_shape[0], : mask_shape[1], offset : mask_shape[2] + offset, : mask_shape[3]] = (
+ mask_slice
+ )
 
  if (
  self.config._attn_implementation == "sdpa"
@@ -1655,9 +1655,9 @@ def _dbrx_update_causal_mask_legacy(
  offset = 0
  mask_shape = attention_mask.shape
  mask_slice = (attention_mask.eq(0.0)).to(dtype=dtype) * min_dtype
- causal_mask[
- : mask_shape[0], : mask_shape[1], offset : mask_shape[2] + offset, : mask_shape[3]
- ] = mask_slice
+ causal_mask[: mask_shape[0], : mask_shape[1], offset : mask_shape[2] + offset, : mask_shape[3]] = (
+ mask_slice
+ )
 
  if (
  self.config._attn_implementation == "sdpa"

optimum/intel/__init__.py

@@ -100,6 +100,7 @@
  "OVStableDiffusionXLPipeline",
  "OVStableDiffusionXLImg2ImgPipeline",
  "OVLatentConsistencyModelPipeline",
+ "OVStableDiffusionControlNetPipeline",
  ]
 else:
  _import_structure["openvino"].extend(
@@ -110,6 +111,7 @@
  "OVStableDiffusionXLPipeline",
  "OVStableDiffusionXLImg2ImgPipeline",
  "OVLatentConsistencyModelPipeline",
+ "OVStableDiffusionControlNetPipeline",
  ]
  )
 
@@ -233,6 +235,7 @@
  except OptionalDependencyNotAvailable:
  from .utils.dummy_openvino_and_diffusers_objects import (
  OVLatentConsistencyModelPipeline,
+ OVStableDiffusionControlNetPipeline,
  OVStableDiffusionImg2ImgPipeline,
  OVStableDiffusionInpaintPipeline,
  OVStableDiffusionPipeline,
@@ -242,6 +245,7 @@
  else:
  from .openvino import (
  OVLatentConsistencyModelPipeline,
+ OVStableDiffusionControlNetPipeline,
  OVStableDiffusionImg2ImgPipeline,
  OVStableDiffusionInpaintPipeline,
  OVStableDiffusionPipeline,

optimum/intel/openvino/__init__.py

@@ -71,6 +71,7 @@
 if is_diffusers_available():
  from .modeling_diffusion import (
  OVLatentConsistencyModelPipeline,
+ OVStableDiffusionControlNetPipeline,
  OVStableDiffusionImg2ImgPipeline,
  OVStableDiffusionInpaintPipeline,
  OVStableDiffusionPipeline,