    def forward(self, x, x_mask, cond, diffusion_step, spk_query_emb):
        """
        x: (B, 128, T)
        x_mask: (B, T), mask is 0
        cond: (B, T, 512)
        diffusion_step: (B,)
        spk_query_emb: (B, 32, 512)
        """
        cond = self.cond_ln(cond)
        cond_input = cond.transpose(1, 2)

        x_input = self.in_proj(x)

        x_input = F.relu(x_input)

        diffusion_step = self.diffusion_embedding(diffusion_step).to(x.dtype)
        diffusion_step = self.mlp(diffusion_step)

        skip = []
        for _, layer in enumerate(self.layers):
            x_input, skip_connection = layer(
                x_input, x_mask, cond_input, diffusion_step, spk_query_emb
            )
            skip.append(skip_connection)

        x_input = torch.sum(torch.stack(skip), dim=0) / math.sqrt(self.num_layers)

        x_out = self.skip_proj(x_input)

        x_out = F.relu(x_out)

        x_out = self.out_proj(x_out)  # (B, 128, T)

        return x_out
x: (B, 128, T)
x_mask: (B, T), mask is 0
cond: (B, T, 512)
diffusion_step: (B,)
spk_query_emb: (B, 32, 512)
models/tts/naturalspeech2/wavenet.py