Change to use thread pool executor for multiprocess

bert_gen.py

@@ -1,5 +1,5 @@
 import argparse
-from multiprocessing import Pool
+from concurrent.futures import ThreadPoolExecutor
 
 import torch
 import torch.multiprocessing as mp
@@ -73,13 +73,13 @@ if __name__ == "__main__":
 
     if len(lines) != 0:
         num_processes = args.num_processes
-        with Pool(processes=num_processes) as pool:
-            for _ in tqdm(
-                pool.imap_unordered(process_line, zip(lines, add_blank)),
-                total=len(lines),
-                file=SAFE_STDOUT,
-            ):
-                # 这里是缩进的代码块，表示循环体
-                pass  # 使用pass语句作为占位符
+        with ThreadPoolExecutor(max_workers=num_processes) as executor:
+            _ = list(
+                tqdm(
+                    executor.map(process_line, zip(lines, add_blank)),
+                    total=len(lines),
+                    file=SAFE_STDOUT,
+                )
+            )
 
     logger.info(f"bert.pt is generated! total: {len(lines)} bert.pt files.")

colab.ipynb

@@ -42,7 +42,7 @@
         "%cd Style-Bert-VITS2/\n",
         "!pip install -r requirements.txt\n",
         "!apt install libcublas11\n",
-        "!python initialize.py"
+        "!python initialize.py --skip_jvnv"
       ]
     },
     {

resample.py

@@ -1,6 +1,6 @@
 import argparse
 import os
-from multiprocessing import Pool, cpu_count
+from concurrent.futures import ThreadPoolExecutor
 
 import librosa
 import pyloudnorm as pyln
@@ -110,13 +110,22 @@ if __name__ == "__main__":
         logger.error(f"No wav files found in {args.in_dir}")
         raise ValueError(f"No wav files found in {args.in_dir}")
 
-    pool = Pool(processes=processes)
-    for _ in tqdm(
-        pool.imap_unordered(process, tasks), file=SAFE_STDOUT, total=len(tasks)
-    ):
-        pass
+    # pool = Pool(processes=processes)
+    # for _ in tqdm(
+    #     pool.imap_unordered(process, tasks), file=SAFE_STDOUT, total=len(tasks)
+    # ):
+    #     pass
 
-    pool.close()
-    pool.join()
+    # pool.close()
+    # pool.join()
+
+    with ThreadPoolExecutor(max_workers=processes) as executor:
+        _ = list(
+            tqdm(
+                executor.map(process, tasks),
+                total=len(tasks),
+                file=SAFE_STDOUT,
+            )
+        )
 
     logger.info("Resampling Done!")