| import os
|
| import logging
|
| import functools
|
| import torch
|
|
|
| logger = logging.getLogger(__name__)
|
|
|
| try:
|
| from fairseq import checkpoint_utils
|
| FAIRSEQ_AVAILABLE = True
|
| except Exception:
|
| FAIRSEQ_AVAILABLE = False
|
|
|
|
|
| def _patch_torch_load():
|
| """Patch torch.load to default weights_only=False for fairseq compatibility (PyTorch 2.6+)."""
|
| _original = torch.load
|
|
|
| @functools.wraps(_original)
|
| def _patched(*args, **kwargs):
|
| kwargs.setdefault("weights_only", False)
|
| return _original(*args, **kwargs)
|
|
|
| return _original, _patched
|
|
|
|
|
| def get_index_path_from_model(sid):
|
| return next(
|
| (
|
| f
|
| for f in [
|
| os.path.join(root, name)
|
| for root, _, files in os.walk(os.getenv("index_root"), topdown=False)
|
| for name in files
|
| if name.endswith(".index") and "trained" not in name
|
| ]
|
| if sid.split(".")[0] in f
|
| ),
|
| "",
|
| )
|
|
|
|
|
| def load_hubert(config):
|
| if FAIRSEQ_AVAILABLE:
|
| _original, _patched = _patch_torch_load()
|
| torch.load = _patched
|
| try:
|
| models, _, _ = checkpoint_utils.load_model_ensemble_and_task(
|
| ["assets/hubert/hubert_base.pt"],
|
| suffix="",
|
| )
|
| finally:
|
| torch.load = _original
|
| hubert_model = models[0]
|
| hubert_model = hubert_model.to(config.device)
|
| if config.is_half:
|
| hubert_model = hubert_model.half()
|
| else:
|
| hubert_model = hubert_model.float()
|
| return hubert_model.eval()
|
|
|
| try:
|
| import torchaudio
|
|
|
| class HubertWrapper:
|
| def __init__(self, model):
|
| self.model = model
|
| self.final_proj = getattr(model, "final_proj", torch.nn.Identity())
|
|
|
| def extract_features(self, source, padding_mask=None, output_layer=None):
|
| feats, _ = self.model.extract_features(source)
|
| if output_layer is None:
|
| idx = -1
|
| else:
|
| idx = min(output_layer - 1, len(feats) - 1)
|
| return (feats[idx], None)
|
|
|
| def to(self, device):
|
| self.model = self.model.to(device)
|
| return self
|
|
|
| def half(self):
|
| self.model = self.model.half()
|
| return self
|
|
|
| def float(self):
|
| self.model = self.model.float()
|
| return self
|
|
|
| def eval(self):
|
| self.model.eval()
|
| return self
|
|
|
| model = torchaudio.pipelines.HUBERT_BASE.get_model()
|
| hubert_model = HubertWrapper(model).to(config.device)
|
| if config.is_half:
|
| hubert_model = hubert_model.half()
|
| else:
|
| hubert_model = hubert_model.float()
|
| return hubert_model.eval()
|
| except Exception as e:
|
| raise RuntimeError(
|
| "HuBERT 模型加载失败，请检查 fairseq 和 torchaudio 是否已安装"
|
| ) from e
|
|
|
