Added extract-small-model-cli. Now users can process their checkpoints without having to run up the GUI. Alot quicker.

This commit is contained in:
Mangio621
2023-05-10 23:24:50 +10:00
parent b34c376509
commit ffe350349c
2 changed files with 24 additions and 4 deletions

View File

@@ -0,0 +1,19 @@
import sys
from train.process_ckpt import extract_small_model
path = str(sys.argv[1])
name = str(sys.argv[2])
sample_rate = str(sys.argv[3])
if_f0 = int(sys.argv[4])
if(len(sys.argv) == 5):
info = ""
else:
info = str(sys.argv[5])
print("Mangio-RVC-Fork Small Model Extraction: Performing extraction...")
extraction = extract_small_model(path, name, sample_rate, if_f0, info)
print("Mangio-RVC-Fork Small Model Extraction: %s" % (extraction))
if(extraction == "Success."):
print("Mangio-RVC-Fork Small Model Extraction: Placed %s.pth into ./weights" % (name))

View File

@@ -1,4 +1,5 @@
# Fork Feature Mangio RVC Fork. Train the feature index (faiss) through the cli
# Please note that the train_index function should probably be exported into a separate file that both the web ui uses and the cli version uses to prevent code duplication.
import os, sys, warnings, shutil, numpy as np
import faiss
@@ -38,10 +39,10 @@ def train_index(exp_dir1):
np.save("%s/total_fea.npy" % exp_dir, big_npy)
# n_ivf = big_npy.shape[0] // 39
n_ivf = min(int(16 * np.sqrt(big_npy.shape[0])), big_npy.shape[0] // 39)
yield "%s,%s \n" % (big_npy.shape, n_ivf)
yield "%s,%s" % (big_npy.shape, n_ivf)
index = faiss.index_factory(256, "IVF%s,Flat" % n_ivf)
# index = faiss.index_factory(256, "IVF%s,PQ128x4fs,RFlat"%n_ivf)
yield "training the index...\n"
yield "training the index..."
index_ivf = faiss.extract_index_ivf(index) #
# index_ivf.nprobe = int(np.power(n_ivf,0.3))
index_ivf.nprobe = 1
@@ -51,7 +52,7 @@ def train_index(exp_dir1):
"%s/trained_IVF%s_Flat_nprobe_%s.index" % (exp_dir, n_ivf, index_ivf.nprobe),
)
# faiss.write_index(index, '%s/trained_IVF%s_Flat_FastScan.index'%(exp_dir,n_ivf))
yield "adding the index... \n"
yield "adding the index..."
batch_size_add = 8192
for i in range(0, big_npy.shape[0], batch_size_add):
index.add(big_npy[i : i + batch_size_add])
@@ -61,7 +62,7 @@ def train_index(exp_dir1):
)
# faiss.write_index(index, '%s/added_IVF%s_Flat_FastScan.index'%(exp_dir,n_ivf))
# infos.append("成功构建索引added_IVF%s_Flat_FastScan.index"%(n_ivf))
yield "Done! added_IVF%s_Flat_nprobe_%s.index \n" % (n_ivf, index_ivf.nprobe)
yield "Done! added_IVF%s_Flat_nprobe_%s.index" % (n_ivf, index_ivf.nprobe)
train_output = train_index(model_name)
if(train_output == fail_msg):