Flax controlnet (#2727)

* add contronet flax

---------

Co-authored-by: yiyixuxu <yixu310@gmail,com>

Author: yiyixuxu

docs/source/en/api/models.mdx

@@ -99,3 +99,9 @@ The models are built on the base class ['ModelMixin'] that is a `torch.nn.module
 
 ## FlaxAutoencoderKL
 [[autodoc]] FlaxAutoencoderKL
+
+## FlaxControlNetOutput
+[[autodoc]] models.controlnet_flax.FlaxControlNetOutput
+
+## FlaxControlNetModel
+[[autodoc]] FlaxControlNetModel

docs/source/en/api/pipelines/stable_diffusion/controlnet.mdx

@@ -272,3 +272,9 @@ All checkpoints can be found under the authors' namespace [lllyasviel](https://h
 	- disable_vae_slicing
 	- enable_xformers_memory_efficient_attention
 	- disable_xformers_memory_efficient_attention
+
+## FlaxStableDiffusionControlNetPipeline
+[[autodoc]] FlaxStableDiffusionControlNetPipeline
+	- all
+	- __call__
+

src/diffusers/__init__.py

@@ -188,6 +188,7 @@
 except OptionalDependencyNotAvailable:
     from .utils.dummy_flax_objects import *  # noqa F403
 else:
+    from .models.controlnet_flax import FlaxControlNetModel
     from .models.modeling_flax_utils import FlaxModelMixin
     from .models.unet_2d_condition_flax import FlaxUNet2DConditionModel
     from .models.vae_flax import FlaxAutoencoderKL
@@ -211,6 +212,7 @@
     from .utils.dummy_flax_and_transformers_objects import *  # noqa F403
 else:
     from .pipelines import (
+        FlaxStableDiffusionControlNetPipeline,
         FlaxStableDiffusionImg2ImgPipeline,
         FlaxStableDiffusionInpaintPipeline,
         FlaxStableDiffusionPipeline,

src/diffusers/models/__init__.py

@@ -30,5 +30,6 @@
     from .vq_model import VQModel
 
 if is_flax_available():
+    from .controlnet_flax import FlaxControlNetModel
     from .unet_2d_condition_flax import FlaxUNet2DConditionModel
     from .vae_flax import FlaxAutoencoderKL