Update core.py

core.py

@@ -1,4 +1,5 @@
 import os
+import spaces
 import sys
 import json
 import argparse
@@ -50,6 +51,7 @@ def get_config():
 
 
 # Infer
+@spaces.GPU(duration=120)
 def run_infer_script(
     pitch: int,
     filter_radius: int,
@@ -101,6 +103,7 @@ def run_infer_script(
 
 
 # Batch infer
+@spaces.GPU(duration=200)
 def run_batch_infer_script(
     pitch: int,
     filter_radius: int,
@@ -167,6 +170,7 @@ def run_batch_infer_script(
 
 
 # TTS
+@spaces.GPU(duration=120)
 def run_tts_script(
     tts_text: str,
     tts_voice: str,
@@ -242,6 +246,7 @@ def run_tts_script(
 
 
 # Preprocess
+@spaces.GPU(duration=360)
 def run_preprocess_script(
     model_name: str, dataset_path: str, sample_rate: int, cpu_cores: int
 ):
@@ -268,6 +273,7 @@ def run_preprocess_script(
 
 
 # Extract
+@spaces.GPU(duration=360)
 def run_extract_script(
     model_name: str,
     rvc_version: str,
@@ -280,13 +286,16 @@ def run_extract_script(
     embedder_model: str,
     embedder_model_custom: str = None,
 ):
-
+    config = get_config()
     model_path = os.path.join(logs_path, model_name)
-    extract = os.path.join("rvc", "train", "extract", "extract.py")
+    pitch_extractor = os.path.join("rvc", "train", "extract", "pitch_extractor.py")
+    embedding_extractor = os.path.join(
+        "rvc", "train", "extract", "embedding_extractor.py"
+    )
 
     command_1 = [
         python,
-        extract,
+        pitch_extractor,
         *map(
             str,
             [
@@ -295,14 +304,26 @@ def run_extract_script(
                 hop_length,
                 cpu_cores,
                 gpu,
+            ],
+        ),
+    ]
+
+    command_2 = [
+        python,
+        embedding_extractor,
+        *map(
+            str,
+            [
+                model_path,
                 rvc_version,
+                gpu,
                 embedder_model,
                 embedder_model_custom,
             ],
         ),
     ]
-
     subprocess.run(command_1)
+    subprocess.run(command_2)
 
     generate_config(rvc_version, sample_rate, model_path)
     generate_filelist(pitch_guidance, model_path, rvc_version, sample_rate)
@@ -310,6 +331,7 @@ def run_extract_script(
 
 
 # Train
+@spaces.GPU(duration=360)
 def run_train_script(
     model_name: str,
     rvc_version: str,
@@ -325,7 +347,6 @@ def run_train_script(
     overtraining_threshold: int,
     pretrained: bool,
     sync_graph: bool,
-    index_algorithm: str,
     cache_data_in_gpu: bool,
     custom_pretrained: bool = False,
     g_pretrained_path: str = None,
@@ -375,19 +396,19 @@ def run_train_script(
         ),
     ]
     subprocess.run(command)
-    run_index_script(model_name, rvc_version, index_algorithm)
+    run_index_script(model_name, rvc_version)
     return f"Model {model_name} trained successfully."
 
 
 # Index
-def run_index_script(model_name: str, rvc_version: str, index_algorithm: str):
+@spaces.GPU
+def run_index_script(model_name: str, rvc_version: str):
     index_script_path = os.path.join("rvc", "train", "process", "extract_index.py")
     command = [
         python,
         index_script_path,
         os.path.join(logs_path, model_name),
         rvc_version,
-        index_algorithm,
     ]
 
     subprocess.run(command)
@@ -395,6 +416,7 @@ def run_index_script(model_name: str, rvc_version: str, index_algorithm: str):
 
 
 # Model extract
+@spaces.GPU
 def run_model_extract_script(
     pth_path: str,
     model_name: str,
@@ -411,6 +433,7 @@ def run_model_extract_script(
 
 
 # Model information
+@spaces.GPU
 def run_model_information_script(pth_path: str):
     print(model_information(pth_path))
 
@@ -424,6 +447,7 @@ def run_model_blender_script(
 
 
 # Tensorboard
+@spaces.GPU
 def run_tensorboard_script():
     launch_tensorboard_pipeline()
 
@@ -1193,14 +1217,6 @@ def parse_arguments():
         choices=["v1", "v2"],
         default="v2",
     )
-    index_parser.add_argument(
-        "--index_algorithm",
-        type=str,
-        choices=["Auto", "Faiss", "KMeans"],
-        help="Choose the method for generating the index file.",
-        default="Auto",
-        required=False,
-    )
 
     # Parser for 'model_extract' mode
     model_extract_parser = subparsers.add_parser(
@@ -1473,7 +1489,6 @@ def main():
             run_index_script(
                 model_name=args.model_name,
                 rvc_version=args.rvc_version,
-                index_algorithm=args.index_algorithm,
             )
         elif args.mode == "model_extract":
             run_model_extract_script(
@@ -1527,4 +1542,4 @@ def main():
 
 
 if __name__ == "__main__":
-    main()
+    main()