From f1429242b2f2108049fa0ed14f0b4efe0ad5c2fb Mon Sep 17 00:00:00 2001
From: litagin02 <litagin02@gmail.com>
Date: Wed, 7 Feb 2024 19:54:38 +0900
Subject: [PATCH] Change to use thread pool executor for multiprocess

---
 bert_gen.py | 18 +++++++++---------
 colab.ipynb |  2 +-
 resample.py | 25 +++++++++++++++++--------
 3 files changed, 27 insertions(+), 18 deletions(-)

diff --git a/bert_gen.py b/bert_gen.py
index bde90f9..a5f7c25 100644
--- a/bert_gen.py
+++ b/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.")
diff --git a/colab.ipynb b/colab.ipynb
index fb1a401..370ca15 100644
--- a/colab.ipynb
+++ b/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"
       ]
     },
     {
diff --git a/resample.py b/resample.py
index f34ca71..5aff5ef 100644
--- a/resample.py
+++ b/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!")