Method encode

acestep/models/ace_step_transformer.py:376–412

(
        self,
        encoder_text_hidden_states: Optional[torch.Tensor] = None,
        text_attention_mask: Optional[torch.LongTensor] = None,
        speaker_embeds: Optional[torch.FloatTensor] = None,
        lyric_token_idx: Optional[torch.LongTensor] = None,
        lyric_mask: Optional[torch.LongTensor] = None,
    )

        return prompt_prenet_out

    def encode(
        self,
        encoder_text_hidden_states: Optional[torch.Tensor] = None,
        text_attention_mask: Optional[torch.LongTensor] = None,
        speaker_embeds: Optional[torch.FloatTensor] = None,
        lyric_token_idx: Optional[torch.LongTensor] = None,
        lyric_mask: Optional[torch.LongTensor] = None,
    ):

        bs = encoder_text_hidden_states.shape[0]
        device = encoder_text_hidden_states.device

        # speaker embedding
        encoder_spk_hidden_states = self.speaker_embedder(speaker_embeds).unsqueeze(1)
        speaker_mask = torch.ones(bs, 1, device=device)

        # genre embedding
        encoder_text_hidden_states = self.genre_embedder(encoder_text_hidden_states)

        # lyric
        encoder_lyric_hidden_states = self.forward_lyric_encoder(
            lyric_token_idx=lyric_token_idx,
            lyric_mask=lyric_mask,
        )

        encoder_hidden_states = torch.cat(
            [
                encoder_spk_hidden_states,
                encoder_text_hidden_states,
                encoder_lyric_hidden_states,
            ],
            dim=1,
        )
        encoder_hidden_mask = torch.cat(
            [speaker_mask, text_attention_mask, lyric_mask], dim=1
        )
        return encoder_hidden_states, encoder_hidden_mask

    def decode(
        self,