Class PretrainedWhisperEncoder

fireredtts2/codec/whisper.py:333–371

# Pretrained encoder from whisper-large-v3
class PretrainedWhisperEncoder(WhisperEncoder):
    @classmethod
    def from_pretrained(cls, pretrained_path: str = None):
        encoder = cls(
            in_dim=128,
            embed_dim=1280,
            num_layers=32,
            num_heads=20,
            ffn_dim=5120,
            attn_dropout=0.0,
            max_positions=1500,
        )
        if pretrained_path is not None:
            ckpt = torch.load(pretrained_path, map_location="cpu")
            encoder.load_state_dict(ckpt)
        encoder.eval()
        # Disable grad
        for p in encoder.parameters():
            p.requires_grad_(False)
        # Add Mel extractor
        encoder.feature_extractor = WhisperMelExtractor(
            num_mels=128,
            sampling_rate=16000,
            hop_length=160,
            n_fft=400,
            fmin=0,
            fmax=8000,
        )
        return encoder

    @torch.inference_mode()
    def forward(self, audio16k: torch.Tensor, audio16k_length: torch.Tensor):
        # Extract mel
        mel, mel_length = self.feature_extractor(audio16k, audio16k_length)
        # Forward model
        semantic_feats, semantic_length = super().forward(
            mel, mel_length, apply_position=True
        )
        return semantic_feats, semantic_length


class WhisperAcousticEncoder(WhisperEncoder):