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_")
# 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.vits_model_only_d import  Vits_models_only_decoder
#from VitsModelSplit.vits_model import VitsModel
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
# )
# class model_onxx:
#     def __init__(self):
#         self.model=None
#         self.n_onxx=""
#         self.storage_dir = "uploads"  
#         pass
     

 


#     def download_file(self,file_path):
#         ff= gr.File(value=file_path, visible=True) 
#         file_url = ff.value['url']  
#         return file_url
#     def function_change(self,n_model,token,n_onxx,choice):
#         if choice=="decoder":

#             V=self.convert_to_onnx_only_decoder(n_model,token,n_onxx)
#         elif choice=="all only decoder":
#             V=self.convert_to_onnx_only_decoder(n_model,token,n_onxx)
#         else:
#             V=self.convert_to_onnx_only_decoder(n_model,token,n_onxx)
#         return V

#     def install_model(self,n_model,token,n_onxx):
#         self.n_onxx=n_onxx
#         self.model= VitsModel.from_pretrained(n_model,token=token)
#         return self.model
#     def convert_model_decoder_onxx(self,n_model,token,namemodelonxx):
#         self.model= VitsModel.from_pretrained(n_model,token=token)
#         x=f"/tmp/{namemodelonxx}.onnx"
#         return x
#     def convert_to_onnx_only_decoder(self,n_model,token,namemodelonxx):
#           model=VitsModel.from_pretrained(n_model,token=token)
#           x=f"/tmp/{namemodelonxx}.onnx"
        
#           vocab_size = model.text_encoder.embed_tokens.weight.size(0)
#           example_input = torch.randint(0, vocab_size, (1, 100), dtype=torch.long)
#           torch.onnx.export(
#               model,  # The model to be exported
#               example_input,  # Example input for the model
#               x,# The filename for the exported ONNX model
#               opset_version=11,  # Use an appropriate ONNX opset version
#               input_names=['input'],  # Name of the input layer
#               output_names=['output'],  # Name of the output layer
#               dynamic_axes={
#                   'input': {0: 'batch_size', 1: 'sequence_length'},  # Dynamic axes for variable-length inputs
#                   'output': {0: 'batch_size'}
#               }
#            )
#           return x
           
#     def convert_to_onnx_all(self,n_model,token ,namemodelonxx):

#           model=VitsModel.from_pretrained(n_model,token=token)
#           x=f"dowload_file/{namemodelonxx}.onnx"

#           vocab_size = model.text_encoder.embed_tokens.weight.size(0)
#           example_input = torch.randint(0, vocab_size, (1, 100), dtype=torch.long)
#           torch.onnx.export(
#               model,  # The model to be exported
#               example_input,  # Example input for the model
#               x,  # The filename for the exported ONNX model
#               opset_version=11,  # Use an appropriate ONNX opset version
#               input_names=['input'],  # Name of the input layer
#               output_names=['output'],  # Name of the output layer
#               dynamic_axes={
#                   'input': {0: 'batch_size', 1: 'sequence_length'},  # Dynamic axes for variable-length inputs
#                   'output': {0: 'batch_size'}
#               }
#            )
#           return x
#     def starrt(self):
#         #with gr.Blocks() as demo:
#             with gr.Row():
#                   with gr.Column():
#                     text_n_model=gr.Textbox(label="name model")
#                     text_n_token=gr.Textbox(label="token")
#                     text_n_onxx=gr.Textbox(label="name model onxx")
#                     choice = gr.Dropdown(choices=["decoder", "all anoly decoder", "All"], label="My Dropdown")

#                   with gr.Column():

#                     btn=gr.Button("convert")
#                     label=gr.Label("return  name model onxx")
#                     btn.click(self.function_change,[text_n_model,text_n_token,text_n_onxx,choice],[gr.File(label="Download File")])
#                     #choice.change(fn=function_change, inputs=choice, outputs=label)
#         #return demo
# c=model_onxx()

#3333333333333333333333333333
class OnnxModelConverter:
    def __init__(self):
        self.model = None
    def download_file(self,file_path):
        ff= gr.File(value=file_path, visible=True) 
        file_url = ff.value['url']  
        return file_url

    def convert(self, model_name, token, onnx_filename, conversion_type):
        """
        Main function to handle different types of model conversions.

        Args:
            model_name (str): Name of the model to convert.
            token (str): Access token for loading the model.
            onnx_filename (str): Desired filename for the ONNX output.
            conversion_type (str): Type of conversion ('decoder', 'only_decoder', or 'full_model').

        Returns:
            str: The path to the generated ONNX file.
        """
        if conversion_type == "decoder":
            return self.convert_decoder(model_name, token, onnx_filename)
        elif conversion_type == "only_decoder":
            return self.convert_only_decoder(model_name, token, onnx_filename)
        elif conversion_type == "full_model":
            return self.convert_full_model(model_name, token, onnx_filename)
        else:
            raise ValueError("Invalid conversion type. Choose from 'decoder', 'only_decoder', or 'full_model'.")

    def convert_decoder(self, model_name, token, onnx_filename):
        """
        Converts only the decoder part of the Vits model to ONNX format.

        Args:
            model_name (str): Name of the model to convert.
            token (str): Access token for loading the model.
            onnx_filename (str): Desired filename for the ONNX output.

        Returns:
            str: The path to the generated ONNX file.
        """
        model = VitsModel.from_pretrained(model_name, token=token)
        onnx_file = f"/tmp/{onnx_filename}.onnx"
        vocab_size = model.text_encoder.embed_tokens.weight.size(0)
        example_input = torch.randint(0, vocab_size, (1, 100), dtype=torch.long)

        torch.onnx.export(
            model,
            example_input,
            onnx_file,
            opset_version=11,
            input_names=['input'],
            output_names=['output'],
            dynamic_axes={'input': {0: 'batch_size', 1: 'sequence_length'}, 'output': {0: 'batch_size'}}
        )
        
        return onnx_file


    def convert_only_decoder(self, model_name, token, onnx_filename):
        """
        Converts only the decoder part of the Vits model to ONNX format.

        Args:
            model_name (str): Name of the model to convert.
            token (str): Access token for loading the model.
            onnx_filename (str): Desired filename for the ONNX output.

        Returns:
            str: The path to the generated ONNX file.
        """
        model = Vits_models_only_decoder.from_pretrained(model_name, token=token)
        onnx_file = f"/tmp/{onnx_filename}.onnx"

        vocab_size = model.text_encoder.embed_tokens.weight.size(0)
        example_input = torch.randint(0, vocab_size, (1, 100), dtype=torch.long)

        torch.onnx.export(
            model,
            example_input,
            onnx_file,
            opset_version=11,
            input_names=['input'],
            output_names=['output'],
            dynamic_axes={'input': {0: 'batch_size', 1: 'sequence_length'}, 'output': {0: 'batch_size'}}
        )

        return onnx_file

    def convert_full_model(self, model_name, token, onnx_filename):
        """
        Converts the full Vits model (including encoder and decoder) to ONNX format.

        Args:
            model_name (str): Name of the model to convert.
            token (str): Access token for loading the model.
            onnx_filename (str): Desired filename for the ONNX output.

        Returns:
            str: The path to the generated ONNX file.
        """
        model = VitsModel.from_pretrained(model_name, token=token)
        onnx_file = f"/tmp/{onnx_filename}.onnx"

        vocab_size = model.text_encoder.embed_tokens.weight.size(0)
        example_input = torch.randint(0, vocab_size, (1, 100), dtype=torch.long)

        torch.onnx.export(
            model,
            example_input,
            onnx_file,
            opset_version=11,
            input_names=['input'],
            output_names=['output'],
            dynamic_axes={'input': {0: 'batch_size', 1: 'sequence_length'}, 'output': {0: 'batch_size'}}
        )

        return onnx_file
    def starrt(self):
        with gr.Blocks() as demo:
            with gr.Row():
                  with gr.Column():
                    text_n_model=gr.Textbox(label="name model")
                    text_n_token=gr.Textbox(label="token")
                    text_n_onxx=gr.Textbox(label="name model onxx")
                    choice = gr.Dropdown(choices=["decoder", "only_decoder", "full_model"], label="My Dropdown")
                    
                  with gr.Column():
 
                    btn=gr.Button("convert")
                    label=gr.Label("return  name model onxx")
                    btn.click(self.convert,[text_n_model,text_n_token,text_n_onxx,choice],[gr.File(label="Download File")])
                    #choice.change(fn=function_change, inputs=choice, outputs=label)
        return demo
c=OnnxModelConverter()
###############################################################
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:
    KC=0
    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
        self.models=(self.model,discriminator)

        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_modell(self,namemodel):
         self.model=VitsModel.from_pretrained(namemodel,token=self.token).to(self.device)
         return "true"
    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
        training_args.fp16=True


        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
        # model,discriminator=self.models

        # optimizer = torch.optim.AdamW(
        #     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),
         ('databatchs',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
                    )
def loadd_d():
    token=os.environ.get("key_")
    #model=VitsModel.from_pretrained(n_model,token=token)
    return token
@spaces.GPU(duration=30)
def run_train_epoch(num):
   TrinerModelVITS.KC+=1
   if num >0:
     pro.init_training()
     for i in range(num):
    # model.train(True)
        return  pro.run_train_epoch() +f'- kc={TrinerModelVITS.KC}'
   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'
def read_modell(n_model):
    #model22=Vits_models_only_decoder.from_pretrained(n_model,token)#.to("cuda")
    return token
with gr.Blocks() as interface:
    with gr.Accordion("get token", open=False):
        btn_init = gr.Button("run")
        label=gr.Label("hhh")
        btn_init.click(loadd_d,[],[label])
    with gr.Accordion("onxx ", open=False):
         c.starrt()
    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')