fix: index_root searching

close #1147

.env

@@ -4,5 +4,5 @@ no_proxy = localhost, 127.0.0.1, ::1
 # You can change the location of the model, etc. by changing here
 weight_root = assets/weights
 weight_uvr5_root = assets/uvr5_weights
-index_root = output
+index_root = logs
 rmvpe_root = assets/rmvpe

configs/config.py

@@ -160,7 +160,7 @@ class Config:
             x_center = 30
             x_max = 32
         if self.dml:
-            print("use DirectML instead")
+            print("Use DirectML instead")
             if (
                 os.path.exists(
                     "runtime\Lib\site-packages\onnxruntime\capi\DirectML.dll"
@@ -188,7 +188,7 @@ class Config:
             self.is_half = False
         else:
             if self.instead:
-                print(f"use {self.instead} instead")
+                print(f"Use {self.instead} instead")
             if (
                 os.path.exists(
                     "runtime\Lib\site-packages\onnxruntime\capi\onnxruntime_providers_cuda.dll"

gui_v1.py

@@ -356,7 +356,7 @@ if __name__ == "__main__":
                     )
                 if event == "start_vc" and self.flag_vc == False:
                     if self.set_values(values) == True:
-                        print("using_cuda:" + str(torch.cuda.is_available()))
+                        print("Use CUDA:" + str(torch.cuda.is_available()))
                         self.start_vc()
                         settings = {
                             "pth_path": values["pth_path"],
@@ -623,7 +623,7 @@ if __name__ == "__main__":
                 sola_offset = sola_offset.item()
             else:
                 sola_offset = torch.argmax(cor_nom[0, 0] / cor_den[0, 0])
-            print("sola offset: " + str(int(sola_offset)))
+            print("sola_offset =" + str(int(sola_offset)))
             self.output_wav[:] = infer_wav[sola_offset : sola_offset + self.block_frame]
             self.output_wav[: self.crossfade_frame] *= self.fade_in_window
             self.output_wav[: self.crossfade_frame] += self.sola_buffer[:]
@@ -663,7 +663,7 @@ if __name__ == "__main__":
                     outdata[:] = self.output_wav[:].repeat(2, 1).t().cpu().numpy()
             total_time = time.perf_counter() - start_time
             self.window["infer_time"].update(int(total_time * 1000))
-            print("infer time:" + str(total_time))
+            print("Infer time:" + str(total_time))
 
         def get_devices(self, update: bool = True):
             """获取设备列表"""
@@ -716,9 +716,9 @@ if __name__ == "__main__":
             sd.default.device[1] = output_device_indices[
                 output_devices.index(output_device)
             ]
-            print("input device:" + str(sd.default.device[0]) + ":" + str(input_device))
+            print("Input device:" + str(sd.default.device[0]) + ":" + str(input_device))
             print(
-                "output device:" + str(sd.default.device[1]) + ":" + str(output_device)
+                "Output device:" + str(sd.default.device[1]) + ":" + str(output_device)
             )
 
     gui = GUI()

infer-web.py

@@ -124,7 +124,7 @@ class ToolButton(gr.Button, gr.components.FormComponent):
 
 weight_root = os.getenv("weight_root")
 weight_uvr5_root = os.getenv("weight_uvr5_root")
-index_root = "logs"
+index_root = os.getenv("index_root")
 
 names = []
 for name in os.listdir(weight_root):
@@ -556,14 +556,14 @@ def click_train(
     shuffle(opt)
     with open("%s/filelist.txt" % exp_dir, "w") as f:
         f.write("\n".join(opt))
-    print("write filelist done")
+    print("Write filelist done")
     # 生成config#无需生成config
     # cmd = python_cmd + " train_nsf_sim_cache_sid_load_pretrain.py -e mi-test -sr 40k -f0 1 -bs 4 -g 0 -te 10 -se 5 -pg pretrained/f0G40k.pth -pd pretrained/f0D40k.pth -l 1 -c 0"
-    print("use gpus:", gpus16)
+    print("Use gpus:", gpus16)
     if pretrained_G14 == "":
-        print("no pretrained Generator")
+        print("No pretrained Generator")
     if pretrained_D15 == "":
-        print("no pretrained Discriminator")
+        print("No pretrained Discriminator")
     if gpus16:
         cmd = get_quoted_python_cmd() + ' infer/modules/train/train.py -e "%s" -sr %s -f0 %s -bs %s -g %s -te %s -se %s %s %s -l %s -c %s -sw %s -v %s' % (
             exp_dir1,

infer/lib/infer_pack/onnx_inference.py

@@ -6,7 +6,7 @@ import soundfile
 
 class ContentVec:
     def __init__(self, vec_path="pretrained/vec-768-layer-12.onnx", device=None):
-        print("load model(s) from {}".format(vec_path))
+        print("Load model(s) from {}".format(vec_path))
         if device == "cpu" or device is None:
             providers = ["CPUExecutionProvider"]
         elif device == "cuda":

infer/lib/train/mel_processing.py

@@ -51,9 +51,9 @@ def spectrogram_torch(y, n_fft, sampling_rate, hop_size, win_size, center=False)
     """
     # Validation
     if torch.min(y) < -1.07:
-        print("min value is ", torch.min(y))
+        print("spectrogram_torch min value is ", torch.min(y))
     if torch.max(y) > 1.07:
-        print("max value is ", torch.max(y))
+        print("spectrogram_torch max value is ", torch.max(y))
 
     # Window - Cache if needed
     global hann_window

infer/lib/train/utils.py

@@ -34,7 +34,7 @@ def load_checkpoint_d(checkpoint_path, combd, sbd, optimizer=None, load_opt=1):
                 new_state_dict[k] = saved_state_dict[k]
                 if saved_state_dict[k].shape != state_dict[k].shape:
                     print(
-                        "shape-%s-mismatch|need-%s|get-%s"
+                        "shape-%s-mismatch. need: %s, get: %s"
                         % (k, state_dict[k].shape, saved_state_dict[k].shape)
                     )  #
                     raise KeyError

infer/modules/train/extract/extract_f0_print.py

@@ -85,7 +85,7 @@ class FeatureInput(object):
             if hasattr(self, "model_rmvpe") == False:
                 from infer.lib.rmvpe import RMVPE
 
-                print("loading rmvpe model")
+                print("Loading rmvpe model")
                 self.model_rmvpe = RMVPE(
                     "assets/rmvpe/rmvpe.pt", is_half=False, device="cpu"
                 )

infer/modules/train/extract/extract_f0_rmvpe.py

@@ -48,7 +48,7 @@ class FeatureInput(object):
             if hasattr(self, "model_rmvpe") == False:
                 from infer.lib.rmvpe import RMVPE
 
-                print("loading rmvpe model")
+                print("Loading rmvpe model")
                 self.model_rmvpe = RMVPE(
                     "assets/rmvpe/rmvpe.pt", is_half=is_half, device="cuda"
                 )

infer/modules/train/extract/extract_f0_rmvpe_dml.py

@@ -46,7 +46,7 @@ class FeatureInput(object):
             if hasattr(self, "model_rmvpe") == False:
                 from infer.lib.rmvpe import RMVPE
 
-                print("loading rmvpe model")
+                print("Loading rmvpe model")
                 self.model_rmvpe = RMVPE(
                     "assets/rmvpe/rmvpe.pt", is_half=False, device=device
                 )

infer/modules/uvr5/mdxnet.py

@@ -101,7 +101,7 @@ class Predictor:
                 "CPUExecutionProvider",
             ],
         )
-        print("onnx load done")
+        print("ONNX load done")
 
     def demix(self, mix):
         samples = mix.shape[-1]

infer/modules/uvr5/modules.py

@@ -90,7 +90,7 @@ def uvr(model_name, inp_root, save_root_vocal, paths, save_root_ins, agg, format
                 del pre_fun
         except:
             traceback.print_exc()
-        print("clean_empty_cache")
         if torch.cuda.is_available():
             torch.cuda.empty_cache()
+            print("Executed torch.cuda.empty_cache()")
     yield "\n".join(infos)

infer/modules/vc/modules.py

@@ -31,7 +31,7 @@ class VC:
 
     def get_vc(self, sid, *to_return_protect):
         person = f'{os.getenv("weight_root")}/{sid}'
-        print(f"loading {person}")
+        print(f"Loading: {person}")
 
         self.cpt = torch.load(person, map_location="cpu")
         self.tgt_sr = self.cpt["config"][-1]
@@ -77,6 +77,7 @@ class VC:
         self.pipeline = Pipeline(self.tgt_sr, self.config)
         n_spk = self.cpt["config"][-3]
         index = {"value": get_index_path_from_model(sid), "__type__": "update"}
+        print("Select index:", index["value"])
 
         return (
             (

infer/modules/vc/pipeline.py

@@ -140,7 +140,7 @@ class Pipeline(object):
                 from infer.lib.rmvpe import RMVPE
 
                 print(
-                    "loading rmvpe model,%s" % "%s/rmvpe.pt" % os.environ["rmvpe_root"]
+                    "Loading rmvpe model,%s" % "%s/rmvpe.pt" % os.environ["rmvpe_root"]
                 )
                 self.model_rmvpe = RMVPE(
                     "%s/rmvpe.pt" % os.environ["rmvpe_root"],
@@ -152,7 +152,7 @@ class Pipeline(object):
         if "privateuseone" in str(self.device):  # clean ortruntime memory
             del self.model_rmvpe.model
             del self.model_rmvpe
-            print("cleaning ortruntime memory")
+            print("Cleaning ortruntime memory")
 
         f0 *= pow(2, f0_up_key / 12)
         # with open("test.txt","w")as f:f.write("\n".join([str(i)for i in f0.tolist()]))
@@ -262,17 +262,12 @@ class Pipeline(object):
             feats = feats.to(feats0.dtype)
         p_len = torch.tensor([p_len], device=self.device).long()
         with torch.no_grad():
-            if pitch is not None and pitchf is not None:
-                audio1 = (
-                    (net_g.infer(feats, p_len, pitch, pitchf, sid)[0][0, 0])
-                    .data.cpu()
-                    .float()
-                    .numpy()
-                )
-            else:
-                audio1 = (
-                    (net_g.infer(feats, p_len, sid)[0][0, 0]).data.cpu().float().numpy()
-                )
+            hasp = pitch is not None and pitchf is not None
+            arg = (feats, p_len, pitch, pitchf, sid) if hasp else (feats, p_len, sid)
+            audio1 = (
+                (net_g.infer(*arg)[0][0, 0]).data.cpu().float().numpy()
+            )
+            del hasp, arg
         del feats, p_len, padding_mask
         if torch.cuda.is_available():
             torch.cuda.empty_cache()

infer/modules/vc/utils.py

@@ -2,33 +2,20 @@ import os
 
 from fairseq import checkpoint_utils
 
-### don't modify the code before you test it
-# def get_index_path_from_model(sid):
-#     return next(
-#         (
-#             f
-#             for f in [
-#                 os.path.join(root, name)
-#                 for root, dirs, 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 get_index_path_from_model(sid):
-    sel_index_path = ""
-    name = os.path.join("logs", sid.split(".")[0], "")
-    # print(name)
-    for f in index_paths:
-        if name in f:
-            # print("selected index path:", f)
-            sel_index_path = f
-            break
-    return sel_index_path
+    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):

tools/app.py

@@ -23,7 +23,7 @@ vc = VC(config)
 
 weight_root = os.getenv("weight_root")
 weight_uvr5_root = os.getenv("weight_uvr5_root")
-index_root = "logs"
+index_root = os.getenv("index_root")
 names = []
 hubert_model = None
 for name in os.listdir(weight_root):

tools/calc_rvc_model_similarity.py

@@ -55,7 +55,7 @@ def main(path, root):
     torch.manual_seed(114514)
     model_a = torch.load(path, map_location="cpu")["weight"]
 
-    print("query:\t\t%s\t%s" % (path, model_hash(path)))
+    print("Query:\t\t%s\t%s" % (path, model_hash(path)))
 
     map_attn_a = {}
     map_rand_input = {}
@@ -83,7 +83,7 @@ def main(path, root):
             sims.append(sim)
 
         print(
-            "reference:\t%s\t%s\t%s"
+            "Reference:\t%s\t%s\t%s"
             % (path, model_hash(path), f"{torch.mean(torch.stack(sims)) * 1e2:.2f}%")
         )
 

tools/infer/infer-pm-index256.py

@@ -34,7 +34,7 @@ from scipy.io import wavfile
 
 device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 model_path = r"E:\codes\py39\vits_vc_gpu_train\assets\hubert\hubert_base.pt"  #
-print("load model(s) from {}".format(model_path))
+print("Load model(s) from {}".format(model_path))
 models, saved_cfg, task = checkpoint_utils.load_model_ensemble_and_task(
     [model_path],
     suffix="",

tools/infer/train-index-v2.py

@@ -50,14 +50,14 @@ np.save("tools/infer/big_src_feature_mi.npy", big_npy)
 # big_npy=np.load("/bili-coeus/jupyter/jupyterhub-liujing04/vits_ch/inference_f0/big_src_feature_mi.npy")
 n_ivf = min(int(16 * np.sqrt(big_npy.shape[0])), big_npy.shape[0] // 39)
 index = faiss.index_factory(768, "IVF%s,Flat" % n_ivf)  # mi
-print("training")
+print("Training...")
 index_ivf = faiss.extract_index_ivf(index)  #
 index_ivf.nprobe = 1
 index.train(big_npy)
 faiss.write_index(
     index, "tools/infer/trained_IVF%s_Flat_baseline_src_feat_v2.index" % (n_ivf)
 )
-print("adding")
+print("Adding...")
 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])

tools/infer/train-index.py

@@ -20,12 +20,12 @@ np.save("infer/big_src_feature_mi.npy", big_npy)
 # big_npy=np.load("/bili-coeus/jupyter/jupyterhub-liujing04/vits_ch/inference_f0/big_src_feature_mi.npy")
 print(big_npy.shape)
 index = faiss.index_factory(256, "IVF512,Flat")  # mi
-print("training")
+print("Training...")
 index_ivf = faiss.extract_index_ivf(index)  #
 index_ivf.nprobe = 9
 index.train(big_npy)
 faiss.write_index(index, "infer/trained_IVF512_Flat_mi_baseline_src_feat.index")
-print("adding")
+print("Adding...")
 index.add(big_npy)
 faiss.write_index(index, "infer/added_IVF512_Flat_mi_baseline_src_feat.index")
 """

tools/rvc_for_realtime.py

@@ -67,7 +67,7 @@ class RVC:
             if index_rate != 0:
                 self.index = faiss.read_index(index_path)
                 self.big_npy = self.index.reconstruct_n(0, self.index.ntotal)
-                print("index search enabled")
+                print("Index search enabled")
             self.index_path = index_path
             self.index_rate = index_rate
             models, _, _ = fairseq.checkpoint_utils.load_model_ensemble_and_task(
@@ -120,7 +120,7 @@ class RVC:
         if new_index_rate != 0 and self.index_rate == 0:
             self.index = faiss.read_index(self.index_path)
             self.big_npy = self.index.reconstruct_n(0, self.index.ntotal)
-            print("index search enabled")
+            print("Index search enabled")
         self.index_rate = new_index_rate
 
     def get_f0_post(self, f0):
@@ -237,7 +237,7 @@ class RVC:
         if hasattr(self, "model_rmvpe") == False:
             from infer.lib.rmvpe import RMVPE
 
-            print("loading rmvpe model")
+            print("Loading rmvpe model")
             self.model_rmvpe = RMVPE(
                 # "rmvpe.pt", is_half=self.is_half if self.device.type!="privateuseone" else False, device=self.device if self.device.type!="privateuseone"else "cpu"####dml时强制对rmvpe用cpu跑
                 #  "rmvpe.pt", is_half=False, device=self.device####dml配置
@@ -295,10 +295,10 @@ class RVC:
                     + (1 - self.index_rate) * feats[0][-leng_replace_head:]
                 )
             else:
-                print("index search FAIL or disabled")
+                print("Index search FAILED or disabled")
         except:
             traceback.print_exc()
-            print("index search FAIL")
+            print("Index search FAILED")
         feats = F.interpolate(feats.permute(0, 2, 1), scale_factor=2).permute(0, 2, 1)
         t3 = ttime()
         if self.if_f0 == 1:
@@ -338,5 +338,5 @@ class RVC:
                     .float()
                 )
         t5 = ttime()
-        print("time->fea-index-f0-model:", t2 - t1, t3 - t2, t4 - t3, t5 - t4)
+        print("Spent time: fea =", t2 - t1, ", index =", t3 - t2, ", f0 =", t4 - t3, ", model =", t5 - t4)
         return infered_audio