fix(pre_hubert): fix vram estimiation

voicepaw · Apr 3, 2023 · 17fa86c · 17fa86c
1 parent 229f9ba
commit 17fa86c
Showing 1 changed file with 5 additions and 2 deletions.
diff --git a/src/so_vits_svc_fork/preprocessing/preprocess_hubert_f0.py b/src/so_vits_svc_fork/preprocessing/preprocess_hubert_f0.py
@@ -18,7 +18,8 @@
 from .preprocess_utils import check_hubert_min_duration
 
 LOG = getLogger(__name__)
-HUBERT_MEMORY = 1250
+HUBERT_MEMORY = 1600
+HUBERT_MEMORY_CREPE = 2600
 
 
 def _process_one(
@@ -99,7 +100,9 @@ def preprocess_hubert_f0(
     hps = utils.get_hparams(config_path)
     if n_jobs is None:
         memory = get_total_gpu_memory("free")
-        n_jobs = memory // HUBERT_MEMORY
+        n_jobs = memory // (
+            HUBERT_MEMORY_CREPE if f0_method == "crepe" else HUBERT_MEMORY
+        )
         LOG.info(f"n_jobs automatically set to {n_jobs}, memory: {memory} MiB")
 
     filepaths = list(input_dir.rglob("*.wav"))