Update demo.py

demo.py

@@ -1,8 +1,106 @@
 import os
 import shutil
 from os import listdir
-import gradio as gr
 from colorama import Fore
+import os
+import shutil
+import numpy as np
+import faiss
+from pathlib import Path
+from sklearn.cluster import MiniBatchKMeans
+import traceback
+import gradio as gr
+
+# Function to preprocess data
+def preprocess_data(model_name, dataset_folder):
+    logs_path = f'/content/RVC/logs/{model_name}'
+    temp_DG_path = '/content/temp_DG'
+
+    if os.path.exists(logs_path):
+        print("Model already exists, This will be resume training.")
+        os.makedirs(temp_DG_path, exist_ok=True)
+
+        # Move files for resuming training
+        for item in os.listdir(logs_path):
+            item_path = os.path.join(logs_path, item)
+            if os.path.isfile(item_path) and (item.startswith('D_') or item.startswith('G_')) and item.endswith('.pth'):
+                shutil.copy(item_path, temp_DG_path)
+
+        for item in os.listdir(logs_path):
+            item_path = os.path.join(logs_path, item)
+            if os.path.isfile(item_path):
+                os.remove(item_path)
+            elif os.path.isdir(item_path):
+                shutil.rmtree(item_path)
+
+        for file_name in os.listdir(temp_DG_path):
+            shutil.move(os.path.join(temp_DG_path, file_name), logs_path)
+
+        shutil.rmtree(temp_DG_path)
+
+    if len(os.listdir(dataset_folder)) < 1:
+        return "Error: Dataset folder is empty."
+
+    os.makedirs(f'./logs/{model_name}', exist_ok=True)
+    !python infer/modules/train/preprocess.py {dataset_folder} 32000 2 ./logs/{model_name} False 3.0
+
+    with open(f'./logs/{model_name}/preprocess.log', 'r') as f:
+        log_content = f.read()
+
+    if 'end preprocess' in log_content:
+        return "Success: Data preprocessing complete."
+    else:
+        return "Error preprocessing data. Check your dataset folder."
+
+# Function to extract F0 feature
+def extract_f0_feature(model_name, f0method):
+    if f0method != "rmvpe_gpu":
+        !python infer/modules/train/extract/extract_f0_print.py ./logs/{model_name} 2 {f0method}
+    else:
+        !python infer/modules/train/extract/extract_f0_rmvpe.py 1 0 0 ./logs/{model_name} True
+
+    with open(f'./logs/{model_name}/extract_f0_feature.log', 'r') as f:
+        log_content = f.read()
+
+    if 'all-feature-done' in log_content:
+        return "Success: F0 feature extraction complete."
+    else:
+        return "Error extracting F0 feature."
+
+# Function to train index
+def train_index(exp_dir1, version19):
+    exp_dir = f"logs/{exp_dir1}"
+    os.makedirs(exp_dir, exist_ok=True)
+    feature_dir = f"{exp_dir}/3_feature768" if version19 == "v2" else f"{exp_dir}/3_feature256"
+
+    if not os.path.exists(feature_dir) or len(os.listdir(feature_dir)) == 0:
+        return "Please run feature extraction first."
+
+    npys = [np.load(f"{feature_dir}/{name}") for name in sorted(os.listdir(feature_dir))]
+    big_npy = np.concatenate(npys, axis=0)
+    big_npy_idx = np.arange(big_npy.shape[0])
+    np.random.shuffle(big_npy_idx)
+    big_npy = big_npy[big_npy_idx]
+
+    if big_npy.shape[0] > 2e5:
+        big_npy = MiniBatchKMeans(n_clusters=10000, batch_size=256, init="random").fit(big_npy).cluster_centers_
+
+    np.save(f"{exp_dir}/total_fea.npy", big_npy)
+    n_ivf = min(int(16 * np.sqrt(big_npy.shape[0])), big_npy.shape[0] // 39)
+
+    index = faiss.index_factory(768 if version19 == "v2" else 256, f"IVF{n_ivf},Flat")
+    index.train(big_npy)
+    faiss.write_index(index, f"{exp_dir}/trained_IVF{n_ivf}_Flat_nprobe_1_{exp_dir1}_{version19}.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])
+
+    faiss.write_index(index, f"{exp_dir}/added_IVF{n_ivf}_Flat_nprobe_1_{exp_dir1}_{version19}.index")
+    return f"Indexing completed. Index saved to {exp_dir}/added_IVF{n_ivf}_Flat_nprobe_1_{exp_dir1}_{version19}.index"
+
+
+
 
 def run_inference(model_name, pitch, input_path, f0_method, save_as, index_rate, volume_normalization, consonant_protection):
     # Setting paths for model and index files