Update app.py

app.py

@@ -10,6 +10,526 @@ tokenizer = AutoTokenizer.from_pretrained("wasmdashai/vits-ar-sa-huba",token=tok
 #device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 model_vits=VitsModel.from_pretrained("wasmdashai/vits-ar-sa-huba",token=token)#.to(device)
 
+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")
+model=VitsModel.from_pretrained("facebook/mms-tts-eng").to(device)
+# model1= VitsModel.from_pretrained("/content/drive/MyDrive/vitsM/OneBatch/S6/MMMMM-dash-azd60").to("cuda")
+# model= VitsModel.from_pretrained("/content/drive/MyDrive/vitsM/TO/sp3/core/vend").to("cuda")
+# model=VitsModel.from_pretrained("/content/drive/MyDrive/vitsM/heppa/EndCore3/v0").to("cuda")
+# model.discriminator=model1.discriminator
+# model.duration_predictor=model1.duration_predictor
+
+model.setMfA(monotonic_align.maximum_path)
+# tokenizer = AutoTokenizer.from_pretrained("facebook/mms-tts-ara",cache_dir="./")
+feature_extractor = VitsFeatureExtractor()
+parser = HfArgumentParser((ModelArguments, DataTrainingArguments, VITSTrainingArguments))
+json_file = os.path.abspath('VitsModelSplit/finetune_config_ara.json')
+model_args, data_args, training_args = parser.parse_json_file(json_file = json_file)
+sgl=get_state_grad_loss(mel=True,
+                       # generator=False,
+                    #    discriminator=False,
+                        duration=False)
+
+training_args.num_train_epochs=1000
+training_args.fp16=True
+training_args.eval_steps=300
+# 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}
+
+
+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
+
+
+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
+     if training_args.do_eval:
+        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 trainer_to_cuda(self,
+          ctrain_datasets = None,
+          eval_dataset = None,
+          full_generation_dataset = None,
+          feature_extractor = VitsFeatureExtractor(),
+          training_args = None,
+          full_generation_sample_index= 0,
+          project_name = "Posterior_Decoder_Finetuning",
+          wandbKey = "782b6a6e82bbb5a5348de0d3c7d40d1e76351e79",
+          is_used_text_encoder=True,
+          is_used_posterior_encode=True,
+          dict_state_grad_loss=None,
+          nk=1,
+          path_save_model='./',
+                    maf=None,
+          n_back_save_model=3000,
+                    start_speeker=0,
+                    end_speeker=1,
+                    n_epoch=0,
+
+
+
+          ):
+
+
+    # os.makedirs(training_args.output_dir,exist_ok=True)
+    # logger = logging.getLogger(f"{__name__} Training")
+    # log_level = training_args.get_process_log_level()
+    # logger.setLevel(log_level)
+
+    # # wandb.login(key= wandbKey)
+    # # wandb.init(project= project_name,config = training_args.to_dict())
+    if dict_state_grad_loss is None:
+        dict_state_grad_loss=get_state_grad_loss()
+    global global_step
+
+
+
+    set_seed(training_args.seed)
+    scaler = GradScaler(enabled=training_args.fp16)
+    self.config.save_pretrained(training_args.output_dir)
+    len_db=len(ctrain_datasets)
+    self.full_generation_sample = full_generation_dataset[full_generation_sample_index]
+
+    # init optimizer, lr_scheduler
+    for disc in self.discriminator.discriminators:
+          disc.apply_weight_norm()
+    self.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.flow.flows:
+        torch.nn.utils.weight_norm(flow.conv_pre)
+        torch.nn.utils.weight_norm(flow.conv_post)
+
+    discriminator=self.discriminator
+    self.discriminator=None
+
+    optimizer = torch.optim.AdamW(
+        self.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)
+
+
+    logger.info("***** Running training *****")
+    logger.info(f"  Num Epochs = {training_args.num_train_epochs}")
+
+
+    #.......................loop training............................
+
+
+
+    for epoch in range(training_args.num_train_epochs):
+        train_losses_sum = 0
+        loss_gen=0
+        loss_des=0
+        loss_durationsall=0
+        loss_melall=0
+        loss_klall=0
+        loss_fmapsall=0
+        lr_scheduler.step()
+
+        disc_lr_scheduler.step()
+        train_dataset,speaker_id=ctrain_datasets[epoch%len_db]
+        print(f"  Num Epochs = {int((epoch+n_epoch)/len_db)}, speaker_id DB ={speaker_id}")
+        num_div_proc=int(len(train_dataset)/10)
+        print('       -process traning : [',end='')
+
+
+
+
+
+
+
+        for step, batch in enumerate(train_dataset):
+            # if speaker_id==None:
+            #     if step<3 :continue
+
+            #     if step>200:break
+
+
+            batch=covert_cuda_batch(batch)
+            displayloss={}
+
+            with autocast(enabled=training_args.fp16):
+              speaker_embeddings=get_embed_speaker(self,batch["speaker_id"] if speaker_id ==None else 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
+                    )
+
+                    dk={"step_loss_disc": loss_disc.detach().item(),
+                        "step_loss_real_disc": loss_real_disc.detach().item(),
+                        "step_loss_fake_disc": loss_fake_disc.detach().item()}
+                    displayloss['dict_loss_discriminator']=dk
+                    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
+                  displayloss['loss_gen'] = loss_gen.detach().item()
+              #   loss_gen.backward(retain_graph=True)
+                  loss_fmaps=loss_fmaps * training_args.weight_fmaps
+                  displayloss['loss_fmaps'] = loss_fmaps.detach().item()
+              #   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
+
+
+
+
+
+
+            # optimizer.step()
+
+
+
+
+            # print(f"TRAINIG - batch {step}, waveform {(batch['waveform'].shape)}, lr {lr_scheduler.get_last_lr()[0]}... ")
+            # print(f"display loss function  enable  :{displayloss}")
+
+            global_step +=1
+            if step%num_div_proc==0:
+                   print('==',end='')
+
+            # validation
+
+            do_eval = training_args.do_eval and (global_step % training_args.eval_steps == 0)
+            if do_eval:
+                speaker_id_c=int(torch.randint(start_speeker,end_speeker,size=(1,))[0])
+                logger.info("Running validation... ")
+                eval_losses_sum = 0
+                cc=0;
+                for step, batch in enumerate(eval_dataset):
+                    break
+                    if cc>2: break
+                    cc+=1
+                    with torch.no_grad():
+                        model_outputs = self.forward(
+                            input_ids=batch["input_ids"],
+                        attention_mask=batch["attention_mask"],
+                        labels=batch["labels"],
+                        labels_attention_mask=batch["labels_attention_mask"],
+                        speaker_id=batch["speaker_id"],
+
+
+                        return_dict=True,
+
+                            )
+
+                    mel_scaled_labels = batch["mel_scaled_input_features"]
+                    mel_scaled_target = self.slice_segments(mel_scaled_labels, model_outputs.ids_slice,self.segment_size)
+                    mel_scaled_generation = feature_extractor._torch_extract_fbank_features(model_outputs.waveform.squeeze(1))[1]
+                    loss = loss_mel.detach().item()
+                    eval_losses_sum +=loss
+
+                    loss_mel = torch.nn.functional.l1_loss(mel_scaled_target, mel_scaled_generation)
+                    print(f"VALIDATION - batch {step}, waveform {(batch['waveform'].shape)}, step_loss_mel {loss} ... ")
+
+
+
+                with torch.no_grad():
+                    full_generation_sample = self.full_generation_sample
+                    full_generation =self.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({
+                "eval_losses": eval_losses_sum,
+                "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_gen/step})
+        wandb.log({"loss_des":loss_des/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='')
+
+
+
+
+          # self.save_pretrained(path_save_model)
+
+
+    self.discriminator=discriminator
+    for disc in self.discriminator.discriminators:
+          disc.remove_weight_norm()
+    self.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 self.flow.flows:
+            torch.nn.utils.remove_weight_norm(flow.conv_pre)
+            torch.nn.utils.remove_weight_norm(flow.conv_post)
+
+    self.save_pretrained(path_save_model)
+
+    logger.info("Running final full generations samples... ")
+
+
+
+    logger.info("***** Training / Inference Done *****")
 def   modelspeech(texts):