ControlNetModel

The ControlNet model was introduced in Adding Conditional Control to Text-to-Image Diffusion Models by Lvmin Zhang, Anyi Rao, Maneesh Agrawala. It provides a greater degree of control over text-to-image generation by conditioning the model on additional inputs such as edge maps, depth maps, segmentation maps, and keypoints for pose detection.

The abstract from the paper is:

We present ControlNet, a neural network architecture to add spatial conditioning controls to large, pretrained text-to-image diffusion models. ControlNet locks the production-ready large diffusion models, and reuses their deep and robust encoding layers pretrained with billions of images as a strong backbone to learn a diverse set of conditional controls. The neural architecture is connected with "zero convolutions" (zero-initialized convolution layers) that progressively grow the parameters from zero and ensure that no harmful noise could affect the finetuning. We test various conditioning controls, eg, edges, depth, segmentation, human pose, etc, with Stable Diffusion, using single or multiple conditions, with or without prompts. We show that the training of ControlNets is robust with small (<50k) and large (>1m) datasets. Extensive results show that ControlNet may facilitate wider applications to control image diffusion models.

Loading from the original format

By default the ControlNetModel should be loaded with from_pretrained, but it can also be loaded from the original format using [from_single_file] as follows:

from mindone.diffusers import StableDiffusionControlNetPipeline, ControlNetModel

url = "https://huggingface.co/lllyasviel/ControlNet-v1-1/blob/main/control_v11p_sd15_canny.pth"  # can also be a local path
controlnet = ControlNetModel.from_single_file(url)

url = "https://huggingface.co/runwayml/stable-diffusion-v1-5/blob/main/v1-5-pruned.safetensors"  # can also be a local path
pipe = StableDiffusionControlNetPipeline.from_single_file(url, controlnet=controlnet)

mindone.diffusers.ControlNetModel

Bases: ModelMixin, ConfigMixin, FromOriginalModelMixin

A ControlNet model.

PARAMETER	DESCRIPTION
in_channels	The number of channels in the input sample. TYPE: `int`, defaults to 4 DEFAULT: 4
flip_sin_to_cos	Whether to flip the sin to cos in the time embedding. TYPE: `bool`, defaults to `True` DEFAULT: True
freq_shift	The frequency shift to apply to the time embedding. TYPE: `int`, defaults to 0 DEFAULT: 0
down_block_types	The tuple of downsample blocks to use. TYPE: `tuple[str]`, defaults to `("CrossAttnDownBlock2D", "CrossAttnDownBlock2D", "CrossAttnDownBlock2D", "DownBlock2D")` DEFAULT: ('CrossAttnDownBlock2D', 'CrossAttnDownBlock2D', 'CrossAttnDownBlock2D', 'DownBlock2D')
only_cross_attention	TYPE: `Union[bool, Tuple[bool]]`, defaults to `False` DEFAULT: False
block_out_channels	The tuple of output channels for each block. TYPE: `tuple[int]`, defaults to `(320, 640, 1280, 1280)` DEFAULT: (320, 640, 1280, 1280)
layers_per_block	The number of layers per block. TYPE: `int`, defaults to 2 DEFAULT: 2
downsample_padding	The padding to use for the downsampling convolution. TYPE: `int`, defaults to 1 DEFAULT: 1
mid_block_scale_factor	The scale factor to use for the mid block. TYPE: `float`, defaults to 1 DEFAULT: 1
act_fn	The activation function to use. TYPE: `str`, defaults to "silu" DEFAULT: 'silu'
norm_num_groups	The number of groups to use for the normalization. If None, normalization and activation layers is skipped in post-processing. TYPE: `int`, *optional*, defaults to 32 DEFAULT: 32
norm_eps	The epsilon to use for the normalization. TYPE: `float`, defaults to 1e-5 DEFAULT: 1e-05
cross_attention_dim	The dimension of the cross attention features. TYPE: `int`, defaults to 1280 DEFAULT: 1280
transformer_layers_per_block	The number of transformer blocks of type [~models.attention.BasicTransformerBlock]. Only relevant for [~models.unet_2d_blocks.CrossAttnDownBlock2D], [~models.unet_2d_blocks.CrossAttnUpBlock2D], [~models.unet_2d_blocks.UNetMidBlock2DCrossAttn]. TYPE: `int` or `Tuple[int]`, *optional*, defaults to 1 DEFAULT: 1
encoder_hid_dim	If encoder_hid_dim_type is defined, encoder_hidden_states will be projected from encoder_hid_dim dimension to cross_attention_dim. TYPE: `int`, *optional*, defaults to None DEFAULT: None
encoder_hid_dim_type	If given, the encoder_hidden_states and potentially other embeddings are down-projected to text embeddings of dimension cross_attention according to encoder_hid_dim_type. TYPE: `str`, *optional*, defaults to `None` DEFAULT: None
attention_head_dim	The dimension of the attention heads. TYPE: `Union[int, Tuple[int]]`, defaults to 8 DEFAULT: 8
use_linear_projection	TYPE: `bool`, defaults to `False` DEFAULT: False
class_embed_type	The type of class embedding to use which is ultimately summed with the time embeddings. Choose from None, "timestep", "identity", "projection", or "simple_projection". TYPE: `str`, *optional*, defaults to `None` DEFAULT: None
addition_embed_type	Configures an optional embedding which will be summed with the time embeddings. Choose from None or "text". "text" will use the TextTimeEmbedding layer. TYPE: `str`, *optional*, defaults to `None` DEFAULT: None
num_class_embeds	Input dimension of the learnable embedding matrix to be projected to time_embed_dim, when performing class conditioning with class_embed_type equal to None. TYPE: `int`, *optional*, defaults to 0 DEFAULT: None
upcast_attention	TYPE: `bool`, defaults to `False` DEFAULT: False
resnet_time_scale_shift	Time scale shift config for ResNet blocks (see ResnetBlock2D). Choose from default or scale_shift. TYPE: `str`, defaults to `"default"` DEFAULT: 'default'
projection_class_embeddings_input_dim	The dimension of the class_labels input when class_embed_type="projection". Required when class_embed_type="projection". TYPE: `int`, *optional*, defaults to `None` DEFAULT: None
controlnet_conditioning_channel_order	The channel order of conditional image. Will convert to rgb if it's bgr. TYPE: `str`, defaults to `"rgb"` DEFAULT: 'rgb'
conditioning_embedding_out_channels	The tuple of output channel for each block in the conditioning_embedding layer. TYPE: `tuple[int]`, *optional*, defaults to `(16, 32, 96, 256)` DEFAULT: (16, 32, 96, 256)
global_pool_conditions	TODO(Patrick) - unused parameter. TYPE: `bool`, defaults to `False` DEFAULT: False
addition_embed_type_num_heads	The number of heads to use for the TextTimeEmbedding layer. TYPE: `int`, defaults to 64 DEFAULT: 64

Source code in mindone/diffusers/models/controlnet.py

class ControlNetModel(ModelMixin, ConfigMixin, FromOriginalModelMixin):
    """
    A ControlNet model.

    Args:
        in_channels (`int`, defaults to 4):
            The number of channels in the input sample.
        flip_sin_to_cos (`bool`, defaults to `True`):
            Whether to flip the sin to cos in the time embedding.
        freq_shift (`int`, defaults to 0):
            The frequency shift to apply to the time embedding.
        down_block_types (`tuple[str]`, defaults to `("CrossAttnDownBlock2D", "CrossAttnDownBlock2D", "CrossAttnDownBlock2D", "DownBlock2D")`):
            The tuple of downsample blocks to use.
        only_cross_attention (`Union[bool, Tuple[bool]]`, defaults to `False`):
        block_out_channels (`tuple[int]`, defaults to `(320, 640, 1280, 1280)`):
            The tuple of output channels for each block.
        layers_per_block (`int`, defaults to 2):
            The number of layers per block.
        downsample_padding (`int`, defaults to 1):
            The padding to use for the downsampling convolution.
        mid_block_scale_factor (`float`, defaults to 1):
            The scale factor to use for the mid block.
        act_fn (`str`, defaults to "silu"):
            The activation function to use.
        norm_num_groups (`int`, *optional*, defaults to 32):
            The number of groups to use for the normalization. If None, normalization and activation layers is skipped
            in post-processing.
        norm_eps (`float`, defaults to 1e-5):
            The epsilon to use for the normalization.
        cross_attention_dim (`int`, defaults to 1280):
            The dimension of the cross attention features.
        transformer_layers_per_block (`int` or `Tuple[int]`, *optional*, defaults to 1):
            The number of transformer blocks of type [`~models.attention.BasicTransformerBlock`]. Only relevant for
            [`~models.unet_2d_blocks.CrossAttnDownBlock2D`], [`~models.unet_2d_blocks.CrossAttnUpBlock2D`],
            [`~models.unet_2d_blocks.UNetMidBlock2DCrossAttn`].
        encoder_hid_dim (`int`, *optional*, defaults to None):
            If `encoder_hid_dim_type` is defined, `encoder_hidden_states` will be projected from `encoder_hid_dim`
            dimension to `cross_attention_dim`.
        encoder_hid_dim_type (`str`, *optional*, defaults to `None`):
            If given, the `encoder_hidden_states` and potentially other embeddings are down-projected to text
            embeddings of dimension `cross_attention` according to `encoder_hid_dim_type`.
        attention_head_dim (`Union[int, Tuple[int]]`, defaults to 8):
            The dimension of the attention heads.
        use_linear_projection (`bool`, defaults to `False`):
        class_embed_type (`str`, *optional*, defaults to `None`):
            The type of class embedding to use which is ultimately summed with the time embeddings. Choose from None,
            `"timestep"`, `"identity"`, `"projection"`, or `"simple_projection"`.
        addition_embed_type (`str`, *optional*, defaults to `None`):
            Configures an optional embedding which will be summed with the time embeddings. Choose from `None` or
            "text". "text" will use the `TextTimeEmbedding` layer.
        num_class_embeds (`int`, *optional*, defaults to 0):
            Input dimension of the learnable embedding matrix to be projected to `time_embed_dim`, when performing
            class conditioning with `class_embed_type` equal to `None`.
        upcast_attention (`bool`, defaults to `False`):
        resnet_time_scale_shift (`str`, defaults to `"default"`):
            Time scale shift config for ResNet blocks (see `ResnetBlock2D`). Choose from `default` or `scale_shift`.
        projection_class_embeddings_input_dim (`int`, *optional*, defaults to `None`):
            The dimension of the `class_labels` input when `class_embed_type="projection"`. Required when
            `class_embed_type="projection"`.
        controlnet_conditioning_channel_order (`str`, defaults to `"rgb"`):
            The channel order of conditional image. Will convert to `rgb` if it's `bgr`.
        conditioning_embedding_out_channels (`tuple[int]`, *optional*, defaults to `(16, 32, 96, 256)`):
            The tuple of output channel for each block in the `conditioning_embedding` layer.
        global_pool_conditions (`bool`, defaults to `False`):
            TODO(Patrick) - unused parameter.
        addition_embed_type_num_heads (`int`, defaults to 64):
            The number of heads to use for the `TextTimeEmbedding` layer.
    """

    _supports_gradient_checkpointing = True

    @register_to_config
    def __init__(
        self,
        in_channels: int = 4,
        conditioning_channels: int = 3,
        flip_sin_to_cos: bool = True,
        freq_shift: int = 0,
        down_block_types: Tuple[str, ...] = (
            "CrossAttnDownBlock2D",
            "CrossAttnDownBlock2D",
            "CrossAttnDownBlock2D",
            "DownBlock2D",
        ),
        mid_block_type: Optional[str] = "UNetMidBlock2DCrossAttn",
        only_cross_attention: Union[bool, Tuple[bool]] = False,
        block_out_channels: Tuple[int, ...] = (320, 640, 1280, 1280),
        layers_per_block: int = 2,
        downsample_padding: int = 1,
        mid_block_scale_factor: float = 1,
        act_fn: str = "silu",
        norm_num_groups: Optional[int] = 32,
        norm_eps: float = 1e-5,
        cross_attention_dim: int = 1280,
        transformer_layers_per_block: Union[int, Tuple[int, ...]] = 1,
        encoder_hid_dim: Optional[int] = None,
        encoder_hid_dim_type: Optional[str] = None,
        attention_head_dim: Union[int, Tuple[int, ...]] = 8,
        num_attention_heads: Optional[Union[int, Tuple[int, ...]]] = None,
        use_linear_projection: bool = False,
        class_embed_type: Optional[str] = None,
        addition_embed_type: Optional[str] = None,
        addition_time_embed_dim: Optional[int] = None,
        num_class_embeds: Optional[int] = None,
        upcast_attention: bool = False,
        resnet_time_scale_shift: str = "default",
        projection_class_embeddings_input_dim: Optional[int] = None,
        controlnet_conditioning_channel_order: str = "rgb",
        conditioning_embedding_out_channels: Optional[Tuple[int, ...]] = (16, 32, 96, 256),
        global_pool_conditions: bool = False,
        addition_embed_type_num_heads: int = 64,
    ):
        super().__init__()

        # If `num_attention_heads` is not defined (which is the case for most models)
        # it will default to `attention_head_dim`. This looks weird upon first reading it and it is.
        # The reason for this behavior is to correct for incorrectly named variables that were introduced
        # when this library was created. The incorrect naming was only discovered much later in
        # https://github.com/huggingface/diffusers/issues/2011#issuecomment-1547958131
        # Changing `attention_head_dim` to `num_attention_heads` for 40,000+ configurations is too backwards breaking
        # which is why we correct for the naming here.
        num_attention_heads = num_attention_heads or attention_head_dim

        # Check inputs
        if len(block_out_channels) != len(down_block_types):
            raise ValueError(
                f"Must provide the same number of `block_out_channels` as `down_block_types`. "
                f"`block_out_channels`: {block_out_channels}. `down_block_types`: {down_block_types}."
            )

        if not isinstance(only_cross_attention, bool) and len(only_cross_attention) != len(down_block_types):
            raise ValueError(
                f"Must provide the same number of `only_cross_attention` as `down_block_types`."
                f"`only_cross_attention`: {only_cross_attention}. `down_block_types`: {down_block_types}."
            )

        if not isinstance(num_attention_heads, int) and len(num_attention_heads) != len(down_block_types):
            raise ValueError(
                f"Must provide the same number of `num_attention_heads` as `down_block_types`."
                f"`num_attention_heads`: {num_attention_heads}. `down_block_types`: {down_block_types}."
            )

        if isinstance(transformer_layers_per_block, int):
            transformer_layers_per_block = [transformer_layers_per_block] * len(down_block_types)

        # input
        conv_in_kernel = 3
        conv_in_padding = (conv_in_kernel - 1) // 2
        self.conv_in = nn.Conv2d(
            in_channels,
            block_out_channels[0],
            kernel_size=conv_in_kernel,
            pad_mode="pad",
            padding=conv_in_padding,
            has_bias=True,
        )

        # time
        time_embed_dim = block_out_channels[0] * 4
        self.time_proj = Timesteps(block_out_channels[0], flip_sin_to_cos, freq_shift)
        timestep_input_dim = block_out_channels[0]
        self.time_embedding = TimestepEmbedding(
            timestep_input_dim,
            time_embed_dim,
            act_fn=act_fn,
        )

        if encoder_hid_dim_type is None and encoder_hid_dim is not None:
            encoder_hid_dim_type = "text_proj"
            self.register_to_config(encoder_hid_dim_type=encoder_hid_dim_type)
            logger.info("encoder_hid_dim_type defaults to 'text_proj' as `encoder_hid_dim` is defined.")

        if encoder_hid_dim is None and encoder_hid_dim_type is not None:
            raise ValueError(
                f"`encoder_hid_dim` has to be defined when `encoder_hid_dim_type` is set to {encoder_hid_dim_type}."
            )

        if encoder_hid_dim_type == "text_proj":
            self.encoder_hid_proj = nn.Dense(encoder_hid_dim, cross_attention_dim)
        elif encoder_hid_dim_type == "text_image_proj":
            # image_embed_dim DOESN'T have to be `cross_attention_dim`. To not clutter the __init__ too much
            # they are set to `cross_attention_dim` here as this is exactly the required dimension for the currently only use
            # case when `addition_embed_type == "text_image_proj"` (Kandinsky 2.1)`
            self.encoder_hid_proj = TextImageProjection(
                text_embed_dim=encoder_hid_dim,
                image_embed_dim=cross_attention_dim,
                cross_attention_dim=cross_attention_dim,
            )

        elif encoder_hid_dim_type is not None:
            raise ValueError(
                f"encoder_hid_dim_type: {encoder_hid_dim_type} must be None, 'text_proj' or 'text_image_proj'."
            )
        else:
            self.encoder_hid_proj = None

        # class embedding
        if class_embed_type is None and num_class_embeds is not None:
            self.class_embedding = nn.Embedding(num_class_embeds, time_embed_dim)
        elif class_embed_type == "timestep":
            self.class_embedding = TimestepEmbedding(timestep_input_dim, time_embed_dim)
        elif class_embed_type == "identity":
            self.class_embedding = nn.Identity()  # time_embed_dim, time_embed_dim
        elif class_embed_type == "projection":
            if projection_class_embeddings_input_dim is None:
                raise ValueError(
                    "`class_embed_type`: 'projection' requires `projection_class_embeddings_input_dim` be set"
                )
            # The projection `class_embed_type` is the same as the timestep `class_embed_type` except
            # 1. the `class_labels` inputs are not first converted to sinusoidal embeddings
            # 2. it projects from an arbitrary input dimension.
            #
            # Note that `TimestepEmbedding` is quite general, being mainly linear layers and activations.
            # When used for embedding actual timesteps, the timesteps are first converted to sinusoidal embeddings.
            # As a result, `TimestepEmbedding` can be passed arbitrary vectors.
            self.class_embedding = TimestepEmbedding(projection_class_embeddings_input_dim, time_embed_dim)
        else:
            self.class_embedding = None

        if addition_embed_type == "text":
            if encoder_hid_dim is not None:
                text_time_embedding_from_dim = encoder_hid_dim
            else:
                text_time_embedding_from_dim = cross_attention_dim

            self.add_embedding = TextTimeEmbedding(
                text_time_embedding_from_dim, time_embed_dim, num_heads=addition_embed_type_num_heads
            )
        elif addition_embed_type == "text_image":
            # text_embed_dim and image_embed_dim DON'T have to be `cross_attention_dim`. To not clutter the __init__ too much
            # they are set to `cross_attention_dim` here as this is exactly the required dimension for the currently only use
            # case when `addition_embed_type == "text_image"` (Kandinsky 2.1)`
            self.add_embedding = TextImageTimeEmbedding(
                text_embed_dim=cross_attention_dim, image_embed_dim=cross_attention_dim, time_embed_dim=time_embed_dim
            )
        elif addition_embed_type == "text_time":
            self.add_time_proj = Timesteps(addition_time_embed_dim, flip_sin_to_cos, freq_shift)
            self.add_embedding = TimestepEmbedding(projection_class_embeddings_input_dim, time_embed_dim)

        elif addition_embed_type is not None:
            raise ValueError(f"addition_embed_type: {addition_embed_type} must be None, 'text' or 'text_image'.")

        # control net conditioning embedding
        self.controlnet_cond_embedding = ControlNetConditioningEmbedding(
            conditioning_embedding_channels=block_out_channels[0],
            block_out_channels=conditioning_embedding_out_channels,
            conditioning_channels=conditioning_channels,
        )

        self.down_blocks = []
        self.controlnet_down_blocks = []

        if isinstance(only_cross_attention, bool):
            only_cross_attention = [only_cross_attention] * len(down_block_types)

        if isinstance(attention_head_dim, int):
            attention_head_dim = (attention_head_dim,) * len(down_block_types)

        if isinstance(num_attention_heads, int):
            num_attention_heads = (num_attention_heads,) * len(down_block_types)

        # down
        output_channel = block_out_channels[0]

        controlnet_block = nn.Conv2d(
            output_channel, output_channel, kernel_size=1, has_bias=True, weight_init="zeros", bias_init="zeros"
        )
        self.controlnet_down_blocks.append(controlnet_block)

        for i, down_block_type in enumerate(down_block_types):
            input_channel = output_channel
            output_channel = block_out_channels[i]
            is_final_block = i == len(block_out_channels) - 1

            down_block = get_down_block(
                down_block_type,
                num_layers=layers_per_block,
                transformer_layers_per_block=transformer_layers_per_block[i],
                in_channels=input_channel,
                out_channels=output_channel,
                temb_channels=time_embed_dim,
                add_downsample=not is_final_block,
                resnet_eps=norm_eps,
                resnet_act_fn=act_fn,
                resnet_groups=norm_num_groups,
                cross_attention_dim=cross_attention_dim,
                num_attention_heads=num_attention_heads[i],
                attention_head_dim=attention_head_dim[i] if attention_head_dim[i] is not None else output_channel,
                downsample_padding=downsample_padding,
                use_linear_projection=use_linear_projection,
                only_cross_attention=only_cross_attention[i],
                upcast_attention=upcast_attention,
                resnet_time_scale_shift=resnet_time_scale_shift,
            )
            self.down_blocks.append(down_block)

            for _ in range(layers_per_block):
                controlnet_block = nn.Conv2d(
                    output_channel, output_channel, kernel_size=1, has_bias=True, weight_init="zeros", bias_init="zeros"
                )
                self.controlnet_down_blocks.append(controlnet_block)

            if not is_final_block:
                controlnet_block = nn.Conv2d(
                    output_channel, output_channel, kernel_size=1, has_bias=True, weight_init="zeros", bias_init="zeros"
                )
                self.controlnet_down_blocks.append(controlnet_block)

        self.down_blocks = nn.CellList(self.down_blocks)
        self.controlnet_down_blocks = nn.CellList(self.controlnet_down_blocks)

        # mid
        mid_block_channel = block_out_channels[-1]

        controlnet_block = nn.Conv2d(
            mid_block_channel, mid_block_channel, kernel_size=1, has_bias=True, weight_init="zeros", bias_init="zeros"
        )
        self.controlnet_mid_block = controlnet_block

        if mid_block_type == "UNetMidBlock2DCrossAttn":
            self.mid_block = UNetMidBlock2DCrossAttn(
                transformer_layers_per_block=transformer_layers_per_block[-1],
                in_channels=mid_block_channel,
                temb_channels=time_embed_dim,
                resnet_eps=norm_eps,
                resnet_act_fn=act_fn,
                output_scale_factor=mid_block_scale_factor,
                resnet_time_scale_shift=resnet_time_scale_shift,
                cross_attention_dim=cross_attention_dim,
                num_attention_heads=num_attention_heads[-1],
                resnet_groups=norm_num_groups,
                use_linear_projection=use_linear_projection,
                upcast_attention=upcast_attention,
            )
        elif mid_block_type == "UNetMidBlock2D":
            self.mid_block = UNetMidBlock2D(
                in_channels=block_out_channels[-1],
                temb_channels=time_embed_dim,
                num_layers=0,
                resnet_eps=norm_eps,
                resnet_act_fn=act_fn,
                output_scale_factor=mid_block_scale_factor,
                resnet_groups=norm_num_groups,
                resnet_time_scale_shift=resnet_time_scale_shift,
                add_attention=False,
            )
        else:
            raise ValueError(f"unknown mid_block_type : {mid_block_type}")

    @classmethod
    def from_unet(
        cls,
        unet: UNet2DConditionModel,
        controlnet_conditioning_channel_order: str = "rgb",
        conditioning_embedding_out_channels: Optional[Tuple[int, ...]] = (16, 32, 96, 256),
        load_weights_from_unet: bool = True,
        conditioning_channels: int = 3,
    ):
        r"""
        Instantiate a [`ControlNetModel`] from [`UNet2DConditionModel`].

        Parameters:
            unet (`UNet2DConditionModel`):
                The UNet model weights to copy to the [`ControlNetModel`]. All configuration options are also copied
                where applicable.
        """
        transformer_layers_per_block = (
            unet.config.transformer_layers_per_block if "transformer_layers_per_block" in unet.config else 1
        )
        encoder_hid_dim = unet.config.encoder_hid_dim if "encoder_hid_dim" in unet.config else None
        encoder_hid_dim_type = unet.config.encoder_hid_dim_type if "encoder_hid_dim_type" in unet.config else None
        addition_embed_type = unet.config.addition_embed_type if "addition_embed_type" in unet.config else None
        addition_time_embed_dim = (
            unet.config.addition_time_embed_dim if "addition_time_embed_dim" in unet.config else None
        )

        controlnet = cls(
            encoder_hid_dim=encoder_hid_dim,
            encoder_hid_dim_type=encoder_hid_dim_type,
            addition_embed_type=addition_embed_type,
            addition_time_embed_dim=addition_time_embed_dim,
            transformer_layers_per_block=transformer_layers_per_block,
            in_channels=unet.config.in_channels,
            flip_sin_to_cos=unet.config.flip_sin_to_cos,
            freq_shift=unet.config.freq_shift,
            down_block_types=unet.config.down_block_types,
            only_cross_attention=unet.config.only_cross_attention,
            block_out_channels=unet.config.block_out_channels,
            layers_per_block=unet.config.layers_per_block,
            downsample_padding=unet.config.downsample_padding,
            mid_block_scale_factor=unet.config.mid_block_scale_factor,
            act_fn=unet.config.act_fn,
            norm_num_groups=unet.config.norm_num_groups,
            norm_eps=unet.config.norm_eps,
            cross_attention_dim=unet.config.cross_attention_dim,
            attention_head_dim=unet.config.attention_head_dim,
            num_attention_heads=unet.config.num_attention_heads,
            use_linear_projection=unet.config.use_linear_projection,
            class_embed_type=unet.config.class_embed_type,
            num_class_embeds=unet.config.num_class_embeds,
            upcast_attention=unet.config.upcast_attention,
            resnet_time_scale_shift=unet.config.resnet_time_scale_shift,
            projection_class_embeddings_input_dim=unet.config.projection_class_embeddings_input_dim,
            mid_block_type=unet.config.mid_block_type,
            controlnet_conditioning_channel_order=controlnet_conditioning_channel_order,
            conditioning_embedding_out_channels=conditioning_embedding_out_channels,
            conditioning_channels=conditioning_channels,
        )

        if load_weights_from_unet:
            ms.load_param_into_net(controlnet.conv_in, unet.conv_in.parameters_dict(), strict_load=True)
            ms.load_param_into_net(controlnet.time_proj, unet.time_proj.parameters_dict(), strict_load=True)
            ms.load_param_into_net(controlnet.time_embedding, unet.time_embedding.parameters_dict(), strict_load=True)

            if controlnet.class_embedding:
                ms.load_param_into_net(
                    controlnet.class_embedding, unet.class_embedding.parameters_dict(), strict_load=True
                )

            if hasattr(controlnet, "add_embedding"):
                ms.load_param_into_net(controlnet.add_embedding, unet.add_embedding.parameters_dict())

            ms.load_param_into_net(controlnet.down_blocks, unet.down_blocks.parameters_dict(), strict_load=True)
            ms.load_param_into_net(controlnet.mid_block, unet.mid_block.parameters_dict(), strict_load=True)

        return controlnet

    @property
    # Copied from diffusers.models.unets.unet_2d_condition.UNet2DConditionModel.attn_processors
    def attn_processors(self) -> Dict[str, AttentionProcessor]:
        r"""
        Returns:
            `dict` of attention processors: A dictionary containing all attention processors used in the model with
            indexed by its weight name.
        """
        # set recursively
        processors = {}

        def fn_recursive_add_processors(name: str, module: nn.Cell, processors: Dict[str, AttentionProcessor]):
            if hasattr(module, "get_processor"):
                processors[f"{name}.processor"] = module.get_processor()

            for sub_name, child in module.name_cells().items():
                fn_recursive_add_processors(f"{name}.{sub_name}", child, processors)

            return processors

        for name, module in self.name_cells().items():
            fn_recursive_add_processors(name, module, processors)

        return processors

    # Copied from diffusers.models.unets.unet_2d_condition.UNet2DConditionModel.set_attn_processor
    def set_attn_processor(self, processor: Union[AttentionProcessor, Dict[str, AttentionProcessor]]):
        r"""
        Sets the attention processor to use to compute attention.
        Parameters:
            processor (`dict` of `AttentionProcessor` or only `AttentionProcessor`):
                The instantiated processor class or a dictionary of processor classes that will be set as the processor
                for **all** `Attention` layers.
                If `processor` is a dict, the key needs to define the path to the corresponding cross attention
                processor. This is strongly recommended when setting trainable attention processors.
        """
        count = len(self.attn_processors.keys())

        if isinstance(processor, dict) and len(processor) != count:
            raise ValueError(
                f"A dict of processors was passed, but the number of processors {len(processor)} does not match the"
                f" number of attention layers: {count}. Please make sure to pass {count} processor classes."
            )

        def fn_recursive_attn_processor(name: str, module: nn.Cell, processor):
            if hasattr(module, "set_processor"):
                if not isinstance(processor, dict):
                    module.set_processor(processor)
                else:
                    module.set_processor(processor.pop(f"{name}.processor"))

            for sub_name, child in module.name_cells().items():
                fn_recursive_attn_processor(f"{name}.{sub_name}", child, processor)

        for name, module in self.name_cells().items():
            fn_recursive_attn_processor(name, module, processor)

    # Copied from diffusers.models.unets.unet_2d_condition.UNet2DConditionModel.set_default_attn_processor
    def set_default_attn_processor(self):
        """
        Disables custom attention processors and sets the default attention implementation.
        """
        if all(proc.__class__ in CROSS_ATTENTION_PROCESSORS for proc in self.attn_processors.values()):
            processor = AttnProcessor()
        else:
            raise ValueError(
                f"Cannot call `set_default_attn_processor` when attention processors are of type {next(iter(self.attn_processors.values()))}"
            )

        self.set_attn_processor(processor)

    def _set_gradient_checkpointing(self, module, value: bool = False) -> None:
        if isinstance(module, (CrossAttnDownBlock2D, DownBlock2D)):
            module.gradient_checkpointing = value

    def construct(
        self,
        sample: ms.Tensor,
        timestep: Union[ms.Tensor, float, int],
        encoder_hidden_states: ms.Tensor,
        controlnet_cond: ms.Tensor,
        conditioning_scale: float = 1.0,
        class_labels: Optional[ms.Tensor] = None,
        timestep_cond: Optional[ms.Tensor] = None,
        attention_mask: Optional[ms.Tensor] = None,
        added_cond_kwargs: Optional[Dict[str, ms.Tensor]] = None,
        cross_attention_kwargs: Optional[Dict[str, Any]] = None,
        guess_mode: bool = False,
        return_dict: bool = False,
    ) -> Union[ControlNetOutput, Tuple[Tuple[ms.Tensor, ...], ms.Tensor]]:
        """
        The [`ControlNetModel`] forward method.

        Args:
            sample (`ms.Tensor`):
                The noisy input tensor.
            timestep (`Union[ms.Tensor, float, int]`):
                The number of timesteps to denoise an input.
            encoder_hidden_states (`ms.Tensor`):
                The encoder hidden states.
            controlnet_cond (`ms.Tensor`):
                The conditional input tensor of shape `(batch_size, sequence_length, hidden_size)`.
            conditioning_scale (`float`, defaults to `1.0`):
                The scale factor for ControlNet outputs.
            class_labels (`ms.Tensor`, *optional*, defaults to `None`):
                Optional class labels for conditioning. Their embeddings will be summed with the timestep embeddings.
            timestep_cond (`ms.Tensor`, *optional*, defaults to `None`):
                Additional conditional embeddings for timestep. If provided, the embeddings will be summed with the
                timestep_embedding passed through the `self.time_embedding` layer to obtain the final timestep
                embeddings.
            attention_mask (`ms.Tensor`, *optional*, defaults to `None`):
                An attention mask of shape `(batch, key_tokens)` is applied to `encoder_hidden_states`. If `1` the mask
                is kept, otherwise if `0` it is discarded. Mask will be converted into a bias, which adds large
                negative values to the attention scores corresponding to "discard" tokens.
            added_cond_kwargs (`dict`):
                Additional conditions for the Stable Diffusion XL UNet.
            cross_attention_kwargs (`dict[str]`, *optional*, defaults to `None`):
                A kwargs dictionary that if specified is passed along to the `AttnProcessor`.
            guess_mode (`bool`, defaults to `False`):
                In this mode, the ControlNet encoder tries its best to recognize the input content of the input even if
                you remove all prompts. A `guidance_scale` between 3.0 and 5.0 is recommended.
            return_dict (`bool`, defaults to `False`):
                Whether or not to return a [`~models.controlnet.ControlNetOutput`] instead of a plain tuple.

        Returns:
            [`~models.controlnet.ControlNetOutput`] **or** `tuple`:
                If `return_dict` is `True`, a [`~models.controlnet.ControlNetOutput`] is returned, otherwise a tuple is
                returned where the first element is the sample tensor.
        """
        # check channel order
        channel_order = self.config["controlnet_conditioning_channel_order"]

        if channel_order == "rgb":
            # in rgb order by default
            ...
        elif channel_order == "bgr":
            controlnet_cond = ops.flip(controlnet_cond, dims=[1])
        else:
            raise ValueError(f"unknown `controlnet_conditioning_channel_order`: {channel_order}")

        # prepare attention_mask
        if attention_mask is not None:
            attention_mask = (1 - attention_mask.to(sample.dtype)) * -10000.0
            attention_mask = attention_mask.unsqueeze(1)

        # 1. time
        timesteps = timestep
        if not ops.is_tensor(timesteps):
            # TODO: this requires sync between CPU and GPU. So try to pass timesteps as tensors if you can
            # This would be a good case for the `match` statement (Python 3.10+)
            if isinstance(timestep, float):
                dtype = ms.float64
            else:
                dtype = ms.int64
            timesteps = ms.Tensor([timesteps], dtype=dtype)
        elif len(timesteps.shape) == 0:
            timesteps = timesteps[None]

        # broadcast to batch dimension in a way that's compatible with ONNX/Core ML
        if timesteps.shape[0] == 1:
            timesteps = timesteps.tile((sample.shape[0],))

        t_emb = self.time_proj(timesteps)

        # timesteps does not contain any weights and will always return f32 tensors
        # but time_embedding might actually be running in fp16. so we need to cast here.
        # there might be better ways to encapsulate this.
        t_emb = t_emb.to(dtype=sample.dtype)

        emb = self.time_embedding(t_emb, timestep_cond)
        aug_emb = None

        if self.class_embedding is not None:
            if class_labels is None:
                raise ValueError("class_labels should be provided when num_class_embeds > 0")

            if self.config["class_embed_type"] == "timestep":
                class_labels = self.time_proj(class_labels)

                # `Timesteps` does not contain any weights and will always return f32 tensors
                # there might be better ways to encapsulate this.
                class_labels = class_labels.to(dtype=sample.dtype)

            class_emb = self.class_embedding(class_labels).to(dtype=sample.dtype)
            emb = emb + class_emb

        if self.config["addition_embed_type"] is not None:
            if self.config["addition_embed_type"] == "text":
                aug_emb = self.add_embedding(encoder_hidden_states)

            elif self.config["addition_embed_type"] == "text_time":
                if "text_embeds" not in added_cond_kwargs:
                    raise ValueError(
                        f"{self.__class__} has the config param `addition_embed_type` set to 'text_time' "
                        f"which requires the keyword argument `text_embeds` to be passed in `added_cond_kwargs`"
                    )
                text_embeds = added_cond_kwargs.get("text_embeds")
                if "time_ids" not in added_cond_kwargs:
                    raise ValueError(
                        f"{self.__class__} has the config param `addition_embed_type` set to 'text_time' "
                        f"which requires the keyword argument `time_ids` to be passed in `added_cond_kwargs`"
                    )
                time_ids = added_cond_kwargs.get("time_ids")
                time_embeds = self.add_time_proj(time_ids.flatten())
                # `Timesteps` does not contain any weights and will always return f32 tensors
                # there might be better ways to encapsulate this.
                time_embeds = time_embeds.to(emb.dtype)
                time_embeds = time_embeds.reshape((text_embeds.shape[0], -1))

                add_embeds = ops.concat([text_embeds, time_embeds], axis=-1)
                add_embeds = add_embeds.to(emb.dtype)
                aug_emb = self.add_embedding(add_embeds)

        emb = emb + aug_emb if aug_emb is not None else emb

        # 2. pre-process
        sample = self.conv_in(sample)

        controlnet_cond = self.controlnet_cond_embedding(controlnet_cond)
        sample = sample + controlnet_cond

        # 3. down
        down_block_res_samples = (sample,)
        for downsample_block in self.down_blocks:
            if downsample_block.has_cross_attention:
                sample, res_samples = downsample_block(
                    hidden_states=sample,
                    temb=emb,
                    encoder_hidden_states=encoder_hidden_states,
                    attention_mask=attention_mask,
                    cross_attention_kwargs=cross_attention_kwargs,
                )
            else:
                sample, res_samples = downsample_block(hidden_states=sample, temb=emb)

            down_block_res_samples += res_samples

        # 4. mid
        if self.mid_block is not None:
            if self.mid_block.has_cross_attention:
                sample = self.mid_block(
                    sample,
                    emb,
                    encoder_hidden_states=encoder_hidden_states,
                    attention_mask=attention_mask,
                    cross_attention_kwargs=cross_attention_kwargs,
                )
            else:
                sample = self.mid_block(sample, emb)

        # 5. Control net blocks
        controlnet_down_block_res_samples = ()

        for down_block_res_sample, controlnet_block in zip(down_block_res_samples, self.controlnet_down_blocks):
            down_block_res_sample = controlnet_block(down_block_res_sample)
            controlnet_down_block_res_samples = controlnet_down_block_res_samples + (down_block_res_sample,)

        down_block_res_samples = controlnet_down_block_res_samples

        mid_block_res_sample = self.controlnet_mid_block(sample)

        # 6. scaling
        if guess_mode and not self.config["global_pool_conditions"]:
            scales = ops.logspace(-1.0, 0.0, len(down_block_res_samples) + 1)  # 0.1 to 1.0
            scales = scales * conditioning_scale
            # Cast scale to sample.dtype manually as torch do the same automatically for scaler tensors
            down_block_res_samples = [
                sample * scale.to(sample.dtype) for sample, scale in zip(down_block_res_samples, scales)
            ]
            mid_block_res_sample = mid_block_res_sample * scales[-1]  # last one
        else:
            down_block_res_samples = [sample * conditioning_scale for sample in down_block_res_samples]
            mid_block_res_sample = mid_block_res_sample * conditioning_scale

        if self.config["global_pool_conditions"]:
            down_block_res_samples = [
                ops.mean(sample, axis=(2, 3), keep_dims=True) for sample in down_block_res_samples
            ]
            mid_block_res_sample = ops.mean(mid_block_res_sample, axis=(2, 3), keep_dims=True)

        if not return_dict:
            return (down_block_res_samples, mid_block_res_sample)

        return ControlNetOutput(
            down_block_res_samples=down_block_res_samples, mid_block_res_sample=mid_block_res_sample
        )

mindone.diffusers.ControlNetModel.attn_processors: Dict[str, AttentionProcessor] property

RETURNS	DESCRIPTION
Dict[str, AttentionProcessor]	dict of attention processors: A dictionary containing all attention processors used in the model with
Dict[str, AttentionProcessor]	indexed by its weight name.

mindone.diffusers.ControlNetModel.construct(sample, timestep, encoder_hidden_states, controlnet_cond, conditioning_scale=1.0, class_labels=None, timestep_cond=None, attention_mask=None, added_cond_kwargs=None, cross_attention_kwargs=None, guess_mode=False, return_dict=False)

The [ControlNetModel] forward method.

PARAMETER	DESCRIPTION
sample	The noisy input tensor. TYPE: `ms.Tensor`
timestep	The number of timesteps to denoise an input. TYPE: `Union[ms.Tensor, float, int]`
encoder_hidden_states	The encoder hidden states. TYPE: `ms.Tensor`
controlnet_cond	The conditional input tensor of shape (batch_size, sequence_length, hidden_size). TYPE: `ms.Tensor`
conditioning_scale	The scale factor for ControlNet outputs. TYPE: `float`, defaults to `1.0` DEFAULT: 1.0
class_labels	Optional class labels for conditioning. Their embeddings will be summed with the timestep embeddings. TYPE: `ms.Tensor`, *optional*, defaults to `None` DEFAULT: None
timestep_cond	Additional conditional embeddings for timestep. If provided, the embeddings will be summed with the timestep_embedding passed through the self.time_embedding layer to obtain the final timestep embeddings. TYPE: `ms.Tensor`, *optional*, defaults to `None` DEFAULT: None
attention_mask	An attention mask of shape (batch, key_tokens) is applied to encoder_hidden_states. If 1 the mask is kept, otherwise if 0 it is discarded. Mask will be converted into a bias, which adds large negative values to the attention scores corresponding to "discard" tokens. TYPE: `ms.Tensor`, *optional*, defaults to `None` DEFAULT: None
added_cond_kwargs	Additional conditions for the Stable Diffusion XL UNet. TYPE: `dict` DEFAULT: None
cross_attention_kwargs	A kwargs dictionary that if specified is passed along to the AttnProcessor. TYPE: `dict[str]`, *optional*, defaults to `None` DEFAULT: None
guess_mode	In this mode, the ControlNet encoder tries its best to recognize the input content of the input even if you remove all prompts. A guidance_scale between 3.0 and 5.0 is recommended. TYPE: `bool`, defaults to `False` DEFAULT: False
return_dict	Whether or not to return a [~models.controlnet.ControlNetOutput] instead of a plain tuple. TYPE: `bool`, defaults to `False` DEFAULT: False

RETURNS	DESCRIPTION
Union[ControlNetOutput, Tuple[Tuple[Tensor, ...], Tensor]]	[~models.controlnet.ControlNetOutput] or tuple: If return_dict is True, a [~models.controlnet.ControlNetOutput] is returned, otherwise a tuple is returned where the first element is the sample tensor.

Source code in mindone/diffusers/models/controlnet.py

def construct(
    self,
    sample: ms.Tensor,
    timestep: Union[ms.Tensor, float, int],
    encoder_hidden_states: ms.Tensor,
    controlnet_cond: ms.Tensor,
    conditioning_scale: float = 1.0,
    class_labels: Optional[ms.Tensor] = None,
    timestep_cond: Optional[ms.Tensor] = None,
    attention_mask: Optional[ms.Tensor] = None,
    added_cond_kwargs: Optional[Dict[str, ms.Tensor]] = None,
    cross_attention_kwargs: Optional[Dict[str, Any]] = None,
    guess_mode: bool = False,
    return_dict: bool = False,
) -> Union[ControlNetOutput, Tuple[Tuple[ms.Tensor, ...], ms.Tensor]]:
    """
    The [`ControlNetModel`] forward method.

    Args:
        sample (`ms.Tensor`):
            The noisy input tensor.
        timestep (`Union[ms.Tensor, float, int]`):
            The number of timesteps to denoise an input.
        encoder_hidden_states (`ms.Tensor`):
            The encoder hidden states.
        controlnet_cond (`ms.Tensor`):
            The conditional input tensor of shape `(batch_size, sequence_length, hidden_size)`.
        conditioning_scale (`float`, defaults to `1.0`):
            The scale factor for ControlNet outputs.
        class_labels (`ms.Tensor`, *optional*, defaults to `None`):
            Optional class labels for conditioning. Their embeddings will be summed with the timestep embeddings.
        timestep_cond (`ms.Tensor`, *optional*, defaults to `None`):
            Additional conditional embeddings for timestep. If provided, the embeddings will be summed with the
            timestep_embedding passed through the `self.time_embedding` layer to obtain the final timestep
            embeddings.
        attention_mask (`ms.Tensor`, *optional*, defaults to `None`):
            An attention mask of shape `(batch, key_tokens)` is applied to `encoder_hidden_states`. If `1` the mask
            is kept, otherwise if `0` it is discarded. Mask will be converted into a bias, which adds large
            negative values to the attention scores corresponding to "discard" tokens.
        added_cond_kwargs (`dict`):
            Additional conditions for the Stable Diffusion XL UNet.
        cross_attention_kwargs (`dict[str]`, *optional*, defaults to `None`):
            A kwargs dictionary that if specified is passed along to the `AttnProcessor`.
        guess_mode (`bool`, defaults to `False`):
            In this mode, the ControlNet encoder tries its best to recognize the input content of the input even if
            you remove all prompts. A `guidance_scale` between 3.0 and 5.0 is recommended.
        return_dict (`bool`, defaults to `False`):
            Whether or not to return a [`~models.controlnet.ControlNetOutput`] instead of a plain tuple.

    Returns:
        [`~models.controlnet.ControlNetOutput`] **or** `tuple`:
            If `return_dict` is `True`, a [`~models.controlnet.ControlNetOutput`] is returned, otherwise a tuple is
            returned where the first element is the sample tensor.
    """
    # check channel order
    channel_order = self.config["controlnet_conditioning_channel_order"]

    if channel_order == "rgb":
        # in rgb order by default
        ...
    elif channel_order == "bgr":
        controlnet_cond = ops.flip(controlnet_cond, dims=[1])
    else:
        raise ValueError(f"unknown `controlnet_conditioning_channel_order`: {channel_order}")

    # prepare attention_mask
    if attention_mask is not None:
        attention_mask = (1 - attention_mask.to(sample.dtype)) * -10000.0
        attention_mask = attention_mask.unsqueeze(1)

    # 1. time
    timesteps = timestep
    if not ops.is_tensor(timesteps):
        # TODO: this requires sync between CPU and GPU. So try to pass timesteps as tensors if you can
        # This would be a good case for the `match` statement (Python 3.10+)
        if isinstance(timestep, float):
            dtype = ms.float64
        else:
            dtype = ms.int64
        timesteps = ms.Tensor([timesteps], dtype=dtype)
    elif len(timesteps.shape) == 0:
        timesteps = timesteps[None]

    # broadcast to batch dimension in a way that's compatible with ONNX/Core ML
    if timesteps.shape[0] == 1:
        timesteps = timesteps.tile((sample.shape[0],))

    t_emb = self.time_proj(timesteps)

    # timesteps does not contain any weights and will always return f32 tensors
    # but time_embedding might actually be running in fp16. so we need to cast here.
    # there might be better ways to encapsulate this.
    t_emb = t_emb.to(dtype=sample.dtype)

    emb = self.time_embedding(t_emb, timestep_cond)
    aug_emb = None

    if self.class_embedding is not None:
        if class_labels is None:
            raise ValueError("class_labels should be provided when num_class_embeds > 0")

        if self.config["class_embed_type"] == "timestep":
            class_labels = self.time_proj(class_labels)

            # `Timesteps` does not contain any weights and will always return f32 tensors
            # there might be better ways to encapsulate this.
            class_labels = class_labels.to(dtype=sample.dtype)

        class_emb = self.class_embedding(class_labels).to(dtype=sample.dtype)
        emb = emb + class_emb

    if self.config["addition_embed_type"] is not None:
        if self.config["addition_embed_type"] == "text":
            aug_emb = self.add_embedding(encoder_hidden_states)

        elif self.config["addition_embed_type"] == "text_time":
            if "text_embeds" not in added_cond_kwargs:
                raise ValueError(
                    f"{self.__class__} has the config param `addition_embed_type` set to 'text_time' "
                    f"which requires the keyword argument `text_embeds` to be passed in `added_cond_kwargs`"
                )
            text_embeds = added_cond_kwargs.get("text_embeds")
            if "time_ids" not in added_cond_kwargs:
                raise ValueError(
                    f"{self.__class__} has the config param `addition_embed_type` set to 'text_time' "
                    f"which requires the keyword argument `time_ids` to be passed in `added_cond_kwargs`"
                )
            time_ids = added_cond_kwargs.get("time_ids")
            time_embeds = self.add_time_proj(time_ids.flatten())
            # `Timesteps` does not contain any weights and will always return f32 tensors
            # there might be better ways to encapsulate this.
            time_embeds = time_embeds.to(emb.dtype)
            time_embeds = time_embeds.reshape((text_embeds.shape[0], -1))

            add_embeds = ops.concat([text_embeds, time_embeds], axis=-1)
            add_embeds = add_embeds.to(emb.dtype)
            aug_emb = self.add_embedding(add_embeds)

    emb = emb + aug_emb if aug_emb is not None else emb

    # 2. pre-process
    sample = self.conv_in(sample)

    controlnet_cond = self.controlnet_cond_embedding(controlnet_cond)
    sample = sample + controlnet_cond

    # 3. down
    down_block_res_samples = (sample,)
    for downsample_block in self.down_blocks:
        if downsample_block.has_cross_attention:
            sample, res_samples = downsample_block(
                hidden_states=sample,
                temb=emb,
                encoder_hidden_states=encoder_hidden_states,
                attention_mask=attention_mask,
                cross_attention_kwargs=cross_attention_kwargs,
            )
        else:
            sample, res_samples = downsample_block(hidden_states=sample, temb=emb)

        down_block_res_samples += res_samples

    # 4. mid
    if self.mid_block is not None:
        if self.mid_block.has_cross_attention:
            sample = self.mid_block(
                sample,
                emb,
                encoder_hidden_states=encoder_hidden_states,
                attention_mask=attention_mask,
                cross_attention_kwargs=cross_attention_kwargs,
            )
        else:
            sample = self.mid_block(sample, emb)

    # 5. Control net blocks
    controlnet_down_block_res_samples = ()

    for down_block_res_sample, controlnet_block in zip(down_block_res_samples, self.controlnet_down_blocks):
        down_block_res_sample = controlnet_block(down_block_res_sample)
        controlnet_down_block_res_samples = controlnet_down_block_res_samples + (down_block_res_sample,)

    down_block_res_samples = controlnet_down_block_res_samples

    mid_block_res_sample = self.controlnet_mid_block(sample)

    # 6. scaling
    if guess_mode and not self.config["global_pool_conditions"]:
        scales = ops.logspace(-1.0, 0.0, len(down_block_res_samples) + 1)  # 0.1 to 1.0
        scales = scales * conditioning_scale
        # Cast scale to sample.dtype manually as torch do the same automatically for scaler tensors
        down_block_res_samples = [
            sample * scale.to(sample.dtype) for sample, scale in zip(down_block_res_samples, scales)
        ]
        mid_block_res_sample = mid_block_res_sample * scales[-1]  # last one
    else:
        down_block_res_samples = [sample * conditioning_scale for sample in down_block_res_samples]
        mid_block_res_sample = mid_block_res_sample * conditioning_scale

    if self.config["global_pool_conditions"]:
        down_block_res_samples = [
            ops.mean(sample, axis=(2, 3), keep_dims=True) for sample in down_block_res_samples
        ]
        mid_block_res_sample = ops.mean(mid_block_res_sample, axis=(2, 3), keep_dims=True)

    if not return_dict:
        return (down_block_res_samples, mid_block_res_sample)

    return ControlNetOutput(
        down_block_res_samples=down_block_res_samples, mid_block_res_sample=mid_block_res_sample
    )

mindone.diffusers.ControlNetModel.from_unet(unet, controlnet_conditioning_channel_order='rgb', conditioning_embedding_out_channels=(16, 32, 96, 256), load_weights_from_unet=True, conditioning_channels=3) classmethod

Instantiate a [ControlNetModel] from [UNet2DConditionModel].

PARAMETER	DESCRIPTION
unet	The UNet model weights to copy to the [ControlNetModel]. All configuration options are also copied where applicable. TYPE: `UNet2DConditionModel`

Source code in mindone/diffusers/models/controlnet.py

@classmethod
def from_unet(
    cls,
    unet: UNet2DConditionModel,
    controlnet_conditioning_channel_order: str = "rgb",
    conditioning_embedding_out_channels: Optional[Tuple[int, ...]] = (16, 32, 96, 256),
    load_weights_from_unet: bool = True,
    conditioning_channels: int = 3,
):
    r"""
    Instantiate a [`ControlNetModel`] from [`UNet2DConditionModel`].

    Parameters:
        unet (`UNet2DConditionModel`):
            The UNet model weights to copy to the [`ControlNetModel`]. All configuration options are also copied
            where applicable.
    """
    transformer_layers_per_block = (
        unet.config.transformer_layers_per_block if "transformer_layers_per_block" in unet.config else 1
    )
    encoder_hid_dim = unet.config.encoder_hid_dim if "encoder_hid_dim" in unet.config else None
    encoder_hid_dim_type = unet.config.encoder_hid_dim_type if "encoder_hid_dim_type" in unet.config else None
    addition_embed_type = unet.config.addition_embed_type if "addition_embed_type" in unet.config else None
    addition_time_embed_dim = (
        unet.config.addition_time_embed_dim if "addition_time_embed_dim" in unet.config else None
    )

    controlnet = cls(
        encoder_hid_dim=encoder_hid_dim,
        encoder_hid_dim_type=encoder_hid_dim_type,
        addition_embed_type=addition_embed_type,
        addition_time_embed_dim=addition_time_embed_dim,
        transformer_layers_per_block=transformer_layers_per_block,
        in_channels=unet.config.in_channels,
        flip_sin_to_cos=unet.config.flip_sin_to_cos,
        freq_shift=unet.config.freq_shift,
        down_block_types=unet.config.down_block_types,
        only_cross_attention=unet.config.only_cross_attention,
        block_out_channels=unet.config.block_out_channels,
        layers_per_block=unet.config.layers_per_block,
        downsample_padding=unet.config.downsample_padding,
        mid_block_scale_factor=unet.config.mid_block_scale_factor,
        act_fn=unet.config.act_fn,
        norm_num_groups=unet.config.norm_num_groups,
        norm_eps=unet.config.norm_eps,
        cross_attention_dim=unet.config.cross_attention_dim,
        attention_head_dim=unet.config.attention_head_dim,
        num_attention_heads=unet.config.num_attention_heads,
        use_linear_projection=unet.config.use_linear_projection,
        class_embed_type=unet.config.class_embed_type,
        num_class_embeds=unet.config.num_class_embeds,
        upcast_attention=unet.config.upcast_attention,
        resnet_time_scale_shift=unet.config.resnet_time_scale_shift,
        projection_class_embeddings_input_dim=unet.config.projection_class_embeddings_input_dim,
        mid_block_type=unet.config.mid_block_type,
        controlnet_conditioning_channel_order=controlnet_conditioning_channel_order,
        conditioning_embedding_out_channels=conditioning_embedding_out_channels,
        conditioning_channels=conditioning_channels,
    )

    if load_weights_from_unet:
        ms.load_param_into_net(controlnet.conv_in, unet.conv_in.parameters_dict(), strict_load=True)
        ms.load_param_into_net(controlnet.time_proj, unet.time_proj.parameters_dict(), strict_load=True)
        ms.load_param_into_net(controlnet.time_embedding, unet.time_embedding.parameters_dict(), strict_load=True)

        if controlnet.class_embedding:
            ms.load_param_into_net(
                controlnet.class_embedding, unet.class_embedding.parameters_dict(), strict_load=True
            )

        if hasattr(controlnet, "add_embedding"):
            ms.load_param_into_net(controlnet.add_embedding, unet.add_embedding.parameters_dict())

        ms.load_param_into_net(controlnet.down_blocks, unet.down_blocks.parameters_dict(), strict_load=True)
        ms.load_param_into_net(controlnet.mid_block, unet.mid_block.parameters_dict(), strict_load=True)

    return controlnet

mindone.diffusers.ControlNetModel.set_attn_processor(processor)

Sets the attention processor to use to compute attention. Parameters: processor (dict of AttentionProcessor or only AttentionProcessor): The instantiated processor class or a dictionary of processor classes that will be set as the processor for all Attention layers. If processor is a dict, the key needs to define the path to the corresponding cross attention processor. This is strongly recommended when setting trainable attention processors.

Source code in mindone/diffusers/models/controlnet.py

def set_attn_processor(self, processor: Union[AttentionProcessor, Dict[str, AttentionProcessor]]):
    r"""
    Sets the attention processor to use to compute attention.
    Parameters:
        processor (`dict` of `AttentionProcessor` or only `AttentionProcessor`):
            The instantiated processor class or a dictionary of processor classes that will be set as the processor
            for **all** `Attention` layers.
            If `processor` is a dict, the key needs to define the path to the corresponding cross attention
            processor. This is strongly recommended when setting trainable attention processors.
    """
    count = len(self.attn_processors.keys())

    if isinstance(processor, dict) and len(processor) != count:
        raise ValueError(
            f"A dict of processors was passed, but the number of processors {len(processor)} does not match the"
            f" number of attention layers: {count}. Please make sure to pass {count} processor classes."
        )

    def fn_recursive_attn_processor(name: str, module: nn.Cell, processor):
        if hasattr(module, "set_processor"):
            if not isinstance(processor, dict):
                module.set_processor(processor)
            else:
                module.set_processor(processor.pop(f"{name}.processor"))

        for sub_name, child in module.name_cells().items():
            fn_recursive_attn_processor(f"{name}.{sub_name}", child, processor)

    for name, module in self.name_cells().items():
        fn_recursive_attn_processor(name, module, processor)

mindone.diffusers.ControlNetModel.set_default_attn_processor()

Disables custom attention processors and sets the default attention implementation.

Source code in mindone/diffusers/models/controlnet.py

def set_default_attn_processor(self):
    """
    Disables custom attention processors and sets the default attention implementation.
    """
    if all(proc.__class__ in CROSS_ATTENTION_PROCESSORS for proc in self.attn_processors.values()):
        processor = AttnProcessor()
    else:
        raise ValueError(
            f"Cannot call `set_default_attn_processor` when attention processors are of type {next(iter(self.attn_processors.values()))}"
        )

    self.set_attn_processor(processor)

mindone.diffusers.models.controlnet.ControlNetOutput dataclass

Bases: BaseOutput

The output of [ControlNetModel].

PARAMETER	DESCRIPTION
down_block_res_samples	A tuple of downsample activations at different resolutions for each downsampling block. Each tensor should be of shape (batch_size, channel * resolution, height //resolution, width // resolution). Output can be used to condition the original UNet's downsampling activations. TYPE: `tuple[ms.Tensor]`
mid_down_block_re_sample	The activation of the middle block (the lowest sample resolution). Each tensor should be of shape (batch_size, channel * lowest_resolution, height // lowest_resolution, width // lowest_resolution). Output can be used to condition the original UNet's middle block activation. TYPE: `ms.Tensor`

Source code in mindone/diffusers/models/controlnet.py

@dataclass
class ControlNetOutput(BaseOutput):
    """
    The output of [`ControlNetModel`].

    Args:
        down_block_res_samples (`tuple[ms.Tensor]`):
            A tuple of downsample activations at different resolutions for each downsampling block. Each tensor should
            be of shape `(batch_size, channel * resolution, height //resolution, width // resolution)`. Output can be
            used to condition the original UNet's downsampling activations.
        mid_down_block_re_sample (`ms.Tensor`):
            The activation of the middle block (the lowest sample resolution). Each tensor should be of shape
            `(batch_size, channel * lowest_resolution, height // lowest_resolution, width // lowest_resolution)`.
            Output can be used to condition the original UNet's middle block activation.
    """

    down_block_res_samples: Tuple[ms.Tensor]
    mid_block_res_sample: ms.Tensor