Update app.py

app.py

@@ -407,7 +407,47 @@ class TrinerModelVITS:
         self.len_dataset=len(self.DataSets['train'])
         self.load_model()
         self.init_wandb()
+        self.training_args=load_training_args(self.path_training_args)
+        training_args= self.training_args
         scaler = GradScaler(enabled=True)
+        for disc in self.model.discriminator.discriminators:
+            disc.apply_weight_norm()
+        self.model.decoder.apply_weight_norm()
+        # torch.nn.utils.weight_norm(self.decoder.conv_pre)
+        # torch.nn.utils.weight_norm(self.decoder.conv_post)
+        for flow in self.model.flow.flows:
+            torch.nn.utils.weight_norm(flow.conv_pre)
+            torch.nn.utils.weight_norm(flow.conv_post)
+
+        discriminator = self.model.discriminator
+        self.model.discriminator = None
+
+        optimizer = torch.optim.AdamW(
+            self.model.parameters(),
+            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)