源文雨 8d5a77dbe9
Revert "Format code (#1154)" (#1156)
This reverts commit eefb7dbe1d253c13321912a7d10ea186248dcfe1.
2023-09-01 15:17:51 +08:00

33 lines
850 B
Python

import os
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 sid.split(".")[0] in f
),
"",
)
def load_hubert(config):
models, _, _ = checkpoint_utils.load_model_ensemble_and_task(
["assets/hubert/hubert_base.pt"],
suffix="",
)
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()