fix(preprocess): pass device

Author: magic-akari

preprocess_hubert_f0.py

@@ -103,7 +103,7 @@ def process_one(filename, hmodel,f0p,rank,diff=False,mel_extractor=None):
         if not os.path.exists(aug_vol_path):
             np.save(aug_vol_path,aug_vol.to('cpu').numpy())
 
-def process_batch(file_chunk, f0p, diff=False, mel_extractor=None):
+def process_batch(file_chunk, f0p, diff=False, mel_extractor=None, device="cpu"):
     logger.info("Loading speech encoder for content...")
     rank = mp.current_process()._identity
     rank = rank[0] if len(rank) > 0 else 0
@@ -116,19 +116,20 @@ def process_batch(file_chunk, f0p, diff=False, mel_extractor=None):
     for filename in tqdm(file_chunk):
         process_one(filename, hmodel, f0p, gpu_id, diff, mel_extractor)
 
-def parallel_process(filenames, num_processes, f0p, diff, mel_extractor):
+def parallel_process(filenames, num_processes, f0p, diff, mel_extractor, device):
     with ProcessPoolExecutor(max_workers=num_processes) as executor:
         tasks = []
         for i in range(num_processes):
             start = int(i * len(filenames) / num_processes)
             end = int((i + 1) * len(filenames) / num_processes)
             file_chunk = filenames[start:end]
-            tasks.append(executor.submit(process_batch, file_chunk, f0p, diff, mel_extractor))
+            tasks.append(executor.submit(process_batch, file_chunk, f0p, diff, mel_extractor, device=device))
         for task in tqdm(tasks):
             task.result()
 
 if __name__ == "__main__":
     parser = argparse.ArgumentParser()
+    parser.add_argument('-d', '--device', type=str, default=None)
     parser.add_argument(
         "--in_dir", type=str, default="dataset/44k", help="path to input dir"
     )
@@ -143,15 +144,19 @@ def parallel_process(filenames, num_processes, f0p, diff, mel_extractor):
     )
     args = parser.parse_args()
     f0p = args.f0_predictor
-    print(speech_encoder)
-    logger.info("Using " + speech_encoder + " SpeechEncoder")
-    logger.info("Using " + f0p + "f0 extractor")
-    logger.info("Using diff Mode:")
-    print(args.use_diff)
+    device = args.device
+    if device is None:
+        device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
+
+    logger.info("Using device: " + device)
+    logger.info("Using SpeechEncoder: " + speech_encoder)
+    logger.info("Using extractor:" + f0p)
+    logger.info("Using diff Mode: " + args.use_diff)
+
     if args.use_diff:
         print("use_diff")
         print("Loading Mel Extractor...")
-        mel_extractor = Vocoder(dconfig.vocoder.type, dconfig.vocoder.ckpt, device = "cuda:0")
+        mel_extractor = Vocoder(dconfig.vocoder.type, dconfig.vocoder.ckpt, device=device)
         print("Loaded Mel Extractor.")
     else:
         mel_extractor = None
@@ -162,5 +167,5 @@ def parallel_process(filenames, num_processes, f0p, diff, mel_extractor):
     num_processes = args.num_processes
     if num_processes == 0:
         num_processes = os.cpu_count()
-    
-    parallel_process(filenames, num_processes, f0p, args.use_diff, mel_extractor)
+
+    parallel_process(filenames, num_processes, f0p, args.use_diff, mel_extractor, device)