Upload 22 files

auto_slicer.py

@@ -0,0 +1,107 @@
+import os
+from typing import final
+import numpy as np
+import librosa
+import soundfile as sf
+from modules.slicer2 import Slicer
+
+class AutoSlicer:
+    def __init__(self):
+        self.slicer_params = {
+            "threshold": -40,
+            "min_length": 5000,
+            "min_interval": 300,
+            "hop_size": 10,
+            "max_sil_kept": 500,
+        }
+        self.original_min_interval = self.slicer_params["min_interval"]
+
+    def auto_slice(self, filename, input_dir, output_dir, max_sec):
+        audio, sr = librosa.load(os.path.join(input_dir, filename), sr=None, mono=False)
+        slicer = Slicer(sr=sr, **self.slicer_params)
+        chunks = slicer.slice(audio)
+        files_to_delete = []
+        for i, chunk in enumerate(chunks):
+            if len(chunk.shape) > 1:
+                chunk = chunk.T
+            output_filename = f"{os.path.splitext(filename)[0]}_{i}"
+            output_filename = "".join(c for c in output_filename if c.isascii() or c == "_") + ".wav"
+            output_filepath = os.path.join(output_dir, output_filename)
+            sf.write(output_filepath, chunk, sr)
+            #Check and re-slice audio that more than max_sec.
+            while True:
+                new_audio, sr = librosa.load(output_filepath, sr=None, mono=False)
+                if librosa.get_duration(y=new_audio, sr=sr) <= max_sec:
+                    break
+                self.slicer_params["min_interval"] = self.slicer_params["min_interval"] // 2
+                if self.slicer_params["min_interval"] >= self.slicer_params["hop_size"]:
+                    new_chunks = Slicer(sr=sr, **self.slicer_params).slice(new_audio)
+                    for j, new_chunk in enumerate(new_chunks):
+                        if len(new_chunk.shape) > 1:
+                            new_chunk = new_chunk.T
+                        new_output_filename = f"{os.path.splitext(output_filename)[0]}_{j}.wav"
+                        sf.write(os.path.join(output_dir, new_output_filename), new_chunk, sr)
+                    files_to_delete.append(output_filepath)
+                else:
+                    break
+            self.slicer_params["min_interval"] = self.original_min_interval
+        for file_path in files_to_delete:
+            if os.path.exists(file_path):
+                os.remove(file_path)
+
+    def merge_short(self, output_dir, max_sec, min_sec):
+        short_files = []
+        for filename in os.listdir(output_dir):
+            filepath = os.path.join(output_dir, filename)
+            if filename.endswith(".wav"):
+                audio, sr = librosa.load(filepath, sr=None, mono=False)
+                duration = librosa.get_duration(y=audio, sr=sr)
+                if duration < min_sec:
+                    short_files.append((filepath, audio, duration))
+        short_files.sort(key=lambda x: x[2], reverse=True)
+        merged_audio = []
+        current_duration = 0
+        for filepath, audio, duration in short_files:
+            if current_duration + duration <= max_sec:
+                merged_audio.append(audio)
+                current_duration += duration
+                os.remove(filepath)
+            else:
+                if merged_audio:
+                    output_audio = np.concatenate(merged_audio, axis=-1)
+                    if len(output_audio.shape) > 1:
+                        output_audio = output_audio.T
+                    output_filename = f"merged_{len(os.listdir(output_dir))}.wav"
+                    sf.write(os.path.join(output_dir, output_filename), output_audio, sr)
+                    merged_audio = [audio]
+                    current_duration = duration
+                    os.remove(filepath)
+        if merged_audio and current_duration >= min_sec:
+            output_audio = np.concatenate(merged_audio, axis=-1)
+            if len(output_audio.shape) > 1:
+                output_audio = output_audio.T
+            output_filename = f"merged_{len(os.listdir(output_dir))}.wav"
+            sf.write(os.path.join(output_dir, output_filename), output_audio, sr)
+    
+    def slice_count(self, input_dir, output_dir):
+        orig_duration = final_duration = 0
+        for file in os.listdir(input_dir):
+            if file.endswith(".wav"):
+                _audio, _sr = librosa.load(os.path.join(input_dir, file), sr=None, mono=False)
+                orig_duration += librosa.get_duration(y=_audio, sr=_sr)
+        wav_files = [file for file in os.listdir(output_dir) if file.endswith(".wav")]
+        num_files = len(wav_files)
+        max_duration = -1
+        min_duration = float("inf")
+        for file in wav_files:
+            file_path = os.path.join(output_dir, file)
+            audio, sr = librosa.load(file_path, sr=None, mono=False)
+            duration = librosa.get_duration(y=audio, sr=sr)
+            final_duration += float(duration)
+            if duration > max_duration:
+                max_duration = float(duration)
+            if duration < min_duration:
+                min_duration = float(duration)
+        return num_files, max_duration, min_duration, orig_duration, final_duration
+
+

flask_api.py

@@ -7,7 +7,7 @@ import torchaudio
 from flask import Flask, request, send_file
 from flask_cors import CORS
 
-from inference.infer_tool import Svc, RealTimeVC
+from inference.infer_tool import RealTimeVC, Svc
 
 app = Flask(__name__)
 

inference/infer_tool_webui.py

@@ -0,0 +1,537 @@
+import gc
+import hashlib
+import io
+import json
+import logging
+import os
+import pickle
+import time
+from pathlib import Path
+
+import librosa
+import numpy as np
+
+# import onnxruntime
+import soundfile
+import torch
+import torchaudio
+from tqdm import tqdm
+
+import cluster
+import utils
+from diffusion.unit2mel import load_model_vocoder
+from inference import slicer
+from models import SynthesizerTrn
+
+logging.getLogger('matplotlib').setLevel(logging.WARNING)
+
+
+def read_temp(file_name):
+    if not os.path.exists(file_name):
+        with open(file_name, "w") as f:
+            f.write(json.dumps({"info": "temp_dict"}))
+        return {}
+    else:
+        try:
+            with open(file_name, "r") as f:
+                data = f.read()
+            data_dict = json.loads(data)
+            if os.path.getsize(file_name) > 50 * 1024 * 1024:
+                f_name = file_name.replace("\\", "/").split("/")[-1]
+                print(f"clean {f_name}")
+                for wav_hash in list(data_dict.keys()):
+                    if int(time.time()) - int(data_dict[wav_hash]["time"]) > 14 * 24 * 3600:
+                        del data_dict[wav_hash]
+        except Exception as e:
+            print(e)
+            print(f"{file_name} error,auto rebuild file")
+            data_dict = {"info": "temp_dict"}
+        return data_dict
+
+
+def write_temp(file_name, data):
+    with open(file_name, "w") as f:
+        f.write(json.dumps(data))
+
+
+def timeit(func):
+    def run(*args, **kwargs):
+        t = time.time()
+        res = func(*args, **kwargs)
+        print('executing \'%s\' costed %.3fs' % (func.__name__, time.time() - t))
+        return res
+
+    return run
+
+
+def format_wav(audio_path):
+    if Path(audio_path).suffix == '.wav':
+        return
+    raw_audio, raw_sample_rate = librosa.load(audio_path, mono=True, sr=None)
+    soundfile.write(Path(audio_path).with_suffix(".wav"), raw_audio, raw_sample_rate)
+
+
+def get_end_file(dir_path, end):
+    file_lists = []
+    for root, dirs, files in os.walk(dir_path):
+        files = [f for f in files if f[0] != '.']
+        dirs[:] = [d for d in dirs if d[0] != '.']
+        for f_file in files:
+            if f_file.endswith(end):
+                file_lists.append(os.path.join(root, f_file).replace("\\", "/"))
+    return file_lists
+
+
+def get_md5(content):
+    return hashlib.new("md5", content).hexdigest()
+
+def fill_a_to_b(a, b):
+    if len(a) < len(b):
+        for _ in range(0, len(b) - len(a)):
+            a.append(a[0])
+
+def mkdir(paths: list):
+    for path in paths:
+        if not os.path.exists(path):
+            os.mkdir(path)
+
+def pad_array(arr, target_length):
+    current_length = arr.shape[0]
+    if current_length >= target_length:
+        return arr
+    else:
+        pad_width = target_length - current_length
+        pad_left = pad_width // 2
+        pad_right = pad_width - pad_left
+        padded_arr = np.pad(arr, (pad_left, pad_right), 'constant', constant_values=(0, 0))
+        return padded_arr
+    
+def split_list_by_n(list_collection, n, pre=0):
+    for i in range(0, len(list_collection), n):
+        yield list_collection[i-pre if i-pre>=0 else i: i + n]
+
+
+class F0FilterException(Exception):
+    pass
+
+class Svc(object):
+    def __init__(self, net_g_path, config_path,
+                 device=None,
+                 cluster_model_path="logs/44k/kmeans_10000.pt",
+                 nsf_hifigan_enhance = False,
+                 diffusion_model_path="logs/44k/diffusion/model_0.pt",
+                 diffusion_config_path="configs/diffusion.yaml",
+                 shallow_diffusion = False,
+                 only_diffusion = False,
+                 spk_mix_enable = False,
+                 feature_retrieval = False
+                 ):
+        self.net_g_path = net_g_path
+        self.only_diffusion = only_diffusion
+        self.shallow_diffusion = shallow_diffusion
+        self.feature_retrieval = feature_retrieval
+        if device is None:
+            self.dev = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+        else:
+            self.dev = torch.device(device)
+        self.net_g_ms = None
+        if not self.only_diffusion:
+            self.hps_ms = utils.get_hparams_from_file(config_path,True)
+            self.target_sample = self.hps_ms.data.sampling_rate
+            self.hop_size = self.hps_ms.data.hop_length
+            self.spk2id = self.hps_ms.spk
+            self.unit_interpolate_mode = self.hps_ms.data.unit_interpolate_mode if self.hps_ms.data.unit_interpolate_mode is not None else 'left'
+            self.vol_embedding = self.hps_ms.model.vol_embedding if self.hps_ms.model.vol_embedding is not None else False
+            self.speech_encoder = self.hps_ms.model.speech_encoder if self.hps_ms.model.speech_encoder is not None else 'vec768l12'
+ 
+        self.nsf_hifigan_enhance = nsf_hifigan_enhance
+        if self.shallow_diffusion or self.only_diffusion:
+            if os.path.exists(diffusion_model_path) and os.path.exists(diffusion_model_path):
+                self.diffusion_model,self.vocoder,self.diffusion_args = load_model_vocoder(diffusion_model_path,self.dev,config_path=diffusion_config_path)
+                if self.only_diffusion:
+                    self.target_sample = self.diffusion_args.data.sampling_rate
+                    self.hop_size = self.diffusion_args.data.block_size
+                    self.spk2id = self.diffusion_args.spk
+                    self.speech_encoder = self.diffusion_args.data.encoder
+                    self.unit_interpolate_mode = self.diffusion_args.data.unit_interpolate_mode if self.diffusion_args.data.unit_interpolate_mode is not None else 'left'
+                if spk_mix_enable:
+                    self.diffusion_model.init_spkmix(len(self.spk2id))
+            else:
+                print("No diffusion model or config found. Shallow diffusion mode will False")
+                self.shallow_diffusion = self.only_diffusion = False
+                
+        # load hubert and model
+        if not self.only_diffusion:
+            self.load_model(spk_mix_enable)
+            self.hubert_model = utils.get_speech_encoder(self.speech_encoder,device=self.dev)
+            self.volume_extractor = utils.Volume_Extractor(self.hop_size)
+        else:
+            self.hubert_model = utils.get_speech_encoder(self.diffusion_args.data.encoder,device=self.dev)
+            self.volume_extractor = utils.Volume_Extractor(self.diffusion_args.data.block_size)
+            
+        if os.path.exists(cluster_model_path):
+            if self.feature_retrieval:
+                with open(cluster_model_path,"rb") as f:
+                    self.cluster_model = pickle.load(f)
+                self.big_npy = None
+                self.now_spk_id = -1
+            else:
+                self.cluster_model = cluster.get_cluster_model(cluster_model_path)
+        else:
+            self.feature_retrieval=False
+
+        if self.shallow_diffusion :
+            self.nsf_hifigan_enhance = False
+        if self.nsf_hifigan_enhance:
+            from modules.enhancer import Enhancer
+            self.enhancer = Enhancer('nsf-hifigan', 'pretrain/nsf_hifigan/model',device=self.dev)
+            
+    def load_model(self, spk_mix_enable=False):
+        # get model configuration
+        self.net_g_ms = SynthesizerTrn(
+            self.hps_ms.data.filter_length // 2 + 1,
+            self.hps_ms.train.segment_size // self.hps_ms.data.hop_length,
+            **self.hps_ms.model)
+        _ = utils.load_checkpoint(self.net_g_path, self.net_g_ms, None)
+        self.dtype = list(self.net_g_ms.parameters())[0].dtype
+        if "half" in self.net_g_path and torch.cuda.is_available():
+            _ = self.net_g_ms.half().eval().to(self.dev)
+        else:
+            _ = self.net_g_ms.eval().to(self.dev)
+        if spk_mix_enable:
+            self.net_g_ms.EnableCharacterMix(len(self.spk2id), self.dev)
+
+    def get_unit_f0(self, wav, tran, cluster_infer_ratio, speaker, f0_filter ,f0_predictor,cr_threshold=0.05):
+
+        f0_predictor_object = utils.get_f0_predictor(f0_predictor,hop_length=self.hop_size,sampling_rate=self.target_sample,device=self.dev,threshold=cr_threshold)
+        
+        f0, uv = f0_predictor_object.compute_f0_uv(wav)
+        if f0_filter and sum(f0) == 0:
+            raise F0FilterException("No voice detected")
+        f0 = torch.FloatTensor(f0).to(self.dev)
+        uv = torch.FloatTensor(uv).to(self.dev)
+
+        f0 = f0 * 2 ** (tran / 12)
+        f0 = f0.unsqueeze(0)
+        uv = uv.unsqueeze(0)
+
+        wav16k = librosa.resample(wav, orig_sr=self.target_sample, target_sr=16000)
+        wav16k = torch.from_numpy(wav16k).to(self.dev)
+        c = self.hubert_model.encoder(wav16k)
+        c = utils.repeat_expand_2d(c.squeeze(0), f0.shape[1],self.unit_interpolate_mode)
+
+        if cluster_infer_ratio !=0:
+            if self.feature_retrieval:
+                speaker_id = self.spk2id.get(speaker)
+                if speaker_id is None:
+                    raise RuntimeError("The name you entered is not in the speaker list!")
+                if not speaker_id and type(speaker) is int:
+                    if len(self.spk2id.__dict__) >= speaker:
+                        speaker_id = speaker
+                feature_index = self.cluster_model[speaker_id]
+                feat_np = c.transpose(0,1).cpu().numpy()
+                if self.big_npy is None or self.now_spk_id != speaker_id:
+                   self.big_npy = feature_index.reconstruct_n(0, feature_index.ntotal)
+                   self.now_spk_id = speaker_id
+                print("starting feature retrieval...")
+                score, ix = feature_index.search(feat_np, k=8)
+                weight = np.square(1 / score)
+                weight /= weight.sum(axis=1, keepdims=True)
+                npy = np.sum(self.big_npy[ix] * np.expand_dims(weight, axis=2), axis=1)
+                c = cluster_infer_ratio * npy + (1 - cluster_infer_ratio) * feat_np
+                c = torch.FloatTensor(c).to(self.dev).transpose(0,1)
+                print("end feature retrieval...")
+            else:
+                cluster_c = cluster.get_cluster_center_result(self.cluster_model, c.cpu().numpy().T, speaker).T
+                cluster_c = torch.FloatTensor(cluster_c).to(self.dev)
+                c = cluster_infer_ratio * cluster_c + (1 - cluster_infer_ratio) * c
+
+        c = c.unsqueeze(0)
+        return c, f0, uv
+
+    def infer(self, speaker, tran, raw_path,
+              cluster_infer_ratio=0,
+              auto_predict_f0=False,
+              noice_scale=0.4,
+              f0_filter=False,
+              f0_predictor='pm',
+              enhancer_adaptive_key = 0,
+              cr_threshold = 0.05,
+              k_step = 100,
+              frame = 0,
+              spk_mix = False,
+              second_encoding = False,
+              loudness_envelope_adjustment = 1
+              ):
+        wav, sr = librosa.load(raw_path, sr=self.target_sample)
+        if spk_mix:
+            c, f0, uv = self.get_unit_f0(wav, tran, 0, None, f0_filter,f0_predictor,cr_threshold=cr_threshold)
+            n_frames = f0.size(1)
+            sid = speaker[:, frame:frame+n_frames].transpose(0,1)
+        else:
+            speaker_id = self.spk2id.get(speaker)
+            if not speaker_id and type(speaker) is int:
+                if len(self.spk2id.__dict__) >= speaker:
+                    speaker_id = speaker
+            if speaker_id is None:
+                raise RuntimeError("The name you entered is not in the speaker list!")
+            sid = torch.LongTensor([int(speaker_id)]).to(self.dev).unsqueeze(0)
+            c, f0, uv = self.get_unit_f0(wav, tran, cluster_infer_ratio, speaker, f0_filter,f0_predictor,cr_threshold=cr_threshold)
+            n_frames = f0.size(1)
+        c = c.to(self.dtype)
+        f0 = f0.to(self.dtype)
+        uv = uv.to(self.dtype)
+        with torch.no_grad():
+            start = time.time()
+            vol = None
+            if not self.only_diffusion:
+                vol = self.volume_extractor.extract(torch.FloatTensor(wav).to(self.dev)[None,:])[None,:].to(self.dev) if self.vol_embedding else None
+                audio,f0 = self.net_g_ms.infer(c, f0=f0, g=sid, uv=uv, predict_f0=auto_predict_f0, noice_scale=noice_scale,vol=vol)
+                audio = audio[0,0].data.float()
+                audio_mel = self.vocoder.extract(audio[None,:],self.target_sample) if self.shallow_diffusion else None
+            else:
+                audio = torch.FloatTensor(wav).to(self.dev)
+                audio_mel = None
+            if self.dtype != torch.float32:
+                c = c.to(torch.float32)
+                f0 = f0.to(torch.float32)
+                uv = uv.to(torch.float32)
+            if self.only_diffusion or self.shallow_diffusion:
+                vol = self.volume_extractor.extract(audio[None,:])[None,:,None].to(self.dev) if vol is None else vol[:,:,None]
+                if self.shallow_diffusion and second_encoding:
+                    audio16k = librosa.resample(audio.detach().cpu().numpy(), orig_sr=self.target_sample, target_sr=16000)
+                    audio16k = torch.from_numpy(audio16k).to(self.dev)
+                    c = self.hubert_model.encoder(audio16k)
+                    c = utils.repeat_expand_2d(c.squeeze(0), f0.shape[1],self.unit_interpolate_mode)
+                f0 = f0[:,:,None]
+                c = c.transpose(-1,-2)
+                audio_mel = self.diffusion_model(
+                c, 
+                f0, 
+                vol, 
+                spk_id = sid, 
+                spk_mix_dict = None,
+                gt_spec=audio_mel,
+                infer=True, 
+                infer_speedup=self.diffusion_args.infer.speedup, 
+                method=self.diffusion_args.infer.method,
+                k_step=k_step)
+                audio = self.vocoder.infer(audio_mel, f0).squeeze()
+            if self.nsf_hifigan_enhance:
+                audio, _ = self.enhancer.enhance(
+                                    audio[None,:], 
+                                    self.target_sample, 
+                                    f0[:,:,None], 
+                                    self.hps_ms.data.hop_length, 
+                                    adaptive_key = enhancer_adaptive_key)
+            if loudness_envelope_adjustment != 1:
+                audio = utils.change_rms(wav,self.target_sample,audio,self.target_sample,loudness_envelope_adjustment)
+            use_time = time.time() - start
+            print("vits use time:{}".format(use_time))
+        return audio, audio.shape[-1], n_frames
+
+    def clear_empty(self):
+        # clean up vram
+        torch.cuda.empty_cache()
+
+    def unload_model(self):
+        # unload model
+        self.net_g_ms = self.net_g_ms.to("cpu")
+        del self.net_g_ms
+        if hasattr(self,"enhancer"): 
+            self.enhancer.enhancer = self.enhancer.enhancer.to("cpu")
+            del self.enhancer.enhancer
+            del self.enhancer
+        gc.collect()
+
+    def slice_inference(self,
+                        raw_audio_path,
+                        spk,
+                        tran,
+                        slice_db,
+                        cluster_infer_ratio,
+                        auto_predict_f0,
+                        noice_scale,
+                        pad_seconds=0.5,
+                        clip_seconds=0,
+                        lg_num=0,
+                        lgr_num =0.75,
+                        f0_predictor='pm',
+                        enhancer_adaptive_key = 0,
+                        cr_threshold = 0.05,
+                        k_step = 100,
+                        use_spk_mix = False,
+                        second_encoding = False,
+                        loudness_envelope_adjustment = 1
+                        ):
+        if use_spk_mix:
+            if len(self.spk2id) == 1:
+                spk = self.spk2id.keys()[0]
+                use_spk_mix = False
+        wav_path = Path(raw_audio_path).with_suffix('.wav')
+        chunks = slicer.cut(wav_path, db_thresh=slice_db)
+        audio_data, audio_sr = slicer.chunks2audio(wav_path, chunks)
+        per_size = int(clip_seconds*audio_sr)
+        lg_size = int(lg_num*audio_sr)
+        lg_size_r = int(lg_size*lgr_num)
+        lg_size_c_l = (lg_size-lg_size_r)//2
+        lg_size_c_r = lg_size-lg_size_r-lg_size_c_l
+        lg = np.linspace(0,1,lg_size_r) if lg_size!=0 else 0
+
+        if use_spk_mix:
+            assert len(self.spk2id) == len(spk)
+            audio_length = 0
+            for (slice_tag, data) in audio_data:
+                aud_length = int(np.ceil(len(data) / audio_sr * self.target_sample))
+                if slice_tag:
+                    audio_length += aud_length // self.hop_size
+                    continue
+                if per_size != 0:
+                    datas = split_list_by_n(data, per_size,lg_size)
+                else:
+                    datas = [data]
+                for k,dat in enumerate(datas):
+                    pad_len = int(audio_sr * pad_seconds)
+                    per_length = int(np.ceil(len(dat) / audio_sr * self.target_sample))
+                    a_length = per_length + 2 * pad_len
+                    audio_length += a_length // self.hop_size
+            audio_length += len(audio_data)
+            spk_mix_tensor = torch.zeros(size=(len(spk), audio_length)).to(self.dev)
+            for i in range(len(spk)):
+                last_end = None
+                for mix in spk[i]:
+                    if mix[3]<0. or mix[2]<0.:
+                        raise RuntimeError("mix value must higer Than zero!")
+                    begin = int(audio_length * mix[0])
+                    end = int(audio_length * mix[1])
+                    length = end - begin
+                    if length<=0:                        
+                        raise RuntimeError("begin Must lower Than end!")
+                    step = (mix[3] - mix[2])/length
+                    if last_end is not None:
+                        if last_end != begin:
+                            raise RuntimeError("[i]EndTime Must Equal [i+1]BeginTime!")
+                    last_end = end
+                    if step == 0.:
+                        spk_mix_data = torch.zeros(length).to(self.dev) + mix[2]
+                    else:
+                        spk_mix_data = torch.arange(mix[2],mix[3],step).to(self.dev)
+                    if(len(spk_mix_data)<length):
+                        num_pad = length - len(spk_mix_data)
+                        spk_mix_data = torch.nn.functional.pad(spk_mix_data, [0, num_pad], mode="reflect").to(self.dev)
+                    spk_mix_tensor[i][begin:end] = spk_mix_data[:length]
+
+            spk_mix_ten = torch.sum(spk_mix_tensor,dim=0).unsqueeze(0).to(self.dev)
+            # spk_mix_tensor[0][spk_mix_ten<0.001] = 1.0
+            for i, x in enumerate(spk_mix_ten[0]):
+                if x == 0.0:
+                    spk_mix_ten[0][i] = 1.0
+                    spk_mix_tensor[:,i] = 1.0 / len(spk)
+            spk_mix_tensor = spk_mix_tensor / spk_mix_ten
+            if not ((torch.sum(spk_mix_tensor,dim=0) - 1.)<0.0001).all():
+                raise RuntimeError("sum(spk_mix_tensor) not equal 1")
+            spk = spk_mix_tensor
+
+        global_frame = 0
+        audio = []
+        for (slice_tag, data) in tqdm(audio_data):
+            print(f'#=====segment start, {round(len(data) / audio_sr, 3)}s======')
+            # padd
+            length = int(np.ceil(len(data) / audio_sr * self.target_sample))
+            if slice_tag:
+                print('jump empty segment')
+                _audio = np.zeros(length)
+                audio.extend(list(pad_array(_audio, length)))
+                global_frame += length // self.hop_size
+                continue
+            if per_size != 0:
+                datas = split_list_by_n(data, per_size,lg_size)
+            else:
+                datas = [data]
+            for k,dat in enumerate(datas):
+                per_length = int(np.ceil(len(dat) / audio_sr * self.target_sample)) if clip_seconds!=0 else length
+                if clip_seconds!=0: 
+                    print(f'###=====segment clip start, {round(len(dat) / audio_sr, 3)}s======')
+                # padd
+                pad_len = int(audio_sr * pad_seconds)
+                dat = np.concatenate([np.zeros([pad_len]), dat, np.zeros([pad_len])])
+                raw_path = io.BytesIO()
+                soundfile.write(raw_path, dat, audio_sr, format="wav")
+                raw_path.seek(0)
+                out_audio, out_sr, out_frame = self.infer(spk, tran, raw_path,
+                                                    cluster_infer_ratio=cluster_infer_ratio,
+                                                    auto_predict_f0=auto_predict_f0,
+                                                    noice_scale=noice_scale,
+                                                    f0_predictor = f0_predictor,
+                                                    enhancer_adaptive_key = enhancer_adaptive_key,
+                                                    cr_threshold = cr_threshold,
+                                                    k_step = k_step,
+                                                    frame = global_frame,
+                                                    spk_mix = use_spk_mix,
+                                                    second_encoding = second_encoding,
+                                                    loudness_envelope_adjustment = loudness_envelope_adjustment
+                                                    )
+                global_frame += out_frame
+                _audio = out_audio.cpu().numpy()
+                pad_len = int(self.target_sample * pad_seconds)
+                _audio = _audio[pad_len:-pad_len]
+                _audio = pad_array(_audio, per_length)
+                if lg_size!=0 and k!=0:
+                    lg1 = audio[-(lg_size_r+lg_size_c_r):-lg_size_c_r] if lgr_num != 1 else audio[-lg_size:]
+                    lg2 = _audio[lg_size_c_l:lg_size_c_l+lg_size_r]  if lgr_num != 1 else _audio[0:lg_size]
+                    lg_pre = lg1*(1-lg)+lg2*lg
+                    audio = audio[0:-(lg_size_r+lg_size_c_r)] if lgr_num != 1 else audio[0:-lg_size]
+                    audio.extend(lg_pre)
+                    _audio = _audio[lg_size_c_l+lg_size_r:] if lgr_num != 1 else _audio[lg_size:]
+                audio.extend(list(_audio))
+        return np.array(audio)
+
+class RealTimeVC:
+    def __init__(self):
+        self.last_chunk = None
+        self.last_o = None
+        self.chunk_len = 16000  # chunk length
+        self.pre_len = 3840  # cross fade length, multiples of 640
+
+    # Input and output are 1-dimensional numpy waveform arrays
+
+    def process(self, svc_model, speaker_id, f_pitch_change, input_wav_path,
+                cluster_infer_ratio=0,
+                auto_predict_f0=False,
+                noice_scale=0.4,
+                f0_filter=False):
+
+        import maad
+        audio, sr = torchaudio.load(input_wav_path)
+        audio = audio.cpu().numpy()[0]
+        temp_wav = io.BytesIO()
+        if self.last_chunk is None:
+            input_wav_path.seek(0)
+
+            audio, sr = svc_model.infer(speaker_id, f_pitch_change, input_wav_path,
+                                        cluster_infer_ratio=cluster_infer_ratio,
+                                        auto_predict_f0=auto_predict_f0,
+                                        noice_scale=noice_scale,
+                                        f0_filter=f0_filter)
+            
+            audio = audio.cpu().numpy()
+            self.last_chunk = audio[-self.pre_len:]
+            self.last_o = audio
+            return audio[-self.chunk_len:]
+        else:
+            audio = np.concatenate([self.last_chunk, audio])
+            soundfile.write(temp_wav, audio, sr, format="wav")
+            temp_wav.seek(0)
+
+            audio, sr = svc_model.infer(speaker_id, f_pitch_change, temp_wav,
+                                        cluster_infer_ratio=cluster_infer_ratio,
+                                        auto_predict_f0=auto_predict_f0,
+                                        noice_scale=noice_scale,
+                                        f0_filter=f0_filter)
+
+            audio = audio.cpu().numpy()
+            ret = maad.util.crossfade(self.last_o, audio, self.pre_len)
+            self.last_chunk = audio[-self.pre_len:]
+            self.last_o = audio
+            return ret[self.chunk_len:2 * self.chunk_len]
+            

preprocess_flist_config.py

@@ -1,11 +1,13 @@
-import os
 import argparse
+import json
+import os
 import re
+import wave
+from random import shuffle
 
 from tqdm import tqdm
-from random import shuffle
-import json
-import wave
+
+import diffusion.logger.utils as du
 
 config_template = json.load(open("configs_template/config_template.json"))
 
@@ -26,6 +28,8 @@ if __name__ == "__main__":
     parser.add_argument("--train_list", type=str, default="./filelists/train.txt", help="path to train list")
     parser.add_argument("--val_list", type=str, default="./filelists/val.txt", help="path to val list")
     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")
     args = parser.parse_args()
     
     train = []
@@ -41,8 +45,8 @@ if __name__ == "__main__":
         for file in wavs:
             if not file.endswith("wav"):
                 continue
-            #if not pattern.match(file):
-            #    print(f"warning：文件名{file}中包含非字母数字下划线，可能会导致错误。（也可能不会）")
+            if not pattern.match(file):
+                print(f"warning：文件名{file}中包含非字母数字下划线，可能会导致错误。（也可能不会）")
             if get_wav_duration(file) < 0.3:
                 print("skip too short audio:", file)
                 continue
@@ -67,9 +71,34 @@ if __name__ == "__main__":
             wavpath = fname
             f.write(wavpath + "\n")
 
+
+    d_config_template = du.load_config("configs_template/diffusion_template.yaml")
+    d_config_template["model"]["n_spk"] = spk_id
+    d_config_template["data"]["encoder"] = args.speech_encoder
+    d_config_template["spk"] = spk_dict
+    
     config_template["spk"] = spk_dict
     config_template["model"]["n_speakers"] = spk_id
-	
+    config_template["model"]["speech_encoder"] = args.speech_encoder
+    
+    if args.speech_encoder == "vec768l12" or args.speech_encoder == "dphubert" or args.speech_encoder == "wavlmbase+":
+        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
+        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
+        d_config_template["data"]["encoder_out_channels"] = 1024
+    elif args.speech_encoder == "whisper-ppg-large":
+        config_template["model"]["ssl_dim"] = config_template["model"]["filter_channels"] = config_template["model"]["gin_channels"] = 1280
+        d_config_template["data"]["encoder_out_channels"] = 1280
+        
+    if args.vol_aug:
+        config_template["train"]["vol_aug"] = config_template["model"]["vol_embedding"] = True
+
     print("Writing configs/config.json")
     with open("configs/config.json", "w") as f:
         json.dump(config_template, f, indent=2)
+    print("Writing configs/diffusion.yaml")
+    du.save_config("configs/diffusion.yaml",d_config_template)

preprocess_hubert_f0.py

@@ -1,43 +1,54 @@
-import math
+import argparse
+import logging
 import multiprocessing
 import os
-import argparse
+import random
+from concurrent.futures import ProcessPoolExecutor
+from glob import glob
 from random import shuffle
 
+import librosa
+import numpy as np
 import torch
-from glob import glob
 from tqdm import tqdm
-from modules.mel_processing import spectrogram_torch
 
+import diffusion.logger.utils as du
 import utils
-import logging
+from diffusion.vocoder import Vocoder
+from modules.mel_processing import spectrogram_torch
 
 logging.getLogger("numba").setLevel(logging.WARNING)
-import librosa
-import numpy as np
+logging.getLogger("matplotlib").setLevel(logging.WARNING)
 
 hps = utils.get_hparams_from_file("configs/config.json")
+dconfig = du.load_config("configs/diffusion.yaml")
 sampling_rate = hps.data.sampling_rate
 hop_length = hps.data.hop_length
+speech_encoder = hps["model"]["speech_encoder"]
 
 
-def process_one(filename, hmodel):
+def process_one(filename, hmodel,f0p,diff=False,mel_extractor=None):
     # print(filename)
     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):
-        device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
         wav16k = librosa.resample(wav, orig_sr=sampling_rate, target_sr=16000)
         wav16k = torch.from_numpy(wav16k).to(device)
-        c = utils.get_hubert_content(hmodel, wav_16k_tensor=wav16k)
+        c = hmodel.encoder(wav16k)
         torch.save(c.cpu(), soft_path)
 
     f0_path = filename + ".f0.npy"
     if not os.path.exists(f0_path):
-        f0 = utils.compute_f0_dio(
-            wav, sampling_rate=sampling_rate, hop_length=hop_length
+        f0_predictor = utils.get_f0_predictor(f0p,sampling_rate=sampling_rate, hop_length=hop_length,device=None,threshold=0.05)
+        f0,uv = f0_predictor.compute_f0_uv(
+            wav
         )
-        np.save(f0_path, f0)
+        np.save(f0_path, np.asanyarray((f0,uv),dtype=object))
+
 
     spec_path = filename.replace(".wav", ".spec.pt")
     if not os.path.exists(spec_path):
@@ -45,7 +56,6 @@ def process_one(filename, hmodel):
         # The following code can't be replaced by torch.FloatTensor(wav)
         # because load_wav_to_torch return a tensor that need to be normalized
 
-        audio, sr = utils.load_wav_to_torch(filename)
         if sr != hps.data.sampling_rate:
             raise ValueError(
                 "{} SR doesn't match target {} SR".format(
@@ -53,8 +63,7 @@ def process_one(filename, hmodel):
                 )
             )
 
-        audio_norm = audio / hps.data.max_wav_value
-        audio_norm = audio_norm.unsqueeze(0)
+        #audio_norm = audio / hps.data.max_wav_value
 
         spec = spectrogram_torch(
             audio_norm,
@@ -67,35 +76,88 @@ def process_one(filename, hmodel):
         spec = torch.squeeze(spec, 0)
         torch.save(spec, spec_path)
 
-
-def process_batch(filenames):
-    print("Loading hubert for content...")
+    if diff or hps.model.vol_embedding:
+        volume_path = filename + ".vol.npy"
+        volume_extractor = utils.Volume_Extractor(hop_length)
+        if not os.path.exists(volume_path):
+            volume = volume_extractor.extract(audio_norm)
+            np.save(volume_path, volume.to('cpu').numpy())
+
+    if diff:
+        mel_path = filename + ".mel.npy"
+        if not os.path.exists(mel_path) and mel_extractor is not None:
+            mel_t = mel_extractor.extract(audio_norm.to(device), sampling_rate)
+            mel = mel_t.squeeze().to('cpu').numpy()
+            np.save(mel_path, mel)
+        aug_mel_path = filename + ".aug_mel.npy"
+        aug_vol_path = filename + ".aug_vol.npy"
+        max_amp = float(torch.max(torch.abs(audio_norm))) + 1e-5
+        max_shift = min(1, np.log10(1/max_amp))
+        log10_vol_shift = random.uniform(-1, max_shift)
+        keyshift = random.uniform(-5, 5)
+        if mel_extractor is not None:
+            aug_mel_t = mel_extractor.extract(audio_norm * (10 ** log10_vol_shift), sampling_rate, keyshift = keyshift)
+        aug_mel = aug_mel_t.squeeze().to('cpu').numpy()
+        aug_vol = volume_extractor.extract(audio_norm * (10 ** log10_vol_shift))
+        if not os.path.exists(aug_mel_path):
+            np.save(aug_mel_path,np.asanyarray((aug_mel,keyshift),dtype=object))
+        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_hubert_model().to(device)
-    print("Loaded hubert.")
-    for filename in tqdm(filenames):
-        process_one(filename, hmodel)
+    hmodel = utils.get_speech_encoder(speech_encoder, device=device)
+    print("Loaded speech encoder.")
 
+    for filename in tqdm(file_chunk):
+        process_one(filename, hmodel, f0p, diff, mel_extractor)
+
+def parallel_process(filenames, num_processes, f0p, diff, mel_extractor):
+    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))
+
+        for task in tqdm(tasks):
+            task.result()
 
 if __name__ == "__main__":
     parser = argparse.ArgumentParser()
     parser.add_argument(
         "--in_dir", type=str, default="dataset/44k", help="path to input dir"
     )
-
+    parser.add_argument(
+        '--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)'
+    )
+    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
+    print(speech_encoder)
+    print(f0p)
+    print(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)
+        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)
 
-    num_processes = 1
-    chunk_size = int(math.ceil(len(filenames) / num_processes))
-    chunks = [
-        filenames[i : i + chunk_size] for i in range(0, len(filenames), chunk_size)
-    ]
-    print([len(c) for c in chunks])
-    processes = [
-        multiprocessing.Process(target=process_batch, args=(chunk,)) for chunk in chunks
-    ]
-    for p in processes:
-        p.start()
+    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)

resample.py

@@ -1,48 +1,98 @@
-import os
 import argparse
+import concurrent.futures
+import os
+from concurrent.futures import ProcessPoolExecutor
+from multiprocessing import cpu_count
+
 import librosa
 import numpy as np
-from multiprocessing import Pool, cpu_count
 from scipy.io import wavfile
 from tqdm import tqdm
 
 
+def load_wav(wav_path):
+    return librosa.load(wav_path, sr=None)
+
+
+def trim_wav(wav, top_db=40):
+    return librosa.effects.trim(wav, top_db=top_db)
+
+
+def normalize_peak(wav, threshold=1.0):
+    peak = np.abs(wav).max()
+    if peak > threshold:
+        wav = 0.98 * wav / peak
+    return wav
+
+
+def resample_wav(wav, sr, target_sr):
+    return librosa.resample(wav, orig_sr=sr, target_sr=target_sr)
+
+
+def save_wav_to_path(wav, save_path, sr):
+    wavfile.write(
+        save_path,
+        sr,
+        (wav * np.iinfo(np.int16).max).astype(np.int16)
+    )
+
+
 def process(item):
     spkdir, wav_name, args = item
-    # speaker 's5', 'p280', 'p315' are excluded,
     speaker = spkdir.replace("\\", "/").split("/")[-1]
+
     wav_path = os.path.join(args.in_dir, speaker, wav_name)
     if os.path.exists(wav_path) and '.wav' in wav_path:
         os.makedirs(os.path.join(args.out_dir2, speaker), exist_ok=True)
-        wav, sr = librosa.load(wav_path, sr=None)
-        wav, _ = librosa.effects.trim(wav, top_db=20)
-        peak = np.abs(wav).max()
-        if peak > 1.0:
-            wav = 0.98 * wav / peak
-        wav2 = librosa.resample(wav, orig_sr=sr, target_sr=args.sr2)
-        wav2 /= max(wav2.max(), -wav2.min())
-        save_name = wav_name
-        save_path2 = os.path.join(args.out_dir2, speaker, save_name)
-        wavfile.write(
-            save_path2,
-            args.sr2,
-            (wav2 * np.iinfo(np.int16).max).astype(np.int16)
-        )
+
+        wav, sr = load_wav(wav_path)
+        wav, _ = trim_wav(wav)
+        wav = normalize_peak(wav)
+        resampled_wav = resample_wav(wav, sr, args.sr2)
+
+        if not args.skip_loudnorm:
+            resampled_wav /= np.max(np.abs(resampled_wav))
+
+        save_path2 = os.path.join(args.out_dir2, speaker, wav_name)
+        save_wav_to_path(resampled_wav, save_path2, args.sr2)
 
 
+"""
+def process_all_speakers():
+    process_count = 30 if os.cpu_count() > 60 else (os.cpu_count() - 2 if os.cpu_count() > 4 else 1)
+
+    with ThreadPoolExecutor(max_workers=process_count) as executor:
+        for speaker in speakers:
+            spk_dir = os.path.join(args.in_dir, speaker)
+            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)):
+                    pass
+"""
+# multi process
+
+
+def process_all_speakers():
+    process_count = 30 if os.cpu_count() > 60 else (os.cpu_count() - 2 if os.cpu_count() > 4 else 1)
+    with ProcessPoolExecutor(max_workers=process_count) as executor:
+        for speaker in speakers:
+            spk_dir = os.path.join(args.in_dir, speaker)
+            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)):
+                    pass
+
 
 if __name__ == "__main__":
     parser = argparse.ArgumentParser()
     parser.add_argument("--sr2", type=int, default=44100, help="sampling rate")
     parser.add_argument("--in_dir", type=str, default="./dataset_raw", help="path to source dir")
     parser.add_argument("--out_dir2", type=str, default="./dataset/44k", help="path to target dir")
+    parser.add_argument("--skip_loudnorm", action="store_true", help="Skip loudness matching if you have done it")
     args = parser.parse_args()
-    processs = cpu_count()-2 if cpu_count() >4 else 1
-    pool = Pool(processes=processs)
-
-    for speaker in os.listdir(args.in_dir):
-        spk_dir = os.path.join(args.in_dir, speaker)
-        if os.path.isdir(spk_dir):
-            print(spk_dir)
-            for _ in tqdm(pool.imap_unordered(process, [(spk_dir, i, args) for i in os.listdir(spk_dir) if i.endswith("wav")])):
-                pass
+
+    print(f"CPU count: {cpu_count()}")
+    speakers = os.listdir(args.in_dir)
+    process_all_speakers()

spkmix.py

@@ -0,0 +1,11 @@
+# 角色混合轨道 编写规则：
+# 角色ID : [[起始时间1, 终止时间1, 起始数值1, 起始数值1], [起始时间2, 终止时间2, 起始数值2, 起始数值2]]
+# 起始时间和前一个的终止时间必须相同，第一个起始时间必须为0，最后一个终止时间必须为1 （时间的范围为0-1）
+# 全部角色必须填写，不使用的角色填[[0., 1., 0., 0.]]即可
+# 融合数值可以随便填，在指定的时间段内从起始数值线性变化为终止数值，内部会自动确保线性组合为1，可以放心使用
+
+spk_mix_map = {
+    0 : [[0., 0.5, 1, 0.5], [0.5, 1, 0.5, 1]],
+    1 : [[0., 0.35, 1, 0.5], [0.35, 0.75, 0.75, 1], [0.75, 1, 0.45, 1]],
+    2 : [[0., 0.35, 1, 0.5], [0.35, 0.75, 0.75, 1], [0.75, 1, 0.45, 1]]
+}

train_diff.py

@@ -0,0 +1,76 @@
+import argparse
+
+import torch
+from torch.optim import lr_scheduler
+
+from diffusion.data_loaders import get_data_loaders
+from diffusion.logger import utils
+from diffusion.solver import train
+from diffusion.unit2mel import Unit2Mel
+from diffusion.vocoder import Vocoder
+
+
+def parse_args(args=None, namespace=None):
+    """Parse command-line arguments."""
+    parser = argparse.ArgumentParser()
+    parser.add_argument(
+        "-c",
+        "--config",
+        type=str,
+        required=True,
+        help="path to the config file")
+    return parser.parse_args(args=args, namespace=namespace)
+
+
+if __name__ == '__main__':
+    # parse commands
+    cmd = parse_args()
+    
+    # load config
+    args = utils.load_config(cmd.config)
+    print(' > config:', cmd.config)
+    print(' >    exp:', args.env.expdir)
+    
+    # load vocoder
+    vocoder = Vocoder(args.vocoder.type, args.vocoder.ckpt, device=args.device)
+    
+    # load model
+    model = Unit2Mel(
+                args.data.encoder_out_channels, 
+                args.model.n_spk,
+                args.model.use_pitch_aug,
+                vocoder.dimension,
+                args.model.n_layers,
+                args.model.n_chans,
+                args.model.n_hidden,
+                args.model.timesteps,
+                args.model.k_step_max
+                )
+    
+    print(f' > INFO: now model timesteps is {model.timesteps}, and k_step_max is {model.k_step_max}')
+    
+    # load parameters
+    optimizer = torch.optim.AdamW(model.parameters())
+    initial_global_step, model, optimizer = utils.load_model(args.env.expdir, model, optimizer, device=args.device)
+    for param_group in optimizer.param_groups:
+        param_group['initial_lr'] = args.train.lr
+        param_group['lr'] = args.train.lr * (args.train.gamma ** max(((initial_global_step-2)//args.train.decay_step),0) )
+        param_group['weight_decay'] = args.train.weight_decay
+    scheduler = lr_scheduler.StepLR(optimizer, step_size=args.train.decay_step, gamma=args.train.gamma,last_epoch=initial_global_step-2)
+    
+    # device
+    if args.device == 'cuda':
+        torch.cuda.set_device(args.env.gpu_id)
+    model.to(args.device)
+    
+    for state in optimizer.state.values():
+        for k, v in state.items():
+            if torch.is_tensor(v):
+                state[k] = v.to(args.device)
+                    
+    # datas
+    loader_train, loader_valid = get_data_loaders(args, whole_audio=False)
+    
+    # run
+    train(args, initial_global_step, model, optimizer, scheduler, vocoder, loader_train, loader_valid)
+    

train_index.py

@@ -0,0 +1,30 @@
+import argparse
+import os
+import pickle
+
+import utils
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser()
+    parser.add_argument(
+        "--root_dir", type=str, default="dataset/44k", help="path to root dir"
+    )
+    parser.add_argument('-c', '--config', type=str, default="./configs/config.json",
+                    help='JSON file for configuration')
+    parser.add_argument(
+        "--output_dir", type=str, default="logs/44k", help="path to output dir"
+    )
+
+    args = parser.parse_args()
+
+    hps = utils.get_hparams_from_file(args.config)
+    spk_dic = hps.spk
+    result = {}
+    
+    for k,v in spk_dic.items():
+        print(f"now, index {k} feature...")
+        index = utils.train_index(k,args.root_dir)
+        result[v] = index
+
+    with open(os.path.join(args.output_dir,"feature_and_index.pkl"),"wb") as f:
+        pickle.dump(result,f)

utils.py

@@ -1,22 +1,21 @@
-import os
-import glob
-import re
-import sys
 import argparse
-import logging
+import glob
 import json
+import logging
+import os
+import re
 import subprocess
-import warnings
-import random
-import functools
+import sys
+import traceback
+from multiprocessing import cpu_count
+
+import faiss
 import librosa
 import numpy as np
-from scipy.io.wavfile import read
 import torch
+from scipy.io.wavfile import read
+from sklearn.cluster import MiniBatchKMeans
 from torch.nn import functional as F
-from modules.commons import sequence_mask
-import faiss
-import tqdm
 
 MATPLOTLIB_FLAG = False
 
@@ -97,7 +96,10 @@ def get_f0_predictor(f0_predictor,hop_length,sampling_rate,**kargs):
         f0_predictor_object = HarvestF0Predictor(hop_length=hop_length,sampling_rate=sampling_rate)
     elif f0_predictor == "dio":
         from modules.F0Predictor.DioF0Predictor import DioF0Predictor
-        f0_predictor_object = DioF0Predictor(hop_length=hop_length,sampling_rate=sampling_rate)
+        f0_predictor_object = DioF0Predictor(hop_length=hop_length,sampling_rate=sampling_rate) 
+    elif f0_predictor == "rmvpe":
+        from modules.F0Predictor.RMVPEF0Predictor import RMVPEF0Predictor
+        f0_predictor_object = RMVPEF0Predictor(hop_length=hop_length,sampling_rate=sampling_rate,dtype=torch.float32 ,device=kargs["device"],threshold=kargs["threshold"])
     else:
         raise Exception("Unknown f0 predictor")
     return f0_predictor_object
@@ -130,6 +132,18 @@ def get_speech_encoder(speech_encoder,device=None,**kargs):
     elif speech_encoder == "whisper-ppg":
         from vencoder.WhisperPPG import WhisperPPG
         speech_encoder_object = WhisperPPG(device = device)
+    elif speech_encoder == "cnhubertlarge":
+        from vencoder.CNHubertLarge import CNHubertLarge
+        speech_encoder_object = CNHubertLarge(device = device)
+    elif speech_encoder == "dphubert":
+        from vencoder.DPHubert import DPHubert
+        speech_encoder_object = DPHubert(device = device)
+    elif speech_encoder == "whisper-ppg-large":
+        from vencoder.WhisperPPGLarge import WhisperPPGLarge
+        speech_encoder_object = WhisperPPGLarge(device = device)
+    elif speech_encoder == "wavlmbase+":
+        from vencoder.WavLMBasePlus import WavLMBasePlus
+        speech_encoder_object = WavLMBasePlus(device = device)
     else:
         raise Exception("Unknown speech encoder")
     return speech_encoder_object 
@@ -142,6 +156,7 @@ def load_checkpoint(checkpoint_path, model, optimizer=None, skip_optimizer=False
     if optimizer is not None and not skip_optimizer and checkpoint_dict['optimizer'] is not None:
         optimizer.load_state_dict(checkpoint_dict['optimizer'])
     saved_state_dict = checkpoint_dict['model']
+    model = model.to(list(saved_state_dict.values())[0].dtype)
     if hasattr(model, 'module'):
         state_dict = model.module.state_dict()
     else:
@@ -153,10 +168,11 @@ def load_checkpoint(checkpoint_path, model, optimizer=None, skip_optimizer=False
             # print("load", k)
             new_state_dict[k] = saved_state_dict[k]
             assert saved_state_dict[k].shape == v.shape, (saved_state_dict[k].shape, v.shape)
-        except:
-            print("error, %s is not in the checkpoint" % k)
-            logger.info("%s is not in the checkpoint" % k)
-            new_state_dict[k] = v
+        except Exception:
+            if "enc_q" not in k or "emb_g" not in k:
+              print("error, %s is not in the checkpoint" % k)
+              logger.info("%s is not in the checkpoint" % k)
+              new_state_dict[k] = v
     if hasattr(model, 'module'):
         model.module.load_state_dict(new_state_dict)
     else:
@@ -189,15 +205,20 @@ def clean_checkpoints(path_to_models='logs/44k/', n_ckpts_to_keep=2, sort_by_tim
                         False -> lexicographically delete ckpts
   """
   ckpts_files = [f for f in os.listdir(path_to_models) if os.path.isfile(os.path.join(path_to_models, f))]
-  name_key = (lambda _f: int(re.compile('._(\d+)\.pth').match(_f).group(1)))
-  time_key = (lambda _f: os.path.getmtime(os.path.join(path_to_models, _f)))
+  def name_key(_f):
+      return int(re.compile("._(\\d+)\\.pth").match(_f).group(1))
+  def time_key(_f):
+      return os.path.getmtime(os.path.join(path_to_models, _f))
   sort_key = time_key if sort_by_time else name_key
-  x_sorted = lambda _x: sorted([f for f in ckpts_files if f.startswith(_x) and not f.endswith('_0.pth')], key=sort_key)
+  def x_sorted(_x):
+      return sorted([f for f in ckpts_files if f.startswith(_x) and not f.endswith("_0.pth")], key=sort_key)
   to_del = [os.path.join(path_to_models, fn) for fn in
             (x_sorted('G')[:-n_ckpts_to_keep] + x_sorted('D')[:-n_ckpts_to_keep])]
-  del_info = lambda fn: logger.info(f".. Free up space by deleting ckpt {fn}")
-  del_routine = lambda x: [os.remove(x), del_info(x)]
-  rs = [del_routine(fn) for fn in to_del]
+  def del_info(fn):
+      return logger.info(f".. Free up space by deleting ckpt {fn}")
+  def del_routine(x):
+      return [os.remove(x), del_info(x)]
+  [del_routine(fn) for fn in to_del]
 
 def summarize(writer, global_step, scalars={}, histograms={}, images={}, audios={}, audio_sampling_rate=22050):
   for k, v in scalars.items():
@@ -325,11 +346,11 @@ def get_hparams_from_dir(model_dir):
   return hparams
 
 
-def get_hparams_from_file(config_path):
+def get_hparams_from_file(config_path, infer_mode = False):
   with open(config_path, "r") as f:
     data = f.read()
   config = json.loads(data)
-  hparams =HParams(**config)
+  hparams =HParams(**config) if not infer_mode else InferHParams(**config)
   return hparams
 
 
@@ -368,7 +389,13 @@ def get_logger(model_dir, filename="train.log"):
   return logger
 
 
-def repeat_expand_2d(content, target_len):
+def repeat_expand_2d(content, target_len, mode = 'left'):
+    # content : [h, t]
+    return repeat_expand_2d_left(content, target_len) if mode == 'left' else repeat_expand_2d_other(content, target_len, mode)
+
+
+
+def repeat_expand_2d_left(content, target_len):
     # content : [h, t]
 
     src_len = content.shape[-1]
@@ -385,6 +412,14 @@ def repeat_expand_2d(content, target_len):
     return target
 
 
+# mode : 'nearest'| 'linear'| 'bilinear'| 'bicubic'| 'trilinear'| 'area'
+def repeat_expand_2d_other(content, target_len, mode = 'nearest'):
+    # content : [h, t]
+    content = content[None,:,:]
+    target = F.interpolate(content,size=target_len,mode=mode)[0]
+    return target
+
+
 def mix_model(model_paths,mix_rate,mode):
   mix_rate = torch.FloatTensor(mix_rate)/100
   model_tem = torch.load(model_paths[0])
@@ -420,6 +455,7 @@ def change_rms(data1, sr1, data2, sr2, rate):  # 1是输入音频，2是输出
     return data2
 
 def train_index(spk_name,root_dir = "dataset/44k/"):  #from: RVC https://github.com/RVC-Project/Retrieval-based-Voice-Conversion-WebUI
+    n_cpu = cpu_count()
     print("The feature index is constructing.")
     exp_dir = os.path.join(root_dir,spk_name)
     listdir_res = []
@@ -436,6 +472,25 @@ def train_index(spk_name,root_dir = "dataset/44k/"):  #from: RVC https://github.
     big_npy_idx = np.arange(big_npy.shape[0])
     np.random.shuffle(big_npy_idx)
     big_npy = big_npy[big_npy_idx]
+    if big_npy.shape[0] > 2e5:
+        # if(1):
+        info = "Trying doing kmeans %s shape to 10k centers." % big_npy.shape[0]
+        print(info)
+        try:
+            big_npy = (
+                MiniBatchKMeans(
+                    n_clusters=10000,
+                    verbose=True,
+                    batch_size=256 * n_cpu,
+                    compute_labels=False,
+                    init="random",
+                )
+                .fit(big_npy)
+                .cluster_centers_
+            )
+        except Exception:
+            info = traceback.format_exc()
+            print(info)
     n_ivf = min(int(16 * np.sqrt(big_npy.shape[0])), big_npy.shape[0] // 39)
     index = faiss.index_factory(big_npy.shape[1] , "IVF%s,Flat" % n_ivf)
     index_ivf = faiss.extract_index_ivf(index)  #
@@ -486,6 +541,18 @@ class HParams():
   def get(self,index):
     return self.__dict__.get(index)
 
+  
+class InferHParams(HParams):
+  def __init__(self, **kwargs):
+    for k, v in kwargs.items():
+      if type(v) == dict:
+        v = InferHParams(**v)
+      self[k] = v
+
+  def __getattr__(self,index):
+    return self.get(index)
+
+
 class Volume_Extractor:
     def __init__(self, hop_size = 512):
         self.hop_size = hop_size
@@ -496,6 +563,6 @@ class Volume_Extractor:
         n_frames = int(audio.size(-1) // self.hop_size)
         audio2 = audio ** 2
         audio2 = torch.nn.functional.pad(audio2, (int(self.hop_size // 2), int((self.hop_size + 1) // 2)), mode = 'reflect')
-        volume = torch.FloatTensor([torch.mean(audio2[:,int(n * self.hop_size) : int((n + 1) * self.hop_size)]) for n in range(n_frames)])
+        volume = torch.nn.functional.unfold(audio2[:,None,None,:],(1,self.hop_size),stride=self.hop_size)[:,:,:n_frames].mean(dim=1)[0]
         volume = torch.sqrt(volume)
-        return volume
+        return volume