| | |
| | |
| | |
| | |
| | |
| | |
| |
|
| | 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) |
| |
|
