| |
| |
| |
| |
| |
| |
|
|
| import logging |
|
|
| import fairseq |
| import torch |
| import torch.nn.functional as F |
|
|
| from fairseq.data.audio.audio_utils import get_features_or_waveform |
|
|
| logger = logging.getLogger("dump_feature") |
|
|
|
|
| class HubertFeatureReader(object): |
| def __init__(self, ckpt_path: str, layer: int, device: str, max_chunk=1600000): |
| ( |
| model, |
| cfg, |
| task, |
| ) = fairseq.checkpoint_utils.load_model_ensemble_and_task([ckpt_path]) |
|
|
| self.device = device |
| logger.info(f"device = {self.device}") |
|
|
| self.model = model[0].eval().to(self.device) |
| self.task = task |
| self.layer = layer |
| self.max_chunk = max_chunk |
| logger.info(f"TASK CONFIG:\n{self.task.cfg}") |
| logger.info(f" max_chunk = {self.max_chunk}") |
|
|
| def read_audio(self, path, ref_len=None): |
| wav = get_features_or_waveform(path, need_waveform=True, use_sample_rate=self.task.cfg.sample_rate) |
| if wav.ndim == 2: |
| wav = wav.mean(-1) |
| assert wav.ndim == 1, wav.ndim |
| 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): |
| x = self.read_audio(path, ref_len=ref_len) |
| with torch.no_grad(): |
| x = torch.from_numpy(x).float().to(self.device) |
| if self.task.cfg.normalize: |
| x = F.layer_norm(x, x.shape) |
| x = x.view(1, -1) |
|
|
| feat = [] |
| for start in range(0, x.size(1), self.max_chunk): |
| x_chunk = x[:, start: start + self.max_chunk] |
| feat_chunk, _ = self.model.extract_features( |
| source=x_chunk, |
| padding_mask=None, |
| mask=False, |
| output_layer=self.layer, |
| ) |
| feat.append(feat_chunk) |
| return torch.cat(feat, 1).squeeze(0) |
|
|
