    def extract_content_features(self, wav_input):
        """
        Extract content features and quantize using KMeans clustering.

        Args:
            audio_data: tensor (batch_size, T)

        Returns:
            quantize: tensor (batch_size, T, 768)
        """
        # Extract features using HuBERT
        wav_input = F.pad(wav_input, (40, 40), "reflect")
        embed = self.model(
            wav_input,
            features_only=True,
            mask=False,
            output_layer=9,  # adjust based on the hubert model used
        )["x"]

        # make the hop size chaged from 200 to 320
        embed = embed.permute((0, 2, 1))
        embed = F.interpolate(embed, scale_factor=1.6, mode="nearest")
        embed = embed.permute((0, 2, 1))

        batched_cluster_centers = repeat(
            self.cluster_centers, "c d -> b c d", b=embed.shape[0]
        )

        dists = -torch.cdist(embed, batched_cluster_centers, p=2)
        clusters = dists.argmax(dim=-1)  # (batch, seq_len)
        quantize = F.embedding(clusters, self.cluster_centers)

        return clusters, quantize
Extract content features and quantize using KMeans clustering.

Args:
audio_data: tensor (batch_size, T)

Returns:
quantize: tensor (batch_size, T, 768)
processors/content_extractor.py