class Kandinsky3UNet(ModelMixin, ConfigMixin):
    @register_to_config
    def __init__(
        self,
        in_channels: int = 4,
        time_embedding_dim: int = 1536,
        groups: int = 32,
        attention_head_dim: int = 64,
        layers_per_block: Union[int, Tuple[int]] = 3,
        block_out_channels: Tuple[int] = (384, 768, 1536, 3072),
        cross_attention_dim: Union[int, Tuple[int]] = 4096,
        encoder_hid_dim: int = 4096,
    ):
        super().__init__()

        # TODO(Yiyi): Give better name and put into config for the following 4 parameters
        expansion_ratio = 4
        compression_ratio = 2
        add_cross_attention = (False, True, True, True)
        add_self_attention = (False, True, True, True)

        out_channels = in_channels
        init_channels = block_out_channels[0] // 2
        self.time_proj = Timesteps(init_channels, flip_sin_to_cos=False, downscale_freq_shift=1)

        self.time_embedding = TimestepEmbedding(
            init_channels,
            time_embedding_dim,
        )

        self.add_time_condition = Kandinsky3AttentionPooling(
            time_embedding_dim, cross_attention_dim, attention_head_dim
        )

        self.conv_in = nn.Conv2d(in_channels, init_channels, kernel_size=3, padding=1)

        self.encoder_hid_proj = Kandinsky3EncoderProj(encoder_hid_dim, cross_attention_dim)

        hidden_dims = [init_channels] + list(block_out_channels)
        in_out_dims = list(zip(hidden_dims[:-1], hidden_dims[1:]))
        text_dims = [cross_attention_dim if is_exist else None for is_exist in add_cross_attention]
        num_blocks = len(block_out_channels) * [layers_per_block]
        layer_params = [num_blocks, text_dims, add_self_attention]
        rev_layer_params = map(reversed, layer_params)

        cat_dims = []
        self.num_levels = len(in_out_dims)
        self.down_blocks = nn.ModuleList([])
        for level, ((in_dim, out_dim), res_block_num, text_dim, self_attention) in enumerate(
            zip(in_out_dims, *layer_params)
        ):
            down_sample = level != (self.num_levels - 1)
            cat_dims.append(out_dim if level != (self.num_levels - 1) else 0)
            self.down_blocks.append(
                Kandinsky3DownSampleBlock(
                    in_dim,
                    out_dim,
                    time_embedding_dim,
                    text_dim,
                    res_block_num,
                    groups,
                    attention_head_dim,
                    expansion_ratio,
                    compression_ratio,
                    down_sample,
                    self_attention,
                )
            )

        self.up_blocks = nn.ModuleList([])
        for level, ((out_dim, in_dim), res_block_num, text_dim, self_attention) in enumerate(
            zip(reversed(in_out_dims), *rev_layer_params)
        ):
            up_sample = level != 0
            self.up_blocks.append(
                Kandinsky3UpSampleBlock(
                    in_dim,
                    cat_dims.pop(),
                    out_dim,
                    time_embedding_dim,
                    text_dim,
                    res_block_num,
                    groups,
                    attention_head_dim,
                    expansion_ratio,
                    compression_ratio,
                    up_sample,
                    self_attention,
                )
            )

        self.conv_norm_out = nn.GroupNorm(groups, init_channels)
        self.conv_act_out = nn.SiLU()
        self.conv_out = nn.Conv2d(init_channels, out_channels, kernel_size=3, padding=1)

    @property
    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, "set_processor"):
                processors[f"{name}.processor"] = module.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

    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)

    def set_default_attn_processor(self):
        """
        Disables custom attention processors and sets the default attention implementation.
        """
        self.set_attn_processor(AttnProcessor())

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

    def forward(self, sample, timestep, encoder_hidden_states=None, encoder_attention_mask=None, return_dict=True):
        if encoder_attention_mask is not None:
            encoder_attention_mask = (1 - encoder_attention_mask.to(sample.dtype)) * -10000.0
            encoder_attention_mask = encoder_attention_mask.unsqueeze(1)

        if not torch.is_tensor(timestep):
            dtype = torch.float32 if isinstance(timestep, float) else torch.int32
            timestep = torch.tensor([timestep], dtype=dtype, device=sample.device)
        elif len(timestep.shape) == 0:
            timestep = timestep[None].to(sample.device)

        # broadcast to batch dimension in a way that's compatible with ONNX/Core ML
        timestep = timestep.expand(sample.shape[0])
        time_embed_input = self.time_proj(timestep).to(sample.dtype)
        time_embed = self.time_embedding(time_embed_input)

        encoder_hidden_states = self.encoder_hid_proj(encoder_hidden_states)

        if encoder_hidden_states is not None:
            time_embed = self.add_time_condition(time_embed, encoder_hidden_states, encoder_attention_mask)

        hidden_states = []
        sample = self.conv_in(sample)
        for level, down_sample in enumerate(self.down_blocks):
            sample = down_sample(sample, time_embed, encoder_hidden_states, encoder_attention_mask)
            if level != self.num_levels - 1:
                hidden_states.append(sample)

        for level, up_sample in enumerate(self.up_blocks):
            if level != 0:
                sample = torch.cat([sample, hidden_states.pop()], dim=1)
            sample = up_sample(sample, time_embed, encoder_hidden_states, encoder_attention_mask)

        sample = self.conv_norm_out(sample)
        sample = self.conv_act_out(sample)
        sample = self.conv_out(sample)

        if not return_dict:
            return (sample,)
        return Kandinsky3UNetOutput(sample=sample)
/Users/nielsrogge/Documents/python_projecten/diffusers/src/diffusers/models/unets/unet_kandinsky3.py