Upload 179 files

preprocess_flist_config.py

@@ -5,12 +5,11 @@ import re
 import wave
 from random import shuffle
 
+from loguru import logger
 from tqdm import tqdm
 
 import diffusion.logger.utils as du
 
-config_template = json.load(open("configs_template/config_template.json"))
-
 pattern = re.compile(r'^[\.a-zA-Z0-9_\/]+$')
 
 def get_wav_duration(file_path):
@@ -30,13 +29,16 @@ if __name__ == "__main__":
     parser.add_argument("--source_dir", type=str, default="./dataset/44k", help="path to source dir")
     parser.add_argument("--speech_encoder", type=str, default="vec768l12", help="choice a speech encoder|'vec768l12','vec256l9','hubertsoft','whisper-ppg','cnhubertlarge','dphubert','whisper-ppg-large','wavlmbase+'")
     parser.add_argument("--vol_aug", action="store_true", help="Whether to use volume embedding and volume augmentation")
+    parser.add_argument("--tiny", action="store_true", help="Whether to train sovits tiny")
     args = parser.parse_args()
     
+    config_template =  json.load(open("configs_template/config_tiny_template.json")) if args.tiny else json.load(open("configs_template/config_template.json"))
     train = []
     val = []
     idx = 0
     spk_dict = {}
     spk_id = 0
+
     for speaker in tqdm(os.listdir(args.source_dir)):
         spk_dict[speaker] = spk_id
         spk_id += 1
@@ -46,9 +48,9 @@ if __name__ == "__main__":
             if not file.endswith("wav"):
                 continue
             if not pattern.match(file):
-                print(f"warning：文件名{file}中包含非字母数字下划线，可能会导致错误。（也可能不会）")
+                logger.warning(f"文件名{file}中包含非字母数字下划线，可能会导致错误。（也可能不会）")
             if get_wav_duration(file) < 0.3:
-                print("skip too short audio:", file)
+                logger.info("Skip too short audio:" + file)
                 continue
             new_wavs.append(file)
         wavs = new_wavs
@@ -59,13 +61,13 @@ if __name__ == "__main__":
     shuffle(train)
     shuffle(val)
             
-    print("Writing", args.train_list)
+    logger.info("Writing" + args.train_list)
     with open(args.train_list, "w") as f:
         for fname in tqdm(train):
             wavpath = fname
             f.write(wavpath + "\n")
         
-    print("Writing", args.val_list)
+    logger.info("Writing" + args.val_list)
     with open(args.val_list, "w") as f:
         for fname in tqdm(val):
             wavpath = fname
@@ -85,7 +87,7 @@ if __name__ == "__main__":
         config_template["model"]["ssl_dim"] = config_template["model"]["filter_channels"] = config_template["model"]["gin_channels"] = 768
         d_config_template["data"]["encoder_out_channels"] = 768
     elif args.speech_encoder == "vec256l9" or args.speech_encoder == 'hubertsoft':
-        config_template["model"]["ssl_dim"] = config_template["model"]["filter_channels"] = config_template["model"]["gin_channels"] = 256
+        config_template["model"]["ssl_dim"] = config_template["model"]["gin_channels"] = 256
         d_config_template["data"]["encoder_out_channels"] = 256
     elif args.speech_encoder == "whisper-ppg" or args.speech_encoder == 'cnhubertlarge':
         config_template["model"]["ssl_dim"] = config_template["model"]["filter_channels"] = config_template["model"]["gin_channels"] = 1024
@@ -97,8 +99,11 @@ if __name__ == "__main__":
     if args.vol_aug:
         config_template["train"]["vol_aug"] = config_template["model"]["vol_embedding"] = True
 
-    print("Writing configs/config.json")
+    if args.tiny:
+        config_template["model"]["filter_channels"] = 512
+
+    logger.info("Writing to configs/config.json")
     with open("configs/config.json", "w") as f:
         json.dump(config_template, f, indent=2)
-    print("Writing configs/diffusion.yaml")
+    logger.info("Writing to configs/diffusion.yaml")
     du.save_config("configs/diffusion.yaml",d_config_template)

preprocess_hubert_f0.py

@@ -1,6 +1,5 @@
 import argparse
 import logging
-import multiprocessing
 import os
 import random
 from concurrent.futures import ProcessPoolExecutor
@@ -10,6 +9,8 @@ from random import shuffle
 import librosa
 import numpy as np
 import torch
+import torch.multiprocessing as mp
+from loguru import logger
 from tqdm import tqdm
 
 import diffusion.logger.utils as du
@@ -27,13 +28,10 @@ hop_length = hps.data.hop_length
 speech_encoder = hps["model"]["speech_encoder"]
 
 
-def process_one(filename, hmodel,f0p,diff=False,mel_extractor=None):
-    # print(filename)
+def process_one(filename, hmodel, f0p, device, diff=False, mel_extractor=None):
     wav, sr = librosa.load(filename, sr=sampling_rate)
     audio_norm = torch.FloatTensor(wav)
     audio_norm = audio_norm.unsqueeze(0)
-    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
-
     soft_path = filename + ".soft.pt"
     if not os.path.exists(soft_path):
         wav16k = librosa.resample(wav, orig_sr=sampling_rate, target_sr=16000)
@@ -104,29 +102,34 @@ def process_one(filename, hmodel,f0p,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):
-    print("Loading speech encoder for content...")
-    device = "cuda" if torch.cuda.is_available() else "cpu"
-    hmodel = utils.get_speech_encoder(speech_encoder, device=device)
-    print("Loaded speech encoder.")
 
+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
+    if torch.cuda.is_available():
+        gpu_id = rank % torch.cuda.device_count()
+        device = torch.device(f"cuda:{gpu_id}")
+    logger.info(f"Rank {rank} uses device {device}")
+    hmodel = utils.get_speech_encoder(speech_encoder, device=device)
+    logger.info(f"Loaded speech encoder for rank {rank}")
     for filename in tqdm(file_chunk):
-        process_one(filename, hmodel, f0p, diff, mel_extractor)
+        process_one(filename, hmodel, f0p, device, 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"
     )
@@ -134,30 +137,36 @@ if __name__ == "__main__":
         '--use_diff',action='store_true', help='Whether to use the diffusion model'
     )
     parser.add_argument(
-        '--f0_predictor', type=str, default="dio", help='Select F0 predictor, can select crepe,pm,dio,harvest,rmvpe, default pm(note: crepe is original F0 using mean filter)'
+        '--f0_predictor', type=str, default="dio", help='Select F0 predictor, can select crepe,pm,dio,harvest,rmvpe,fcpe|default: pm(note: crepe is original F0 using mean filter)'
     )
     parser.add_argument(
         '--num_processes', type=int, default=1, help='You are advised to set the number of processes to the same as the number of CPU cores'
     )
-    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
     args = parser.parse_args()
     f0p = args.f0_predictor
+    device = args.device
+    if device is None:
+        device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
+
     print(speech_encoder)
-    print(f0p)
-    print(args.use_diff)
+    logger.info("Using device: ", device)
+    logger.info("Using SpeechEncoder: " + speech_encoder)
+    logger.info("Using extractor: " + f0p)
+    logger.info("Using diff Mode: " + str( args.use_diff))
+
     if args.use_diff:
         print("use_diff")
         print("Loading Mel Extractor...")
-        mel_extractor = Vocoder(dconfig.vocoder.type, dconfig.vocoder.ckpt, device = device)
+        mel_extractor = Vocoder(dconfig.vocoder.type, dconfig.vocoder.ckpt, device=device)
         print("Loaded Mel Extractor.")
     else:
         mel_extractor = None
     filenames = glob(f"{args.in_dir}/*/*.wav", recursive=True)  # [:10]
     shuffle(filenames)
-    multiprocessing.set_start_method("spawn", force=True)
+    mp.set_start_method("spawn", force=True)
 
     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)

resample.py

@@ -6,8 +6,8 @@ from multiprocessing import cpu_count
 
 import librosa
 import numpy as np
+from rich.progress import track
 from scipy.io import wavfile
-from tqdm import tqdm
 
 
 def load_wav(wav_path):
@@ -81,7 +81,7 @@ def process_all_speakers():
             if os.path.isdir(spk_dir):
                 print(spk_dir)
                 futures = [executor.submit(process, (spk_dir, i, args)) for i in os.listdir(spk_dir) if i.endswith("wav")]
-                for _ in tqdm(concurrent.futures.as_completed(futures), total=len(futures)):
+                for _ in track(concurrent.futures.as_completed(futures), total=len(futures), description="resampling:"):
                     pass