app.py

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_")
tokenizer = AutoTokenizer.from_pretrained("wasmdashai/vits-ar-sa-huba",token=token)
#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")


# 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(model,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.model[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.epoch_count=0
        self.global_step=0
       
        

    def init_Starting(self):
        self.training_args=load_training_args(self.path_training_args)
        self.stute_mode=False
       
        self.load_dataset(train_dataset_dirs,eval_dataset_dir,full_generation_dir)
        self.len_dataset=len(self.DataSets['train'])
    def init_training(self):
        
        self.load_model()
        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")
        wandb.init(project= 'HugfaceTraining',config = self.training_args.to_dict())
   
    
    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 = os.path.join(dataset_dir,'eval'),
                                                            full_generation_dir = os.path.join(dataset_dir,'full_generation'),
                                                                     device=self.device)
       self.DataSets={'train':ctrain_datasets,'eval':eval_dataset,'full_generation':full_generation_dataset}


    



    

    def initialize_training_components(self):

        
    
        self.training_args=training_args

        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
      

    
    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)

        self.input_save_pretrained(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'),
                                                                     device=device
                    )

@spaces.GPU
def run_train_epoch(num):
   for i in range(10):
    # model.train(True)
        yield pro.run_train_epoch()
       
@spaces.GPU
def init_training():
    pro.init_training()
    return pro.dir_model,'init_training'

@spaces.GPU
def init_Starting():
    pro.init_Starting()
    return 'init_Starting'
@spaces.GPU
def init_wandb():
    pro.init_wandb()
    return 'init_wandb'
    
@spaces.GPU
def save_pretrained(path):
    pro.save_pretrained(path)
    pro.init_training()
    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_training,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')