import spaces import gradio as gr GK=0 from transformers import AutoTokenizer import torch import os from VitsModelSplit.vits_model2 import VitsModel,get_state_grad_loss import VitsModelSplit.monotonic_align as monotonic_align token=os.environ.get("key_") # import VitsModelSplit.monotonic_align as monotonic_align from IPython.display import clear_output from transformers import set_seed import wandb import logging import copy import torch import numpy as np import torch from datasets import DatasetDict,Dataset import os from VitsModelSplit.vits_model2 import VitsModel,get_state_grad_loss from VitsModelSplit.PosteriorDecoderModel import PosteriorDecoderModel from VitsModelSplit.feature_extraction import VitsFeatureExtractor from transformers import AutoTokenizer, HfArgumentParser, set_seed from VitsModelSplit.Arguments import DataTrainingArguments, ModelArguments, VITSTrainingArguments from VitsModelSplit.dataset_features_collector import FeaturesCollectionDataset from torch.cuda.amp import autocast, GradScaler device = torch.device("cuda" if torch.cuda.is_available() else "cpu") # sgl=get_state_grad_loss(k1=True,#generator=False, # discriminator=False, # duration=False # ) Lst=['input_ids', 'attention_mask', 'waveform', 'labels', 'labels_attention_mask', 'mel_scaled_input_features'] def covert_cuda_batch(d): return d for key in Lst: d[key]=d[key].cuda(non_blocking=True) # for key in d['text_encoder_output']: # d['text_encoder_output'][key]=d['text_encoder_output'][key].cuda(non_blocking=True) # for key in d['posterior_encode_output']: # d['posterior_encode_output'][key]=d['posterior_encode_output'][key].cuda(non_blocking=True) return d def generator_loss(disc_outputs): total_loss = 0 gen_losses = [] for disc_output in disc_outputs: disc_output = disc_output loss = torch.mean((1 - disc_output) ** 2) gen_losses.append(loss) total_loss += loss return total_loss, gen_losses def discriminator_loss(disc_real_outputs, disc_generated_outputs): loss = 0 real_losses = 0 generated_losses = 0 for disc_real, disc_generated in zip(disc_real_outputs, disc_generated_outputs): real_loss = torch.mean((1 - disc_real) ** 2) generated_loss = torch.mean(disc_generated**2) loss += real_loss + generated_loss real_losses += real_loss generated_losses += generated_loss return loss, real_losses, generated_losses def feature_loss(feature_maps_real, feature_maps_generated): loss = 0 for feature_map_real, feature_map_generated in zip(feature_maps_real, feature_maps_generated): for real, generated in zip(feature_map_real, feature_map_generated): real = real.detach() loss += torch.mean(torch.abs(real - generated)) return loss * 2 def kl_loss(z_p, logs_q, m_p, logs_p, z_mask): """ z_p, logs_q: [b, h, t_t] m_p, logs_p: [b, h, t_t] """ z_p = z_p.float() logs_q = logs_q.float() m_p = m_p.float() logs_p = logs_p.float() z_mask = z_mask.float() kl = logs_p - logs_q - 0.5 kl += 0.5 * ((z_p - m_p)**2) * torch.exp(-2. * logs_p) kl = torch.sum(kl * z_mask) l = kl / torch.sum(z_mask) return l #............................................. # def kl_loss(prior_latents, posterior_log_variance, prior_means, prior_log_variance, labels_mask): # kl = prior_log_variance - posterior_log_variance - 0.5 # kl += 0.5 * ((prior_latents - prior_means) ** 2) * torch.exp(-2.0 * prior_log_variance) # kl = torch.sum(kl * labels_mask) # loss = kl / torch.sum(labels_mask) # return loss def get_state_grad_loss(k1=True, mel=True, duration=True, generator=True, discriminator=True): return {'k1':k1,'mel':mel,'duration':duration,'generator':generator,'discriminator':discriminator} @spaces.GPU def clip_grad_value_(parameters, clip_value, norm_type=2): if isinstance(parameters, torch.Tensor): parameters = [parameters] parameters = list(filter(lambda p: p.grad is not None, parameters)) norm_type = float(norm_type) if clip_value is not None: clip_value = float(clip_value) total_norm = 0 for p in parameters: param_norm = p.grad.data.norm(norm_type) total_norm += param_norm.item() ** norm_type if clip_value is not None: p.grad.data.clamp_(min=-clip_value, max=clip_value) total_norm = total_norm ** (1. / norm_type) return total_norm @spaces.GPU def get_embed_speaker(self,speaker_id): if self.config.num_speakers > 1 and speaker_id is not None: if isinstance(speaker_id, int): speaker_id = torch.full(size=(1,), fill_value=speaker_id, device=self.device) elif isinstance(speaker_id, (list, tuple, np.ndarray)): speaker_id = torch.tensor(speaker_id, device=self.device) if not ((0 <= speaker_id).all() and (speaker_id < self.config.num_speakers).all()).item(): raise ValueError(f"Set `speaker_id` in the range 0-{self.config.num_speakers - 1}.") return self.embed_speaker(speaker_id).unsqueeze(-1) else: return None def get_data_loader(train_dataset_dirs,eval_dataset_dir,full_generation_dir,device): ctrain_datasets=[] for dataset_dir ,id_sp in train_dataset_dirs: train_dataset = FeaturesCollectionDataset(dataset_dir = os.path.join(dataset_dir,'train'), device = device ) ctrain_datasets.append((train_dataset,id_sp)) eval_dataset = None eval_dataset = FeaturesCollectionDataset(dataset_dir = eval_dataset_dir, device = device ) full_generation_dataset = FeaturesCollectionDataset(dataset_dir = full_generation_dir, device = device) return ctrain_datasets,eval_dataset,full_generation_dataset global_step=0 def train_step(batch,models=[],optimizers=[], training_args=None,tools=[]): self,discriminator=models optimizer,disc_optimizer,scaler=optimizers feature_extractor,maf,dict_state_grad_loss=tools with autocast(enabled=training_args.fp16): speaker_embeddings=get_embed_speaker(self,batch["speaker_id"]) waveform,ids_slice,log_duration,prior_latents,posterior_log_variances,prior_means,prior_log_variances,labels_padding_mask = self.forward_train( input_ids=batch["input_ids"], attention_mask=batch["attention_mask"], labels=batch["labels"], labels_attention_mask=batch["labels_attention_mask"], text_encoder_output =None , posterior_encode_output=None , return_dict=True, monotonic_alignment_function=maf, speaker_embeddings=speaker_embeddings ) mel_scaled_labels = batch["mel_scaled_input_features"] mel_scaled_target = self.slice_segments(mel_scaled_labels, ids_slice,self.segment_size) mel_scaled_generation = feature_extractor._torch_extract_fbank_features(waveform.squeeze(1))[1] target_waveform = batch["waveform"].transpose(1, 2) target_waveform = self.slice_segments( target_waveform, ids_slice * feature_extractor.hop_length, self.config.segment_size ) discriminator_target, fmaps_target = discriminator(target_waveform) discriminator_candidate, fmaps_candidate = discriminator(waveform.detach()) with autocast(enabled=False): if dict_state_grad_loss['discriminator']: loss_disc, loss_real_disc, loss_fake_disc = discriminator_loss( discriminator_target, discriminator_candidate ) loss_dd = loss_disc# + loss_real_disc + loss_fake_disc # loss_dd.backward() disc_optimizer.zero_grad() scaler.scale(loss_dd).backward() scaler.unscale_(disc_optimizer ) grad_norm_d = clip_grad_value_(discriminator.parameters(), None) scaler.step(disc_optimizer) loss_des=grad_norm_d with autocast(enabled=training_args.fp16): # backpropagate discriminator_target, fmaps_target = discriminator(target_waveform) discriminator_candidate, fmaps_candidate = discriminator(waveform.detach()) with autocast(enabled=False): if dict_state_grad_loss['k1']: loss_kl = kl_loss( prior_latents, posterior_log_variances, prior_means, prior_log_variances, labels_padding_mask, ) loss_kl=loss_kl*training_args.weight_kl loss_klall=loss_kl.detach().item() #if displayloss['loss_kl']>=0: # loss_kl.backward() if dict_state_grad_loss['mel']: loss_mel = torch.nn.functional.l1_loss(mel_scaled_target, mel_scaled_generation)*training_args.weight_mel loss_melall= loss_mel.detach().item() # train_losses_sum = train_losses_sum + displayloss['loss_mel'] # if displayloss['loss_mel']>=0: # loss_mel.backward() if dict_state_grad_loss['duration']: loss_duration=torch.sum(log_duration)*training_args.weight_duration loss_durationsall=loss_duration.detach().item() # if displayloss['loss_duration']>=0: # loss_duration.backward() if dict_state_grad_loss['generator']: loss_fmaps = feature_loss(fmaps_target, fmaps_candidate) loss_gen, losses_gen = generator_loss(discriminator_candidate) loss_gen=loss_gen * training_args.weight_gen # loss_gen.backward(retain_graph=True) loss_fmaps=loss_fmaps * training_args.weight_fmaps # loss_fmaps.backward(retain_graph=True) total_generator_loss = ( loss_duration + loss_mel + loss_kl + loss_fmaps + loss_gen ) # total_generator_loss.backward() optimizer.zero_grad() scaler.scale(total_generator_loss).backward() scaler.unscale_(optimizer) grad_norm_g = clip_grad_value_(self.parameters(), None) scaler.step(optimizer) scaler.update() loss_gen=grad_norm_g return loss_gen,loss_des,loss_durationsall,loss_melall,loss_klall def train_epoch(obtrainer,index_db=0,epoch=0,idspeakers=[],full_generation_sample_index=-1): train_losses_sum = 0 loss_genall=0 loss_desall=0 loss_durationsall=0 loss_melall=0 loss_klall=0 loss_fmapsall=0 start_speeker,end_speeker=idspeakers datatrain=obtrainer.DataSets['train'][index_db] lr_scheduler,disc_lr_scheduler=obtrainer.lr_schedulers lr_scheduler.step() disc_lr_scheduler.step() train_dataset,speaker_id=datatrain print(f" Num Epochs = {epoch}, speaker_id DB ={speaker_id}") num_div_proc=int(len(train_dataset)/10)+1 print(' -process traning : [',end='') full_generation_sample =obtrainer.DataSets['full_generation'][full_generation_sample_index] for step, batch in enumerate(train_dataset): loss_gen,loss_des,loss_durationsa,loss_mela,loss_kl=train_step(batch, models=obtrainer.models, optimizers=obtrainer.optimizers, training_args=obtrainer.training_args, tools=obtrainer.tools) loss_genall+=loss_gen loss_desall+=loss_des loss_durationsall+=loss_durationsa loss_melall+=loss_mela loss_klall+=loss_kl obtrainer.global_step +=1 if step%num_div_proc==0: print('==',end='') # validation do_eval = obtrainer.training_args.do_eval and (obtrainer.global_step % obtrainer.training_args.eval_steps == 0) if do_eval: speaker_id_c=int(torch.randint(start_speeker,end_speeker,size=(1,))[0]) model=obtrainer.models[0] with torch.no_grad(): full_generation =model.forward( input_ids =full_generation_sample["input_ids"], attention_mask=full_generation_sample["attention_mask"], speaker_id=speaker_id_c ) full_generation_waveform = full_generation.waveform.cpu().numpy() wandb.log({ "full generations samples": [ wandb.Audio(w.reshape(-1), caption=f"Full generation sample {epoch}", sample_rate=16000) for w in full_generation_waveform],}) step+=1 # wandb.log({"train_losses":loss_melall}) wandb.log({"loss_gen":loss_genall/step}) wandb.log({"loss_des":loss_desall/step}) wandb.log({"loss_duration":loss_durationsall/step}) wandb.log({"loss_mel":loss_melall/step}) wandb.log({f"loss_kl_db{speaker_id}":loss_klall/step}) print(']',end='') def load_training_args(path): parser = HfArgumentParser((ModelArguments, DataTrainingArguments, VITSTrainingArguments)) json_file = os.path.abspath(path) model_args, data_args, training_args = parser.parse_json_file(json_file = json_file) return training_args def load_tools(): feature_extractor = VitsFeatureExtractor() dict_state_grad_loss=get_state_grad_loss() return feature_extractor,monotonic_align.maximum_path,dict_state_grad_loss class TrinerModelVITS: def __init__(self,dir_model="", path_training_args="", train_dataset_dirs=[], eval_dataset_dir="", full_generation_dir="", token="", device="cpu"): self.device=device self.dir_model=dir_model self.path_training_args=path_training_args self.stute_mode=False self.token=token self.load_dataset(train_dataset_dirs,eval_dataset_dir,full_generation_dir) self.epoch_count=0 self.global_step=0 self.len_dataset=len(self.DataSets['train']) self.load_model() self.init_wandb() self.training_args=load_training_args(self.path_training_args) training_args= self.training_args scaler = GradScaler(enabled=True) for disc in self.model.discriminator.discriminators: disc.apply_weight_norm() self.model.decoder.apply_weight_norm() # torch.nn.utils.weight_norm(self.decoder.conv_pre) # torch.nn.utils.weight_norm(self.decoder.conv_post) for flow in self.model.flow.flows: torch.nn.utils.weight_norm(flow.conv_pre) torch.nn.utils.weight_norm(flow.conv_post) discriminator = self.model.discriminator self.model.discriminator = None optimizer = torch.optim.AdamW( self.model.parameters(), 2e-4, betas=[0.8, 0.99], # eps=training_args.adam_epsilon, ) # Hack to be able to train on multiple device disc_optimizer = torch.optim.AdamW( discriminator.parameters(), 2e-4, betas=[0.8, 0.99], # eps=training_args.adam_epsilon, ) lr_scheduler = torch.optim.lr_scheduler.ExponentialLR( optimizer,gamma=0.999875, last_epoch=-1 ) disc_lr_scheduler = torch.optim.lr_scheduler.ExponentialLR( disc_optimizer, gamma=0.999875,last_epoch=-1 ) self.models=(self.model,discriminator) self.optimizers=(optimizer,disc_optimizer,scaler) self.lr_schedulers=(lr_scheduler,disc_lr_scheduler) self.tools=load_tools() self.stute_mode=True print(self.lr_schedulers) def init_Starting(self): print('init_Starting') self.training_args=load_training_args(self.path_training_args) self.stute_mode=False print('end training_args') def init_training(self): self.initialize_training_components() # self.epoch_count=0 def load_model(self): self.model=VitsModel.from_pretrained(self.dir_model,token=self.token).to(self.device) self.model.setMfA(monotonic_align.maximum_path) def init_wandb(self): wandb.login(key= "782b6a6e82bbb5a5348de0d3c7d40d1e76351e79") #config = self.training_args.to_dict() wandb.init(project= 'HugfaceTraining') def load_dataset(self,train_dataset_dirs,eval_dataset_dir,full_generation_dir): ctrain_datasets,eval_dataset,full_generation_dataset=get_data_loader(train_dataset_dirs = train_dataset_dirs, eval_dataset_dir =eval_dataset_dir , full_generation_dir =full_generation_dir , device=self.device) self.DataSets={'train':ctrain_datasets,'eval':eval_dataset,'full_generation':full_generation_dataset} def initialize_training_components(self): self.training_args=load_training_args(self.path_training_args) training_args= self.training_args training_args.weight_kl=1 training_args.d_learning_rate=2e-4 training_args.learning_rate=2e-4 training_args.weight_mel=45 training_args.num_train_epochs=4 training_args.eval_steps=1000 # set_seed(training_args.seed) # scaler = GradScaler(enabled=training_args.fp16) # # Initialize optimizer, lr_scheduler # for disc in self.model.discriminator.discriminators: # disc.apply_weight_norm() # self.model.decoder.apply_weight_norm() # # torch.nn.utils.weight_norm(self.decoder.conv_pre) # # torch.nn.utils.weight_norm(self.decoder.conv_post) # for flow in self.model.flow.flows: # torch.nn.utils.weight_norm(flow.conv_pre) # torch.nn.utils.weight_norm(flow.conv_post) # discriminator = self.model.discriminator # self.model.discriminator = None # optimizer = torch.optim.AdamW( # self.model.parameters(), # training_args.learning_rate, # betas=[training_args.adam_beta1, training_args.adam_beta2], # eps=training_args.adam_epsilon, # ) # # Hack to be able to train on multiple device # disc_optimizer = torch.optim.AdamW( # discriminator.parameters(), # training_args.d_learning_rate, # betas=[training_args.d_adam_beta1, training_args.d_adam_beta2], # eps=training_args.adam_epsilon, # ) # lr_scheduler = torch.optim.lr_scheduler.ExponentialLR( # optimizer, gamma=training_args.lr_decay, last_epoch=-1 # ) # disc_lr_scheduler = torch.optim.lr_scheduler.ExponentialLR( # disc_optimizer, gamma=training_args.lr_decay, last_epoch=-1 # ) # self.models=(self.model,discriminator) # self.optimizers=(optimizer,disc_optimizer,scaler) # self.lr_schedulers=(lr_scheduler,disc_lr_scheduler) # self.tools=load_tools() # self.stute_mode=True # print(self.lr_schedulers) def save_pretrained(self,path_save_model): model,discriminator=self.models model.discriminator=discriminator for disc in model.discriminator.discriminators: disc.remove_weight_norm() model.decoder.remove_weight_norm() # torch.nn.utils.remove_weight_norm(self.decoder.conv_pre) # torch.nn.utils.remove_weight_norm(self.decoder.conv_post) for flow in model.flow.flows: torch.nn.utils.remove_weight_norm(flow.conv_pre) torch.nn.utils.remove_weight_norm(flow.conv_post) model.push_to_hub(path_save_model,token=self.token) def run_train_epoch(self): index_db=self.epoch_count%self.len_dataset train_epoch(self,index_db=index_db,epoch=self.epoch_count,idspeakers=(0,1),full_generation_sample_index=-1) self.epoch_count+=1 return f'epoch_count:{self.epoch_count},global_step:{self.global_step},index_db"{index_db}' # return (self.model,discriminator),(optimizer, disc_optimizer), (lr_scheduler, disc_lr_scheduler) # logger.info("***** Training / Inference Done *****") def modelspeech(texts): inputs = tokenizer(texts, return_tensors="pt")#.cuda() wav = model_vits(input_ids=inputs["input_ids"]).waveform#.detach() # display(Audio(wav, rate=model.config.sampling_rate)) return model_vits.config.sampling_rate,wav#remove_noise_nr(wav) dataset_dir='ABThag-db' train_dataset_dirs=[ # ('/content/drive/MyDrive/vitsM/DATA/fahd_db',0), # ('/content/drive/MyDrive/vitsM/DATA/fahd_db',0), # ('/content/drive/MyDrive/vitsM/DB2KKKK',1), # ('/content/drive/MyDrive/vitsM/DATA/Db_Amgd_50_Bitch10',0), # ('/content/drive/MyDrive/vitsM/DB2KKKK',1), # # ('/content/drive/MyDrive/vitsM/DATA/Db_Amgd_50_Bitch10',0), # ('/content/drive/MyDrive/vitsM/DATA/DBWfaa-Bitch:8-Count:60',0), # ('/content/drive/MyDrive/vitsM/DATA/Wafa/b10r',0), # ('/content/drive/MyDrive/vitsM/DATA/Wafa/b16r',0), # ('/content/drive/MyDrive/vitsM/DATA/Wafa/b4',0), # ('/content/drive/MyDrive/vitsM/DATA/fahd_db',None), # ('/content/drive/MyDrive/vitsM/DATA/wafa-db',None), # ('/content/drive/MyDrive/vitsM/DATA/wafa-db',4), # ('/content/drive/MyDrive/vitsM/DATA/DB-ABThag-Bitch:5-Count-37',4), # ('/content/drive/MyDrive/vitsM/DB-300-k',6), ('ABThag-db',0), #('/content/drive/MyDrive/dataset_ljBatchs',0), ] dir_model='wasmdashai/vits-ar-huba-fine' pro=TrinerModelVITS(dir_model=dir_model, path_training_args='VitsModelSplit/finetune_config_ara.json', train_dataset_dirs = train_dataset_dirs, eval_dataset_dir = os.path.join(dataset_dir,'eval'), full_generation_dir = os.path.join(dataset_dir,'full_generation'), token=token, device=device ) @spaces.GPU(duration=5) def run_train_epoch(num): if num >0: pro.init_training() for i in range(num): # model.train(True) yield pro.run_train_epoch() else: pro.save_pretrained(pro.dir_model) pro.load_model() return 'save model ' @spaces.GPU def init_training(): pro.init_training() return pro.dir_model,'init_training' @spaces.GPU def init_Starting(): device = torch.device("cuda" if torch.cuda.is_available() else "cpu") return 'init_Starting' @spaces.GPU def init_wandb(): pro.init_wandb() return 'init_wandb' def save_pretrained(path): pro.save_pretrained(path) pro.load_model() return 'save_pretrained' with gr.Blocks() as interface: with gr.Accordion("init_Starting ", open=False): btn_init = gr.Button("init start") output_init = gr.Textbox(label="init") btn_init.click(fn=init_Starting,inputs=[],outputs=[output_init]) with gr.Accordion("init_wandb ", open=False): btn_init_wandb = gr.Button("nit_wandb") output_initbtn_init_wandb = gr.Textbox(label="init") btn_init_wandb.click(fn=init_wandb,inputs=[],outputs=[output_initbtn_init_wandb]) with gr.Accordion("init_training ", open=False): btn_init_train = gr.Button("init init_train") output_btn_init_train = gr.Textbox(label="init") # btn_init_train.click(fn=init_training,inputs=[],outputs=[output_btn_init_train]) with gr.Accordion("run_train_epoch ", open=False): btn_run_train_epoch = gr.Button("run_train_epoch") input_run_train_epoch = gr.Number(label="number _train_epoch") output_run_train_epoch = gr.Textbox(label="run_train_epoch") btn_run_train_epoch.click(fn=run_train_epoch,inputs=[input_run_train_epoch],outputs=[output_run_train_epoch]) with gr.Accordion("save_pretrained ", open=False): btn_save_pretrained = gr.Button("save_pretrained") input_save_pretrained = gr.Textbox(label="save_pretrained") output_save_pretrained = gr.Textbox(label="save_pretrained") btn_save_pretrained.click(fn=save_pretrained,inputs=[input_save_pretrained],outputs=[output_save_pretrained]) btn_init_train.click(fn=init_training,inputs=[],outputs=[input_save_pretrained,output_btn_init_train]) interface.launch() print('loadeed')