| | |
| | |
| | |
| | |
| | |
| | |
| | |
| |
|
| | import io |
| | import logging |
| | import os |
| | from typing import Optional, Union |
| |
|
| | import soundfile as sf |
| | import torch |
| | from whisper import _MODELS, _download, _ALIGNMENT_HEADS, available_models |
| | from whisper.audio import log_mel_spectrogram |
| | from whisper.model import ModelDimensions |
| |
|
| | from whisper_model import Whisper_ |
| |
|
| | logger = logging.getLogger("dump_feature") |
| |
|
| |
|
| | def load_model( |
| | name: str, |
| | device: Optional[Union[str, torch.device]] = None, |
| | download_root: str = None, |
| | in_memory: bool = False, |
| | ) -> Whisper_: |
| | """ |
| | Reference: https://github.com/openai/whisper/blob/main/whisper/__init__.py#L97 |
| | But we will load a `Whisper_` model for feature extraction. |
| | |
| | Parameters |
| | ---------- |
| | name : str |
| | one of the official model names listed by `whisper.available_models()`, or |
| | path to a model checkpoint containing the model dimensions and the model state_dict. |
| | device : Union[str, torch.device] |
| | the PyTorch device to put the model into |
| | download_root: str |
| | path to download the model files; by default, it uses "~/.cache/whisper" |
| | in_memory: bool |
| | whether to preload the model weights into host memory |
| | |
| | Returns |
| | ------- |
| | model : Whisper |
| | The Whisper ASR model instance |
| | """ |
| |
|
| | if device is None: |
| | device = "cuda" if torch.cuda.is_available() else "cpu" |
| | if download_root is None: |
| | default = os.path.join(os.path.expanduser("~"), ".cache") |
| | download_root = os.path.join(os.getenv("XDG_CACHE_HOME", default), "whisper") |
| |
|
| | if name in _MODELS: |
| | checkpoint_file = _download(_MODELS[name], download_root, in_memory) |
| | alignment_heads = _ALIGNMENT_HEADS[name] |
| | elif os.path.isfile(name): |
| | checkpoint_file = open(name, "rb").read() if in_memory else name |
| | alignment_heads = None |
| | else: |
| | raise RuntimeError( |
| | f"Model {name} not found; available models = {available_models()}" |
| | ) |
| |
|
| | with ( |
| | io.BytesIO(checkpoint_file) if in_memory else open(checkpoint_file, "rb") |
| | ) as fp: |
| | checkpoint = torch.load(fp, map_location=device) |
| | del checkpoint_file |
| |
|
| | dims = ModelDimensions(**checkpoint["dims"]) |
| | model = Whisper_(dims) |
| | model.load_state_dict(checkpoint["model_state_dict"]) |
| |
|
| | if alignment_heads is not None: |
| | model.set_alignment_heads(alignment_heads) |
| |
|
| | return model.to(device) |
| |
|
| |
|
| | class WhisperFeatureReader(object): |
| | def __init__(self, root, ckpt, layer, device): |
| | self.device = device |
| | logger.info(f"device = {self.device}") |
| |
|
| | self.model: Whisper_ = load_model(name=ckpt, device=self.device, download_root=root).eval() |
| | self.model.decoder = None |
| | self.layer = layer |
| |
|
| | def read_audio(self, path, ref_len=None): |
| | wav, sample_rate = sf.read(path) |
| | assert sample_rate == 16000, sample_rate |
| | if ref_len is not None and abs(ref_len - len(wav)) > 160: |
| | logger.warning(f"ref {ref_len} != read {len(wav)} ({path})") |
| | return wav |
| |
|
| | def get_feats(self, path, ref_len=None): |
| | wav = self.read_audio(path, ref_len) |
| | audio_length = len(wav) |
| | with torch.no_grad(): |
| | mel = log_mel_spectrogram(torch.from_numpy(wav).float().to(self.device)) |
| | hidden = self.model.extract_features(mel.unsqueeze(0), target_layer=self.layer) |
| | feature_length = audio_length // 320 |
| | hidden = hidden[0, :feature_length] |
| | return hidden.contiguous() |
| |
|
