mirror of
https://github.com/RVC-Project/Retrieval-based-Voice-Conversion-WebUI.git
synced 2026-08-29 10:09:32 +02:00
250 lines
8.3 KiB
Python
250 lines
8.3 KiB
Python
import os
|
||
import platform
|
||
import sys
|
||
import traceback
|
||
import glob
|
||
import json
|
||
|
||
import faiss
|
||
import numpy as np
|
||
from sklearn.cluster import MiniBatchKMeans
|
||
from i18n.i18n import I18nAuto
|
||
from tools.progress import should_report
|
||
|
||
|
||
i18n = I18nAuto()
|
||
|
||
|
||
exp_name = sys.argv[1]
|
||
version = sys.argv[2]
|
||
outside_index_root = sys.argv[3]
|
||
n_cpu = int(sys.argv[4])
|
||
index_mode = sys.argv[5] if len(sys.argv) > 5 else "auto"
|
||
exp_dir = os.path.join("logs", exp_name)
|
||
feature_dir = os.path.join(
|
||
exp_dir, "3_feature256" if version == "v1" else "3_feature768"
|
||
)
|
||
log_path = os.path.join(exp_dir, "train_index.log")
|
||
os.makedirs(exp_dir, exist_ok=True)
|
||
|
||
|
||
def log(message):
|
||
print(message, flush=True)
|
||
with open(log_path, "a", encoding="utf8") as f:
|
||
f.write(str(message) + "\n")
|
||
|
||
|
||
with open(log_path, "w", encoding="utf8"):
|
||
pass
|
||
|
||
|
||
def newest_index(pattern):
|
||
paths = [path for path in glob.glob(pattern) if os.path.isfile(path)]
|
||
return max(paths, key=os.path.getmtime) if paths else ""
|
||
|
||
|
||
def speaker_scope(message, speaker_id):
|
||
if speaker_id is None:
|
||
return message
|
||
return "%s：%s | %s" % (i18n("说话人ID（0~109）"), speaker_id, message)
|
||
|
||
|
||
def link_added_index(added_path, speaker_id=None):
|
||
added_name = os.path.basename(added_path)
|
||
try:
|
||
os.makedirs(outside_index_root, exist_ok=True)
|
||
source = os.path.abspath(added_path)
|
||
outside_root = os.path.abspath(outside_index_root)
|
||
if os.path.commonpath([source, outside_root]) == outside_root:
|
||
log(speaker_scope(i18n("[索引训练] 外部索引链接已存在：%s") % source, speaker_id))
|
||
return
|
||
target = os.path.abspath(
|
||
os.path.join(outside_index_root, "%s_%s" % (exp_name, added_name))
|
||
)
|
||
if os.path.lexists(target):
|
||
try:
|
||
if os.path.samefile(source, target):
|
||
log(speaker_scope(i18n("[索引训练] 外部索引链接已存在：%s") % target, speaker_id))
|
||
return
|
||
except (FileNotFoundError, OSError):
|
||
pass
|
||
os.unlink(target)
|
||
if platform.system() == "Windows":
|
||
os.link(source, target)
|
||
else:
|
||
os.symlink(source, target)
|
||
log(speaker_scope(i18n("[索引训练] 已链接索引到外部目录：%s") % outside_index_root, speaker_id))
|
||
except Exception:
|
||
log(
|
||
speaker_scope(
|
||
i18n("[索引训练][失败] 无法链接索引到外部目录：%s\n%s")
|
||
% (outside_index_root, traceback.format_exc()),
|
||
speaker_id,
|
||
)
|
||
)
|
||
|
||
|
||
if not os.path.isdir(feature_dir) or not os.listdir(feature_dir):
|
||
log(i18n("[索引训练][失败] 请先进行特征提取"))
|
||
raise SystemExit(1)
|
||
|
||
manifest_path = os.path.join(exp_dir, "multispeaker_manifest.json")
|
||
manifest_by_key = {}
|
||
if index_mode != "single" and os.path.isfile(manifest_path):
|
||
try:
|
||
with open(manifest_path, "r", encoding="utf8") as file:
|
||
manifest = json.load(file)
|
||
entries = manifest.get("entries", []) if isinstance(manifest, dict) else []
|
||
manifest_by_key = {
|
||
str(entry["output_key"]): int(entry["speaker_id"])
|
||
for entry in entries
|
||
}
|
||
except (KeyError, TypeError, ValueError, OSError, json.JSONDecodeError):
|
||
manifest_by_key = {}
|
||
|
||
feature_paths = [
|
||
os.path.join(feature_dir, name)
|
||
for name in sorted(os.listdir(feature_dir))
|
||
if name.lower().endswith(".npy")
|
||
]
|
||
feature_groups = {}
|
||
if manifest_by_key:
|
||
for path in feature_paths:
|
||
stem = os.path.splitext(os.path.basename(path))[0]
|
||
output_key = stem if stem in manifest_by_key else stem.rsplit("_", 1)[0]
|
||
speaker_id = manifest_by_key.get(output_key)
|
||
if speaker_id is not None:
|
||
feature_groups.setdefault(speaker_id, []).append(path)
|
||
else:
|
||
feature_groups[None] = feature_paths
|
||
|
||
if not feature_groups or not any(feature_groups.values()):
|
||
log(i18n("[索引训练][失败] 请先进行特征提取"))
|
||
raise SystemExit(1)
|
||
|
||
|
||
def train_one_speaker(speaker_id, paths):
|
||
scope = lambda message: speaker_scope(message, speaker_id)
|
||
features = [np.load(path) for path in paths]
|
||
big_npy = np.concatenate(features, 0)
|
||
big_npy = big_npy[np.random.permutation(big_npy.shape[0])]
|
||
suffix = "" if speaker_id is None else "_spkid%s" % speaker_id
|
||
total_path = os.path.join(exp_dir, "total_fea%s.npy" % suffix)
|
||
np.save(total_path, big_npy)
|
||
|
||
trained_pattern = os.path.join(
|
||
exp_dir,
|
||
"trained_IVF*_Flat_nprobe_*_%s_%s%s.index"
|
||
% (exp_name, version, suffix),
|
||
)
|
||
added_pattern = os.path.join(
|
||
exp_dir,
|
||
"added_IVF*_Flat_nprobe_*_%s_%s%s.index"
|
||
% (exp_name, version, suffix),
|
||
)
|
||
existing_trained_path = newest_index(trained_pattern)
|
||
existing_added_path = newest_index(added_pattern)
|
||
if not existing_added_path:
|
||
existing_added_path = newest_index(
|
||
os.path.join(
|
||
outside_index_root,
|
||
"%s_*added_IVF*_Flat_nprobe_*_%s_%s%s.index"
|
||
% (exp_name, exp_name, version, suffix),
|
||
)
|
||
)
|
||
if existing_added_path:
|
||
if existing_trained_path:
|
||
log(
|
||
scope(
|
||
i18n("[索引训练][跳过] trained索引已存在：%s")
|
||
% os.path.basename(existing_trained_path)
|
||
)
|
||
)
|
||
log(
|
||
scope(
|
||
i18n("[索引训练][跳过] added索引已存在：%s")
|
||
% os.path.basename(existing_added_path)
|
||
)
|
||
)
|
||
link_added_index(existing_added_path, speaker_id)
|
||
return
|
||
|
||
if big_npy.shape[0] > 200000:
|
||
log(
|
||
scope(
|
||
i18n("[索引训练] 正在将%s条特征聚类为10000个中心")
|
||
% big_npy.shape[0]
|
||
)
|
||
)
|
||
try:
|
||
big_npy = MiniBatchKMeans(
|
||
n_clusters=10000,
|
||
verbose=False,
|
||
batch_size=256 * n_cpu,
|
||
compute_labels=False,
|
||
init="random",
|
||
).fit(big_npy).cluster_centers_
|
||
except Exception:
|
||
log(
|
||
scope(
|
||
i18n("[索引训练][失败] 聚类失败，将使用原始特征继续\n%s")
|
||
% traceback.format_exc()
|
||
)
|
||
)
|
||
|
||
n_ivf = max(1, min(int(16 * np.sqrt(big_npy.shape[0])), big_npy.shape[0] // 39))
|
||
log(scope(i18n("[索引训练] 特征形状：%s | IVF数量：%s") % (big_npy.shape, n_ivf)))
|
||
if existing_trained_path:
|
||
trained_path = existing_trained_path
|
||
index = faiss.read_index(trained_path)
|
||
index_ivf = faiss.extract_index_ivf(index)
|
||
index_ivf.nprobe = 1
|
||
log(
|
||
scope(
|
||
i18n("[索引训练][跳过] trained索引已存在：%s")
|
||
% os.path.basename(trained_path)
|
||
)
|
||
)
|
||
else:
|
||
index = faiss.index_factory(
|
||
256 if version == "v1" else 768, "IVF%s,Flat" % n_ivf
|
||
)
|
||
index_ivf = faiss.extract_index_ivf(index)
|
||
index_ivf.nprobe = 1
|
||
trained_path = os.path.join(
|
||
exp_dir,
|
||
"trained_IVF%s_Flat_nprobe_%s_%s_%s%s.index"
|
||
% (n_ivf, index_ivf.nprobe, exp_name, version, suffix),
|
||
)
|
||
log(scope(i18n("[索引训练] 正在训练索引")))
|
||
index.train(big_npy)
|
||
faiss.write_index(index, trained_path)
|
||
|
||
log(scope(i18n("[索引训练] 正在写入特征向量")))
|
||
starts = list(range(0, big_npy.shape[0], 8192))
|
||
for batch_index, start in enumerate(starts):
|
||
index.add(big_npy[start : start + 8192])
|
||
if should_report(batch_index, len(starts), 10):
|
||
log(
|
||
scope(
|
||
i18n("[索引训练] 写入进度：%s/%s")
|
||
% (batch_index + 1, len(starts))
|
||
)
|
||
)
|
||
|
||
added_name = "added_IVF%s_Flat_nprobe_%s_%s_%s%s.index" % (
|
||
n_ivf,
|
||
index_ivf.nprobe,
|
||
exp_name,
|
||
version,
|
||
suffix,
|
||
)
|
||
added_path = os.path.join(exp_dir, added_name)
|
||
faiss.write_index(index, added_path)
|
||
log(scope(i18n("[索引训练] 成功构建索引：%s") % added_name))
|
||
link_added_index(added_path, speaker_id)
|
||
|
||
|
||
for speaker_id in sorted(feature_groups, key=lambda value: -1 if value is None else value):
|
||
train_one_speaker(speaker_id, feature_groups[speaker_id])
|