class AudioLDM2ProjectionModel(ModelMixin, ConfigMixin):
    """
    A simple linear projection model to map two text embeddings to a shared latent space. It also inserts learned
    embedding vectors at the start and end of each text embedding sequence respectively. Each variable appended with
    `_1` refers to that corresponding to the second text encoder. Otherwise, it is from the first.

    Args:
        text_encoder_dim (`int`):
            Dimensionality of the text embeddings from the first text encoder (CLAP).
        text_encoder_1_dim (`int`):
            Dimensionality of the text embeddings from the second text encoder (T5 or VITS).
        langauge_model_dim (`int`):
            Dimensionality of the text embeddings from the language model (GPT2).
    """

    @register_to_config
    def __init__(
        self,
        text_encoder_dim,
        text_encoder_1_dim,
        langauge_model_dim,
        use_learned_position_embedding=None,
        max_seq_length=None,
    ):
        super().__init__()
        # additional projection layers for each text encoder
        self.projection = nn.Linear(text_encoder_dim, langauge_model_dim)
        self.projection_1 = nn.Linear(text_encoder_1_dim, langauge_model_dim)

        # learnable SOS / EOS token embeddings for each text encoder
        self.sos_embed = nn.Parameter(torch.ones(langauge_model_dim))
        self.eos_embed = nn.Parameter(torch.ones(langauge_model_dim))

        self.sos_embed_1 = nn.Parameter(torch.ones(langauge_model_dim))
        self.eos_embed_1 = nn.Parameter(torch.ones(langauge_model_dim))

        self.use_learned_position_embedding = use_learned_position_embedding

        # learable positional embedding for vits encoder
        if self.use_learned_position_embedding is not None:
            self.learnable_positional_embedding = torch.nn.Parameter(
                torch.zeros((1, text_encoder_1_dim, max_seq_length))
            )

    def forward(
        self,
        hidden_states: Optional[torch.Tensor] = None,
        hidden_states_1: Optional[torch.Tensor] = None,
        attention_mask: Optional[torch.LongTensor] = None,
        attention_mask_1: Optional[torch.LongTensor] = None,
    ):
        hidden_states = self.projection(hidden_states)
        hidden_states, attention_mask = add_special_tokens(
            hidden_states, attention_mask, sos_token=self.sos_embed, eos_token=self.eos_embed
        )

        # Add positional embedding for Vits hidden state
        if self.use_learned_position_embedding is not None:
            hidden_states_1 = (hidden_states_1.permute(0, 2, 1) + self.learnable_positional_embedding).permute(0, 2, 1)

        hidden_states_1 = self.projection_1(hidden_states_1)
        hidden_states_1, attention_mask_1 = add_special_tokens(
            hidden_states_1, attention_mask_1, sos_token=self.sos_embed_1, eos_token=self.eos_embed_1
        )

        # concatenate clap and t5 text encoding
        hidden_states = torch.cat([hidden_states, hidden_states_1], dim=1)

        # concatenate attention masks
        if attention_mask is None and attention_mask_1 is not None:
            attention_mask = attention_mask_1.new_ones((hidden_states[:2]))
        elif attention_mask is not None and attention_mask_1 is None:
            attention_mask_1 = attention_mask.new_ones((hidden_states_1[:2]))

        if attention_mask is not None and attention_mask_1 is not None:
            attention_mask = torch.cat([attention_mask, attention_mask_1], dim=-1)
        else:
            attention_mask = None

        return AudioLDM2ProjectionModelOutput(
            hidden_states=hidden_states,
            attention_mask=attention_mask,
        )
/Users/nielsrogge/Documents/python_projecten/diffusers/src/diffusers/pipelines/audioldm2/modeling_audioldm2.py