diff --git a/train.py b/train.py index a8cc54e..79da223 100644 --- a/train.py +++ b/train.py @@ -1,6 +1,8 @@ import os +import sys os.environ['HF_HUB_CACHE'] = './checkpoints/hf_cache' import torch +import torch.multiprocessing as mp import random import librosa import yaml @@ -9,14 +11,12 @@ import torchaudio.compliance.kaldi as kaldi import glob from tqdm import tqdm +import shutil from modules.commons import recursive_munch, build_model, load_checkpoint from optimizers import build_optimizer from data.ft_dataset import build_ft_dataloader from hf_utils import load_custom_model_from_hf -import shutil - - class Trainer: @@ -79,23 +79,22 @@ def __init__(self, # initialize optimizers after preparing models for compatibility with FSDP self.optimizer = build_optimizer({key: self.model[key] for key in self.model}, - lr=float(scheduler_params['base_lr'])) + lr=float(scheduler_params['base_lr'])) if pretrained_ckpt_path is None: - # find latest checkpoint with name pattern of 'T2V_epoch_*_step_*.pth' + # find latest checkpoint available_checkpoints = glob.glob(os.path.join(self.log_dir, "DiT_epoch_*_step_*.pth")) if len(available_checkpoints) > 0: - # find the checkpoint that has the highest step number latest_checkpoint = max( available_checkpoints, key=lambda x: int(x.split("_")[-1].split(".")[0]) ) earliest_checkpoint = min( available_checkpoints, key=lambda x: int(x.split("_")[-1].split(".")[0]) ) - # delete the earliest checkpoint + # delete the earliest checkpoint if we have more than 2 if ( - earliest_checkpoint != latest_checkpoint - and len(available_checkpoints) > 2 + earliest_checkpoint != latest_checkpoint + and len(available_checkpoints) > 2 ): os.remove(earliest_checkpoint) print(f"Removed {earliest_checkpoint}") @@ -108,16 +107,18 @@ def __init__(self, latest_checkpoint = pretrained_ckpt_path if os.path.exists(latest_checkpoint): - self.model, self.optimizer, self.epoch, self.iters = load_checkpoint(self.model, self.optimizer, latest_checkpoint, - load_only_params=True, - ignore_modules=[], - is_distributed=False) + self.model, self.optimizer, self.epoch, self.iters = load_checkpoint( + self.model, self.optimizer, latest_checkpoint, + load_only_params=True, + ignore_modules=[], + is_distributed=False + ) print(f"Loaded checkpoint from {latest_checkpoint}") else: self.epoch, self.iters = 0, 0 - print("Failed to load any checkpoint, this implies you are training from scratch.") + print("Failed to load any checkpoint, training from scratch.") + def build_sv_model(self, device, config): - # speaker verification model from modules.campplus.DTDNN import CAMPPlus self.campplus_model = CAMPPlus(feat_dim=80, embedding_size=192) campplus_sd_path = load_custom_model_from_hf("funasr/campplus", "campplus_cn_common.bin", config_filename=None) @@ -126,16 +127,17 @@ def build_sv_model(self, device, config): self.campplus_model.eval() self.campplus_model.to(device) self.sv_fn = self.campplus_model + def build_f0_fn(self, device, config): from modules.rmvpe import RMVPE model_path = load_custom_model_from_hf("lj1995/VoiceConversionWebUI", "rmvpe.pt", None) self.rmvpe = RMVPE(model_path, is_half=False, device=device) self.f0_fn = self.rmvpe + def build_converter(self, device, config): - # speaker perturbation model from modules.openvoice.api import ToneColorConverter ckpt_converter, config_converter = load_custom_model_from_hf("myshell-ai/OpenVoiceV2", "converter/checkpoint.pth", "converter/config.json") - self.tone_color_converter = ToneColorConverter(config_converter, device=device,) + self.tone_color_converter = ToneColorConverter(config_converter, device=device) self.tone_color_converter.load_ckpt(ckpt_converter) self.tone_color_converter.model.eval() se_db_path = load_custom_model_from_hf("Plachta/Seed-VC", "se_db.pt", None) @@ -146,9 +148,7 @@ def build_vocoder(self, device, config): vocoder_name = config['model_params']['vocoder'].get('name', None) if vocoder_type == 'bigvgan': from modules.bigvgan import bigvgan - bigvgan_name = vocoder_name - self.bigvgan_model = bigvgan.BigVGAN.from_pretrained(bigvgan_name, use_cuda_kernel=False) - # remove weight norm in the model and set to eval mode + self.bigvgan_model = bigvgan.BigVGAN.from_pretrained(vocoder_name, use_cuda_kernel=False) self.bigvgan_model.remove_weight_norm() self.bigvgan_model = self.bigvgan_model.eval().to(device) vocoder_fn = self.bigvgan_model @@ -158,7 +158,7 @@ def build_vocoder(self, device, config): hift_config = yaml.safe_load(open('configs/hifigan.yml', 'r')) hift_path = load_custom_model_from_hf("FunAudioLLM/CosyVoice-300M", 'hift.pt', None) self.hift_gen = HiFTGenerator(**hift_config['hift'], - f0_predictor=ConvRNNF0Predictor(**hift_config['f0_predictor'])) + f0_predictor=ConvRNNF0Predictor(**hift_config['f0_predictor'])) self.hift_gen.load_state_dict(torch.load(hift_path, map_location='cpu')) self.hift_gen.eval() self.hift_gen.to(device) @@ -168,21 +168,25 @@ def build_vocoder(self, device, config): self.vocoder_fn = vocoder_fn def build_semantic_fn(self, device, config): - # speech tokenizer speech_tokenizer_type = config['model_params']['speech_tokenizer'].get('type', 'cosyvoice') if speech_tokenizer_type == 'whisper': from transformers import AutoFeatureExtractor, WhisperModel whisper_model_name = config['model_params']['speech_tokenizer']['name'] self.whisper_model = WhisperModel.from_pretrained(whisper_model_name).to(device) self.whisper_feature_extractor = AutoFeatureExtractor.from_pretrained(whisper_model_name) + # remove decoder to save memory del self.whisper_model.decoder + def semantic_fn(waves_16k): - ori_inputs = self.whisper_feature_extractor([w16k.cpu().numpy() for w16k in waves_16k], - return_tensors="pt", - return_attention_mask=True, - sampling_rate=16000,) + ori_inputs = self.whisper_feature_extractor( + [w16k.cpu().numpy() for w16k in waves_16k], + return_tensors="pt", + return_attention_mask=True, + sampling_rate=16000, + ) ori_input_features = self.whisper_model._mask_input_features( - ori_inputs.input_features, attention_mask=ori_inputs.attention_mask).to(device) + ori_inputs.input_features, attention_mask=ori_inputs.attention_mask + ).to(device) with torch.no_grad(): ori_outputs = self.whisper_model.encoder( ori_input_features.to(self.whisper_model.encoder.dtype), @@ -194,6 +198,7 @@ def semantic_fn(waves_16k): S_ori = ori_outputs.last_hidden_state.to(torch.float32) S_ori = S_ori[:, :waves_16k.size(-1) // 320 + 1] return S_ori + elif speech_tokenizer_type == 'xlsr': from transformers import ( Wav2Vec2FeatureExtractor, @@ -209,15 +214,14 @@ def semantic_fn(waves_16k): self.wav2vec_model = self.wav2vec_model.half() def semantic_fn(waves_16k): - ori_waves_16k_input_list = [ - waves_16k[bib].cpu().numpy() - for bib in range(len(waves_16k)) - ] - ori_inputs = self.wav2vec_feature_extractor(ori_waves_16k_input_list, - return_tensors="pt", - return_attention_mask=True, - padding=True, - sampling_rate=16000).to(device) + ori_waves_16k_input_list = [waves_16k[bib].cpu().numpy() for bib in range(len(waves_16k))] + ori_inputs = self.wav2vec_feature_extractor( + ori_waves_16k_input_list, + return_tensors="pt", + return_attention_mask=True, + padding=True, + sampling_rate=16000 + ).to(device) with torch.no_grad(): ori_outputs = self.wav2vec_model( ori_inputs.input_values.half(), @@ -246,11 +250,12 @@ def train_one_step(self, batch): se_batch = self.tone_color_converter.extract_se(waves_22k, wave_lengths_22k) ref_se_idx = torch.randint(0, len(self.se_db), (B,)) - ref_se = self.se_db[ref_se_idx] - ref_se = ref_se.to(self.device) + ref_se = self.se_db[ref_se_idx].to(self.device) # convert - converted_waves_22k = self.tone_color_converter.convert(waves_22k, wave_lengths_22k, se_batch, ref_se).squeeze(1) + converted_waves_22k = self.tone_color_converter.convert( + waves_22k, wave_lengths_22k, se_batch, ref_se + ).squeeze(1) if self.sr != 22050: converted_waves = torchaudio.functional.resample(converted_waves_22k, 22050, self.sr) @@ -260,6 +265,7 @@ def train_one_step(self, batch): waves_16k = torchaudio.functional.resample(waves, self.sr, 16000) wave_lengths_16k = (wave_lengths.float() * 16000 / self.sr).long() converted_waves_16k = torchaudio.functional.resample(converted_waves, self.sr, 16000) + # extract S_alt (perturbed speech tokens) S_ori = self.semantic_fn(waves_16k) S_alt = self.semantic_fn(converted_waves_16k) @@ -268,11 +274,14 @@ def train_one_step(self, batch): F0_ori = self.rmvpe.infer_from_audio_batch(waves_16k) else: F0_ori = None + # interpolate speech token to match acoustic feature length alt_cond, _, alt_codes, alt_commitment_loss, alt_codebook_loss = ( - self.model.length_regulator(S_alt, ylens=target_lengths, f0=F0_ori)) + self.model.length_regulator(S_alt, ylens=target_lengths, f0=F0_ori) + ) ori_cond, _, ori_codes, ori_commitment_loss, ori_codebook_loss = ( - self.model.length_regulator(S_ori, ylens=target_lengths, f0=F0_ori)) + self.model.length_regulator(S_ori, ylens=target_lengths, f0=F0_ori) + ) if alt_commitment_loss is None: alt_commitment_loss = 0 alt_codebook_loss = 0 @@ -281,10 +290,10 @@ def train_one_step(self, batch): # randomly set a length as prompt prompt_len_max = target_lengths - 1 - prompt_len = (torch.rand([B], device=alt_cond.device) * prompt_len_max).floor().to(dtype=torch.long) + prompt_len = (torch.rand([B], device=alt_cond.device) * prompt_len_max).floor().long() prompt_len[torch.rand([B], device=alt_cond.device) < 0.1] = 0 - # for prompt cond token, it must be from ori_cond instead of alt_cond + # for prompt cond token, use ori_cond instead of alt_cond cond = alt_cond.clone() for bib in range(B): cond[bib, :prompt_len[bib]] = ori_cond[bib, :prompt_len[bib]] @@ -295,13 +304,16 @@ def train_one_step(self, batch): cond = cond[:, :common_min_len] target_lengths = torch.clamp(target_lengths, max=common_min_len) x = target - # style vectors are extracted from prompt only to avoid inference time OOD + + # style vectors are extracted from the prompt only feat_list = [] for bib in range(B): - feat = kaldi.fbank(waves_16k[bib:bib + 1, :wave_lengths_16k[bib]], - num_mel_bins=80, - dither=0, - sample_frequency=16000) + feat = kaldi.fbank( + waves_16k[bib:bib + 1, :wave_lengths_16k[bib]], + num_mel_bins=80, + dither=0, + sample_frequency=16000 + ) feat = feat - feat.mean(dim=0, keepdim=True) feat_list.append(feat) y_list = [] @@ -313,31 +325,39 @@ def train_one_step(self, batch): loss, _ = self.model.cfm(x, target_lengths, prompt_len, cond, y) - loss_total = (loss + - (alt_commitment_loss + ori_commitment_loss) * 0.05 + - (ori_codebook_loss + alt_codebook_loss) * 0.15) + loss_total = ( + loss + + (alt_commitment_loss + ori_commitment_loss) * 0.05 + + (ori_codebook_loss + alt_codebook_loss) * 0.15 + ) self.optimizer.zero_grad() loss_total.backward() - grad_norm_g = torch.nn.utils.clip_grad_norm_(self.model.cfm.parameters(), 10.0) - grad_norm_g2 = torch.nn.utils.clip_grad_norm_(self.model.length_regulator.parameters(), 10.0) + torch.nn.utils.clip_grad_norm_(self.model.cfm.parameters(), 10.0) + torch.nn.utils.clip_grad_norm_(self.model.length_regulator.parameters(), 10.0) self.optimizer.step('cfm') self.optimizer.step('length_regulator') self.optimizer.scheduler(key='cfm') self.optimizer.scheduler(key='length_regulator') return loss.detach().item() + def train_one_epoch(self): _ = [self.model[key].train() for key in self.model] for i, batch in enumerate(tqdm(self.train_dataloader)): batch = [b.to(self.device) for b in batch] loss = self.train_one_step(batch) - self.ema_loss = self.ema_loss * self.loss_smoothing_rate + loss * (1 - self.loss_smoothing_rate) if self.iters > 0 else loss + self.ema_loss = ( + self.ema_loss * self.loss_smoothing_rate + loss * (1 - self.loss_smoothing_rate) + if self.iters > 0 else loss + ) if self.iters % self.log_interval == 0: print(f"epoch {self.epoch}, step {self.iters}, loss: {self.ema_loss}") self.iters += 1 + if self.iters >= self.max_steps: break + if self.iters % self.save_interval == 0: print('Saving..') state = { @@ -347,13 +367,15 @@ def train_one_epoch(self): 'iters': self.iters, 'epoch': self.epoch, } - save_path = os.path.join(self.log_dir, 'DiT_epoch_%05d_step_%05d.pth' % (self.epoch, self.iters)) + save_path = os.path.join( + self.log_dir, + f'DiT_epoch_{self.epoch:05d}_step_{self.iters:05d}.pth' + ) torch.save(state, save_path) # find all checkpoints and remove old ones checkpoints = glob.glob(os.path.join(self.log_dir, 'DiT_epoch_*.pth')) if len(checkpoints) > 2: - # sort by step checkpoints.sort(key=lambda x: int(x.split('_')[-1].split('.')[0])) for cp in checkpoints[:-2]: os.remove(cp) @@ -364,15 +386,34 @@ def train(self): for epoch in range(self.n_epochs): self.epoch = epoch self.train_one_epoch() + # Save after each epoch + print('Epoch completed. Saving..') + state = { + 'net': {key: self.model[key].state_dict() for key in self.model}, + 'optimizer': self.optimizer.state_dict(), + 'scheduler': self.optimizer.scheduler_state_dict(), + 'iters': self.iters, + 'epoch': self.epoch, + } + save_path = os.path.join( + self.log_dir, + f'DiT_epoch_{self.epoch:05d}_step_{self.iters:05d}.pth' + ) + torch.save(state, save_path) + print(f"Checkpoint saved at {save_path}") + if self.iters >= self.max_steps: break - print('Saving..') + + print('Saving final model..') state = { 'net': {key: self.model[key].state_dict() for key in self.model}, } os.makedirs(self.log_dir, exist_ok=True) save_path = os.path.join(self.log_dir, 'ft_model.pth') torch.save(state, save_path) + print(f"Final model saved at {save_path}") + def main(args): trainer = Trainer( @@ -387,8 +428,12 @@ def main(args): num_workers=args.num_workers, ) trainer.train() - + if __name__ == '__main__': + if sys.platform == 'win32': + mp.freeze_support() + mp.set_start_method('spawn', force=True) + parser = argparse.ArgumentParser() parser.add_argument('--config', type=str, default='./configs/presets/config_dit_mel_seed_uvit_xlsr_tiny.yml') parser.add_argument('--pretrained-ckpt', type=str, default=None) @@ -400,4 +445,5 @@ def main(args): parser.add_argument('--save-every', type=int, default=500) parser.add_argument('--num-workers', type=int, default=0) args = parser.parse_args() - main(args) \ No newline at end of file + + main(args)