class UNetControlNetXSModel(ModelMixin, ConfigMixin):
    r"""
    A UNet fused with a ControlNet-XS adapter model

    This model inherits from [`ModelMixin`] and [`ConfigMixin`]. Check the superclass documentation for it's generic
    methods implemented for all models (such as downloading or saving).

    `UNetControlNetXSModel` is compatible with StableDiffusion and StableDiffusion-XL. It's default parameters are
    compatible with StableDiffusion.

    It's parameters are either passed to the underlying `UNet2DConditionModel` or used exactly like in
    `ControlNetXSAdapter` . See their documentation for details.
    """

    _supports_gradient_checkpointing = True

    @register_to_config
    def __init__(
        self,
        # unet configs
        sample_size: Optional[int] = 96,
        down_block_types: Tuple[str] = (
            "CrossAttnDownBlock2D",
            "CrossAttnDownBlock2D",
            "CrossAttnDownBlock2D",
            "DownBlock2D",
        ),
        up_block_types: Tuple[str] = ("UpBlock2D", "CrossAttnUpBlock2D", "CrossAttnUpBlock2D", "CrossAttnUpBlock2D"),
        block_out_channels: Tuple[int] = (320, 640, 1280, 1280),
        norm_num_groups: Optional[int] = 32,
        cross_attention_dim: Union[int, Tuple[int]] = 1024,
        transformer_layers_per_block: Union[int, Tuple[int]] = 1,
        num_attention_heads: Union[int, Tuple[int]] = 8,
        addition_embed_type: Optional[str] = None,
        addition_time_embed_dim: Optional[int] = None,
        upcast_attention: bool = True,
        use_linear_projection: bool = True,
        time_cond_proj_dim: Optional[int] = None,
        projection_class_embeddings_input_dim: Optional[int] = None,
        # additional controlnet configs
        time_embedding_mix: float = 1.0,
        ctrl_conditioning_channels: int = 3,
        ctrl_conditioning_embedding_out_channels: Tuple[int] = (16, 32, 96, 256),
        ctrl_conditioning_channel_order: str = "rgb",
        ctrl_learn_time_embedding: bool = False,
        ctrl_block_out_channels: Tuple[int] = (4, 8, 16, 16),
        ctrl_num_attention_heads: Union[int, Tuple[int]] = 4,
        ctrl_max_norm_num_groups: int = 32,
    ):
        super().__init__()

        if time_embedding_mix < 0 or time_embedding_mix > 1:
            raise ValueError("`time_embedding_mix` needs to be between 0 and 1.")
        if time_embedding_mix < 1 and not ctrl_learn_time_embedding:
            raise ValueError("To use `time_embedding_mix` < 1, `ctrl_learn_time_embedding` must be `True`")

        if addition_embed_type is not None and addition_embed_type != "text_time":
            raise ValueError(
                "As `UNetControlNetXSModel` currently only supports StableDiffusion and StableDiffusion-XL, `addition_embed_type` must be `None` or `'text_time'`."
            )

        if not isinstance(transformer_layers_per_block, (list, tuple)):
            transformer_layers_per_block = [transformer_layers_per_block] * len(down_block_types)
        if not isinstance(cross_attention_dim, (list, tuple)):
            cross_attention_dim = [cross_attention_dim] * len(down_block_types)
        if not isinstance(num_attention_heads, (list, tuple)):
            num_attention_heads = [num_attention_heads] * len(down_block_types)
        if not isinstance(ctrl_num_attention_heads, (list, tuple)):
            ctrl_num_attention_heads = [ctrl_num_attention_heads] * len(down_block_types)

        base_num_attention_heads = num_attention_heads

        self.in_channels = 4

        # # Input
        self.base_conv_in = nn.Conv2d(4, block_out_channels[0], kernel_size=3, padding=1)
        self.controlnet_cond_embedding = ControlNetConditioningEmbedding(
            conditioning_embedding_channels=ctrl_block_out_channels[0],
            block_out_channels=ctrl_conditioning_embedding_out_channels,
            conditioning_channels=ctrl_conditioning_channels,
        )
        self.ctrl_conv_in = nn.Conv2d(4, ctrl_block_out_channels[0], kernel_size=3, padding=1)
        self.control_to_base_for_conv_in = make_zero_conv(ctrl_block_out_channels[0], block_out_channels[0])

        # # Time
        time_embed_input_dim = block_out_channels[0]
        time_embed_dim = block_out_channels[0] * 4

        self.base_time_proj = Timesteps(block_out_channels[0], flip_sin_to_cos=True, downscale_freq_shift=0)
        self.base_time_embedding = TimestepEmbedding(
            time_embed_input_dim,
            time_embed_dim,
            cond_proj_dim=time_cond_proj_dim,
        )
        if ctrl_learn_time_embedding:
            self.ctrl_time_embedding = TimestepEmbedding(
                in_channels=time_embed_input_dim, time_embed_dim=time_embed_dim
            )
        else:
            self.ctrl_time_embedding = None

        if addition_embed_type is None:
            self.base_add_time_proj = None
            self.base_add_embedding = None
        else:
            self.base_add_time_proj = Timesteps(addition_time_embed_dim, flip_sin_to_cos=True, downscale_freq_shift=0)
            self.base_add_embedding = TimestepEmbedding(projection_class_embeddings_input_dim, time_embed_dim)

        # # Create down blocks
        down_blocks = []
        base_out_channels = block_out_channels[0]
        ctrl_out_channels = ctrl_block_out_channels[0]
        for i, down_block_type in enumerate(down_block_types):
            base_in_channels = base_out_channels
            base_out_channels = block_out_channels[i]
            ctrl_in_channels = ctrl_out_channels
            ctrl_out_channels = ctrl_block_out_channels[i]
            has_crossattn = "CrossAttn" in down_block_type
            is_final_block = i == len(down_block_types) - 1

            down_blocks.append(
                ControlNetXSCrossAttnDownBlock2D(
                    base_in_channels=base_in_channels,
                    base_out_channels=base_out_channels,
                    ctrl_in_channels=ctrl_in_channels,
                    ctrl_out_channels=ctrl_out_channels,
                    temb_channels=time_embed_dim,
                    norm_num_groups=norm_num_groups,
                    ctrl_max_norm_num_groups=ctrl_max_norm_num_groups,
                    has_crossattn=has_crossattn,
                    transformer_layers_per_block=transformer_layers_per_block[i],
                    base_num_attention_heads=base_num_attention_heads[i],
                    ctrl_num_attention_heads=ctrl_num_attention_heads[i],
                    cross_attention_dim=cross_attention_dim[i],
                    add_downsample=not is_final_block,
                    upcast_attention=upcast_attention,
                    use_linear_projection=use_linear_projection,
                )
            )

        # # Create mid block
        self.mid_block = ControlNetXSCrossAttnMidBlock2D(
            base_channels=block_out_channels[-1],
            ctrl_channels=ctrl_block_out_channels[-1],
            temb_channels=time_embed_dim,
            norm_num_groups=norm_num_groups,
            ctrl_max_norm_num_groups=ctrl_max_norm_num_groups,
            transformer_layers_per_block=transformer_layers_per_block[-1],
            base_num_attention_heads=base_num_attention_heads[-1],
            ctrl_num_attention_heads=ctrl_num_attention_heads[-1],
            cross_attention_dim=cross_attention_dim[-1],
            upcast_attention=upcast_attention,
            use_linear_projection=use_linear_projection,
        )

        # # Create up blocks
        up_blocks = []
        rev_transformer_layers_per_block = list(reversed(transformer_layers_per_block))
        rev_num_attention_heads = list(reversed(base_num_attention_heads))
        rev_cross_attention_dim = list(reversed(cross_attention_dim))

        # The skip connection channels are the output of the conv_in and of all the down subblocks
        ctrl_skip_channels = [ctrl_block_out_channels[0]]
        for i, out_channels in enumerate(ctrl_block_out_channels):
            number_of_subblocks = (
                3 if i < len(ctrl_block_out_channels) - 1 else 2
            )  # every block has 3 subblocks, except last one, which has 2 as it has no downsampler
            ctrl_skip_channels.extend([out_channels] * number_of_subblocks)

        reversed_block_out_channels = list(reversed(block_out_channels))

        out_channels = reversed_block_out_channels[0]
        for i, up_block_type in enumerate(up_block_types):
            prev_output_channel = out_channels
            out_channels = reversed_block_out_channels[i]
            in_channels = reversed_block_out_channels[min(i + 1, len(block_out_channels) - 1)]
            ctrl_skip_channels_ = [ctrl_skip_channels.pop() for _ in range(3)]

            has_crossattn = "CrossAttn" in up_block_type
            is_final_block = i == len(block_out_channels) - 1

            up_blocks.append(
                ControlNetXSCrossAttnUpBlock2D(
                    in_channels=in_channels,
                    out_channels=out_channels,
                    prev_output_channel=prev_output_channel,
                    ctrl_skip_channels=ctrl_skip_channels_,
                    temb_channels=time_embed_dim,
                    resolution_idx=i,
                    has_crossattn=has_crossattn,
                    transformer_layers_per_block=rev_transformer_layers_per_block[i],
                    num_attention_heads=rev_num_attention_heads[i],
                    cross_attention_dim=rev_cross_attention_dim[i],
                    add_upsample=not is_final_block,
                    upcast_attention=upcast_attention,
                    norm_num_groups=norm_num_groups,
                    use_linear_projection=use_linear_projection,
                )
            )

        self.down_blocks = nn.ModuleList(down_blocks)
        self.up_blocks = nn.ModuleList(up_blocks)

        self.base_conv_norm_out = nn.GroupNorm(num_channels=block_out_channels[0], num_groups=norm_num_groups)
        self.base_conv_act = nn.SiLU()
        self.base_conv_out = nn.Conv2d(block_out_channels[0], 4, kernel_size=3, padding=1)

    @classmethod
    def from_unet(
        cls,
        unet: UNet2DConditionModel,
        controlnet: Optional[ControlNetXSAdapter] = None,
        size_ratio: Optional[float] = None,
        ctrl_block_out_channels: Optional[List[float]] = None,
        time_embedding_mix: Optional[float] = None,
        ctrl_optional_kwargs: Optional[Dict] = None,
    ):
        r"""
        Instantiate a [`UNetControlNetXSModel`] from a [`UNet2DConditionModel`] and an optional [`ControlNetXSAdapter`]
        .

        Parameters:
            unet (`UNet2DConditionModel`):
                The UNet model we want to control.
            controlnet (`ControlNetXSAdapter`):
                The ConntrolNet-XS adapter with which the UNet will be fused. If none is given, a new ConntrolNet-XS
                adapter will be created.
            size_ratio (float, *optional*, defaults to `None`):
                Used to contruct the controlnet if none is given. See [`ControlNetXSAdapter.from_unet`] for details.
            ctrl_block_out_channels (`List[int]`, *optional*, defaults to `None`):
                Used to contruct the controlnet if none is given. See [`ControlNetXSAdapter.from_unet`] for details,
                where this parameter is called `block_out_channels`.
            time_embedding_mix (`float`, *optional*, defaults to None):
                Used to contruct the controlnet if none is given. See [`ControlNetXSAdapter.from_unet`] for details.
            ctrl_optional_kwargs (`Dict`, *optional*, defaults to `None`):
                Passed to the `init` of the new controlent if no controlent was given.
        """
        if controlnet is None:
            controlnet = ControlNetXSAdapter.from_unet(
                unet, size_ratio, ctrl_block_out_channels, **ctrl_optional_kwargs
            )
        else:
            if any(
                o is not None for o in (size_ratio, ctrl_block_out_channels, time_embedding_mix, ctrl_optional_kwargs)
            ):
                raise ValueError(
                    "When a controlnet is passed, none of these parameters should be passed: size_ratio, ctrl_block_out_channels, time_embedding_mix, ctrl_optional_kwargs."
                )

        # # get params
        params_for_unet = [
            "sample_size",
            "down_block_types",
            "up_block_types",
            "block_out_channels",
            "norm_num_groups",
            "cross_attention_dim",
            "transformer_layers_per_block",
            "addition_embed_type",
            "addition_time_embed_dim",
            "upcast_attention",
            "use_linear_projection",
            "time_cond_proj_dim",
            "projection_class_embeddings_input_dim",
        ]
        params_for_unet = {k: v for k, v in unet.config.items() if k in params_for_unet}
        # The naming seems a bit confusing and it is, see https://github.com/huggingface/diffusers/issues/2011#issuecomment-1547958131 for why.
        params_for_unet["num_attention_heads"] = unet.config.attention_head_dim

        params_for_controlnet = [
            "conditioning_channels",
            "conditioning_embedding_out_channels",
            "conditioning_channel_order",
            "learn_time_embedding",
            "block_out_channels",
            "num_attention_heads",
            "max_norm_num_groups",
        ]
        params_for_controlnet = {"ctrl_" + k: v for k, v in controlnet.config.items() if k in params_for_controlnet}
        params_for_controlnet["time_embedding_mix"] = controlnet.config.time_embedding_mix

        # # create model
        model = cls.from_config({**params_for_unet, **params_for_controlnet})

        # # load weights
        # from unet
        modules_from_unet = [
            "time_embedding",
            "conv_in",
            "conv_norm_out",
            "conv_out",
        ]
        for m in modules_from_unet:
            getattr(model, "base_" + m).load_state_dict(getattr(unet, m).state_dict())

        optional_modules_from_unet = [
            "add_time_proj",
            "add_embedding",
        ]
        for m in optional_modules_from_unet:
            if hasattr(unet, m) and getattr(unet, m) is not None:
                getattr(model, "base_" + m).load_state_dict(getattr(unet, m).state_dict())

        # from controlnet
        model.controlnet_cond_embedding.load_state_dict(controlnet.controlnet_cond_embedding.state_dict())
        model.ctrl_conv_in.load_state_dict(controlnet.conv_in.state_dict())
        if controlnet.time_embedding is not None:
            model.ctrl_time_embedding.load_state_dict(controlnet.time_embedding.state_dict())
        model.control_to_base_for_conv_in.load_state_dict(controlnet.control_to_base_for_conv_in.state_dict())

        # from both
        model.down_blocks = nn.ModuleList(
            ControlNetXSCrossAttnDownBlock2D.from_modules(b, c)
            for b, c in zip(unet.down_blocks, controlnet.down_blocks)
        )
        model.mid_block = ControlNetXSCrossAttnMidBlock2D.from_modules(unet.mid_block, controlnet.mid_block)
        model.up_blocks = nn.ModuleList(
            ControlNetXSCrossAttnUpBlock2D.from_modules(b, c)
            for b, c in zip(unet.up_blocks, controlnet.up_connections)
        )

        # ensure that the UNetControlNetXSModel is the same dtype as the UNet2DConditionModel
        model.to(unet.dtype)

        return model

    def freeze_unet_params(self) -> None:
        """Freeze the weights of the parts belonging to the base UNet2DConditionModel, and leave everything else unfrozen for fine
        tuning."""
        # Freeze everything
        for param in self.parameters():
            param.requires_grad = True

        # Unfreeze ControlNetXSAdapter
        base_parts = [
            "base_time_proj",
            "base_time_embedding",
            "base_add_time_proj",
            "base_add_embedding",
            "base_conv_in",
            "base_conv_norm_out",
            "base_conv_act",
            "base_conv_out",
        ]
        base_parts = [getattr(self, part) for part in base_parts if getattr(self, part) is not None]
        for part in base_parts:
            for param in part.parameters():
                param.requires_grad = False

        for d in self.down_blocks:
            d.freeze_base_params()
        self.mid_block.freeze_base_params()
        for u in self.up_blocks:
            u.freeze_base_params()

    def _set_gradient_checkpointing(self, module, value=False):
        if hasattr(module, "gradient_checkpointing"):
            module.gradient_checkpointing = value

    @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: torch.nn.Module, processors: Dict[str, AttentionProcessor]):
            if hasattr(module, "get_processor"):
                processors[f"{name}.processor"] = module.get_processor()

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

            return processors

        for name, module in self.named_children():
            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: torch.nn.Module, 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.named_children():
                fn_recursive_attn_processor(f"{name}.{sub_name}", child, processor)

        for name, module in self.named_children():
            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 ADDED_KV_ATTENTION_PROCESSORS for proc in self.attn_processors.values()):
            processor = AttnAddedKVProcessor()
        elif 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)

    # Copied from diffusers.models.unets.unet_2d_condition.UNet2DConditionModel.enable_freeu
    def enable_freeu(self, s1: float, s2: float, b1: float, b2: float):
        r"""Enables the FreeU mechanism from https://arxiv.org/abs/2309.11497.

        The suffixes after the scaling factors represent the stage blocks where they are being applied.

        Please refer to the [official repository](https://github.com/ChenyangSi/FreeU) for combinations of values that
        are known to work well for different pipelines such as Stable Diffusion v1, v2, and Stable Diffusion XL.

        Args:
            s1 (`float`):
                Scaling factor for stage 1 to attenuate the contributions of the skip features. This is done to
                mitigate the "oversmoothing effect" in the enhanced denoising process.
            s2 (`float`):
                Scaling factor for stage 2 to attenuate the contributions of the skip features. This is done to
                mitigate the "oversmoothing effect" in the enhanced denoising process.
            b1 (`float`): Scaling factor for stage 1 to amplify the contributions of backbone features.
            b2 (`float`): Scaling factor for stage 2 to amplify the contributions of backbone features.
        """
        for i, upsample_block in enumerate(self.up_blocks):
            setattr(upsample_block, "s1", s1)
            setattr(upsample_block, "s2", s2)
            setattr(upsample_block, "b1", b1)
            setattr(upsample_block, "b2", b2)

    # Copied from diffusers.models.unets.unet_2d_condition.UNet2DConditionModel.disable_freeu
    def disable_freeu(self):
        """Disables the FreeU mechanism."""
        freeu_keys = {"s1", "s2", "b1", "b2"}
        for i, upsample_block in enumerate(self.up_blocks):
            for k in freeu_keys:
                if hasattr(upsample_block, k) or getattr(upsample_block, k, None) is not None:
                    setattr(upsample_block, k, None)

    # Copied from diffusers.models.unets.unet_2d_condition.UNet2DConditionModel.fuse_qkv_projections
    def fuse_qkv_projections(self):
        """
        Enables fused QKV projections. For self-attention modules, all projection matrices (i.e., query, key, value)
        are fused. For cross-attention modules, key and value projection matrices are fused.

        <Tip warning={true}>

        This API is 🧪 experimental.

        </Tip>
        """
        self.original_attn_processors = None

        for _, attn_processor in self.attn_processors.items():
            if "Added" in str(attn_processor.__class__.__name__):
                raise ValueError("`fuse_qkv_projections()` is not supported for models having added KV projections.")

        self.original_attn_processors = self.attn_processors

        for module in self.modules():
            if isinstance(module, Attention):
                module.fuse_projections(fuse=True)

        self.set_attn_processor(FusedAttnProcessor2_0())

    # Copied from diffusers.models.unets.unet_2d_condition.UNet2DConditionModel.unfuse_qkv_projections
    def unfuse_qkv_projections(self):
        """Disables the fused QKV projection if enabled.

        <Tip warning={true}>

        This API is 🧪 experimental.

        </Tip>

        """
        if self.original_attn_processors is not None:
            self.set_attn_processor(self.original_attn_processors)

    def forward(
        self,
        sample: Tensor,
        timestep: Union[torch.Tensor, float, int],
        encoder_hidden_states: torch.Tensor,
        controlnet_cond: Optional[torch.Tensor] = None,
        conditioning_scale: Optional[float] = 1.0,
        class_labels: Optional[torch.Tensor] = None,
        timestep_cond: Optional[torch.Tensor] = None,
        attention_mask: Optional[torch.Tensor] = None,
        cross_attention_kwargs: Optional[Dict[str, Any]] = None,
        added_cond_kwargs: Optional[Dict[str, torch.Tensor]] = None,
        return_dict: bool = True,
        apply_control: bool = True,
    ) -> Union[ControlNetXSOutput, Tuple]:
        """
        The [`ControlNetXSModel`] forward method.

        Args:
            sample (`Tensor`):
                The noisy input tensor.
            timestep (`Union[torch.Tensor, float, int]`):
                The number of timesteps to denoise an input.
            encoder_hidden_states (`torch.Tensor`):
                The encoder hidden states.
            controlnet_cond (`Tensor`):
                The conditional input tensor of shape `(batch_size, sequence_length, hidden_size)`.
            conditioning_scale (`float`, defaults to `1.0`):
                How much the control model affects the base model outputs.
            class_labels (`torch.Tensor`, *optional*, defaults to `None`):
                Optional class labels for conditioning. Their embeddings will be summed with the timestep embeddings.
            timestep_cond (`torch.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 (`torch.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.
            cross_attention_kwargs (`dict[str]`, *optional*, defaults to `None`):
                A kwargs dictionary that if specified is passed along to the `AttnProcessor`.
            added_cond_kwargs (`dict`):
                Additional conditions for the Stable Diffusion XL UNet.
            return_dict (`bool`, defaults to `True`):
                Whether or not to return a [`~models.controlnets.controlnet.ControlNetOutput`] instead of a plain
                tuple.
            apply_control (`bool`, defaults to `True`):
                If `False`, the input is run only through the base model.

        Returns:
            [`~models.controlnetxs.ControlNetXSOutput`] **or** `tuple`:
                If `return_dict` is `True`, a [`~models.controlnetxs.ControlNetXSOutput`] is returned, otherwise a
                tuple is returned where the first element is the sample tensor.
        """

        # check channel order
        if self.config.ctrl_conditioning_channel_order == "bgr":
            controlnet_cond = torch.flip(controlnet_cond, dims=[1])

        # 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 torch.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+)
            is_mps = sample.device.type == "mps"
            if isinstance(timestep, float):
                dtype = torch.float32 if is_mps else torch.float64
            else:
                dtype = torch.int32 if is_mps else torch.int64
            timesteps = torch.tensor([timesteps], dtype=dtype, device=sample.device)
        elif len(timesteps.shape) == 0:
            timesteps = timesteps[None].to(sample.device)

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

        t_emb = self.base_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)

        if self.config.ctrl_learn_time_embedding and apply_control:
            ctrl_temb = self.ctrl_time_embedding(t_emb, timestep_cond)
            base_temb = self.base_time_embedding(t_emb, timestep_cond)
            interpolation_param = self.config.time_embedding_mix**0.3

            temb = ctrl_temb * interpolation_param + base_temb * (1 - interpolation_param)
        else:
            temb = self.base_time_embedding(t_emb)

        # added time & text embeddings
        aug_emb = None

        if self.config.addition_embed_type is None:
            pass
        elif self.config.addition_embed_type == "text_time":
            # SDXL - style
            if "text_embeds" not in added_cond_kwargs:
                raise ValueError(
                    f"{self.__class__} has the config param `addition_embed_type` set to 'text_time' 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' 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.base_add_time_proj(time_ids.flatten())
            time_embeds = time_embeds.reshape((text_embeds.shape[0], -1))
            add_embeds = torch.concat([text_embeds, time_embeds], dim=-1)
            add_embeds = add_embeds.to(temb.dtype)
            aug_emb = self.base_add_embedding(add_embeds)
        else:
            raise ValueError(
                f"ControlNet-XS currently only supports StableDiffusion and StableDiffusion-XL, so addition_embed_type = {self.config.addition_embed_type} is currently not supported."
            )

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

        # text embeddings
        cemb = encoder_hidden_states

        # Preparation
        h_ctrl = h_base = sample
        hs_base, hs_ctrl = [], []

        # Cross Control
        guided_hint = self.controlnet_cond_embedding(controlnet_cond)

        # 1 - conv in & down

        h_base = self.base_conv_in(h_base)
        h_ctrl = self.ctrl_conv_in(h_ctrl)
        if guided_hint is not None:
            h_ctrl += guided_hint
        if apply_control:
            h_base = h_base + self.control_to_base_for_conv_in(h_ctrl) * conditioning_scale  # add ctrl -> base

        hs_base.append(h_base)
        hs_ctrl.append(h_ctrl)

        for down in self.down_blocks:
            h_base, h_ctrl, residual_hb, residual_hc = down(
                hidden_states_base=h_base,
                hidden_states_ctrl=h_ctrl,
                temb=temb,
                encoder_hidden_states=cemb,
                conditioning_scale=conditioning_scale,
                cross_attention_kwargs=cross_attention_kwargs,
                attention_mask=attention_mask,
                apply_control=apply_control,
            )
            hs_base.extend(residual_hb)
            hs_ctrl.extend(residual_hc)

        # 2 - mid
        h_base, h_ctrl = self.mid_block(
            hidden_states_base=h_base,
            hidden_states_ctrl=h_ctrl,
            temb=temb,
            encoder_hidden_states=cemb,
            conditioning_scale=conditioning_scale,
            cross_attention_kwargs=cross_attention_kwargs,
            attention_mask=attention_mask,
            apply_control=apply_control,
        )

        # 3 - up
        for up in self.up_blocks:
            n_resnets = len(up.resnets)
            skips_hb = hs_base[-n_resnets:]
            skips_hc = hs_ctrl[-n_resnets:]
            hs_base = hs_base[:-n_resnets]
            hs_ctrl = hs_ctrl[:-n_resnets]
            h_base = up(
                hidden_states=h_base,
                res_hidden_states_tuple_base=skips_hb,
                res_hidden_states_tuple_ctrl=skips_hc,
                temb=temb,
                encoder_hidden_states=cemb,
                conditioning_scale=conditioning_scale,
                cross_attention_kwargs=cross_attention_kwargs,
                attention_mask=attention_mask,
                apply_control=apply_control,
            )

        # 4 - conv out
        h_base = self.base_conv_norm_out(h_base)
        h_base = self.base_conv_act(h_base)
        h_base = self.base_conv_out(h_base)

        if not return_dict:
            return (h_base,)

        return ControlNetXSOutput(sample=h_base)
/Users/nielsrogge/Documents/python_projecten/diffusers/src/diffusers/models/controlnets/controlnet_xs.py