VitsModelSplit/vits_model2.py

import numpy as np
import torch
from torch import nn
import math
from typing import Any, Callable, Optional, Tuple, Union
from torch.cuda.amp import autocast, GradScaler

from .vits_config import VitsConfig,VitsPreTrainedModel
from .flow import VitsResidualCouplingBlock
from .duration_predictor import VitsDurationPredictor, VitsStochasticDurationPredictor
from .encoder import VitsTextEncoder
from .decoder import VitsHifiGan
from .posterior_encoder import VitsPosteriorEncoder
from .discriminator import VitsDiscriminator
from .vits_output import VitsModelOutput, VitsTrainingOutput
from  .dataset_features_collector import FeaturesCollectionDataset
from .feature_extraction import VitsFeatureExtractor

import os
import sys
from typing import Optional
import tempfile
from torch.cuda.amp import autocast, GradScaler

from IPython.display import clear_output
from transformers import set_seed
import wandb
import logging
import copy
Lst=['input_ids',
 'attention_mask',
 'waveform',
 'labels',
 'labels_attention_mask',
 'mel_scaled_input_features']

def covert_cuda_batch(d):
  #return d
  for key in Lst:
      d[key]=d[key].cuda(non_blocking=True)
  # for key in d['text_encoder_output']:
  #   d['text_encoder_output'][key]=d['text_encoder_output'][key].cuda(non_blocking=True)
  for key in d['posterior_encode_output']:
    d['posterior_encode_output'][key]=d['posterior_encode_output'][key].cuda(non_blocking=True)

  return d
def generator_loss(disc_outputs):
    total_loss = 0
    gen_losses = []
    for disc_output in disc_outputs:
        disc_output = disc_output
        loss = torch.mean((1 - disc_output) ** 2)
        gen_losses.append(loss)
        total_loss += loss

    return total_loss, gen_losses

def discriminator_loss(disc_real_outputs, disc_generated_outputs):
    loss = 0
    real_losses = 0
    generated_losses = 0
    for disc_real, disc_generated in zip(disc_real_outputs, disc_generated_outputs):
        real_loss = torch.mean((1 - disc_real) ** 2)
        generated_loss = torch.mean(disc_generated**2)
        loss += real_loss + generated_loss
        real_losses += real_loss
        generated_losses += generated_loss

    return loss, real_losses, generated_losses

def feature_loss(feature_maps_real, feature_maps_generated):
    loss = 0
    for feature_map_real, feature_map_generated in zip(feature_maps_real, feature_maps_generated):
        for real, generated in zip(feature_map_real, feature_map_generated):
            real = real.detach()
            loss += torch.mean(torch.abs(real - generated))

    return loss * 2
def kl_loss(z_p, logs_q, m_p, logs_p, z_mask):
  """
  z_p, logs_q: [b, h, t_t]
  m_p, logs_p: [b, h, t_t]
  """
  z_p = z_p.float()
  logs_q = logs_q.float()
  m_p = m_p.float()
  logs_p = logs_p.float()
  z_mask = z_mask.float()

  kl = logs_p - logs_q - 0.5
  kl += 0.5 * ((z_p - m_p)**2) * torch.exp(-2. * logs_p)
  kl = torch.sum(kl * z_mask)
  l = kl / torch.sum(z_mask)
  return l
#.............................................
# def kl_loss(prior_latents, posterior_log_variance, prior_means, prior_log_variance, labels_mask):


#   kl = prior_log_variance - posterior_log_variance - 0.5
#   kl += 0.5 * ((prior_latents - prior_means) ** 2) * torch.exp(-2.0 * prior_log_variance)
#   kl = torch.sum(kl * labels_mask)
#   loss = kl / torch.sum(labels_mask)
#   return loss

def get_state_grad_loss(k1=True,
                             mel=True,
                             duration=True,
                             generator=True,
                             discriminator=True):
                             return {'k1':k1,'mel':mel,'duration':duration,'generator':generator,'discriminator':discriminator}


def clip_grad_value_(parameters, clip_value, norm_type=2):
  if isinstance(parameters, torch.Tensor):
    parameters = [parameters]
  parameters = list(filter(lambda p: p.grad is not None, parameters))
  norm_type = float(norm_type)
  if clip_value is not None:
    clip_value = float(clip_value)

  total_norm = 0
  for p in parameters:
    param_norm = p.grad.data.norm(norm_type)
    total_norm += param_norm.item() ** norm_type
    if clip_value is not None:
      p.grad.data.clamp_(min=-clip_value, max=clip_value)
  total_norm = total_norm ** (1. / norm_type)
  return total_norm


class VitsModel(VitsPreTrainedModel):

    def __init__(self, config: VitsConfig):
        super().__init__(config)

        self.config = config
        self.text_encoder = VitsTextEncoder(config)
        self.flow = VitsResidualCouplingBlock(config)
        self.decoder = VitsHifiGan(config)



        if config.use_stochastic_duration_prediction:
            self.duration_predictor = VitsStochasticDurationPredictor(config)
        else:
            self.duration_predictor = VitsDurationPredictor(config)

        if config.num_speakers > 1:
            self.embed_speaker = nn.Embedding(config.num_speakers, config.speaker_embedding_size)

        # This is used only for training.
        self.posterior_encoder = VitsPosteriorEncoder(config)
        self.discriminator = VitsDiscriminator(config)

        # These parameters control the synthesised speech properties
        self.speaking_rate = config.speaking_rate
        self.noise_scale = config.noise_scale
        self.noise_scale_duration = config.noise_scale_duration
        self.segment_size = self.config.segment_size // self.config.hop_length

        # Initialize weights and apply final processing
        self.post_init()
        self.monotonic_alignment_function=self.monotonic_align_max_path



    #....................................
    def setMfA(self,fn):
      self.monotonic_alignment_function=fn



    def monotonic_align_max_path(self,log_likelihoods, mask):
        # used for training - awfully slow
        # an alternative is proposed in examples/pytorch/text-to-speech/run_vits_finetuning.py
        path = torch.zeros_like(log_likelihoods)

        text_length_maxs = mask.sum(1)[:, 0]
        latent_length_maxs = mask.sum(2)[:, 0]

        indexes = latent_length_maxs - 1

        max_neg_val = -1e9

        for batch_id in range(len(path)):
            index = int(indexes[batch_id].item())
            text_length_max = int(text_length_maxs[batch_id].item())
            latent_length_max = int(latent_length_maxs[batch_id].item())

            for y in range(text_length_max):
                for x in range(max(0, latent_length_max + y - text_length_max), min(latent_length_max, y + 1)):
                    if x == y:
                        v_cur = max_neg_val
                    else:
                        v_cur = log_likelihoods[batch_id, y - 1, x]
                    if x == 0:
                        if y == 0:
                            v_prev = 0.0
                        else:
                            v_prev = max_neg_val
                    else:
                        v_prev = log_likelihoods[batch_id, y - 1, x - 1]
                    log_likelihoods[batch_id, y, x] += max(v_prev, v_cur)

            for y in range(text_length_max - 1, -1, -1):
                path[batch_id, y, index] = 1
                if index != 0 and (
                    index == y or log_likelihoods[batch_id, y - 1, index] < log_likelihoods[batch_id, y - 1, index - 1]
                ):
                    index = index - 1
        return path

    #....................................

    def slice_segments(self,hidden_states, ids_str, segment_size=4):

        batch_size, channels, _ = hidden_states.shape
        # 1d tensor containing the indices to keep
        indices = torch.arange(segment_size).to(ids_str.device)
        # extend the indices to match the shape of hidden_states
        indices = indices.view(1, 1, -1).expand(batch_size, channels, -1)
        # offset indices with ids_str
        indices = indices + ids_str.view(-1, 1, 1)
        # gather indices
        output = torch.gather(hidden_states, dim=2, index=indices)

        return output


    #....................................


    def rand_slice_segments(self,hidden_states, sample_lengths=None, segment_size=4):

        batch_size, _, seq_len = hidden_states.size()
        if sample_lengths is None:
            sample_lengths = seq_len
        ids_str_max = sample_lengths - segment_size + 1
        ids_str = (torch.rand([batch_size]).to(device=hidden_states.device) * ids_str_max).to(dtype=torch.long)
        ret = self.slice_segments(hidden_states, ids_str, segment_size)

        return ret, ids_str

    #....................................

    def resize_speaker_embeddings(
        self,
        new_num_speakers: int,
        speaker_embedding_size: Optional[int] = None,
        pad_to_multiple_of: Optional[int] = 2,
    ):
        if pad_to_multiple_of is not None:
            new_num_speakers = ((new_num_speakers + pad_to_multiple_of - 1) // pad_to_multiple_of) * pad_to_multiple_of

        # first, take care of embed_speaker
        if self.config.num_speakers <= 1:
            if speaker_embedding_size is None:
                raise ValueError(
                    "The current model had no previous speaker embedding, but `speaker_embedding_size` is not specified. Pass `speaker_embedding_size` to this method."
                )
            # create new embedding layer
            new_embeddings = nn.Embedding(
                new_num_speakers,
                speaker_embedding_size,
                device=self.device,
            )
            # initialize all new embeddings
            self._init_weights(new_embeddings)
        else:
            new_embeddings = self._get_resized_embeddings(self.embed_speaker, new_num_speakers)

        self.embed_speaker = new_embeddings

        # then take care of sub-models
        self.flow.resize_speaker_embeddings(speaker_embedding_size)
        for flow in self.flow.flows:
            self._init_weights(flow.wavenet.cond_layer)

        self.decoder.resize_speaker_embedding(speaker_embedding_size)
        self._init_weights(self.decoder.cond)

        self.duration_predictor.resize_speaker_embeddings(speaker_embedding_size)
        self._init_weights(self.duration_predictor.cond)

        self.posterior_encoder.resize_speaker_embeddings(speaker_embedding_size)
        self._init_weights(self.posterior_encoder.wavenet.cond_layer)

        self.config.num_speakers = new_num_speakers
        self.config.speaker_embedding_size = speaker_embedding_size

    #....................................

    def get_input_embeddings(self):
        return self.text_encoder.get_input_embeddings()

    #....................................

    def set_input_embeddings(self, value):
        self.text_encoder.set_input_embeddings(value)

    #....................................

    def apply_weight_norm(self):
        self.decoder.apply_weight_norm()
        self.flow.apply_weight_norm()
        self.posterior_encoder.apply_weight_norm()

    #....................................

    def remove_weight_norm(self):
        self.decoder.remove_weight_norm()
        self.flow.remove_weight_norm()
        self.posterior_encoder.remove_weight_norm()

    #....................................

    def discriminate(self, hidden_states):
        return self.discriminator(hidden_states)

    #....................................

    def get_encoder(self):
        return self.text_encoder

    #....................................

    def _inference_forward(
        self,
        input_ids: Optional[torch.Tensor] = None,
        attention_mask: Optional[torch.Tensor] = None,
        speaker_embeddings: Optional[torch.Tensor] = None,
        output_attentions: Optional[bool] = None,
        output_hidden_states: Optional[bool] = None,
        return_dict: Optional[bool] = None,
        padding_mask: Optional[torch.Tensor] = None,
    ):
        text_encoder_output = self.text_encoder(
            input_ids=input_ids,
            padding_mask=padding_mask,
            attention_mask=attention_mask,
            output_attentions=output_attentions,
            output_hidden_states=output_hidden_states,
            return_dict=return_dict,
        )
        hidden_states = text_encoder_output[0] if not return_dict else text_encoder_output.last_hidden_state
        hidden_states = hidden_states.transpose(1, 2)
        input_padding_mask = padding_mask.transpose(1, 2)

        prior_means = text_encoder_output[1] if not return_dict else text_encoder_output.prior_means
        prior_log_variances = text_encoder_output[2] if not return_dict else text_encoder_output.prior_log_variances

        if self.config.use_stochastic_duration_prediction:
            log_duration = self.duration_predictor(
                hidden_states,
                input_padding_mask,
                speaker_embeddings,
                reverse=True,
                noise_scale=self.noise_scale_duration,
            )
        else:
            log_duration = self.duration_predictor(hidden_states, input_padding_mask, speaker_embeddings)

        length_scale = 1.0 / self.speaking_rate
        duration = torch.ceil(torch.exp(log_duration) * input_padding_mask * length_scale)
        predicted_lengths = torch.clamp_min(torch.sum(duration, [1, 2]), 1).long()


        # Create a padding mask for the output lengths of shape (batch, 1, max_output_length)
        indices = torch.arange(predicted_lengths.max(), dtype=predicted_lengths.dtype, device=predicted_lengths.device)
        output_padding_mask = indices.unsqueeze(0) < predicted_lengths.unsqueeze(1)
        output_padding_mask = output_padding_mask.unsqueeze(1).to(input_padding_mask.dtype)

        # Reconstruct an attention tensor of shape (batch, 1, out_length, in_length)
        attn_mask = torch.unsqueeze(input_padding_mask, 2) * torch.unsqueeze(output_padding_mask, -1)
        batch_size, _, output_length, input_length = attn_mask.shape
        cum_duration = torch.cumsum(duration, -1).view(batch_size * input_length, 1)
        indices = torch.arange(output_length, dtype=duration.dtype, device=duration.device)
        valid_indices = indices.unsqueeze(0) < cum_duration
        valid_indices = valid_indices.to(attn_mask.dtype).view(batch_size, input_length, output_length)
        padded_indices = valid_indices - nn.functional.pad(valid_indices, [0, 0, 1, 0, 0, 0])[:, :-1]
        attn = padded_indices.unsqueeze(1).transpose(2, 3) * attn_mask

        # Expand prior distribution
        prior_means = torch.matmul(attn.squeeze(1), prior_means).transpose(1, 2)
        prior_log_variances = torch.matmul(attn.squeeze(1), prior_log_variances).transpose(1, 2)

        prior_latents = prior_means + torch.randn_like(prior_means) * torch.exp(prior_log_variances) * self.noise_scale
        latents = self.flow(prior_latents, output_padding_mask, speaker_embeddings, reverse=True)

        spectrogram = latents * output_padding_mask
        waveform = self.decoder(spectrogram, speaker_embeddings)
        waveform = waveform.squeeze(1)
        sequence_lengths = predicted_lengths * np.prod(self.config.upsample_rates)

        if not return_dict:
            outputs = (waveform, sequence_lengths, spectrogram) + text_encoder_output[3:]
            return outputs

        return VitsModelOutput(
            waveform=waveform,
            sequence_lengths=sequence_lengths,
            spectrogram=spectrogram,
            hidden_states=text_encoder_output.hidden_states,
            attentions=text_encoder_output.attentions,
        )

    #....................................

    def forward(
        self,
        input_ids: Optional[torch.Tensor] = None,
        attention_mask: Optional[torch.Tensor] = None,
        speaker_id: Optional[int] = None,
        output_attentions: Optional[bool] = None,
        output_hidden_states: Optional[bool] = None,
        return_dict: Optional[bool] = None,
        labels: Optional[torch.FloatTensor] = None,
        labels_attention_mask: Optional[torch.Tensor] = None,
        monotonic_alignment_function: Optional[Callable] = None,
    ) -> Union[Tuple[Any], VitsModelOutput]:

        output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
        output_hidden_states = (
            output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
        )
        return_dict = return_dict if return_dict is not None else self.config.use_return_dict

        monotonic_alignment_function = (
            self.monotonic_align_max_path if monotonic_alignment_function is None else monotonic_alignment_function
        )

        if attention_mask is not None:
            input_padding_mask = attention_mask.unsqueeze(-1).float()
        else:
            input_padding_mask = torch.ones_like(input_ids).unsqueeze(-1).float()

        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}.")
            if not (len(speaker_id) == 1 or len(speaker_id == len(input_ids))):
                raise ValueError(
                    f"You passed {len(speaker_id)} `speaker_id` but you should either pass one speaker id or `batch_size` `speaker_id`."
                )

            speaker_embeddings = self.embed_speaker(speaker_id).unsqueeze(-1)
        else:
            speaker_embeddings = None

        # if inference, return inference forward of VitsModel
        if labels is None:
            return self._inference_forward(
                input_ids,
                attention_mask,
                speaker_embeddings,
                output_attentions,
                output_hidden_states,
                return_dict,
                input_padding_mask,
            )

        if labels_attention_mask is not None:
            labels_padding_mask = labels_attention_mask.unsqueeze(1).float()
        else:
            labels_attention_mask = torch.ones((labels.shape[0], labels.shape[2])).float().to(self.device)
            labels_padding_mask = labels_attention_mask.unsqueeze(1)

        text_encoder_output = self.text_encoder(
            input_ids=input_ids,
            padding_mask=input_padding_mask,
            attention_mask=attention_mask,
            output_attentions=output_attentions,
            output_hidden_states=output_hidden_states,
            return_dict=return_dict,
        )
        hidden_states = text_encoder_output[0] if not return_dict else text_encoder_output.last_hidden_state
        hidden_states = hidden_states.transpose(1, 2)
        input_padding_mask = input_padding_mask.transpose(1, 2)
        prior_means = text_encoder_output[1] if not return_dict else text_encoder_output.prior_means
        prior_log_variances = text_encoder_output[2] if not return_dict else text_encoder_output.prior_log_variances

        latents, posterior_means, posterior_log_variances = self.posterior_encoder(
            labels, labels_padding_mask, speaker_embeddings
        )
        prior_latents = self.flow(latents, labels_padding_mask, speaker_embeddings, reverse=False)

        prior_means, prior_log_variances = prior_means.transpose(1, 2), prior_log_variances.transpose(1, 2)
        with torch.no_grad():
            # negative cross-entropy

            # [batch_size, d, latent_length]
            prior_variances = torch.exp(-2 * prior_log_variances)
            # [batch_size, 1, latent_length]
            neg_cent1 = torch.sum(-0.5 * math.log(2 * math.pi) - prior_log_variances, [1], keepdim=True)
            # [batch_size, text_length, d] x [batch_size, d, latent_length] = [batch_size, text_length, latent_length]
            neg_cent2 = torch.matmul(-0.5 * (prior_latents**2).transpose(1, 2), prior_variances)
            # [batch_size, text_length, d] x [batch_size, d, latent_length] = [batch_size, text_length, latent_length]
            neg_cent3 = torch.matmul(prior_latents.transpose(1, 2), (prior_means * prior_variances))
            # [batch_size, 1, latent_length]
            neg_cent4 = torch.sum(-0.5 * (prior_means**2) * prior_variances, [1], keepdim=True)

            # [batch_size, text_length, latent_length]
            neg_cent = neg_cent1 + neg_cent2 + neg_cent3 + neg_cent4

            attn_mask = torch.unsqueeze(input_padding_mask, 2) * torch.unsqueeze(labels_padding_mask, -1)

            attn = monotonic_alignment_function(neg_cent, attn_mask.squeeze(1)).unsqueeze(1).detach()

        durations = attn.sum(2)

        if self.config.use_stochastic_duration_prediction:
            log_duration = self.duration_predictor(
                hidden_states, input_padding_mask, speaker_embeddings, durations=durations, reverse=False
            )
            log_duration = log_duration / torch.sum(input_padding_mask)
        else:
            log_duration_padded = torch.log(durations + 1e-6) * input_padding_mask
            log_duration = self.duration_predictor(hidden_states, input_padding_mask, speaker_embeddings)
            log_duration = torch.sum((log_duration - log_duration_padded) ** 2, [1, 2]) / torch.sum(input_padding_mask)

        # expand priors
        prior_means = torch.matmul(attn.squeeze(1), prior_means.transpose(1, 2)).transpose(1, 2)
        prior_log_variances = torch.matmul(attn.squeeze(1), prior_log_variances.transpose(1, 2)).transpose(1, 2)

        label_lengths = labels_attention_mask.sum(dim=1)
        latents_slice, ids_slice = self.rand_slice_segments(latents, label_lengths, segment_size=self.segment_size)

        waveform = self.decoder(latents_slice, speaker_embeddings)

        if not return_dict:
            outputs = (
                waveform,
                log_duration,
                attn,
                ids_slice,
                input_padding_mask,
                labels_padding_mask,
                latents,
                prior_latents,
                prior_means,
                prior_log_variances,
                posterior_means,
                posterior_log_variances,
            )
            return outputs

        return VitsTrainingOutput(
            waveform=waveform,
            log_duration=log_duration,
            attn=attn,
            ids_slice=ids_slice,
            input_padding_mask=input_padding_mask,
            labels_padding_mask=labels_padding_mask,
            latents=latents,
            prior_latents=prior_latents,
            prior_means=prior_means,
            prior_log_variances=prior_log_variances,
            posterior_means=posterior_means,
            posterior_log_variances=posterior_log_variances,
        )
    def slice_segments(self,hidden_states, ids_str, segment_size=4):

        batch_size, channels, _ = hidden_states.shape
        # 1d tensor containing the indices to keep
        indices = torch.arange(segment_size).to(ids_str.device)
        # extend the indices to match the shape of hidden_states
        indices = indices.view(1, 1, -1).expand(batch_size, channels, -1)
        # offset indices with ids_str
        indices = indices + ids_str.view(-1, 1, 1)
        # gather indices
        output = torch.gather(hidden_states, dim=2, index=indices)

        return output

    #....................................

    def rand_slice_segments(self,hidden_states, sample_lengths=None, segment_size=4):
        batch_size, _, seq_len = hidden_states.size()
        if sample_lengths is None:
            sample_lengths = seq_len
        ids_str_max = sample_lengths - segment_size + 1
        ids_str = (torch.rand([batch_size]).to(device=hidden_states.device) * ids_str_max).to(dtype=torch.long)
        ret = self.slice_segments(hidden_states, ids_str, segment_size)

        return ret, ids_str

    def forward_k(
        self,
        input_ids: Optional[torch.Tensor] = None,
        attention_mask: Optional[torch.Tensor] = None,
        speaker_id: Optional[int] = None,
        output_attentions: Optional[bool] = None,
        output_hidden_states: Optional[bool] = None,
        return_dict: Optional[bool] = None,
        labels: Optional[torch.FloatTensor] = None,
        labels_attention_mask: Optional[torch.Tensor] = None,
        text_encoder_output=None,
        posterior_encode_output=None,
        monotonic_alignment_function: Optional[Callable] = None,
    ) -> Union[Tuple[Any], VitsModelOutput]:

        output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
        output_hidden_states = (
            output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
        )
        return_dict = return_dict if return_dict is not None else self.config.use_return_dict

        monotonic_alignment_function = (
            self.monotonic_align_max_path if monotonic_alignment_function is None else monotonic_alignment_function
        )

        if attention_mask is not None:
            input_padding_mask = attention_mask.unsqueeze(-1).float()
        else:
            input_padding_mask = torch.ones_like(input_ids).unsqueeze(-1).float()

        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}.")
            if not (len(speaker_id) == 1 or len(speaker_id == len(input_ids))):
                raise ValueError(
                    f"You passed {len(speaker_id)} `speaker_id` but you should either pass one speaker id or `batch_size` `speaker_id`."
                )

            speaker_embeddings = self.embed_speaker(speaker_id).unsqueeze(-1)
        else:
            speaker_embeddings = None

        # if inference, return inference forward of VitsModel
        if labels is None:
            return self._inference_forward(
                input_ids,
                attention_mask,
                speaker_embeddings,
                output_attentions,
                output_hidden_states,
                return_dict,
                input_padding_mask,
            )

        if labels_attention_mask is not None:
            labels_padding_mask = labels_attention_mask.unsqueeze(1).float()
        else:
            labels_attention_mask = torch.ones((labels.shape[0], labels.shape[2])).float().to(self.device)
            labels_padding_mask = labels_attention_mask.unsqueeze(1)
        if text_encoder_output is None:
            text_encoder_output = self.text_encoder(
                input_ids=input_ids,
                padding_mask=input_padding_mask,
                attention_mask=attention_mask,
                output_attentions=output_attentions,
                output_hidden_states=output_hidden_states,
                return_dict=return_dict,
            )
        hidden_states = text_encoder_output[0] if not return_dict else text_encoder_output.last_hidden_state
        hidden_states = hidden_states.transpose(1, 2)
        input_padding_mask = input_padding_mask.transpose(1, 2)
        prior_means = text_encoder_output[1] if not return_dict else text_encoder_output.prior_means
        prior_log_variances = text_encoder_output[2] if not return_dict else text_encoder_output.prior_log_variances
        if posterior_encode_output is None:
            latents, posterior_means, posterior_log_variances = self.posterior_encoder(
                labels, labels_padding_mask, speaker_embeddings
            )
        else:
          latents=posterior_encode_output['posterior_latents']
          posterior_means=posterior_encode_output['posterior_means']
          posterior_log_variances=posterior_encode_output['posterior_log_variances']

        prior_latents = self.flow(latents, labels_padding_mask, speaker_embeddings, reverse=False)

        prior_means, prior_log_variances = prior_means.transpose(1, 2), prior_log_variances.transpose(1, 2)
        with torch.no_grad():
            # negative cross-entropy

            # [batch_size, d, latent_length]
            prior_variances = torch.exp(-2 * prior_log_variances)
            # [batch_size, 1, latent_length]
            neg_cent1 = torch.sum(-0.5 * math.log(2 * math.pi) - prior_log_variances, [1], keepdim=True)
            # [batch_size, text_length, d] x [batch_size, d, latent_length] = [batch_size, text_length, latent_length]
            neg_cent2 = torch.matmul(-0.5 * (prior_latents**2).transpose(1, 2), prior_variances)
            # [batch_size, text_length, d] x [batch_size, d, latent_length] = [batch_size, text_length, latent_length]
            neg_cent3 = torch.matmul(prior_latents.transpose(1, 2), (prior_means * prior_variances))
            # [batch_size, 1, latent_length]
            neg_cent4 = torch.sum(-0.5 * (prior_means**2) * prior_variances, [1], keepdim=True)

            # [batch_size, text_length, latent_length]
            neg_cent = neg_cent1 + neg_cent2 + neg_cent3 + neg_cent4

            attn_mask = torch.unsqueeze(input_padding_mask, 2) * torch.unsqueeze(labels_padding_mask, -1)

            attn = monotonic_alignment_function(neg_cent, attn_mask.squeeze(1)).unsqueeze(1).detach()

        durations = attn.sum(2)

        if self.config.use_stochastic_duration_prediction:
            log_duration = self.duration_predictor(
                hidden_states, input_padding_mask, speaker_embeddings, durations=durations, reverse=False
            )
            log_duration = log_duration / torch.sum(input_padding_mask)
        else:
            log_duration_padded = torch.log(durations + 1e-6) * input_padding_mask
            log_duration = self.duration_predictor(hidden_states, input_padding_mask, speaker_embeddings)
            log_duration = torch.sum((log_duration - log_duration_padded) ** 2, [1, 2]) / torch.sum(input_padding_mask)

        # expand priors
        prior_means = torch.matmul(attn.squeeze(1), prior_means.transpose(1, 2)).transpose(1, 2)
        prior_log_variances = torch.matmul(attn.squeeze(1), prior_log_variances.transpose(1, 2)).transpose(1, 2)

        label_lengths = labels_attention_mask.sum(dim=1)
        latents_slice, ids_slice = self.rand_slice_segments(latents, label_lengths, segment_size=self.segment_size)

        waveform = self.decoder(latents_slice, speaker_embeddings)

        if not return_dict:
            outputs = (
                waveform,
                log_duration,
                attn,
                ids_slice,
                input_padding_mask,
                labels_padding_mask,
                latents,
                prior_latents,
                prior_means,
                prior_log_variances,
                posterior_means,
                posterior_log_variances,
            )
            return outputs

        return VitsTrainingOutput(
            waveform=waveform,
            log_duration=log_duration,
            attn=attn,
            ids_slice=ids_slice,
            input_padding_mask=input_padding_mask,
            labels_padding_mask=labels_padding_mask,
            latents=latents,
            prior_latents=prior_latents,
            prior_means=prior_means,
            prior_log_variances=prior_log_variances,
            posterior_means=posterior_means,
            posterior_log_variances=posterior_log_variances,
        )

    def trainer(self,
              train_dataset_dir = None,
              eval_dataset_dir = None,
              full_generation_dir = 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


              ):


        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()


        set_seed(training_args.seed)
        # Apply Weight Norm Decoder
        # self.apply_weight_norm()
        # Save Config
        self.config.save_pretrained(training_args.output_dir)

        train_dataset = FeaturesCollectionDataset(dataset_dir = train_dataset_dir,
                                                  device = self.device
                                                  )

        eval_dataset = None
        if training_args.do_eval:
            eval_dataset = FeaturesCollectionDataset(dataset_dir = eval_dataset_dir,
                                                     device = self.device
                                                     )

        full_generation_dataset = FeaturesCollectionDataset(dataset_dir = full_generation_dir,
                                                            device = self.device
                                                            )
        self.full_generation_sample = full_generation_dataset[full_generation_sample_index]

        # init optimizer, lr_scheduler

        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(
        #         self.discriminator.parameters(),
        #         training_args.learning_rate,
        #         betas=[training_args.adam_beta1, training_args.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............................

        global_step = 0

        for epoch in range(training_args.num_train_epochs):
            train_losses_sum = 0
            lr_scheduler.step()
          #  disc_lr_scheduler.step()
            print(f"  Num Epochs = {epoch}")
            if epoch%nk==0:
              print('Save checkpoints Model :',int(epoch/nk))
              self.save_pretrained(path_save_model)




            for step, batch in enumerate(train_dataset):

                # forward through model
                # outputs = self.forward(
                #     labels=batch["labels"],
                #     labels_attention_mask=batch["labels_attention_mask"],
                #     speaker_id=batch["speaker_id"]
                #     )
                #if step==10:break

                model_outputs = self.forward_k(
                    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"],
                    text_encoder_output =None  if is_used_text_encoder else batch['text_encoder_output'],
                    posterior_encode_output=None  if is_used_posterior_encode else batch['posterior_encode_output'],
                    return_dict=True,
                    monotonic_alignment_function=maf,
                )

                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]

                target_waveform = batch["waveform"].transpose(1, 2)
                target_waveform = self.slice_segments(
                                    target_waveform,
                                    model_outputs.ids_slice * feature_extractor.hop_length,
                                    self.config.segment_size
                                )
                optimizer.zero_grad()

                displayloss={}
                # backpropagate
                if dict_state_grad_loss['k1']:
                    loss_kl = kl_loss(
                        model_outputs.prior_latents,
                        model_outputs.posterior_log_variances,
                        model_outputs.prior_means,
                        model_outputs.prior_log_variances,
                        model_outputs.labels_padding_mask,
                    )
                    loss_kl=loss_kl*training_args.weight_kl
                    displayloss['loss_kl']=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)
                    displayloss['loss_mel'] = 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(model_outputs.log_duration)*training_args.weight_duration
                    displayloss['loss_duration'] = loss_duration.detach().item()
                  #  if displayloss['loss_duration']>=0:
                     #   loss_duration.backward()

                discriminator_target, fmaps_target = self.discriminator(target_waveform)
                discriminator_candidate, fmaps_candidate = self.discriminator(model_outputs.waveform.detach())
                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()
                discriminator_target, fmaps_target = self.discriminator(target_waveform)

                discriminator_candidate, fmaps_candidate = self.discriminator(model_outputs.waveform.detach())

                if dict_state_grad_loss['generator']:
                    loss_fmaps = feature_loss(fmaps_target, fmaps_candidate)
                    loss_gen, losses_gen = generator_loss(discriminator_candidate)
                    loss_gen=loss_gen * training_args.weight_gen
                    displayloss['loss_gen'] = loss_gen.detach().item()
                 #   loss_gen.backward(retain_graph=True)
                    loss_fmaps=loss_fmaps * training_args.weight_fmaps
                    displayloss['loss_fmaps'] = loss_fmaps.detach().item()
                 #   loss_fmaps.backward(retain_graph=True)
                    total_generator_loss = (
                        loss_duration
                        + loss_mel
                        + loss_kl
                        + loss_fmaps
                        + loss_gen
                    )
                    total_generator_loss.backward()





                optimizer.step()




                print(f"TRAINIG - batch {step}, waveform {(batch['waveform'].shape)}, lr {lr_scheduler.get_last_lr()[0]}... ")
                print(f"display loss function  enable  :{displayloss}")

                global_step +=1

                # validation

                do_eval = training_args.do_eval and (global_step % training_args.eval_steps == 0)
                if do_eval:
                    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,
                            monotonic_alignment_function=None,
                                )

                        mel_scaled_labels = batch["mel_scaled_input_features"]
                        mel_scaled_target = self.slice_segments(mel_scaled_labels, model_outputs.ids_slice,self.segment_size)
                        mel_scaled_generation = feature_extractor._torch_extract_fbank_features(model_outputs.waveform.squeeze(1))[1]
                        loss = loss_mel.detach().item()
                        eval_losses_sum +=loss

                        loss_mel = torch.nn.functional.l1_loss(mel_scaled_target, mel_scaled_generation)
                        print(f"VALIDATION - batch {step}, waveform {(batch['waveform'].shape)}, step_loss_mel {loss} ... ")



                    with torch.no_grad():
                        full_generation_sample = self.full_generation_sample
                        full_generation =self.forward(
                          input_ids =full_generation_sample["input_ids"],
                          attention_mask=full_generation_sample["attention_mask"],
                          speaker_id=full_generation_sample["speaker_id"]
                          )

                        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],})

            wandb.log({"train_losses":train_losses_sum})

        # add weight norms
        # self.remove_weight_norm()

        try:
          torch.save(self.posterior_encoder.state_dict(), os.path.join(training_args.output_dir,"posterior_encoder.pt"))
          torch.save(self.decoder.state_dict(), os.path.join(training_args.output_dir,"decoder.pt"))
        except:pass


        logger.info("Running final full generations samples... ")


        with torch.no_grad():

            full_generation_sample = self.full_generation_sample
            full_generation = self.forward(
                    input_ids=full_generation_sample["labels"],
                    attention_mask=full_generation_sample["labels_attention_mask"],
                    speaker_id=full_generation_sample["speaker_id"]
                    )

            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],
                })


        logger.info("***** Training / Inference Done *****")

    #....................................


    def trainer_to_cuda(self,
              train_dataset_dir = None,
              eval_dataset_dir = None,
              full_generation_dir = 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


              ):


        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()


        set_seed(training_args.seed)
        scaler = GradScaler(enabled=training_args.fp16)

        # Apply Weight Norm Decoder
        # self.apply_weight_norm()
        # Save Config
        self.config.save_pretrained(training_args.output_dir)

        train_dataset = FeaturesCollectionDataset(dataset_dir = train_dataset_dir,
                                                  device = self.device
                                                  )

        eval_dataset = None
        if training_args.do_eval:
            eval_dataset = FeaturesCollectionDataset(dataset_dir = eval_dataset_dir,
                                                     device = self.device
                                                     )

        full_generation_dataset = FeaturesCollectionDataset(dataset_dir = full_generation_dir,
                                                            device = self.device
                                                            )
        self.full_generation_sample = full_generation_dataset[full_generation_sample_index]

        # init optimizer, lr_scheduler
        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............................

        global_step = 0

        for epoch in range(training_args.num_train_epochs):
            train_losses_sum = 0
            lr_scheduler.step()

            disc_lr_scheduler.step()
            print(f"  Num Epochs = {epoch}")
            if (epoch+1)%nk==0:
              clear_output()
              print('Save checkpoints Model :',int(epoch/nk))
              self.discriminator=discriminator

              self.save_pretrained(path_save_model)
              self.discriminator=None




            for step, batch in enumerate(train_dataset):

                # forward through model
                # outputs = self.forward(
                #     labels=batch["labels"],
                #     labels_attention_mask=batch["labels_attention_mask"],
                #     speaker_id=batch["speaker_id"]
                #     )
                #if step==10:break
                batch=covert_cuda_batch(batch)

                with autocast(enabled=training_args.fp16):


                  model_outputs = self.forward_k(
                      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"],
                      text_encoder_output =None  if is_used_text_encoder else batch['text_encoder_output'],
                      posterior_encode_output=None  if is_used_posterior_encode else batch['posterior_encode_output'],
                      return_dict=True,
                      monotonic_alignment_function= maf,
                  )

                  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]

                  target_waveform = batch["waveform"].transpose(1, 2)
                  target_waveform = self.slice_segments(
                                      target_waveform,
                                      model_outputs.ids_slice * feature_extractor.hop_length,
                                      self.config.segment_size
                                  )

                  discriminator_target, fmaps_target = discriminator(target_waveform)
                  discriminator_candidate, fmaps_candidate = discriminator(model_outputs.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)


                with autocast(enabled=training_args.fp16):

                  displayloss={}
                  # backpropagate
                  if dict_state_grad_loss['k1']:
                      loss_kl = kl_loss(
                          model_outputs.prior_latents,
                          model_outputs.posterior_log_variances,
                          model_outputs.prior_means,
                          model_outputs.prior_log_variances,
                          model_outputs.labels_padding_mask,
                      )
                      loss_kl=loss_kl*training_args.weight_kl
                      displayloss['loss_kl']=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)
                      displayloss['loss_mel'] = 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(model_outputs.log_duration)*training_args.weight_duration
                      displayloss['loss_duration'] = loss_duration.detach().item()
                    #  if displayloss['loss_duration']>=0:
                      #   loss_duration.backward()




                  discriminator_target, fmaps_target = discriminator(target_waveform)

                  discriminator_candidate, fmaps_candidate = discriminator(model_outputs.waveform.detach())

                  if dict_state_grad_loss['generator']:
                      loss_fmaps = feature_loss(fmaps_target, fmaps_candidate)
                      loss_gen, losses_gen = generator_loss(discriminator_candidate)
                      loss_gen=loss_gen * training_args.weight_gen
                      displayloss['loss_gen'] = loss_gen.detach().item()
                  #   loss_gen.backward(retain_graph=True)
                      loss_fmaps=loss_fmaps * training_args.weight_fmaps
                      displayloss['loss_fmaps'] = loss_fmaps.detach().item()
                  #   loss_fmaps.backward(retain_graph=True)
                      total_generator_loss = (
                          loss_duration
                          + loss_mel
                          + loss_kl
                          + loss_fmaps
                          + loss_gen
                      )
                     # total_generator_loss.backward()
                optimizer.zero_grad()
                scaler.scale(total_generator_loss).backward()
                scaler.unscale_(optimizer)
                grad_norm_g = clip_grad_value_(self.parameters(), None)
                scaler.step(optimizer)
                scaler.update()






                # optimizer.step()




                print(f"TRAINIG - batch {step}, waveform {(batch['waveform'].shape)}, lr {lr_scheduler.get_last_lr()[0]}... ")
                print(f"display loss function  enable  :{displayloss}")

                global_step +=1

                # validation

                do_eval = training_args.do_eval and (global_step % training_args.eval_steps == 0)
                if do_eval:
                    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,
                            monotonic_alignment_function=None,
                                )

                        mel_scaled_labels = batch["mel_scaled_input_features"]
                        mel_scaled_target = self.slice_segments(mel_scaled_labels, model_outputs.ids_slice,self.segment_size)
                        mel_scaled_generation = feature_extractor._torch_extract_fbank_features(model_outputs.waveform.squeeze(1))[1]
                        loss = loss_mel.detach().item()
                        eval_losses_sum +=loss

                        loss_mel = torch.nn.functional.l1_loss(mel_scaled_target, mel_scaled_generation)
                        print(f"VALIDATION - batch {step}, waveform {(batch['waveform'].shape)}, step_loss_mel {loss} ... ")



                    with torch.no_grad():
                        full_generation_sample = self.full_generation_sample
                        full_generation =self.forward(
                          input_ids =full_generation_sample["input_ids"],
                          attention_mask=full_generation_sample["attention_mask"],
                          speaker_id=full_generation_sample["speaker_id"]
                          )

                        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],})

            wandb.log({"train_losses":train_losses_sum})

        # add weight norms
        # self.remove_weight_norm()

        try:
          torch.save(self.posterior_encoder.state_dict(), os.path.join(training_args.output_dir,"posterior_encoder.pt"))
          torch.save(self.decoder.state_dict(), os.path.join(training_args.output_dir,"decoder.pt"))
        except:pass


        logger.info("Running final full generations samples... ")


        with torch.no_grad():

            full_generation_sample = self.full_generation_sample
            full_generation = self.forward(
                    input_ids=full_generation_sample["labels"],
                    attention_mask=full_generation_sample["labels_attention_mask"],
                    speaker_id=full_generation_sample["speaker_id"]
                    )

            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],
                })


        logger.info("***** Training / Inference Done *****")

    #....................................
    def trainer_to_cuda(self,
              train_dataset_dir = None,
              eval_dataset_dir = None,
              full_generation_dir = 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


              ):


        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()


        set_seed(training_args.seed)
        scaler = GradScaler(enabled=training_args.fp16)

        # Apply Weight Norm Decoder
        # self.apply_weight_norm()
        # Save Config
        self.config.save_pretrained(training_args.output_dir)

        train_dataset = FeaturesCollectionDataset(dataset_dir = train_dataset_dir,
                                                  device = self.device
                                                  )

        eval_dataset = None
        if training_args.do_eval:
            eval_dataset = FeaturesCollectionDataset(dataset_dir = eval_dataset_dir,
                                                     device = self.device
                                                     )

        full_generation_dataset = FeaturesCollectionDataset(dataset_dir = full_generation_dir,
                                                            device = self.device
                                                            )
        self.full_generation_sample = full_generation_dataset[full_generation_sample_index]

        # init optimizer, lr_scheduler

        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(
        #         self.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............................

        global_step = 0

        for epoch in range(training_args.num_train_epochs):
            train_losses_sum = 0
            lr_scheduler.step()

            # disc_lr_scheduler.step()
            print(f"  Num Epochs = {epoch}")
            if (epoch+1)%nk==0:
             clear_output()
             print('Save checkpoints Model :',int(epoch/nk))
             self.save_pretrained(path_save_model)

            for step, batch in enumerate(train_dataset):

                # forward through model
                # outputs = self.forward(
                #     labels=batch["labels"],
                #     labels_attention_mask=batch["labels_attention_mask"],
                #     speaker_id=batch["speaker_id"]
                #     )
                #if step==10:break
                batch=covert_cuda_batch(batch)
                displayloss={}


                with autocast(enabled=training_args.fp16):


                  model_outputs = self.forward_k(
                      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"],
                      text_encoder_output =None  if is_used_text_encoder else batch['text_encoder_output'],
                      posterior_encode_output=None  if is_used_posterior_encode else batch['posterior_encode_output'],
                      return_dict=True,
                      monotonic_alignment_function=maf,
                  )

                  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]

                  target_waveform = batch["waveform"].transpose(1, 2)
                  target_waveform = self.slice_segments(
                                      target_waveform,
                                      model_outputs.ids_slice * feature_extractor.hop_length,
                                      self.config.segment_size
                                  )

                  discriminator_target, fmaps_target = self.discriminator(target_waveform)
                  discriminator_candidate, fmaps_candidate = self.discriminator(model_outputs.waveform.detach())
                  with autocast(enabled=False):
                     if dict_state_grad_loss['discriminator']:
                      #  disc_optimizer.zero_grad()

                        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()
                optimizer.zero_grad()
                # disc_optimizer.zero_grad()
                scaler.scale(loss_dd).backward()
                # scaler.unscale_(disc_optimizer)
                #grad_norm_d = clip_grad_value_(self.discriminator.parameters(), None)
                # scaler.step(disc_optimizer)


                with autocast(enabled=training_args.fp16):




                  # backpropagate




                  discriminator_target, fmaps_target = self.discriminator(target_waveform)

                  discriminator_candidate, fmaps_candidate = self.discriminator(model_outputs.waveform.detach())
                  with autocast(enabled=False):
                    if dict_state_grad_loss['k1']:
                      loss_kl = kl_loss(
                          model_outputs.prior_latents,
                          model_outputs.posterior_log_variances,
                          model_outputs.prior_means,
                          model_outputs.prior_log_variances,
                          model_outputs.labels_padding_mask,
                      )
                      loss_kl=loss_kl*training_args.weight_kl
                      displayloss['loss_kl']=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)
                          displayloss['loss_mel'] = 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(model_outputs.log_duration)*training_args.weight_duration
                          displayloss['loss_duration'] = loss_duration.detach().item()
                        #  if displayloss['loss_duration']>=0:
                          #   loss_duration.backward()

                    if dict_state_grad_loss['generator']:
                        loss_fmaps = feature_loss(fmaps_target, fmaps_candidate)
                        loss_gen, losses_gen = generator_loss(discriminator_candidate)
                        loss_gen=loss_gen * training_args.weight_gen
                        displayloss['loss_gen'] = loss_gen.detach().item()
                    #   loss_gen.backward(retain_graph=True)
                        loss_fmaps=loss_fmaps * training_args.weight_fmaps
                        displayloss['loss_fmaps'] = loss_fmaps.detach().item()
                    #   loss_fmaps.backward(retain_graph=True)
                        total_generator_loss = (
                            loss_duration
                            + loss_mel
                            + loss_kl
                            + loss_fmaps
                            + loss_gen
                        )
                     # total_generator_loss.backward()
                scaler.scale(total_generator_loss).backward()
                scaler.unscale_(optimizer)
                grad_norm_g = clip_grad_value_(self.parameters(), None)
                scaler.step(optimizer)
                scaler.update()






                # optimizer.step()




                print(f"TRAINIG - batch {step},Grad G{grad_norm_g}, waveform {(batch['waveform'].shape)}, lr {lr_scheduler.get_last_lr()[0]}... ")
                print(f"display loss function  enable  :{displayloss}")

                global_step +=1

                # validation

                do_eval = training_args.do_eval and (global_step % training_args.eval_steps == 0)
                if do_eval:
                    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,
                            monotonic_alignment_function=None,
                                )

                        mel_scaled_labels = batch["mel_scaled_input_features"]
                        mel_scaled_target = self.slice_segments(mel_scaled_labels, model_outputs.ids_slice,self.segment_size)
                        mel_scaled_generation = feature_extractor._torch_extract_fbank_features(model_outputs.waveform.squeeze(1))[1]
                        loss = loss_mel.detach().item()
                        eval_losses_sum +=loss

                        loss_mel = torch.nn.functional.l1_loss(mel_scaled_target, mel_scaled_generation)
                        print(f"VALIDATION - batch {step}, waveform {(batch['waveform'].shape)}, step_loss_mel {loss} ... ")



                    with torch.no_grad():
                        full_generation_sample = self.full_generation_sample
                        full_generation =self.forward(
                          input_ids =full_generation_sample["input_ids"],
                          attention_mask=full_generation_sample["attention_mask"],
                          speaker_id=full_generation_sample["speaker_id"]
                          )

                        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],})

            wandb.log({"train_losses":train_losses_sum})

        # add weight norms
        # self.remove_weight_norm()

        try:
          torch.save(self.posterior_encoder.state_dict(), os.path.join(training_args.output_dir,"posterior_encoder.pt"))
          torch.save(self.decoder.state_dict(), os.path.join(training_args.output_dir,"decoder.pt"))
        except:pass


        logger.info("Running final full generations samples... ")


        with torch.no_grad():

            full_generation_sample = self.full_generation_sample
            full_generation = self.forward(
                    input_ids=full_generation_sample["labels"],
                    attention_mask=full_generation_sample["labels_attention_mask"],
                    speaker_id=full_generation_sample["speaker_id"]
                    )

            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],
                })


        logger.info("***** Training / Inference Done *****")

    #....................................

    def trainer_to(self,
              train_dataset_dir = None,
              eval_dataset_dir = None,
              full_generation_dir = 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


              ):


        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()


        set_seed(training_args.seed)
        # Apply Weight Norm Decoder
        # self.apply_weight_norm()
        # Save Config
        self.config.save_pretrained(training_args.output_dir)

        train_dataset = FeaturesCollectionDataset(dataset_dir = train_dataset_dir,
                                                  device = self.device
                                                  )

        eval_dataset = None
        if training_args.do_eval:
            eval_dataset = FeaturesCollectionDataset(dataset_dir = eval_dataset_dir,
                                                     device = self.device
                                                     )

        full_generation_dataset = FeaturesCollectionDataset(dataset_dir = full_generation_dir,
                                                            device = self.device
                                                            )
        self.full_generation_sample = full_generation_dataset[full_generation_sample_index]

        # init optimizer, lr_scheduler

        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(
                self.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............................

        global_step = 0

        for epoch in range(training_args.num_train_epochs):
            train_losses_sum = 0
            lr_scheduler.step()

            disc_lr_scheduler.step()
            print(f"  Num Epochs = {epoch}")
            if epoch%nk==0:
              clear_output()
              print('')
              print('Save checkpoints Model :',int(epoch/nk))
              self.save_pretrained(path_save_model)




            for step, batch in enumerate(train_dataset):

                # forward through model
                # outputs = self.forward(
                #     labels=batch["labels"],
                #     labels_attention_mask=batch["labels_attention_mask"],
                #     speaker_id=batch["speaker_id"]
                #     )
                #if step==10:break
                batch=covert_cuda_batch(batch)

                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"],
                    speaker_id=batch["speaker_id"],
                    text_encoder_output =None , #if is_used_text_encoder else batch['text_encoder_output'],
                    posterior_encode_output=batch['posterior_encode_output'] ,#  if is_used_posterior_encode else ,
                    return_dict=True,
                    monotonic_alignment_function= maf,
                )

                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
                                )



                displayloss={}
                # backpropagate
                #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
                displayloss['loss_kl']=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)
                displayloss['loss_mel'] = 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
                displayloss['loss_duration'] = loss_duration.detach().item()
              #  if displayloss['loss_duration']>=0:
                     #   loss_duration.backward()

                discriminator_target, fmaps_target = self.discriminator(target_waveform)
                discriminator_candidate, fmaps_candidate = self.discriminator(waveform.detach())
                #if dict_state_grad_loss['discriminator']:
                disc_optimizer.zero_grad()

                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.step()


                discriminator_target, fmaps_target = self.discriminator(target_waveform)

                discriminator_candidate, fmaps_candidate = self.discriminator(waveform.detach())
                optimizer.zero_grad()
             #   if dict_state_grad_loss['generator']:
                loss_fmaps = feature_loss(fmaps_target, fmaps_candidate)
                loss_gen, losses_gen = generator_loss(discriminator_candidate)
                loss_gen=loss_gen * training_args.weight_gen
                displayloss['loss_gen'] = loss_gen.detach().item()
              #   loss_gen.backward(retain_graph=True)
                loss_fmaps=loss_fmaps * training_args.weight_fmaps
                displayloss['loss_fmaps'] = loss_fmaps.detach().item()
              #   loss_fmaps.backward(retain_graph=True)
                total_generator_loss = (
                    loss_duration
                    + loss_mel
                    + loss_kl
                    + loss_fmaps
                    + loss_gen
                )
                total_generator_loss.backward()





                optimizer.step()




                print(f"TRAINIG - batch {step}, waveform {(batch['waveform'].shape)}, lr {lr_scheduler.get_last_lr()[0]}... ")
                print(f"display loss function  enable  :{displayloss}")

                global_step +=1

                # validation

                do_eval = training_args.do_eval and (global_step % training_args.eval_steps == 0)
                if do_eval:
                    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,
                            monotonic_alignment_function=None,
                                )

                        mel_scaled_labels = batch["mel_scaled_input_features"]
                        mel_scaled_target = self.slice_segments(mel_scaled_labels, model_outputs.ids_slice,self.segment_size)
                        mel_scaled_generation = feature_extractor._torch_extract_fbank_features(model_outputs.waveform.squeeze(1))[1]
                        loss = loss_mel.detach().item()
                        eval_losses_sum +=loss

                        loss_mel = torch.nn.functional.l1_loss(mel_scaled_target, mel_scaled_generation)
                        print(f"VALIDATION - batch {step}, waveform {(batch['waveform'].shape)}, step_loss_mel {loss} ... ")



                    with torch.no_grad():
                        full_generation_sample = self.full_generation_sample
                        full_generation =self.forward(
                          input_ids =full_generation_sample["input_ids"],
                          attention_mask=full_generation_sample["attention_mask"],
                          speaker_id=full_generation_sample["speaker_id"]
                          )

                        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],})

            wandb.log({"train_losses":train_losses_sum})

        # add weight norms
        # self.remove_weight_norm()

        try:
          torch.save(self.posterior_encoder.state_dict(), os.path.join(training_args.output_dir,"posterior_encoder.pt"))
          torch.save(self.decoder.state_dict(), os.path.join(training_args.output_dir,"decoder.pt"))
        except:pass


        logger.info("Running final full generations samples... ")


        with torch.no_grad():

            full_generation_sample = self.full_generation_sample
            full_generation = self.forward(
                    input_ids=full_generation_sample["labels"],
                    attention_mask=full_generation_sample["labels_attention_mask"],
                    speaker_id=full_generation_sample["speaker_id"]
                    )

            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],
                })


        logger.info("***** Training / Inference Done *****")

    #....................................




    def trainer_to_cuda1(self,
              train_dataset_dir = None,
              eval_dataset_dir = None,
              full_generation_dir = 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


              ):


        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()


        set_seed(training_args.seed)
        scaler = GradScaler(enabled=training_args.fp16)

        # Apply Weight Norm Decoder
        # self.apply_weight_norm()
        # Save Config
        self.config.save_pretrained(training_args.output_dir)

        train_dataset = FeaturesCollectionDataset(dataset_dir = train_dataset_dir,
                                                  device = self.device
                                                  )

        eval_dataset = None
        if training_args.do_eval:
            eval_dataset = FeaturesCollectionDataset(dataset_dir = eval_dataset_dir,
                                                     device = self.device
                                                     )

        full_generation_dataset = FeaturesCollectionDataset(dataset_dir = full_generation_dir,
                                                            device = self.device
                                                            )
        self.full_generation_sample = full_generation_dataset[full_generation_sample_index]

        # init optimizer, lr_scheduler
        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............................

        global_step = 0

        for epoch in range(training_args.num_train_epochs):
            train_losses_sum = 0
            lr_scheduler.step()

            disc_lr_scheduler.step()
            print(f"  Num Epochs = {epoch}")
            if epoch%nk==0:
              clear_output()
              print('Save checkpoints Model :',int(epoch/nk))
              self.discriminator=discriminator

              self.save_pretrained(path_save_model)
              self.discriminator=None




            for step, batch in enumerate(train_dataset):

                # forward through model
                # outputs = self.forward(
                #     labels=batch["labels"],
                #     labels_attention_mask=batch["labels_attention_mask"],
                #     speaker_id=batch["speaker_id"]
                #     )
                #if step==10:break
                batch=covert_cuda_batch(batch)
                displayloss={}

                with autocast(enabled=training_args.fp16):


                  model_outputs = self.forward_k(
                      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"],
                      text_encoder_output =None  if is_used_text_encoder else batch['text_encoder_output'],
                      posterior_encode_output=None  if is_used_posterior_encode else batch['posterior_encode_output'],
                      return_dict=True,
                      monotonic_alignment_function= maf,
                  )

                  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]

                  target_waveform = batch["waveform"].transpose(1, 2)
                  target_waveform = self.slice_segments(
                                      target_waveform,
                                      model_outputs.ids_slice * feature_extractor.hop_length,
                                      self.config.segment_size
                                  )

                  discriminator_target, fmaps_target = discriminator(target_waveform)
                  discriminator_candidate, fmaps_candidate = discriminator(model_outputs.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

                disc_optimizer.zero_grad()
                loss_dd.backward()


                # scaler.scale(loss_dd).backward()
                # scaler.unscale_(disc_optimizer )
                grad_norm_d = clip_grad_value_(discriminator.parameters(), None)
                disc_optimizer.step()


                with autocast(enabled=training_args.fp16):


                  # backpropagate
                  if dict_state_grad_loss['k1']:
                      loss_kl = kl_loss(
                          model_outputs.prior_latents,
                          model_outputs.posterior_log_variances,
                          model_outputs.prior_means,
                          model_outputs.prior_log_variances,
                          model_outputs.labels_padding_mask,
                      )
                      loss_kl=loss_kl*training_args.weight_kl
                      displayloss['loss_kl']=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)
                      displayloss['loss_mel'] = 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(model_outputs.log_duration)*training_args.weight_duration
                      displayloss['loss_duration'] = loss_duration.detach().item()
                    #  if displayloss['loss_duration']>=0:
                      #   loss_duration.backward()




                  discriminator_target, fmaps_target = discriminator(target_waveform)

                  discriminator_candidate, fmaps_candidate = discriminator(model_outputs.waveform.detach())

                  if dict_state_grad_loss['generator']:
                      loss_fmaps = feature_loss(fmaps_target, fmaps_candidate)
                      loss_gen, losses_gen = generator_loss(discriminator_candidate)
                      loss_gen=loss_gen * training_args.weight_gen
                      displayloss['loss_gen'] = loss_gen.detach().item()
                  #   loss_gen.backward(retain_graph=True)
                      loss_fmaps=loss_fmaps * training_args.weight_fmaps
                      displayloss['loss_fmaps'] = loss_fmaps.detach().item()
                  #   loss_fmaps.backward(retain_graph=True)
                      total_generator_loss = (
                          loss_duration
                          + loss_mel
                          + loss_kl
                          + loss_fmaps
                          + loss_gen
                      )

                optimizer.zero_grad()
                total_generator_loss.backward()
                # scaler.scale(total_generator_loss).backward()
                # scaler.unscale_(optimizer)
                grad_norm_g = clip_grad_value_(self.parameters(), None)
                optimizer.step()
                # scaler.update()






                # optimizer.step()




                print(f"TRAINIG - batch {step}, waveform {(batch['waveform'].shape)}, lr {lr_scheduler.get_last_lr()[0]}... ")
                print(f"display loss function  enable  :{displayloss}")

                global_step +=1

                # validation

                do_eval = training_args.do_eval and (global_step % training_args.eval_steps == 0)
                if do_eval:
                    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,
                            monotonic_alignment_function=None,
                                )

                        mel_scaled_labels = batch["mel_scaled_input_features"]
                        mel_scaled_target = self.slice_segments(mel_scaled_labels, model_outputs.ids_slice,self.segment_size)
                        mel_scaled_generation = feature_extractor._torch_extract_fbank_features(model_outputs.waveform.squeeze(1))[1]
                        loss = loss_mel.detach().item()
                        eval_losses_sum +=loss

                        loss_mel = torch.nn.functional.l1_loss(mel_scaled_target, mel_scaled_generation)
                        print(f"VALIDATION - batch {step}, waveform {(batch['waveform'].shape)}, step_loss_mel {loss} ... ")



                    with torch.no_grad():
                        full_generation_sample = self.full_generation_sample
                        full_generation =self.forward(
                          input_ids =full_generation_sample["input_ids"],
                          attention_mask=full_generation_sample["attention_mask"],
                          speaker_id=full_generation_sample["speaker_id"]
                          )

                        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],})

            wandb.log({"train_losses":train_losses_sum})

        # add weight norms
        # self.remove_weight_norm()

        try:
          torch.save(self.posterior_encoder.state_dict(), os.path.join(training_args.output_dir,"posterior_encoder.pt"))
          torch.save(self.decoder.state_dict(), os.path.join(training_args.output_dir,"decoder.pt"))
        except:pass


        logger.info("Running final full generations samples... ")


        with torch.no_grad():

            full_generation_sample = self.full_generation_sample
            full_generation = self.forward(
                    input_ids=full_generation_sample["labels"],
                    attention_mask=full_generation_sample["labels_attention_mask"],
                    speaker_id=full_generation_sample["speaker_id"]
                    )

            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],
                })


        logger.info("***** Training / Inference Done *****")

    def forward_train(
        self,
        input_ids: Optional[torch.Tensor] = None,
        attention_mask: Optional[torch.Tensor] = None,
        speaker_id: Optional[int] = None,
        output_attentions: Optional[bool] = None,
        output_hidden_states: Optional[bool] = None,
        return_dict: Optional[bool] = None,
        labels: Optional[torch.FloatTensor] = None,
        labels_attention_mask: Optional[torch.Tensor] = None,
        text_encoder_output=None,
        posterior_encode_output=None,
        monotonic_alignment_function: Optional[Callable] = None,
        speaker_embeddings=None
    ) -> Union[Tuple[Any], VitsModelOutput]:

        #output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
        output_hidden_states = (
            output_hidden_states# if output_hidden_states is not None else self.config.output_hidden_states
        )
       # return_dict = return_dict if return_dict is not None else self.config.use_return_dict


       # if attention_mask is not None:
        input_padding_mask = attention_mask.unsqueeze(-1).float()
        #else:
         #   input_padding_mask = torch.ones_like(input_ids).unsqueeze(-1).float()

        # speaker_embeddings=None
       # if labels_attention_mask is not None:
        labels_padding_mask = labels_attention_mask.unsqueeze(1).float()
        # else:
        #     labels_attention_mask = torch.ones((labels.shape[0], labels.shape[2])).float().to(self.device)
        #     labels_padding_mask = labels_attention_mask.unsqueeze(1)
        if text_encoder_output is None:
          text_encoder_output = self.text_encoder(
                  input_ids=input_ids,
                  padding_mask=input_padding_mask,
                  attention_mask=attention_mask,
                  output_attentions=output_attentions,
                  output_hidden_states=output_hidden_states,
                  return_dict=return_dict,
              )
        #hidden_states = text_encoder_output[0] #if not return_dict else text_encoder_output.last_hidden_state
        hidden_states = text_encoder_output[0].transpose(1, 2)
        input_padding_mask = input_padding_mask.transpose(1, 2)
        prior_means = text_encoder_output[1].transpose(1, 2) #if not return_dict else text_encoder_output.prior_means
        prior_log_variances = text_encoder_output[2].transpose(1, 2) #if not return_dict else text_encoder_output.prior_log_variances

        if posterior_encode_output is None:
             latents, posterior_means, posterior_log_variances = self.posterior_encoder(
                labels, labels_padding_mask, speaker_embeddings
            )
        else:
             latents=posterior_encode_output['posterior_latents']
             posterior_means=posterior_encode_output['posterior_means']
             posterior_log_variances=posterior_encode_output['posterior_log_variances']

        prior_latents = self.flow(latents, labels_padding_mask, speaker_embeddings, reverse=False)

        # prior_means, prior_log_variances = prior_means.transpose(1, 2), prior_log_variances.transpose(1, 2)
        with torch.no_grad():
            # negative cross-entropy

            # [batch_size, d, latent_length]
            prior_variances = torch.exp(-2 * prior_log_variances)
            # [batch_size, 1, latent_length]
            neg_cent1 = torch.sum(-0.5 * math.log(2 * math.pi) - prior_log_variances, [1], keepdim=True)
            # [batch_size, text_length, d] x [batch_size, d, latent_length] = [batch_size, text_length, latent_length]
            neg_cent2 = torch.matmul(-0.5 * (prior_latents**2).transpose(1, 2), prior_variances)
            # [batch_size, text_length, d] x [batch_size, d, latent_length] = [batch_size, text_length, latent_length]
            neg_cent3 = torch.matmul(prior_latents.transpose(1, 2), (prior_means * prior_variances))
            # [batch_size, 1, latent_length]
            neg_cent4 = torch.sum(-0.5 * (prior_means**2) * prior_variances, [1], keepdim=True)

            # [batch_size, text_length, latent_length]
            neg_cent = neg_cent1 + neg_cent2 + neg_cent3 + neg_cent4

            attn_mask = torch.unsqueeze(input_padding_mask, 2) * torch.unsqueeze(labels_padding_mask, -1)

            attn = monotonic_alignment_function(neg_cent, attn_mask.squeeze(1)).unsqueeze(1).detach()

        durations = attn.sum(2)

        #if self.config.use_stochastic_duration_prediction:
        log_duration = self.duration_predictor(
                hidden_states, input_padding_mask, speaker_embeddings, durations=durations, reverse=False
            )
        log_duration = log_duration / torch.sum(input_padding_mask)
        # else:
        #     log_duration_padded = torch.log(durations + 1e-6) * input_padding_mask
        #     log_duration = self.duration_predictor(hidden_states, input_padding_mask, speaker_embeddings)
        #     log_duration = torch.sum((log_duration - log_duration_padded) ** 2, [1, 2]) / torch.sum(input_padding_mask)

        # expand priors
        prior_means = torch.matmul(attn.squeeze(1), prior_means.transpose(1, 2)).transpose(1, 2)
        prior_log_variances = torch.matmul(attn.squeeze(1), prior_log_variances.transpose(1, 2)).transpose(1, 2)

        label_lengths = labels_attention_mask.sum(dim=1)
        latents_slice, ids_slice = self.rand_slice_segments(latents, label_lengths, segment_size=self.segment_size)
        waveform = self.decoder(latents_slice, speaker_embeddings)
        return waveform,ids_slice,log_duration,prior_latents,posterior_log_variances,prior_means,prior_log_variances,labels_padding_mask