Update app.py

app.py

@@ -38,8 +38,6 @@ from torch.cuda.amp import autocast, GradScaler
 device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 
 
-feature_extractor = VitsFeatureExtractor()
-
 # sgl=get_state_grad_loss(k1=True,#generator=False,
 #                         discriminator=False,
 #                         duration=False
@@ -182,188 +180,63 @@ def  get_data_loader(train_dataset_dirs,eval_dataset_dir,full_generation_dir,dev
                                                         device = device)
      return ctrain_datasets,eval_dataset,full_generation_dataset
 global_step=0
-@spaces.GPU
-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)+1
-        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']:
+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
                     )
 
-                    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
-
+     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):
+     with autocast(enabled=training_args.fp16):
 
               # backpropagate
 
-
-
-
-
               discriminator_target, fmaps_target = discriminator(target_waveform)
 
               discriminator_candidate, fmaps_candidate = discriminator(waveform.detach())
@@ -377,30 +250,30 @@ def trainer_to_cuda(self,
                      labels_padding_mask,
                   )
                   loss_kl=loss_kl*training_args.weight_kl
-                  loss_klall+=loss_kl.detach().item()
+                  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()
+                    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()
+                    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
@@ -410,111 +283,250 @@ def trainer_to_cuda(self,
                       + 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.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
 
-            # optimizer.step()
+          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()
 
-            # print(f"TRAINIG - batch {step}, waveform {(batch['waveform'].shape)}, lr {lr_scheduler.get_last_lr()[0]}... ")
-            # print(f"display loss function  enable  :{displayloss}")
+                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='')
+      
 
-            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
+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
+       
+        
 
-                    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} ... ")
+    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}
 
 
+    
 
-                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='')
+    
 
+    def initialize_training_components(self):
 
+        
+    
+        self.training_args=training_args
 
+        set_seed(training_args.seed)
+        scaler = GradScaler(enabled=training_args.fp16)
+        
 
-          # self.save_pretrained(path_save_model)
+        # 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
+      
 
-    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)
+    
+    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}'
+          
+          
+          
 
-    self.save_pretrained(path_save_model)
 
-    # logger.info("Running final full generations samples... ")
+        
+        # return (self.model,discriminator),(optimizer, disc_optimizer), (lr_scheduler, disc_lr_scheduler)
 
+  
 
 
     # logger.info("***** Training / Inference Done *****")
@@ -560,85 +572,73 @@ train_dataset_dirs=[
 
 
 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
+                    )
 
-
-
-global_step=0
-wandb.login(key= "782b6a6e82bbb5a5348de0d3c7d40d1e76351e79")
-
-
-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="cuda")
-
-print('load Data')
-wandb.init(project= 'AZ')
-
-print('wandb')
-model=VitsModel.from_pretrained(dir_model,token=token).to("cuda")
-print('loadeed')
 @spaces.GPU
-def greet(text,id):
-    global GK 
-    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)
-    print('start')
-    sgl=get_state_grad_loss(mel=True,
-                           # generator=False,
-                        #    discriminator=False,
-                            duration=False)
-    
-    print(training_args)
-    training_args.num_train_epochs=1000
-    training_args.fp16=True
-    training_args.eval_steps=300
-    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
-    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
-    print(device)
-
-
-    for i in range(10000):
-# model.train(True)
-        print(f'clcye epochs ={i}')
-        yield f'clcye epochs ={i}'
-        
-        model=VitsModel.from_pretrained(dir_model,token=token).to("cuda")
+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'
 
-        #dir_model_save=dir_model+'/vend'
-        
-    
-        trainer_to_cuda(model,
-                        ctrain_datasets = ctrain_datasets,
-                        eval_dataset = eval_dataset,
-                        full_generation_dataset = ctrain_datasets[0][0],
-                        feature_extractor = VitsFeatureExtractor(),
-                        training_args = training_args,
-                        full_generation_sample_index= -1,
-                        project_name = "AZ",
-                        wandbKey = "782b6a6e82bbb5a5348de0d3c7d40d1e76351e79",
-                        is_used_text_encoder=True,
-                        is_used_posterior_encode=True,
-                    #    dict_state_grad_loss=sgl,
-                        nk=50,
-                            path_save_model=dir_model,
-                        maf=monotonic_align.maximum_path,
+@spaces.GPU
+def init_Starting():
+    pro.init_Starting()
+    return 'init_Starting'
+@spaces.GPU
+def init_wandb():
+    pro.init_wandb()
+    return 'init_wandb'
     
-                        n_back_save_model=3000,
-                        start_speeker=0,
-                        end_speeker=1,
-                        n_epoch=i*training_args.num_train_epochs,
+@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])
     
-                                      )
+        
 
+        
 
-demo = gr.Interface(fn=greet, inputs=["text","text"], outputs="text")
-demo.launch()
+interface.launch()
+print('loadeed')