try L2 kv

.gitattributes

@@ -47,3 +47,4 @@ uc_spk_Landscape2Soundscape_Masterpieces_pics/03_Schinkel_WS200-002.jpg filter=l
 uc_spk_Landscape2Soundscape_Masterpieces_pics/11_Liebermann_NG4-94_001.jpg filter=lfs diff=lfs merge=lfs -text
 uc_spk_Landscape2Soundscape_Masterpieces_pics/12_Slevogt_AII1022_001.jpg filter=lfs diff=lfs merge=lfs -text
 assets/audiobook_TTS.docx filter=lfs diff=lfs merge=lfs -text
+assets/ocr.jpg filter=lfs diff=lfs merge=lfs -text

Modules/vits/models.py

@@ -5,136 +5,13 @@ import numpy as np
 import torch
 import torch.utils.checkpoint
 from torch import nn
-from transformers.activations import ACT2FN
+
 from transformers.modeling_attn_mask_utils import _prepare_4d_attention_mask
 from transformers.modeling_outputs import BaseModelOutput, ModelOutput
 from transformers.modeling_utils import PreTrainedModel
 from transformers.configuration_utils import PretrainedConfig
 
 class VitsConfig(PretrainedConfig):
-    r"""
-    This is the configuration class to store the configuration of a [`VitsModel`]. It is used to instantiate a VITS
-    model according to the specified arguments, defining the model architecture. Instantiating a configuration with the
-    defaults will yield a similar configuration to that of the VITS
-    [facebook/mms-tts-eng](https://huggingface.co/facebook/mms-tts-eng) architecture.
-
-    Configuration objects inherit from [`PretrainedConfig`] and can be used to control the model outputs. Read the
-    documentation from [`PretrainedConfig`] for more information.
-
-    Args:
-        vocab_size (`int`, *optional*, defaults to 38):
-            Vocabulary size of the VITS model. Defines the number of different tokens that can be represented by the
-            `inputs_ids` passed to the forward method of [`VitsModel`].
-        hidden_size (`int`, *optional*, defaults to 192):
-            Dimensionality of the text encoder layers.
-        num_hidden_layers (`int`, *optional*, defaults to 6):
-            Number of hidden layers in the Transformer encoder.
-        num_attention_heads (`int`, *optional*, defaults to 2):
-            Number of attention heads for each attention layer in the Transformer encoder.
-        window_size (`int`, *optional*, defaults to 4):
-            Window size for the relative positional embeddings in the attention layers of the Transformer encoder.
-        use_bias (`bool`, *optional*, defaults to `True`):
-            Whether to use bias in the key, query, value projection layers in the Transformer encoder.
-        ffn_dim (`int`, *optional*, defaults to 768):
-            Dimensionality of the "intermediate" (i.e., feed-forward) layer in the Transformer encoder.
-        layerdrop (`float`, *optional*, defaults to 0.1):
-            The LayerDrop probability for the encoder. See the [LayerDrop paper](see https://arxiv.org/abs/1909.11556)
-            for more details.
-        ffn_kernel_size (`int`, *optional*, defaults to 3):
-            Kernel size of the 1D convolution layers used by the feed-forward network in the Transformer encoder.
-        flow_size (`int`, *optional*, defaults to 192):
-            Dimensionality of the flow layers.
-        spectrogram_bins (`int`, *optional*, defaults to 513):
-            Number of frequency bins in the target spectrogram.
-        hidden_act (`str` or `function`, *optional*, defaults to `"relu"`):
-            The non-linear activation function (function or string) in the encoder and pooler. If string, `"gelu"`,
-            `"relu"`, `"selu"` and `"gelu_new"` are supported.
-        hidden_dropout (`float`, *optional*, defaults to 0.1):
-            The dropout probability for all fully connected layers in the embeddings and encoder.
-        attention_dropout (`float`, *optional*, defaults to 0.1):
-            The dropout ratio for the attention probabilities.
-        activation_dropout (`float`, *optional*, defaults to 0.1):
-            The dropout ratio for activations inside the fully connected layer.
-        initializer_range (`float`, *optional*, defaults to 0.02):
-            The standard deviation of the truncated_normal_initializer for initializing all weight matrices.
-        layer_norm_eps (`float`, *optional*, defaults to 1e-05):
-            The epsilon used by the layer normalization layers.
-        use_stochastic_duration_prediction (`bool`, *optional*, defaults to `True`):
-            Whether to use the stochastic duration prediction module or the regular duration predictor.
-        num_speakers (`int`, *optional*, defaults to 1):
-            Number of speakers if this is a multi-speaker model.
-        speaker_embedding_size (`int`, *optional*, defaults to 0):
-            Number of channels used by the speaker embeddings. Is zero for single-speaker models.
-        upsample_initial_channel (`int`, *optional*, defaults to 512):
-            The number of input channels into the HiFi-GAN upsampling network.
-        upsample_rates (`Tuple[int]` or `List[int]`, *optional*, defaults to `[8, 8, 2, 2]`):
-            A tuple of integers defining the stride of each 1D convolutional layer in the HiFi-GAN upsampling network.
-            The length of `upsample_rates` defines the number of convolutional layers and has to match the length of
-            `upsample_kernel_sizes`.
-        upsample_kernel_sizes (`Tuple[int]` or `List[int]`, *optional*, defaults to `[16, 16, 4, 4]`):
-            A tuple of integers defining the kernel size of each 1D convolutional layer in the HiFi-GAN upsampling
-            network. The length of `upsample_kernel_sizes` defines the number of convolutional layers and has to match
-            the length of `upsample_rates`.
-        resblock_kernel_sizes (`Tuple[int]` or `List[int]`, *optional*, defaults to `[3, 7, 11]`):
-            A tuple of integers defining the kernel sizes of the 1D convolutional layers in the HiFi-GAN
-            multi-receptive field fusion (MRF) module.
-        resblock_dilation_sizes (`Tuple[Tuple[int]]` or `List[List[int]]`, *optional*, defaults to `[[1, 3, 5], [1, 3, 5], [1, 3, 5]]`):
-            A nested tuple of integers defining the dilation rates of the dilated 1D convolutional layers in the
-            HiFi-GAN multi-receptive field fusion (MRF) module.
-        leaky_relu_slope (`float`, *optional*, defaults to 0.1):
-            The angle of the negative slope used by the leaky ReLU activation.
-        depth_separable_channels (`int`, *optional*, defaults to 2):
-            Number of channels to use in each depth-separable block.
-        depth_separable_num_layers (`int`, *optional*, defaults to 3):
-            Number of convolutional layers to use in each depth-separable block.
-        duration_predictor_flow_bins (`int`, *optional*, defaults to 10):
-            Number of channels to map using the unonstrained rational spline in the duration predictor model.
-        duration_predictor_tail_bound (`float`, *optional*, defaults to 5.0):
-            Value of the tail bin boundary when computing the unconstrained rational spline in the duration predictor
-            model.
-        duration_predictor_kernel_size (`int`, *optional*, defaults to 3):
-            Kernel size of the 1D convolution layers used in the duration predictor model.
-        duration_predictor_dropout (`float`, *optional*, defaults to 0.5):
-            The dropout ratio for the duration predictor model.
-        duration_predictor_num_flows (`int`, *optional*, defaults to 4):
-            Number of flow stages used by the duration predictor model.
-        duration_predictor_filter_channels (`int`, *optional*, defaults to 256):
-            Number of channels for the convolution layers used in the duration predictor model.
-        prior_encoder_num_flows (`int`, *optional*, defaults to 4):
-            Number of flow stages used by the prior encoder flow model.
-        prior_encoder_num_wavenet_layers (`int`, *optional*, defaults to 4):
-            Number of WaveNet layers used by the prior encoder flow model.
-        posterior_encoder_num_wavenet_layers (`int`, *optional*, defaults to 16):
-            Number of WaveNet layers used by the posterior encoder model.
-        wavenet_kernel_size (`int`, *optional*, defaults to 5):
-            Kernel size of the 1D convolution layers used in the WaveNet model.
-        wavenet_dilation_rate (`int`, *optional*, defaults to 1):
-            Dilation rates of the dilated 1D convolutional layers used in the WaveNet model.
-        wavenet_dropout (`float`, *optional*, defaults to 0.0):
-            The dropout ratio for the WaveNet layers.
-        speaking_rate (`float`, *optional*, defaults to 1.0):
-            Speaking rate. Larger values give faster synthesised speech.
-        noise_scale (`float`, *optional*, defaults to 0.667):
-            How random the speech prediction is. Larger values create more variation in the predicted speech.
-        noise_scale_duration (`float`, *optional*, defaults to 0.8):
-            How random the duration prediction is. Larger values create more variation in the predicted durations.
-        sampling_rate (`int`, *optional*, defaults to 16000):
-            The sampling rate at which the output audio waveform is digitalized expressed in hertz (Hz).
-
-    Example:
-
-    ```python
-    >>> from transformers import VitsModel, VitsConfig
-
-    >>> # Initializing a "facebook/mms-tts-eng" style configuration
-    >>> configuration = VitsConfig()
-
-    >>> # Initializing a model (with random weights) from the "facebook/mms-tts-eng" style configuration
-    >>> model = VitsModel(configuration)
-
-    >>> # Accessing the model configuration
-    >>> configuration = model.config
-    ```"""
 
     model_type = "vits"
 
@@ -151,7 +28,7 @@ class VitsConfig(PretrainedConfig):
         ffn_kernel_size=3,
         flow_size=192,
         spectrogram_bins=513,
-        hidden_act="relu",
+        # hidden_act="relu",
         hidden_dropout=0.1,
         attention_dropout=0.1,
         activation_dropout=0.1,
@@ -180,7 +57,7 @@ class VitsConfig(PretrainedConfig):
         wavenet_kernel_size=5,
         wavenet_dilation_rate=1,
         wavenet_dropout=0.0,
-        speaking_rate=1.0,
+        speaking_rate=1.0,  # unused
         noise_scale=0.667,
         noise_scale_duration=0.8,
         sampling_rate=16_000,
@@ -197,10 +74,8 @@ class VitsConfig(PretrainedConfig):
         self.ffn_kernel_size = ffn_kernel_size
         self.flow_size = flow_size
         self.spectrogram_bins = spectrogram_bins
-        self.hidden_act = hidden_act
-        self.hidden_dropout = hidden_dropout
-        self.attention_dropout = attention_dropout
-        self.activation_dropout = activation_dropout
+
+
         self.initializer_range = initializer_range
         self.layer_norm_eps = layer_norm_eps
         self.use_stochastic_duration_prediction = use_stochastic_duration_prediction
@@ -217,7 +92,7 @@ class VitsConfig(PretrainedConfig):
         self.duration_predictor_flow_bins = duration_predictor_flow_bins
         self.duration_predictor_tail_bound = duration_predictor_tail_bound
         self.duration_predictor_kernel_size = duration_predictor_kernel_size
-        self.duration_predictor_dropout = duration_predictor_dropout
+        
         self.duration_predictor_num_flows = duration_predictor_num_flows
         self.duration_predictor_filter_channels = duration_predictor_filter_channels
         self.prior_encoder_num_flows = prior_encoder_num_flows
@@ -225,8 +100,8 @@ class VitsConfig(PretrainedConfig):
         self.posterior_encoder_num_wavenet_layers = posterior_encoder_num_wavenet_layers
         self.wavenet_kernel_size = wavenet_kernel_size
         self.wavenet_dilation_rate = wavenet_dilation_rate
-        self.wavenet_dropout = wavenet_dropout
-        self.speaking_rate = speaking_rate  # reset during long txt inference for natural variation
+        
+        
         self.noise_scale = noise_scale
         self.noise_scale_duration = noise_scale_duration
         self.sampling_rate = sampling_rate
@@ -241,50 +116,14 @@ class VitsConfig(PretrainedConfig):
 
 
 
-
-# ============================ modeling
-
 @dataclass
 class VitsTextEncoderOutput(ModelOutput):
-    """
-    Describes the outputs for the VITS text encoder model, with potential hidden states and attentions.
-
-    Args:
-        last_hidden_state (`torch.FloatTensor` of shape `(batch_size, sequence_length, hidden_size)`):
-            Sequence of hidden-states at the output of the last layer of the model.
-        prior_means (`torch.FloatTensor` of shape `(batch_size, sequence_length, hidden_size)`):
-            The predicted mean values of the prior distribution for the latent text variables.
-        prior_log_variances (`torch.FloatTensor` of shape `(batch_size, sequence_length, hidden_size)`):
-            The predicted log-variance values of the prior distribution for the latent text variables.
-        hidden_states (`tuple(torch.FloatTensor)`, *optional*, returned when `output_hidden_states=True` is passed or when `config.output_hidden_states=True`):
-            Tuple of `torch.FloatTensor` (one for the output of the embeddings, if the model has an embedding layer, +
-            one for the output of each layer) of shape `(batch_size, sequence_length, hidden_size)`.
-
-            Hidden-states of the model at the output of each layer plus the optional initial embedding outputs.
-        attentions (`tuple(torch.FloatTensor)`, *optional*, returned when `output_attentions=True` is passed or when `config.output_attentions=True`):
-            Tuple of `torch.FloatTensor` (one for each layer) of shape `(batch_size, num_heads, sequence_length,
-            sequence_length)`.
-
-            Attention weights after the attention softmax, used to compute the weighted average in the self-attention
-            heads.
-    """
-
     last_hidden_state: torch.FloatTensor = None
     prior_means: torch.FloatTensor = None
     prior_log_variances: torch.FloatTensor = None
     hidden_states: Optional[Tuple[torch.FloatTensor]] = None
     attentions: Optional[Tuple[torch.FloatTensor]] = None
 
-
-@torch.jit.script
-def fused_add_tanh_sigmoid_multiply(input_a, input_b, num_channels):
-    in_act = input_a + input_b
-    t_act = torch.tanh(in_act[:, :num_channels, :])
-    s_act = torch.sigmoid(in_act[:, num_channels:, :])
-    acts = t_act * s_act
-    return acts
-
-
 def _unconstrained_rational_quadratic_spline(
     inputs,
     unnormalized_widths,
@@ -296,46 +135,11 @@ def _unconstrained_rational_quadratic_spline(
     min_bin_height=1e-3,
     min_derivative=1e-3,
 ):
-    """
-    This transformation represents a monotonically increasing piecewise rational quadratic function. Outside of the
-    `tail_bound`, the transform behaves as an identity function.
-
-    Args:
-        inputs (`torch.FloatTensor` of shape `(batch_size, channels, seq_len)`:
-            Second half of the hidden-states input to the Vits convolutional flow module.
-        unnormalized_widths (`torch.FloatTensor` of shape `(batch_size, channels, seq_len, duration_predictor_flow_bins)`):
-            First `duration_predictor_flow_bins` of the hidden-states from the output of the convolution projection
-            layer in the convolutional flow module
-        unnormalized_heights (`torch.FloatTensor` of shape `(batch_size, channels, seq_len, duration_predictor_flow_bins)`):
-            Second `duration_predictor_flow_bins` of the hidden-states from the output of the convolution projection
-            layer in the convolutional flow module
-        unnormalized_derivatives (`torch.FloatTensor` of shape `(batch_size, channels, seq_len, duration_predictor_flow_bins)`):
-            Third `duration_predictor_flow_bins` of the hidden-states from the output of the convolution projection
-            layer in the convolutional flow module
-        reverse (`bool`, *optional*, defaults to `False`):
-            Whether the model is being run in reverse mode.
-        tail_bound (`float`, *optional* defaults to 5):
-            Upper and lower limit bound for the rational quadratic function. Outside of this `tail_bound`, the
-            transform behaves as an identity function.
-        min_bin_width (`float`, *optional*, defaults to 1e-3):
-            Minimum bin value across the width dimension for the piecewise rational quadratic function.
-        min_bin_height (`float`, *optional*, defaults to 1e-3):
-            Minimum bin value across the height dimension for the piecewise rational quadratic function.
-        min_derivative (`float`, *optional*, defaults to 1e-3):
-            Minimum bin value across the derivatives for the piecewise rational quadratic function.
-    Returns:
-        outputs (`torch.FloatTensor` of shape `(batch_size, channels, seq_len)`:
-            Hidden-states as transformed by the piecewise rational quadratic function with the `tail_bound` limits
-            applied.
-        log_abs_det (`torch.FloatTensor` of shape `(batch_size, channels, seq_len)`:
-            Logarithm of the absolute value of the determinants corresponding to the `outputs` with the `tail_bound`
-            limits applied.
-    """
     inside_interval_mask = (inputs >= -tail_bound) & (inputs <= tail_bound)
     outside_interval_mask = ~inside_interval_mask
 
     outputs = torch.zeros_like(inputs)
-    log_abs_det = torch.zeros_like(inputs)
+    
     constant = np.log(np.exp(1 - min_derivative) - 1)
 
     unnormalized_derivatives = nn.functional.pad(unnormalized_derivatives, pad=(1, 1))
@@ -343,9 +147,9 @@ def _unconstrained_rational_quadratic_spline(
     unnormalized_derivatives[..., -1] = constant
 
     outputs[outside_interval_mask] = inputs[outside_interval_mask]
-    log_abs_det[outside_interval_mask] = 0.0
+    
 
-    outputs[inside_interval_mask], log_abs_det[inside_interval_mask] = _rational_quadratic_spline(
+    outputs[inside_interval_mask] = _rational_quadratic_spline(
         inputs=inputs[inside_interval_mask],
         unnormalized_widths=unnormalized_widths[inside_interval_mask, :],
         unnormalized_heights=unnormalized_heights[inside_interval_mask, :],
@@ -356,7 +160,7 @@ def _unconstrained_rational_quadratic_spline(
         min_bin_height=min_bin_height,
         min_derivative=min_derivative,
     )
-    return outputs, log_abs_det
+    return outputs
 
 
 def _rational_quadratic_spline(
@@ -455,20 +259,21 @@ def _rational_quadratic_spline(
 
     intermediate1 = input_derivatives + input_derivatives_plus_one - 2 * input_delta
     if not reverse:
-        theta = (inputs - input_cumwidths) / input_bin_widths
-        theta_one_minus_theta = theta * (1 - theta)
-
-        numerator = input_heights * (input_delta * theta.pow(2) + input_derivatives * theta_one_minus_theta)
-        denominator = input_delta + intermediate1 * theta_one_minus_theta
-        outputs = input_cumheights + numerator / denominator
-
-        derivative_numerator = input_delta.pow(2) * (
-            input_derivatives_plus_one * theta.pow(2)
-            + 2 * input_delta * theta_one_minus_theta
-            + input_derivatives * (1 - theta).pow(2)
-        )
-        log_abs_det = torch.log(derivative_numerator) - 2 * torch.log(denominator)
-        return outputs, log_abs_det
+        raise ValueError
+        # theta = (inputs - input_cumwidths) / input_bin_widths
+        # theta_one_minus_theta = theta * (1 - theta)
+
+        # numerator = input_heights * (input_delta * theta.pow(2) + input_derivatives * theta_one_minus_theta)
+        # denominator = input_delta + intermediate1 * theta_one_minus_theta
+        # outputs = input_cumheights + numerator / denominator
+
+        # derivative_numerator = input_delta.pow(2) * (
+        #     input_derivatives_plus_one * theta.pow(2)
+        #     + 2 * input_delta * theta_one_minus_theta
+        #     + input_derivatives * (1 - theta).pow(2)
+        # )
+        # log_abs_det = torch.log(derivative_numerator) - 2 * torch.log(denominator)
+        # return outputs, log_abs_det
     else:
         # find the roots of a quadratic equation
         intermediate2 = inputs - input_cumheights
@@ -484,26 +289,26 @@ def _rational_quadratic_spline(
         root = (2 * c) / (-b - torch.sqrt(discriminant))
         outputs = root * input_bin_widths + input_cumwidths
 
-        theta_one_minus_theta = root * (1 - root)
-        denominator = input_delta + intermediate1 * theta_one_minus_theta
-        derivative_numerator = input_delta.pow(2) * (
-            input_derivatives_plus_one * root.pow(2)
-            + 2 * input_delta * theta_one_minus_theta
-            + input_derivatives * (1 - root).pow(2)
-        )
-        log_abs_det = torch.log(derivative_numerator) - 2 * torch.log(denominator)
-        return outputs, -log_abs_det
+        # theta_one_minus_theta = root * (1 - root)
+        # denominator = input_delta + intermediate1 * theta_one_minus_theta
+        # derivative_numerator = input_delta.pow(2) * (
+        #     input_derivatives_plus_one * root.pow(2)
+        #     + 2 * input_delta * theta_one_minus_theta
+        #     + input_derivatives * (1 - root).pow(2)
+        # )
+        # log_abs_det = torch.log(derivative_numerator) - 2 * torch.log(denominator)
+        return outputs #, -log_abs_det
 
 
 class VitsWaveNet(torch.nn.Module):
-    def __init__(self, config: VitsConfig, num_layers: int):
+    def __init__(self, config, num_layers):
         super().__init__()
         self.hidden_size = config.hidden_size
         self.num_layers = num_layers
 
         self.in_layers = torch.nn.ModuleList()
         self.res_skip_layers = torch.nn.ModuleList()
-        self.dropout = nn.Dropout(config.wavenet_dropout)
+        
 
         if hasattr(nn.utils.parametrizations, "weight_norm"):
             weight_norm = nn.utils.parametrizations.weight_norm
@@ -539,22 +344,28 @@ class VitsWaveNet(torch.nn.Module):
 
     def forward(self, inputs, padding_mask, global_conditioning=None):
         outputs = torch.zeros_like(inputs)
-        num_channels_tensor = torch.IntTensor([self.hidden_size])
+        num_channels = torch.IntTensor([self.hidden_size])[0]
 
-        if global_conditioning is not None:
-            global_conditioning = self.cond_layer(global_conditioning)
 
         for i in range(self.num_layers):
-            hidden_states = self.in_layers[i](inputs)
+            in_act = self.in_layers[i](inputs)
 
-            if global_conditioning is not None:
-                cond_offset = i * 2 * self.hidden_size
-                global_states = global_conditioning[:, cond_offset : cond_offset + 2 * self.hidden_size, :]
-            else:
-                global_states = torch.zeros_like(hidden_states)
+            
+            # global_states = torch.zeros_like(hidden_states)  # style ?
 
-            acts = fused_add_tanh_sigmoid_multiply(hidden_states, global_states, num_channels_tensor[0])
-            acts = self.dropout(acts)
+            # acts = fused_add_tanh_sigmoid_multiply(hidden_states, global_states, num_channels_tensor[0])
+            
+            # --
+            # def fused_add_tanh_sigmoid_multiply(input_a, input_b, num_channels):
+            # in_act = input_a #  + input_b
+            t_act = torch.tanh(in_act[:, :num_channels, :])
+            s_act = torch.sigmoid(in_act[:, num_channels:, :])
+            acts = t_act * s_act
+    
+            
+            #
+            
+            
 
             res_skip_acts = self.res_skip_layers[i](acts)
             if i < self.num_layers - 1:
@@ -642,7 +453,7 @@ class HifiGanResidualBlock(nn.Module):
 
 
 class VitsHifiGan(nn.Module):
-    def __init__(self, config: VitsConfig):
+    def __init__(self, config):
         super().__init__()
         self.config = config
         self.num_kernels = len(config.resblock_kernel_sizes)
@@ -695,25 +506,12 @@ class VitsHifiGan(nn.Module):
             layer.remove_weight_norm()
 
     def forward(
-        self, spectrogram: torch.FloatTensor, global_conditioning: Optional[torch.FloatTensor] = None
-    ) -> torch.FloatTensor:
-        r"""
-        Converts a spectrogram into a speech waveform.
-
-        Args:
-            spectrogram (`torch.FloatTensor` of shape `(batch_size, config.spectrogram_bins, sequence_length)`):
-                Tensor containing the spectrograms.
-            global_conditioning (`torch.FloatTensor` of shape `(batch_size, config.speaker_embedding_size, 1)`, *optional*):
-                Tensor containing speaker embeddings, for multispeaker models.
-
-        Returns:
-            `torch.FloatTensor`: Tensor of shape shape `(batch_size, 1, num_frames)` containing the speech waveform.
-        """
+        self,
+        spectrogram,
+        global_conditioning=None):
+        
         hidden_states = self.conv_pre(spectrogram)
 
-        if global_conditioning is not None:
-            hidden_states = hidden_states + self.cond(global_conditioning)
-
         for i in range(self.num_upsamples):
             hidden_states = nn.functional.leaky_relu(hidden_states, self.config.leaky_relu_slope)
             hidden_states = self.upsampler[i](hidden_states)
@@ -730,7 +528,7 @@ class VitsHifiGan(nn.Module):
 
 
 class VitsResidualCouplingLayer(nn.Module):
-    def __init__(self, config: VitsConfig):
+    def __init__(self, config):
         super().__init__()
         self.half_channels = config.flow_size // 2
 
@@ -757,7 +555,7 @@ class VitsResidualCouplingLayer(nn.Module):
 
 
 class VitsResidualCouplingBlock(nn.Module):
-    def __init__(self, config: VitsConfig):
+    def __init__(self, config):
         super().__init__()
         self.flows = nn.ModuleList()
         for _ in range(config.prior_encoder_num_flows):
@@ -776,13 +574,12 @@ class VitsResidualCouplingBlock(nn.Module):
 
 
 class VitsDilatedDepthSeparableConv(nn.Module):
-    def __init__(self, config: VitsConfig, dropout_rate=0.0):
+    def __init__(self, config, dropout_rate=0.0):
         super().__init__()
         kernel_size = config.duration_predictor_kernel_size
         channels = config.hidden_size
         self.num_layers = config.depth_separable_num_layers
 
-        self.dropout = nn.Dropout(dropout_rate)
         self.convs_dilated = nn.ModuleList()
         self.convs_pointwise = nn.ModuleList()
         self.norms_1 = nn.ModuleList()
@@ -815,14 +612,14 @@ class VitsDilatedDepthSeparableConv(nn.Module):
             hidden_states = self.convs_pointwise[i](hidden_states)
             hidden_states = self.norms_2[i](hidden_states.transpose(1, -1)).transpose(1, -1)
             hidden_states = nn.functional.gelu(hidden_states)
-            hidden_states = self.dropout(hidden_states)
+            
             inputs = inputs + hidden_states
 
         return inputs * padding_mask
 
 
 class VitsConvFlow(nn.Module):
-    def __init__(self, config: VitsConfig):
+    def __init__(self, config):
         super().__init__()
         self.filter_channels = config.hidden_size
         self.half_channels = config.depth_separable_channels // 2
@@ -847,7 +644,7 @@ class VitsConvFlow(nn.Module):
         unnormalized_heights = hidden_states[..., self.num_bins : 2 * self.num_bins] / math.sqrt(self.filter_channels)
         unnormalized_derivatives = hidden_states[..., 2 * self.num_bins :]
 
-        second_half, log_abs_det = _unconstrained_rational_quadratic_spline(
+        second_half = _unconstrained_rational_quadratic_spline(
             second_half,
             unnormalized_widths,
             unnormalized_heights,
@@ -857,11 +654,8 @@ class VitsConvFlow(nn.Module):
         )
 
         outputs = torch.cat([first_half, second_half], dim=1) * padding_mask
-        if not reverse:
-            log_determinant = torch.sum(log_abs_det * padding_mask, [1, 2])
-            return outputs, log_determinant
-        else:
-            return outputs, None
+        
+        return outputs, None
 
 
 class VitsElementwiseAffine(nn.Module):
@@ -873,10 +667,11 @@ class VitsElementwiseAffine(nn.Module):
 
     def forward(self, inputs, padding_mask, global_conditioning=None, reverse=False):
         if not reverse:
-            outputs = self.translate + torch.exp(self.log_scale) * inputs
-            outputs = outputs * padding_mask
-            log_determinant = torch.sum(self.log_scale * padding_mask, [1, 2])
-            return outputs, log_determinant
+            raise ValueError
+            # outputs = self.translate + torch.exp(self.log_scale) * inputs
+            # outputs = outputs * padding_mask
+            # log_determinant = torch.sum(self.log_scale * padding_mask, [1, 2])
+            # return outputs, log_determinant
         else:
             outputs = (inputs - self.translate) * torch.exp(-self.log_scale) * padding_mask
             return outputs, None
@@ -890,10 +685,7 @@ class VitsStochasticDurationPredictor(nn.Module):
 
         self.conv_pre = nn.Conv1d(filter_channels, filter_channels, 1)
         self.conv_proj = nn.Conv1d(filter_channels, filter_channels, 1)
-        self.conv_dds = VitsDilatedDepthSeparableConv(
-            config,
-            dropout_rate=config.duration_predictor_dropout,
-        )
+        self.conv_dds = VitsDilatedDepthSeparableConv(config)
 
         if embed_dim != 0:
             self.cond = nn.Conv1d(embed_dim, filter_channels, 1)
@@ -903,69 +695,71 @@ class VitsStochasticDurationPredictor(nn.Module):
         for _ in range(config.duration_predictor_num_flows):
             self.flows.append(VitsConvFlow(config))
 
-        self.post_conv_pre = nn.Conv1d(1, filter_channels, 1)
-        self.post_conv_proj = nn.Conv1d(filter_channels, filter_channels, 1)
-        self.post_conv_dds = VitsDilatedDepthSeparableConv(
-            config,
-            dropout_rate=config.duration_predictor_dropout,
-        )
+        # self.post_conv_pre = nn.Conv1d(1, filter_channels, 1)
+        # self.post_conv_proj = nn.Conv1d(filter_channels, filter_channels, 1)
+        # self.post_conv_dds = VitsDilatedDepthSeparableConv(
+        #     config,
+        #     dropout_rate=config.duration_predictor_dropout,
+        # )
 
-        self.post_flows = nn.ModuleList()
-        self.post_flows.append(VitsElementwiseAffine(config))
-        for _ in range(config.duration_predictor_num_flows):
-            self.post_flows.append(VitsConvFlow(config))
+        # self.post_flows = nn.ModuleList()
+        # self.post_flows.append(VitsElementwiseAffine(config))
+        # for _ in range(config.duration_predictor_num_flows):
+        #     self.post_flows.append(VitsConvFlow(config))
 
     def forward(self, inputs, padding_mask, global_conditioning=None, durations=None, reverse=False, noise_scale=1.0):
         inputs = torch.detach(inputs)
         inputs = self.conv_pre(inputs)
 
         if global_conditioning is not None:
-            global_conditioning = torch.detach(global_conditioning)
-            inputs = inputs + self.cond(global_conditioning)
+            raise ValueError
+            # global_conditioning = torch.detach(global_conditioning)
+            # inputs = inputs + self.cond(global_conditioning)
 
         inputs = self.conv_dds(inputs, padding_mask)
         inputs = self.conv_proj(inputs) * padding_mask
 
         if not reverse:
-            hidden_states = self.post_conv_pre(durations)
-            hidden_states = self.post_conv_dds(hidden_states, padding_mask)
-            hidden_states = self.post_conv_proj(hidden_states) * padding_mask
-
-            random_posterior = (
-                torch.randn(durations.size(0), 2, durations.size(2)).to(device=inputs.device, dtype=inputs.dtype)
-                * padding_mask
-            )
-            log_determinant_posterior_sum = 0
-            latents_posterior = random_posterior
-            for flow in self.post_flows:
-                latents_posterior, log_determinant = flow(
-                    latents_posterior, padding_mask, global_conditioning=inputs + hidden_states
-                )
-                latents_posterior = torch.flip(latents_posterior, [1])
-                log_determinant_posterior_sum += log_determinant
-
-            first_half, second_half = torch.split(latents_posterior, [1, 1], dim=1)
-
-            log_determinant_posterior_sum += torch.sum(
-                (nn.functional.logsigmoid(first_half) + nn.functional.logsigmoid(-first_half)) * padding_mask, [1, 2]
-            )
-            logq = (
-                torch.sum(-0.5 * (math.log(2 * math.pi) + (random_posterior**2)) * padding_mask, [1, 2])
-                - log_determinant_posterior_sum
-            )
-
-            first_half = (durations - torch.sigmoid(first_half)) * padding_mask
-            first_half = torch.log(torch.clamp_min(first_half, 1e-5)) * padding_mask
-            log_determinant_sum = torch.sum(-first_half, [1, 2])
-
-            latents = torch.cat([first_half, second_half], dim=1)
-            for flow in self.flows:
-                latents, log_determinant = flow(latents, padding_mask, global_conditioning=inputs)
-                latents = torch.flip(latents, [1])
-                log_determinant_sum += log_determinant
-
-            nll = torch.sum(0.5 * (math.log(2 * math.pi) + (latents**2)) * padding_mask, [1, 2]) - log_determinant_sum
-            return nll + logq
+            raise ValueError
+            # hidden_states = self.post_conv_pre(durations)
+            # hidden_states = self.post_conv_dds(hidden_states, padding_mask)
+            # hidden_states = self.post_conv_proj(hidden_states) * padding_mask
+
+            # random_posterior = (
+            #     torch.randn(durations.size(0), 2, durations.size(2)).to(device=inputs.device, dtype=inputs.dtype)
+            #     * padding_mask
+            # )
+            # log_determinant_posterior_sum = 0
+            # latents_posterior = random_posterior
+            # for flow in self.post_flows:
+            #     latents_posterior, log_determinant = flow(
+            #         latents_posterior, padding_mask, global_conditioning=inputs + hidden_states
+            #     )
+            #     latents_posterior = torch.flip(latents_posterior, [1])
+            #     log_determinant_posterior_sum += log_determinant
+
+            # first_half, second_half = torch.split(latents_posterior, [1, 1], dim=1)
+
+            # log_determinant_posterior_sum += torch.sum(
+            #     (nn.functional.logsigmoid(first_half) + nn.functional.logsigmoid(-first_half)) * padding_mask, [1, 2]
+            # )
+            # logq = (
+            #     torch.sum(-0.5 * (math.log(2 * math.pi) + (random_posterior**2)) * padding_mask, [1, 2])
+            #     - log_determinant_posterior_sum
+            # )
+
+            # first_half = (durations - torch.sigmoid(first_half)) * padding_mask
+            # first_half = torch.log(torch.clamp_min(first_half, 1e-5)) * padding_mask
+            # log_determinant_sum = torch.sum(-first_half, [1, 2])
+
+            # latents = torch.cat([first_half, second_half], dim=1)
+            # for flow in self.flows:
+            #     latents, log_determinant = flow(latents, padding_mask, global_conditioning=inputs)
+            #     latents = torch.flip(latents, [1])
+            #     log_determinant_sum += log_determinant
+
+            # nll = torch.sum(0.5 * (math.log(2 * math.pi) + (latents**2)) * padding_mask, [1, 2]) - log_determinant_sum
+            # return nll + logq
         else:
             flows = list(reversed(self.flows))
             flows = flows[:-2] + [flows[-1]]  # remove a useless vflow
@@ -982,51 +776,17 @@ class VitsStochasticDurationPredictor(nn.Module):
             return log_duration
 
 
-class VitsDurationPredictor(nn.Module):
-    def __init__(self, config):
-        super().__init__()
-        kernel_size = config.duration_predictor_kernel_size
-        filter_channels = config.duration_predictor_filter_channels
-
-        self.dropout = nn.Dropout(config.duration_predictor_dropout)
-        self.conv_1 = nn.Conv1d(config.hidden_size, filter_channels, kernel_size, padding=kernel_size // 2)
-        self.norm_1 = nn.LayerNorm(filter_channels, eps=config.layer_norm_eps)
-        self.conv_2 = nn.Conv1d(filter_channels, filter_channels, kernel_size, padding=kernel_size // 2)
-        self.norm_2 = nn.LayerNorm(filter_channels, eps=config.layer_norm_eps)
-        self.proj = nn.Conv1d(filter_channels, 1, 1)
-
-        if config.speaker_embedding_size != 0:
-            self.cond = nn.Conv1d(config.speaker_embedding_size, config.hidden_size, 1)
-
-    def forward(self, inputs, padding_mask, global_conditioning=None):
-        inputs = torch.detach(inputs)
-
-        if global_conditioning is not None:
-            global_conditioning = torch.detach(global_conditioning)
-            inputs = inputs + self.cond(global_conditioning)
 
-        inputs = self.conv_1(inputs * padding_mask)
-        inputs = torch.relu(inputs)
-        inputs = self.norm_1(inputs.transpose(1, -1)).transpose(1, -1)
-        inputs = self.dropout(inputs)
-
-        inputs = self.conv_2(inputs * padding_mask)
-        inputs = torch.relu(inputs)
-        inputs = self.norm_2(inputs.transpose(1, -1)).transpose(1, -1)
-        inputs = self.dropout(inputs)
-
-        inputs = self.proj(inputs * padding_mask)
-        return inputs * padding_mask
 
 
 class VitsAttention(nn.Module):
     """Multi-headed attention with relative positional representation."""
 
-    def __init__(self, config: VitsConfig):
+    def __init__(self, config):
         super().__init__()
         self.embed_dim = config.hidden_size
         self.num_heads = config.num_attention_heads
-        self.dropout = config.attention_dropout
+        
         self.window_size = config.window_size
 
         self.head_dim = self.embed_dim // self.num_heads
@@ -1044,10 +804,11 @@ class VitsAttention(nn.Module):
         self.out_proj = nn.Linear(self.embed_dim, self.embed_dim, bias=config.use_bias)
 
         if self.window_size:
+            # Those provide relative pos embs for k/v interpolated from 2*4+1 to 1027 time frames - duration of txt 
             self.emb_rel_k = nn.Parameter(torch.randn(1, self.window_size * 2 + 1, self.head_dim) * self.scaling)
             self.emb_rel_v = nn.Parameter(torch.randn(1, self.window_size * 2 + 1, self.head_dim) * self.scaling)
 
-    def _shape(self, tensor: torch.Tensor, seq_len: int, bsz: int):
+    def _shape(self, tensor, seq_len, bsz):
         return tensor.view(bsz, seq_len, self.num_heads, self.head_dim).transpose(1, 2).contiguous()
 
     def forward(
@@ -1080,55 +841,32 @@ class VitsAttention(nn.Module):
         src_len = key_states.size(1)
         attn_weights = torch.bmm(query_states, key_states.transpose(1, 2))
 
-        if attn_weights.size() != (bsz * self.num_heads, tgt_len, src_len):
-            raise ValueError(
-                f"Attention weights should be of size {(bsz * self.num_heads, tgt_len, src_len)}, but is"
-                f" {attn_weights.size()}"
-            )
+        
 
         if self.window_size is not None:
             # 4
-            key_relative_embeddings = self._get_relative_embeddings(self.emb_rel_k, src_len)
+            # key_relative_embeddings = self._get_relative_embeddings(self.emb_rel_k, src_len)
+            key_relative_embeddings = self._get_relative_embeddings(self.emb_rel_k, src_len)   # try fix k.shape[2] to have consistent voice deu 
+            # print(f'{self.emb_rel_k.shape=} {key_relative_embeddings.shape=}\n\nL855')
             relative_logits = torch.matmul(query_states, key_relative_embeddings.transpose(-2, -1))
+            # -- only here (key)
             rel_pos_bias = self._relative_position_to_absolute_position(relative_logits)
             attn_weights += rel_pos_bias
 
         if attention_mask is not None:
-            if attention_mask.size() != (bsz, 1, tgt_len, src_len):
-                raise ValueError(
-                    f"Attention mask should be of size {(bsz, 1, tgt_len, src_len)}, but is {attention_mask.size()}"
-                )
+
             attn_weights = attn_weights.view(bsz, self.num_heads, tgt_len, src_len) + attention_mask
             attn_weights = attn_weights.view(bsz * self.num_heads, tgt_len, src_len)
-        
-        # Is possible that starting frames of this attentio hold the choice of voice to place the generation in male or female for german
-        #   1. Is plausible to have some pre-append or post-append frames (whose TTS is always male or female )
-        
-        # 
-        # --
-        # ___IN attn 1110__ torch.Size([2, 927, 927])
-        # ___IN attn 1110__ torch.Size([2, 927, 927])
-        # ___IN attn 1110__ torch.Size([2, 927, 927])
-        # ___IN attn 1110__ torch.Size([2, 927, 927])
-        # ___IN attn 1110__ torch.Size([2, 927, 927])
-        # ___IN attn 1110__ torch.Size([2, 927, 927])  # this appears to use always thefull len of bert hidden states 
-        # --
         attn_weights = nn.functional.softmax(attn_weights, dim=-1)
+        attn_output = torch.bmm(attn_weights, 
+                                value_states)
 
-        attn_probs = nn.functional.dropout(attn_weights, p=self.dropout, training=self.training)
 
-        attn_output = torch.bmm(attn_probs, value_states)
-
-        if attn_output.size() != (bsz * self.num_heads, tgt_len, self.head_dim):
-            raise ValueError(
-                f"`attn_output` should be of size {(bsz, self.num_heads, tgt_len, self.head_dim)}, but is"
-                f" {attn_output.size()}"
-            )
 
         if self.window_size is not None:
             # Entering here with self.window_size = 4
             value_relative_embeddings = self._get_relative_embeddings(self.emb_rel_v, src_len)
-            relative_weights = self._absolute_position_to_relative_position(attn_probs)
+            relative_weights = self._absolute_position_to_relative_position(attn_weights)
             rel_pos_bias = torch.matmul(relative_weights, value_relative_embeddings)
             attn_output += rel_pos_bias
 
@@ -1185,12 +923,8 @@ class VitsFeedForward(nn.Module):
         super().__init__()
         self.conv_1 = nn.Conv1d(config.hidden_size, config.ffn_dim, config.ffn_kernel_size)
         self.conv_2 = nn.Conv1d(config.ffn_dim, config.hidden_size, config.ffn_kernel_size)
-        self.dropout = nn.Dropout(config.activation_dropout)
-
-        if isinstance(config.hidden_act, str):
-            self.act_fn = ACT2FN[config.hidden_act]
-        else:
-            self.act_fn = config.hidden_act
+        self.act_fn = nn.ReLU()
+        
 
         if config.ffn_kernel_size > 1:
             pad_left = (config.ffn_kernel_size - 1) // 2
@@ -1209,7 +943,7 @@ class VitsFeedForward(nn.Module):
 
         hidden_states = self.conv_1(hidden_states)
         hidden_states = self.act_fn(hidden_states)
-        hidden_states = self.dropout(hidden_states)
+        
 
         hidden_states = hidden_states * padding_mask
         if self.padding is not None:
@@ -1223,10 +957,10 @@ class VitsFeedForward(nn.Module):
 
 
 class VitsEncoderLayer(nn.Module):
-    def __init__(self, config: VitsConfig):
+    def __init__(self, config):
         super().__init__()
         self.attention = VitsAttention(config)
-        self.dropout = nn.Dropout(config.hidden_dropout)
+        
         self.layer_norm = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps)
         self.feed_forward = VitsFeedForward(config)
         self.final_layer_norm = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps)
@@ -1245,12 +979,12 @@ class VitsEncoderLayer(nn.Module):
             output_attentions=output_attentions,
         )
 
-        hidden_states = self.dropout(hidden_states)
+        
         hidden_states = self.layer_norm(residual + hidden_states)
 
         residual = hidden_states
         hidden_states = self.feed_forward(hidden_states, padding_mask)
-        hidden_states = self.dropout(hidden_states)
+        
         hidden_states = self.final_layer_norm(residual + hidden_states)
 
         outputs = (hidden_states,)
@@ -1262,7 +996,7 @@ class VitsEncoderLayer(nn.Module):
 
 
 class VitsEncoder(nn.Module):
-    def __init__(self, config: VitsConfig):
+    def __init__(self, config):
         super().__init__()
         self.config = config
         self.layers = nn.ModuleList([VitsEncoderLayer(config) for _ in range(config.num_hidden_layers)])
@@ -1277,7 +1011,7 @@ class VitsEncoder(nn.Module):
         output_attentions: Optional[bool] = None,
         output_hidden_states: Optional[bool] = None,
         return_dict: Optional[bool] = None,
-    ) -> Union[Tuple, BaseModelOutput]:
+    ):
         all_hidden_states = () if output_hidden_states else None
         all_self_attentions = () if output_attentions else None
 
@@ -1293,11 +1027,6 @@ class VitsEncoder(nn.Module):
         for encoder_layer in self.layers:
             if output_hidden_states:
                 all_hidden_states = all_hidden_states + (hidden_states,)
-
-            # add LayerDrop (see https://arxiv.org/abs/1909.11556 for description)
-            dropout_probability = np.random.uniform(0, 1)
-
-            skip_the_layer = self.training and (dropout_probability < self.layerdrop)
             
             layer_outputs = encoder_layer(
                 hidden_states,
@@ -1306,21 +1035,12 @@ class VitsEncoder(nn.Module):
                 output_attentions=output_attentions,
             )
             hidden_states = layer_outputs[0]
-
-            if skip_the_layer:
-                layer_outputs = (None, None)
-
+            
             if output_attentions:
                 all_self_attentions = all_self_attentions + (layer_outputs[1],)
 
         hidden_states = hidden_states * padding_mask
 
-        if output_hidden_states:
-            all_hidden_states = all_hidden_states + (hidden_states,)
-
-        if not return_dict:
-            return tuple(v for v in [hidden_states, all_hidden_states, all_self_attentions] if v is not None)
-
         return BaseModelOutput(
             last_hidden_state=hidden_states,
             hidden_states=all_hidden_states,
@@ -1333,18 +1053,18 @@ class VitsTextEncoder(nn.Module):
     Transformer encoder that uses relative positional representation instead of absolute positional encoding.
     """
 
-    def __init__(self, config: VitsConfig):
+    def __init__(self, config):
         super().__init__()
         self.config = config
         self.embed_tokens = nn.Embedding(config.vocab_size, config.hidden_size, config.pad_token_id)
         self.encoder = VitsEncoder(config)  # 6 Layers of VitsAttention
         self.project = nn.Conv1d(config.hidden_size, config.flow_size * 2, kernel_size=1)
 
-    def get_input_embeddings(self):
-        return self.embed_tokens
+    # def get_input_embeddings(self):
+    #     return self.embed_tokens
 
-    def set_input_embeddings(self, value):
-        self.embed_tokens = value
+    # def set_input_embeddings(self, value):
+    #     self.embed_tokens = value
 
     def forward(
         self,
@@ -1354,7 +1074,7 @@ class VitsTextEncoder(nn.Module):
         output_attentions: Optional[bool] = None,
         output_hidden_states: Optional[bool] = None,
         return_dict: Optional[bool] = True,
-    ) -> Union[Tuple[torch.Tensor], VitsTextEncoderOutput]:
+    ):
         hidden_states = self.embed_tokens(input_ids) * math.sqrt(self.config.hidden_size)
 
         encoder_outputs = self.encoder(
@@ -1371,10 +1091,6 @@ class VitsTextEncoder(nn.Module):
         stats = self.project(last_hidden_state.transpose(1, 2)).transpose(1, 2) * padding_mask
         prior_means, prior_log_variances = torch.split(stats, self.config.flow_size, dim=2)
 
-        if not return_dict:
-            outputs = (last_hidden_state, prior_means, prior_log_variances) + encoder_outputs[1:]
-            return outputs
-
         return VitsTextEncoderOutput(
             last_hidden_state=last_hidden_state,
             prior_means=prior_means,
@@ -1416,7 +1132,7 @@ class VitsPreTrainedModel(PreTrainedModel):
 
 
 class VitsModel(VitsPreTrainedModel):
-    def __init__(self, config: VitsConfig):
+    def __init__(self, config):
         super().__init__(config)
         self.config = config
         self.text_encoder = VitsTextEncoder(config)  # has VitsEncoder that includes 6L of VitsAttention
@@ -1426,13 +1142,14 @@ class VitsModel(VitsPreTrainedModel):
         if config.use_stochastic_duration_prediction:
             self.duration_predictor = VitsStochasticDurationPredictor(config)
         else:
-            self.duration_predictor = VitsDurationPredictor(config)
+            raise ValueError
+            # self.duration_predictor = VitsDurationPredictor(config)
 
         if config.num_speakers > 1:
             self.embed_speaker = nn.Embedding(config.num_speakers, config.speaker_embedding_size)
 
-        # 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
 
@@ -1451,33 +1168,9 @@ class VitsModel(VitsPreTrainedModel):
         output_hidden_states = None,
         return_dict = None,
         labels = None,
+        speed=None,
     ):
-        r"""
-        labels (`torch.FloatTensor` of shape `(batch_size, config.spectrogram_bins, sequence_length)`, *optional*):
-            Float values of target spectrogram. Timesteps set to `-100.0` are ignored (masked) for the loss
-            computation.
-
-        Returns:
-
-        Example:
-
-        ```python
-        >>> from transformers import VitsTokenizer, VitsModel, set_seed
-        >>> import torch
-
-        >>> tokenizer = VitsTokenizer.from_pretrained("facebook/mms-tts-eng")
-        >>> model = VitsModel.from_pretrained("facebook/mms-tts-eng")
-
-        >>> inputs = tokenizer(text="Hello - my dog is cute", return_tensors="pt")
-
-        >>> set_seed(555)  # make deterministic
 
-        >>> with torch.no_grad():
-        ...     outputs = model(inputs["input_ids"])
-        >>> outputs.waveform.shape
-        torch.Size([1, 45824])
-        ```
-        """
         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
@@ -1528,7 +1221,7 @@ class VitsModel(VitsPreTrainedModel):
             raise ValueError
             # log_duration = self.duration_predictor(hidden_states, input_padding_mask, speaker_embeddings)
 
-        length_scale = 1.0 / self.speaking_rate
+        length_scale = 1.0 / speed
         duration = torch.ceil(torch.exp(log_duration) * input_padding_mask * length_scale)
         predicted_lengths = torch.clamp_min(torch.sum(duration, [1, 2]), 1).long()
 

Utils/text_utils.py

@@ -87,7 +87,8 @@ def split_into_sentences(text):
     # -- even 400 phonemes sometimes OOM in cuda:4
     sentences = [sub_sent+' ' for s in sentences for sub_sent in textwrap.wrap(s, 300, break_long_words=0)]
     
-    if sentences and not sentences[-1]: sentences = sentences[:-1]
+    # if sentences and not sentences[-1]: 
+    #     sentences = sentences[:-1]
     return sentences
 
 def store_ssml(text=None,

api.py

@@ -16,10 +16,7 @@ from moviepy.video.VideoClip import ImageClip
 from audiocraft.builders import AudioGen
 
 CACHE_DIR = 'flask_cache/'
-PIECE_OF_SOUND_DURATION = 4.74 # seconds 
-sound_generator = AudioGen(
-    duration=PIECE_OF_SOUND_DURATION
-                            ).to('cuda:0').eval()
+sound_generator = AudioGen().to('cuda:0').eval()  # duration chosen in generate()
 
 Path(CACHE_DIR).mkdir(parents=True, exist_ok=True)
 
@@ -58,22 +55,16 @@ def _resize(image, width=None, height=None, inter=cv2.INTER_AREA):
     # return the resized image
     return resized
 
-def overlay(x, soundscape=None):
-    # pre-calculate the n_repeat here then apply torchaudio.resample and repeat insd sound_gen forward()
+def overlay(x,soundscape=None):
     if soundscape is not None:
-        
         background = sound_generator.generate(soundscape,
-            n_repeat=int(len(x) / (PIECE_OF_SOUND_DURATION * 16000)) + 1
-                                        ).detach().cpu().numpy() # bs, 11400 @.74s
+                                              duration=len(x)/24000 + .74, # seconds - TTS @ 24kHz
+                                              ).detach().cpu().numpy() # bs, 11400 @.74s
         
         # blend TTS
         
         x = .5 * x + .5 * background[:len(x)]
-        
-    else:
-        
-        print('sound_background = None')
-        
+    # TTS & AudioGen at 24kHz                
     return x
 
 

audiocraft/builders.py

@@ -11,6 +11,8 @@ from .lm import LMModel
 from .seanet import SEANetDecoder
 from .vq import ResidualVectorQuantizer
 
+N_REPEAT = 7  # num (virtual batch_size) clones of audio sounds
+
 def _shift(x):
     # [bs, samples] shift circular each batch elem of sound
     n = x.shape[1]
@@ -42,29 +44,26 @@ class AudioGen(nn.Module):
     
     # https://huggingface.co/facebook/audiogen-medium
     
-    def __init__(self,
-                 duration=2.24,  # s
-                 ):
+    def __init__(self):
 
         super().__init__()
         self.load_compression_model()
         self.load_lm_model()
-        self.duration = duration
+        
         #  AudioGen = 16KHZ                StyleTTS2 = 24 KHz / MMSTTS = 24 KHz
         self.resample_fn = torchaudio.transforms.Resample(16000, 24000)
-
-    @property
-    def frame_rate(self):
-        return self.compression_model.frame_rate
     
     def generate(self,
                  descriptions,
-                 n_repeat=3):
+                 duration=2.24,  ## seconds of audio
+                 ):
         
         with torch.no_grad():
+            print(duration / N_REPEAT * self.compression_model.frame_rate, 'DURATION TOKENS AudioGen')
+            
             gen_tokens = self.lm.generate(
-                descriptions=[descriptions]*3,
-                max_tokens=int(self.duration * self.frame_rate)) # [bs, 4, 37 * self.lm.n_draw]
+                descriptions=[descriptions] * N_REPEAT,
+                max_tokens=int(duration / N_REPEAT * self.compression_model.frame_rate)) # [bs, 4, 37 * self.lm.n_draw]
             x = self.compression_model.decode(gen_tokens, None)   #[bs, 1, 11840]
 
             x = x[:, 0, :]  # last samples have splash sounds DISCARD 25000 last samples
@@ -75,7 +74,7 @@ class AudioGen(nn.Module):
 
             # batch size = different sounds for same txt
             
-            x = x.repeat(1, n_repeat)
+            x = x.repeat(1, N_REPEAT)
             
             # less periodic - shift every batch elem
             

audiocraft/transformer.py

@@ -32,15 +32,18 @@ class StreamingMultiheadAttention(nn.Module):
     def __init__(self, 
                  embed_dim, 
                  num_heads,
-                 cross_attention = False):
+                 cross_attention = False,
+                 ):
+        
         super().__init__()
+        
         self.cross_attention = cross_attention
         self.embed_dim = embed_dim
         self.k_history = None  # previous k from the previous tokens seen in the current generation - only for selt.attn
         self.v_history = None  # clean up IN LM after finishing GENERATION - Each 1...47 mha has different kv history        
         self.num_heads = num_heads
         self.out_proj = nn.Linear(embed_dim, embed_dim, bias=False)
-        self.register_buffer('in_proj_weight', torch.ones((3 * embed_dim, embed_dim), 
+        self.register_buffer('in_proj_weight', torch.ones((3 * embed_dim, embed_dim),
                                                            dtype=torch.float))
 
     def forward(self,
@@ -62,7 +65,7 @@ class StreamingMultiheadAttention(nn.Module):
             # print(q.shape, k.shape, v.shape, q.sum(), k.sum(), v.sum(),'CROSS A5')
         else:
             # 1st projected makes k,v (instantaneous)
-            # 2nd cat
+            # Here else is self_attention for audio with itself (above is cross attention txt)
             
             
             # HISTORY - DIFFERENT FOR EACH TRANSF LAYER
@@ -75,12 +78,21 @@ class StreamingMultiheadAttention(nn.Module):
 
                 
             if self.k_history is not None:
-                # 
-                # pk.shape=torch.Size([2, 24, 3, 64]) k.shape=torch.Size([2, 24, 1, 64]) CONCAT
-                # has to be 4D with batch 1 due to single condition 3=seqlen
-                # 24 heads 64 dimofh
+                # k_history.shape = torch.Size([2*N_REPEAT, 24, 3, 64])  FOR cfg > k.shape=torch.Size([2, 24, 1, 64])
+                # 24 heads 64 dim
                 self.k_history = torch.cat([self.k_history, k], 2)  # IF ctrl^c here during live demo it is non-atomic k!=v
-                self.v_history = torch.cat([self.v_history, v], 2)
+                self.v_history = torch.cat([self.v_history, v], 2)  # thus it will try to continue with incompatible k/v dims!
+                # Preserve first 4-10 tokens & flush kv
+                if self.k_history.shape[2] > 24:
+
+                    # find LOWEST l2 norm of keys > https://arxiv.org/pdf/2406.11430v4
+                    
+                    low_norm = (self.k_history * self.k_history).mean(3, keepdims=True).sum(1, keepdims=True)   # [bs, 24, T, 64] -> [bs, T]
+                    _, _ix = torch.topk(low_norm, k=10, dim=2, largest=False)  # shows background music due to cfg - looses the txt conditioning if flushed!
+                    _ix = _ix.repeat(1, 24, 1, 64)
+                    # print(_ix.shape)
+                    self.k_history = torch.gather(self.k_history, 2, _ix)
+                    self.v_history = torch.gather(self.v_history, 2, _ix)
 
             else:
                 # init on 1st token (for all 47 transf layers)

demo.py

@@ -64,7 +64,7 @@ def tts_entry(text='»Vom Prof. Friedrich ist noch eine recht schöne große Lan
     else:
 
         # MMS TTS - list of sentences
-        x = msinference.foreign(text=[text],
+        x = msinference.foreign(text=text,
                                 lang=voice,  # voice = 'romanian', 'serbian' 'hungarian'
                                 speed=speed)  # normalisation externally
 
@@ -74,4 +74,4 @@ def tts_entry(text='»Vom Prof. Friedrich ist noch eine recht schöne große Lan
     print(x.shape, 'TTS OK')
     return x
 
-soundfile.write(f'demo.wav', tts_entry(), 24000)
+soundfile.write(f'de00i.wav', tts_entry(), 24000)

msinference.py

@@ -9,7 +9,7 @@ import torchaudio
 import librosa
 from models import ProsodyPredictor, TextEncoder, StyleEncoder, load_F0_models
 from nltk.tokenize import word_tokenize
-
+import textwrap
 # IPA Phonemizer: https://github.com/bootphon/phonemizer
 
 _pad = "$"
@@ -76,10 +76,6 @@ def compute_style(path):
 device = 'cpu'
 if torch.cuda.is_available():
     device = 'cuda'
-elif torch.backends.mps.is_available():
-    # print("MPS would be available but cannot be used rn")
-    pass
-    # device = 'mps'
 
 import phonemizer
 global_phonemizer = phonemizer.backend.EspeakBackend(language='en-us', preserve_punctuation=True,  with_stress=True)
@@ -311,7 +307,8 @@ def has_cyrillic(text):
     # https://stackoverflow.com/questions/48255244/python-check-if-a-string-contains-cyrillic-characters
     return bool(re.search('[\u0400-\u04FF]', text))
 
-def foreign(text=None,   # list of text
+def foreign(text=None,   # split sentences here so we can prepend a txt for german to each sentence to 
+                         # fall on the male voice (Sink attn)
             lang='romanian',
             speed=None):
 
@@ -325,7 +322,7 @@ def foreign(text=None,   # list of text
         
     elif any([i in lang for i in ['ser', 'bosn', 'herzegov', 'montenegr', 'macedon']]):
         
-        if has_cyrillic(text[0]):  # check 0-th sentence if is cyrillic
+        if has_cyrillic(text):  # check 0-th sentence if is cyrillic
             
             lang_code = 'rmc-script_cyrillic'   # romani carpathian (also has latin / cyrillic Vlax)
         
@@ -356,10 +353,19 @@ def foreign(text=None,   # list of text
     
     net_g = VitsModel.from_pretrained(f'facebook/mms-tts-{lang_code}').eval().to(device)
     tokenizer = VitsTokenizer.from_pretrained(f'facebook/mms-tts-{lang_code}')
-    
+
     # CALL MMS TTS VITS
     
-    x = []
+    total_audio = []
+    
+    # Split long sentences if deu to control voice switch - for other languages let text no-split
+    if not isinstance(text, list):
+        if lang_code == 'deu':
+            # Split Very long sentences >500 phoneme - StyleTTS2 crashes # -- even 400 phonemes sometimes OOM in cuda:4
+            # However prosody is nicer on non-split for MMS TTS
+            text = [sub_sent+' ' for sub_sent in textwrap.wrap(text, 300, break_long_words=0)] 
+        else:
+            text = [text]
     
     for _t in text:
 
@@ -373,23 +379,27 @@ def foreign(text=None,   # list of text
         elif lang_code == 'ron':            
 
             _t = _t.replace("ţ", "ț"
-                        ).replace('ț','ts').replace('î', 'u')
+                        ).replace('ț','ts').replace('î', 'u').replace('â','a').replace('ş','s')
 
         # /data/dkounadis/.hf7/hub/models--facebook--mms-tts/snapshots/44cc7fb408064ef9ea6e7c59130d88cac1274671/models/rmc-script_latin/vocab.txt
         inputs = tokenizer(_t, return_tensors="pt")  # input_ids / attention_mask
-
+        
         with torch.no_grad():
-            # -- reset speed
-            net_g.speaking_rate = speed
-            # --
-            x.append(
-                net_g(input_ids=inputs.input_ids.to(device),
-                      attention_mask=inputs.attention_mask.to(device))
-            )
-            print(x[-1].shape)
-        print(f'{speed=}\n\n\n\n_______________________________ {_t}')
             
-    x = torch.cat(x).cpu().numpy()
+            # MMS
+            
+            x = net_g(input_ids=inputs.input_ids.to(device),
+                           attention_mask=inputs.attention_mask.to(device),
+                           speed = .94 + .4 * np.random.rand()  # variable speed / sentence
+                           )[0, :]
+            
+            # crop the 1st audio - is PREFIX text 156000 samples to chose deu voice / VitsAttention()
+            
+            total_audio.append(x)
+            
+        print(f'\n\n_______________________________ {_t} {x.shape=}')
+            
+    x = torch.cat(total_audio).cpu().numpy()
             
     x /= np.abs(x).max() + 1e-7