class FlaxControlNetConditioningEmbedding(nn.Module):
    conditioning_embedding_channels: int
    block_out_channels: Tuple[int, ...] = (16, 32, 96, 256)
    dtype: jnp.dtype = jnp.float32

    def setup(self) -> None:
        self.conv_in = nn.Conv(
            self.block_out_channels[0],
            kernel_size=(3, 3),
            padding=((1, 1), (1, 1)),
            dtype=self.dtype,
        )

        blocks = []
        for i in range(len(self.block_out_channels) - 1):
            channel_in = self.block_out_channels[i]
            channel_out = self.block_out_channels[i + 1]
            conv1 = nn.Conv(
                channel_in,
                kernel_size=(3, 3),
                padding=((1, 1), (1, 1)),
                dtype=self.dtype,
            )
            blocks.append(conv1)
            conv2 = nn.Conv(
                channel_out,
                kernel_size=(3, 3),
                strides=(2, 2),
                padding=((1, 1), (1, 1)),
                dtype=self.dtype,
            )
            blocks.append(conv2)
        self.blocks = blocks

        self.conv_out = nn.Conv(
            self.conditioning_embedding_channels,
            kernel_size=(3, 3),
            padding=((1, 1), (1, 1)),
            kernel_init=nn.initializers.zeros_init(),
            bias_init=nn.initializers.zeros_init(),
            dtype=self.dtype,
        )

    def __call__(self, conditioning: jnp.ndarray) -> jnp.ndarray:
        embedding = self.conv_in(conditioning)
        embedding = nn.silu(embedding)

        for block in self.blocks:
            embedding = block(embedding)
            embedding = nn.silu(embedding)

        embedding = self.conv_out(embedding)

        return embedding
/Users/nielsrogge/Documents/python_projecten/diffusers/src/diffusers/models/controlnets/controlnet_flax.py