45 lines
1.3 KiB
Python
45 lines
1.3 KiB
Python
import os
|
|
import torch
|
|
|
|
from fairseq import checkpoint_utils
|
|
|
|
|
|
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 str(sid).split(".")[0] in f
|
|
),
|
|
"",
|
|
)
|
|
|
|
|
|
def load_hubert(config, hubert_path: str):
|
|
# PyTorch 2.6+ changed weights_only default to True, which breaks fairseq checkpoints
|
|
# Monkey-patch torch.load to use weights_only=False for fairseq
|
|
original_torch_load = torch.load
|
|
|
|
def patched_torch_load(f, map_location=None, *args, **kwargs):
|
|
kwargs.setdefault('weights_only', False)
|
|
return original_torch_load(f, map_location=map_location, *args, **kwargs)
|
|
|
|
torch.load = patched_torch_load
|
|
try:
|
|
models, _, _ = checkpoint_utils.load_model_ensemble_and_task(
|
|
[hubert_path],
|
|
suffix="",
|
|
)
|
|
finally:
|
|
torch.load = original_torch_load
|
|
|
|
hubert_model = models[0]
|
|
hubert_model = hubert_model.to(config.device)
|
|
hubert_model = hubert_model.half() if config.is_half else hubert_model.float()
|
|
return hubert_model.eval()
|