pony123/ChatGLM2-Voice-Cloning

1from multiprocess.pool import ThreadPool2from speaker_encoder.params_data import *3from speaker_encoder.config import librispeech_datasets, anglophone_nationalites4from datetime import datetime5from speaker_encoder import audio6from pathlib import Path7from tqdm import tqdm8import numpy as np9 10 11class DatasetLog:12    """13    Registers metadata about the dataset in a text file.14    """15    def __init__(self, root, name):16        self.text_file = open(Path(root, "Log_%s.txt" % name.replace("/", "_")), "w")17        self.sample_data = dict()18        19        start_time = str(datetime.now().strftime("%A %d %B %Y at %H:%M"))20        self.write_line("Creating dataset %s on %s" % (name, start_time))21        self.write_line("-----")22        self._log_params()23        24    def _log_params(self):25        from speaker_encoder import params_data26        self.write_line("Parameter values:")27        for param_name in (p for p in dir(params_data) if not p.startswith("__")):28            value = getattr(params_data, param_name)29            self.write_line("\t%s: %s" % (param_name, value))30        self.write_line("-----")31    32    def write_line(self, line):33        self.text_file.write("%s\n" % line)34        35    def add_sample(self, **kwargs):36        for param_name, value in kwargs.items():37            if not param_name in self.sample_data:38                self.sample_data[param_name] = []39            self.sample_data[param_name].append(value)40            41    def finalize(self):42        self.write_line("Statistics:")43        for param_name, values in self.sample_data.items():44            self.write_line("\t%s:" % param_name)45            self.write_line("\t\tmin %.3f, max %.3f" % (np.min(values), np.max(values)))46            self.write_line("\t\tmean %.3f, median %.3f" % (np.mean(values), np.median(values)))47        self.write_line("-----")48        end_time = str(datetime.now().strftime("%A %d %B %Y at %H:%M"))49        self.write_line("Finished on %s" % end_time)50        self.text_file.close()51       52        53def _init_preprocess_dataset(dataset_name, datasets_root, out_dir) -> (Path, DatasetLog):54    dataset_root = datasets_root.joinpath(dataset_name)55    if not dataset_root.exists():56        print("Couldn\'t find %s, skipping this dataset." % dataset_root)57        return None, None58    return dataset_root, DatasetLog(out_dir, dataset_name)59 60 61def _preprocess_speaker_dirs(speaker_dirs, dataset_name, datasets_root, out_dir, extension,62                             skip_existing, logger):63    print("%s: Preprocessing data for %d speakers." % (dataset_name, len(speaker_dirs)))64    65    # Function to preprocess utterances for one speaker66    def preprocess_speaker(speaker_dir: Path):67        # Give a name to the speaker that includes its dataset68        speaker_name = "_".join(speaker_dir.relative_to(datasets_root).parts)69        70        # Create an output directory with that name, as well as a txt file containing a 71        # reference to each source file.72        speaker_out_dir = out_dir.joinpath(speaker_name)73        speaker_out_dir.mkdir(exist_ok=True)74        sources_fpath = speaker_out_dir.joinpath("_sources.txt")75        76        # There's a possibility that the preprocessing was interrupted earlier, check if 77        # there already is a sources file.78        if sources_fpath.exists():79            try:80                with sources_fpath.open("r") as sources_file:81                    existing_fnames = {line.split(",")[0] for line in sources_file}82            except:83                existing_fnames = {}84        else:85            existing_fnames = {}86        87        # Gather all audio files for that speaker recursively88        sources_file = sources_fpath.open("a" if skip_existing else "w")89        for in_fpath in speaker_dir.glob("**/*.%s" % extension):90            # Check if the target output file already exists91            out_fname = "_".join(in_fpath.relative_to(speaker_dir).parts)92            out_fname = out_fname.replace(".%s" % extension, ".npy")93            if skip_existing and out_fname in existing_fnames:94                continue95                96            # Load and preprocess the waveform97            wav = audio.preprocess_wav(in_fpath)98            if len(wav) == 0:99                continue100            101            # Create the mel spectrogram, discard those that are too short102            frames = audio.wav_to_mel_spectrogram(wav)103            if len(frames) < partials_n_frames:104                continue105            106            out_fpath = speaker_out_dir.joinpath(out_fname)107            np.save(out_fpath, frames)108            logger.add_sample(duration=len(wav) / sampling_rate)109            sources_file.write("%s,%s\n" % (out_fname, in_fpath))110        111        sources_file.close()112    113    # Process the utterances for each speaker114    with ThreadPool(8) as pool:115        list(tqdm(pool.imap(preprocess_speaker, speaker_dirs), dataset_name, len(speaker_dirs),116                  unit="speakers"))117    logger.finalize()118    print("Done preprocessing %s.\n" % dataset_name)119 120 121# Function to preprocess utterances for one speaker122def __preprocess_speaker(speaker_dir: Path, datasets_root: Path, out_dir: Path, extension: str, skip_existing: bool):123        # Give a name to the speaker that includes its dataset124        speaker_name = "_".join(speaker_dir.relative_to(datasets_root).parts)125        126        # Create an output directory with that name, as well as a txt file containing a 127        # reference to each source file.128        speaker_out_dir = out_dir.joinpath(speaker_name)129        speaker_out_dir.mkdir(exist_ok=True)130        sources_fpath = speaker_out_dir.joinpath("_sources.txt")131        132        # There's a possibility that the preprocessing was interrupted earlier, check if 133        # there already is a sources file.134        # if sources_fpath.exists():135        #     try:136        #         with sources_fpath.open("r") as sources_file:137        #             existing_fnames = {line.split(",")[0] for line in sources_file}138        #     except:139        #         existing_fnames = {}140        # else:141        #     existing_fnames = {}142        existing_fnames = {}143        # Gather all audio files for that speaker recursively144        sources_file = sources_fpath.open("a" if skip_existing else "w")145 146        for in_fpath in speaker_dir.glob("**/*.%s" % extension):147            # Check if the target output file already exists148            out_fname = "_".join(in_fpath.relative_to(speaker_dir).parts)149            out_fname = out_fname.replace(".%s" % extension, ".npy")150            if skip_existing and out_fname in existing_fnames:151                continue152                153            # Load and preprocess the waveform154            wav = audio.preprocess_wav(in_fpath)155            if len(wav) == 0:156                continue157            158            # Create the mel spectrogram, discard those that are too short159            frames = audio.wav_to_mel_spectrogram(wav)160            if len(frames) < partials_n_frames:161                continue162            163            out_fpath = speaker_out_dir.joinpath(out_fname)164            np.save(out_fpath, frames)165            # logger.add_sample(duration=len(wav) / sampling_rate)166            sources_file.write("%s,%s\n" % (out_fname, in_fpath))167        168        sources_file.close()169        return len(wav)170 171def _preprocess_speaker_dirs_vox2(speaker_dirs, dataset_name, datasets_root, out_dir, extension,172                             skip_existing, logger):173    # from multiprocessing import Pool, cpu_count174    from pathos.multiprocessing import ProcessingPool as Pool175    # Function to preprocess utterances for one speaker176    def __preprocess_speaker(speaker_dir: Path):177        # Give a name to the speaker that includes its dataset178        speaker_name = "_".join(speaker_dir.relative_to(datasets_root).parts)179        180        # Create an output directory with that name, as well as a txt file containing a 181        # reference to each source file.182        speaker_out_dir = out_dir.joinpath(speaker_name)183        speaker_out_dir.mkdir(exist_ok=True)184        sources_fpath = speaker_out_dir.joinpath("_sources.txt")185        186        existing_fnames = {}187        # Gather all audio files for that speaker recursively188        sources_file = sources_fpath.open("a" if skip_existing else "w")189        wav_lens = []190        for in_fpath in speaker_dir.glob("**/*.%s" % extension):191            # Check if the target output file already exists192            out_fname = "_".join(in_fpath.relative_to(speaker_dir).parts)193            out_fname = out_fname.replace(".%s" % extension, ".npy")194            if skip_existing and out_fname in existing_fnames:195                continue196                197            # Load and preprocess the waveform198            wav = audio.preprocess_wav(in_fpath)199            if len(wav) == 0:200                continue201            202            # Create the mel spectrogram, discard those that are too short203            frames = audio.wav_to_mel_spectrogram(wav)204            if len(frames) < partials_n_frames:205                continue206            207            out_fpath = speaker_out_dir.joinpath(out_fname)208            np.save(out_fpath, frames)209            # logger.add_sample(duration=len(wav) / sampling_rate)210            sources_file.write("%s,%s\n" % (out_fname, in_fpath))211            wav_lens.append(len(wav))212        sources_file.close()213        return wav_lens214 215    print("%s: Preprocessing data for %d speakers." % (dataset_name, len(speaker_dirs)))216    # Process the utterances for each speaker217    # with ThreadPool(8) as pool:218    #     list(tqdm(pool.imap(preprocess_speaker, speaker_dirs), dataset_name, len(speaker_dirs),219    #               unit="speakers"))220    pool = Pool(processes=20)221    for i, wav_lens in enumerate(pool.map(__preprocess_speaker, speaker_dirs), 1):222        for wav_len in wav_lens:223            logger.add_sample(duration=wav_len / sampling_rate)224        print(f'{i}/{len(speaker_dirs)} \r')225 226    logger.finalize()227    print("Done preprocessing %s.\n" % dataset_name)228 229 230def preprocess_librispeech(datasets_root: Path, out_dir: Path, skip_existing=False):231    for dataset_name in librispeech_datasets["train"]["other"]:232        # Initialize the preprocessing233        dataset_root, logger = _init_preprocess_dataset(dataset_name, datasets_root, out_dir)234        if not dataset_root:235            return 236        237        # Preprocess all speakers238        speaker_dirs = list(dataset_root.glob("*"))239        _preprocess_speaker_dirs(speaker_dirs, dataset_name, datasets_root, out_dir, "flac",240                                 skip_existing, logger)241 242 243def preprocess_voxceleb1(datasets_root: Path, out_dir: Path, skip_existing=False):244    # Initialize the preprocessing245    dataset_name = "VoxCeleb1"246    dataset_root, logger = _init_preprocess_dataset(dataset_name, datasets_root, out_dir)247    if not dataset_root:248        return249 250    # Get the contents of the meta file251    with dataset_root.joinpath("vox1_meta.csv").open("r") as metafile:252        metadata = [line.split("\t") for line in metafile][1:]253    254    # Select the ID and the nationality, filter out non-anglophone speakers255    nationalities = {line[0]: line[3] for line in metadata}256    # keep_speaker_ids = [speaker_id for speaker_id, nationality in nationalities.items() if 257    #                     nationality.lower() in anglophone_nationalites]258    keep_speaker_ids = [speaker_id for speaker_id, nationality in nationalities.items()]                        259    print("VoxCeleb1: using samples from %d (presumed anglophone) speakers out of %d." % 260          (len(keep_speaker_ids), len(nationalities)))261    262    # Get the speaker directories for anglophone speakers only263    speaker_dirs = dataset_root.joinpath("wav").glob("*")264    speaker_dirs = [speaker_dir for speaker_dir in speaker_dirs if265                    speaker_dir.name in keep_speaker_ids]266    print("VoxCeleb1: found %d anglophone speakers on the disk, %d missing (this is normal)." % 267          (len(speaker_dirs), len(keep_speaker_ids) - len(speaker_dirs)))268 269    # Preprocess all speakers270    _preprocess_speaker_dirs(speaker_dirs, dataset_name, datasets_root, out_dir, "wav",271                             skip_existing, logger)272 273 274def preprocess_voxceleb2(datasets_root: Path, out_dir: Path, skip_existing=False):275    # Initialize the preprocessing276    dataset_name = "VoxCeleb2"277    dataset_root, logger = _init_preprocess_dataset(dataset_name, datasets_root, out_dir)278    if not dataset_root:279        return280    281    # Get the speaker directories282    # Preprocess all speakers283    speaker_dirs = list(dataset_root.joinpath("dev", "aac").glob("*"))284    _preprocess_speaker_dirs_vox2(speaker_dirs, dataset_name, datasets_root, out_dir, "m4a",285                             skip_existing, logger)286