Aluode/PerceptionLabPortable
0
1# coding=utf-82# Copyright 2023 The Kakao Enterprise Authors and the HuggingFace Inc. team. All rights reserved.3#4# Licensed under the Apache License, Version 2.0 (the "License");5# you may not use this file except in compliance with the License.6# You may obtain a copy of the License at7#8# http://www.apache.org/licenses/LICENSE-2.09#10# Unless required by applicable law or agreed to in writing, software11# distributed under the License is distributed on an "AS IS" BASIS,12# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.13# See the License for the specific language governing permissions and14# limitations under the License.15"""PyTorch VITS model."""16 17import math18from dataclasses import dataclass19from typing import Any, Optional, Union20 21import numpy as np22import torch23from torch import nn24 25from ...activations import ACT2FN26from ...integrations.deepspeed import is_deepspeed_zero3_enabled27from ...integrations.fsdp import is_fsdp_managed_module28from ...modeling_attn_mask_utils import _prepare_4d_attention_mask29from ...modeling_layers import GradientCheckpointingLayer30from ...modeling_outputs import BaseModelOutput, ModelOutput31from ...modeling_utils import PreTrainedModel32from ...utils import auto_docstring, logging33from .configuration_vits import VitsConfig34 35 36logger = logging.get_logger(__name__)37 38 39@dataclass40@auto_docstring(41 custom_intro="""42 Describes the outputs for the VITS model, with potential hidden states and attentions.43 """44)45class VitsModelOutput(ModelOutput):46 r"""47 waveform (`torch.FloatTensor` of shape `(batch_size, sequence_length)`):48 The final audio waveform predicted by the model.49 sequence_lengths (`torch.FloatTensor` of shape `(batch_size,)`):50 The length in samples of each element in the `waveform` batch.51 spectrogram (`torch.FloatTensor` of shape `(batch_size, sequence_length, num_bins)`):52 The log-mel spectrogram predicted at the output of the flow model. This spectrogram is passed to the Hi-Fi53 GAN decoder model to obtain the final audio waveform.54 """55 56 waveform: Optional[torch.FloatTensor] = None57 sequence_lengths: Optional[torch.FloatTensor] = None58 spectrogram: Optional[tuple[torch.FloatTensor]] = None59 hidden_states: Optional[tuple[torch.FloatTensor]] = None60 attentions: Optional[tuple[torch.FloatTensor]] = None61 62 63@dataclass64@auto_docstring(65 custom_intro="""66 Describes the outputs for the VITS text encoder model, with potential hidden states and attentions.67 """68)69class VitsTextEncoderOutput(ModelOutput):70 r"""71 prior_means (`torch.FloatTensor` of shape `(batch_size, sequence_length, hidden_size)`):72 The predicted mean values of the prior distribution for the latent text variables.73 prior_log_variances (`torch.FloatTensor` of shape `(batch_size, sequence_length, hidden_size)`):74 The predicted log-variance values of the prior distribution for the latent text variables.75 """76 77 last_hidden_state: Optional[torch.FloatTensor] = None78 prior_means: Optional[torch.FloatTensor] = None79 prior_log_variances: Optional[torch.FloatTensor] = None80 hidden_states: Optional[tuple[torch.FloatTensor]] = None81 attentions: Optional[tuple[torch.FloatTensor]] = None82 83 84@torch.jit.script85def fused_add_tanh_sigmoid_multiply(input_a, input_b, num_channels):86 in_act = input_a + input_b87 t_act = torch.tanh(in_act[:, :num_channels, :])88 s_act = torch.sigmoid(in_act[:, num_channels:, :])89 acts = t_act * s_act90 return acts91 92 93def _unconstrained_rational_quadratic_spline(94 inputs,95 unnormalized_widths,96 unnormalized_heights,97 unnormalized_derivatives,98 reverse=False,99 tail_bound=5.0,100 min_bin_width=1e-3,101 min_bin_height=1e-3,102 min_derivative=1e-3,103):104 """105 This transformation represents a monotonically increasing piecewise rational quadratic function. Outside of the106 `tail_bound`, the transform behaves as an identity function.107 108 Args:109 inputs (`torch.FloatTensor` of shape `(batch_size, channels, seq_len)`:110 Second half of the hidden-states input to the Vits convolutional flow module.111 unnormalized_widths (`torch.FloatTensor` of shape `(batch_size, channels, seq_len, duration_predictor_flow_bins)`):112 First `duration_predictor_flow_bins` of the hidden-states from the output of the convolution projection113 layer in the convolutional flow module114 unnormalized_heights (`torch.FloatTensor` of shape `(batch_size, channels, seq_len, duration_predictor_flow_bins)`):115 Second `duration_predictor_flow_bins` of the hidden-states from the output of the convolution projection116 layer in the convolutional flow module117 unnormalized_derivatives (`torch.FloatTensor` of shape `(batch_size, channels, seq_len, duration_predictor_flow_bins)`):118 Third `duration_predictor_flow_bins` of the hidden-states from the output of the convolution projection119 layer in the convolutional flow module120 reverse (`bool`, *optional*, defaults to `False`):121 Whether the model is being run in reverse mode.122 tail_bound (`float`, *optional* defaults to 5):123 Upper and lower limit bound for the rational quadratic function. Outside of this `tail_bound`, the124 transform behaves as an identity function.125 min_bin_width (`float`, *optional*, defaults to 1e-3):126 Minimum bin value across the width dimension for the piecewise rational quadratic function.127 min_bin_height (`float`, *optional*, defaults to 1e-3):128 Minimum bin value across the height dimension for the piecewise rational quadratic function.129 min_derivative (`float`, *optional*, defaults to 1e-3):130 Minimum bin value across the derivatives for the piecewise rational quadratic function.131 Returns:132 outputs (`torch.FloatTensor` of shape `(batch_size, channels, seq_len)`:133 Hidden-states as transformed by the piecewise rational quadratic function with the `tail_bound` limits134 applied.135 log_abs_det (`torch.FloatTensor` of shape `(batch_size, channels, seq_len)`:136 Logarithm of the absolute value of the determinants corresponding to the `outputs` with the `tail_bound`137 limits applied.138 """139 inside_interval_mask = (inputs >= -tail_bound) & (inputs <= tail_bound)140 outside_interval_mask = ~inside_interval_mask141 142 outputs = torch.zeros_like(inputs)143 log_abs_det = torch.zeros_like(inputs)144 constant = np.log(np.exp(1 - min_derivative) - 1)145 146 unnormalized_derivatives = nn.functional.pad(unnormalized_derivatives, pad=(1, 1))147 unnormalized_derivatives[..., 0] = constant148 unnormalized_derivatives[..., -1] = constant149 150 outputs[outside_interval_mask] = inputs[outside_interval_mask]151 log_abs_det[outside_interval_mask] = 0.0152 153 outputs[inside_interval_mask], log_abs_det[inside_interval_mask] = _rational_quadratic_spline(154 inputs=inputs[inside_interval_mask],155 unnormalized_widths=unnormalized_widths[inside_interval_mask, :],156 unnormalized_heights=unnormalized_heights[inside_interval_mask, :],157 unnormalized_derivatives=unnormalized_derivatives[inside_interval_mask, :],158 reverse=reverse,159 tail_bound=tail_bound,160 min_bin_width=min_bin_width,161 min_bin_height=min_bin_height,162 min_derivative=min_derivative,163 )164 return outputs, log_abs_det165 166 167def _rational_quadratic_spline(168 inputs,169 unnormalized_widths,170 unnormalized_heights,171 unnormalized_derivatives,172 reverse,173 tail_bound,174 min_bin_width,175 min_bin_height,176 min_derivative,177):178 """179 This transformation represents a monotonically increasing piecewise rational quadratic function. Unlike the180 function `_unconstrained_rational_quadratic_spline`, the function behaves the same across the `tail_bound`.181 182 Args:183 inputs (`torch.FloatTensor` of shape `(batch_size, channels, seq_len)`:184 Second half of the hidden-states input to the Vits convolutional flow module.185 unnormalized_widths (`torch.FloatTensor` of shape `(batch_size, channels, seq_len, duration_predictor_flow_bins)`):186 First `duration_predictor_flow_bins` of the hidden-states from the output of the convolution projection187 layer in the convolutional flow module188 unnormalized_heights (`torch.FloatTensor` of shape `(batch_size, channels, seq_len, duration_predictor_flow_bins)`):189 Second `duration_predictor_flow_bins` of the hidden-states from the output of the convolution projection190 layer in the convolutional flow module191 unnormalized_derivatives (`torch.FloatTensor` of shape `(batch_size, channels, seq_len, duration_predictor_flow_bins)`):192 Third `duration_predictor_flow_bins` of the hidden-states from the output of the convolution projection193 layer in the convolutional flow module194 reverse (`bool`):195 Whether the model is being run in reverse mode.196 tail_bound (`float`):197 Upper and lower limit bound for the rational quadratic function. Outside of this `tail_bound`, the198 transform behaves as an identity function.199 min_bin_width (`float`):200 Minimum bin value across the width dimension for the piecewise rational quadratic function.201 min_bin_height (`float`):202 Minimum bin value across the height dimension for the piecewise rational quadratic function.203 min_derivative (`float`):204 Minimum bin value across the derivatives for the piecewise rational quadratic function.205 Returns:206 outputs (`torch.FloatTensor` of shape `(batch_size, channels, seq_len)`:207 Hidden-states as transformed by the piecewise rational quadratic function.208 log_abs_det (`torch.FloatTensor` of shape `(batch_size, channels, seq_len)`:209 Logarithm of the absolute value of the determinants corresponding to the `outputs`.210 """211 upper_bound = tail_bound212 lower_bound = -tail_bound213 214 if torch.min(inputs) < lower_bound or torch.max(inputs) > upper_bound:215 raise ValueError("Input to a transform is not within its domain")216 217 num_bins = unnormalized_widths.shape[-1]218 219 if min_bin_width * num_bins > 1.0:220 raise ValueError(f"Minimal bin width {min_bin_width} too large for the number of bins {num_bins}")221 if min_bin_height * num_bins > 1.0:222 raise ValueError(f"Minimal bin height {min_bin_height} too large for the number of bins {num_bins}")223 224 widths = nn.functional.softmax(unnormalized_widths, dim=-1)225 widths = min_bin_width + (1 - min_bin_width * num_bins) * widths226 cumwidths = torch.cumsum(widths, dim=-1)227 cumwidths = nn.functional.pad(cumwidths, pad=(1, 0), mode="constant", value=0.0)228 cumwidths = (upper_bound - lower_bound) * cumwidths + lower_bound229 cumwidths[..., 0] = lower_bound230 cumwidths[..., -1] = upper_bound231 widths = cumwidths[..., 1:] - cumwidths[..., :-1]232 233 derivatives = min_derivative + nn.functional.softplus(unnormalized_derivatives)234 235 heights = nn.functional.softmax(unnormalized_heights, dim=-1)236 heights = min_bin_height + (1 - min_bin_height * num_bins) * heights237 cumheights = torch.cumsum(heights, dim=-1)238 cumheights = nn.functional.pad(cumheights, pad=(1, 0), mode="constant", value=0.0)239 cumheights = (upper_bound - lower_bound) * cumheights + lower_bound240 cumheights[..., 0] = lower_bound241 cumheights[..., -1] = upper_bound242 heights = cumheights[..., 1:] - cumheights[..., :-1]243 244 bin_locations = cumheights if reverse else cumwidths245 bin_locations[..., -1] += 1e-6246 bin_idx = torch.sum(inputs[..., None] >= bin_locations, dim=-1) - 1247 bin_idx = bin_idx[..., None]248 249 input_cumwidths = cumwidths.gather(-1, bin_idx)[..., 0]250 input_bin_widths = widths.gather(-1, bin_idx)[..., 0]251 252 input_cumheights = cumheights.gather(-1, bin_idx)[..., 0]253 delta = heights / widths254 input_delta = delta.gather(-1, bin_idx)[..., 0]255 256 input_derivatives = derivatives.gather(-1, bin_idx)[..., 0]257 input_derivatives_plus_one = derivatives[..., 1:].gather(-1, bin_idx)[..., 0]258 259 input_heights = heights.gather(-1, bin_idx)[..., 0]260 261 intermediate1 = input_derivatives + input_derivatives_plus_one - 2 * input_delta262 if not reverse:263 theta = (inputs - input_cumwidths) / input_bin_widths264 theta_one_minus_theta = theta * (1 - theta)265 266 numerator = input_heights * (input_delta * theta.pow(2) + input_derivatives * theta_one_minus_theta)267 denominator = input_delta + intermediate1 * theta_one_minus_theta268 outputs = input_cumheights + numerator / denominator269 270 derivative_numerator = input_delta.pow(2) * (271 input_derivatives_plus_one * theta.pow(2)272 + 2 * input_delta * theta_one_minus_theta273 + input_derivatives * (1 - theta).pow(2)274 )275 log_abs_det = torch.log(derivative_numerator) - 2 * torch.log(denominator)276 return outputs, log_abs_det277 else:278 # find the roots of a quadratic equation279 intermediate2 = inputs - input_cumheights280 intermediate3 = intermediate2 * intermediate1281 a = input_heights * (input_delta - input_derivatives) + intermediate3282 b = input_heights * input_derivatives - intermediate3283 c = -input_delta * intermediate2284 285 discriminant = b.pow(2) - 4 * a * c286 if not (discriminant >= 0).all():287 raise RuntimeError(f"invalid discriminant {discriminant}")288 289 root = (2 * c) / (-b - torch.sqrt(discriminant))290 outputs = root * input_bin_widths + input_cumwidths291 292 theta_one_minus_theta = root * (1 - root)293 denominator = input_delta + intermediate1 * theta_one_minus_theta294 derivative_numerator = input_delta.pow(2) * (295 input_derivatives_plus_one * root.pow(2)296 + 2 * input_delta * theta_one_minus_theta297 + input_derivatives * (1 - root).pow(2)298 )299 log_abs_det = torch.log(derivative_numerator) - 2 * torch.log(denominator)300 return outputs, -log_abs_det301 302 303class VitsWaveNet(torch.nn.Module):304 def __init__(self, config: VitsConfig, num_layers: int):305 super().__init__()306 self.hidden_size = config.hidden_size307 self.num_layers = num_layers308 309 self.in_layers = torch.nn.ModuleList()310 self.res_skip_layers = torch.nn.ModuleList()311 self.dropout = nn.Dropout(config.wavenet_dropout)312 313 if hasattr(nn.utils.parametrizations, "weight_norm"):314 weight_norm = nn.utils.parametrizations.weight_norm315 else:316 weight_norm = nn.utils.weight_norm317 318 if config.speaker_embedding_size != 0:319 cond_layer = torch.nn.Conv1d(config.speaker_embedding_size, 2 * config.hidden_size * num_layers, 1)320 self.cond_layer = weight_norm(cond_layer, name="weight")321 322 for i in range(num_layers):323 dilation = config.wavenet_dilation_rate**i324 padding = (config.wavenet_kernel_size * dilation - dilation) // 2325 in_layer = torch.nn.Conv1d(326 in_channels=config.hidden_size,327 out_channels=2 * config.hidden_size,328 kernel_size=config.wavenet_kernel_size,329 dilation=dilation,330 padding=padding,331 )332 in_layer = weight_norm(in_layer, name="weight")333 self.in_layers.append(in_layer)334 335 # last one is not necessary336 if i < num_layers - 1:337 res_skip_channels = 2 * config.hidden_size338 else:339 res_skip_channels = config.hidden_size340 341 res_skip_layer = torch.nn.Conv1d(config.hidden_size, res_skip_channels, 1)342 res_skip_layer = weight_norm(res_skip_layer, name="weight")343 self.res_skip_layers.append(res_skip_layer)344 345 def forward(self, inputs, padding_mask, global_conditioning=None):346 outputs = torch.zeros_like(inputs)347 num_channels_tensor = torch.IntTensor([self.hidden_size])348 349 if global_conditioning is not None:350 global_conditioning = self.cond_layer(global_conditioning)351 352 for i in range(self.num_layers):353 hidden_states = self.in_layers[i](inputs)354 355 if global_conditioning is not None:356 cond_offset = i * 2 * self.hidden_size357 global_states = global_conditioning[:, cond_offset : cond_offset + 2 * self.hidden_size, :]358 else:359 global_states = torch.zeros_like(hidden_states)360 361 acts = fused_add_tanh_sigmoid_multiply(hidden_states, global_states, num_channels_tensor[0])362 acts = self.dropout(acts)363 364 res_skip_acts = self.res_skip_layers[i](acts)365 if i < self.num_layers - 1:366 res_acts = res_skip_acts[:, : self.hidden_size, :]367 inputs = (inputs + res_acts) * padding_mask368 outputs = outputs + res_skip_acts[:, self.hidden_size :, :]369 else:370 outputs = outputs + res_skip_acts371 372 return outputs * padding_mask373 374 def remove_weight_norm(self):375 if self.speaker_embedding_size != 0:376 torch.nn.utils.remove_weight_norm(self.cond_layer)377 for layer in self.in_layers:378 torch.nn.utils.remove_weight_norm(layer)379 for layer in self.res_skip_layers:380 torch.nn.utils.remove_weight_norm(layer)381 382 383class VitsPosteriorEncoder(nn.Module):384 def __init__(self, config: VitsConfig):385 super().__init__()386 self.out_channels = config.flow_size387 388 self.conv_pre = nn.Conv1d(config.spectrogram_bins, config.hidden_size, 1)389 self.wavenet = VitsWaveNet(config, num_layers=config.posterior_encoder_num_wavenet_layers)390 self.conv_proj = nn.Conv1d(config.hidden_size, self.out_channels * 2, 1)391 392 def forward(self, inputs, padding_mask, global_conditioning=None):393 inputs = self.conv_pre(inputs) * padding_mask394 inputs = self.wavenet(inputs, padding_mask, global_conditioning)395 stats = self.conv_proj(inputs) * padding_mask396 mean, log_stddev = torch.split(stats, self.out_channels, dim=1)397 sampled = (mean + torch.randn_like(mean) * torch.exp(log_stddev)) * padding_mask398 return sampled, mean, log_stddev399 400 401# Copied from transformers.models.speecht5.modeling_speecht5.HifiGanResidualBlock402class HifiGanResidualBlock(nn.Module):403 def __init__(self, channels, kernel_size=3, dilation=(1, 3, 5), leaky_relu_slope=0.1):404 super().__init__()405 self.leaky_relu_slope = leaky_relu_slope406 407 self.convs1 = nn.ModuleList(408 [409 nn.Conv1d(410 channels,411 channels,412 kernel_size,413 stride=1,414 dilation=dilation[i],415 padding=self.get_padding(kernel_size, dilation[i]),416 )417 for i in range(len(dilation))418 ]419 )420 self.convs2 = nn.ModuleList(421 [422 nn.Conv1d(423 channels,424 channels,425 kernel_size,426 stride=1,427 dilation=1,428 padding=self.get_padding(kernel_size, 1),429 )430 for _ in range(len(dilation))431 ]432 )433 434 def get_padding(self, kernel_size, dilation=1):435 return (kernel_size * dilation - dilation) // 2436 437 def apply_weight_norm(self):438 weight_norm = nn.utils.weight_norm439 if hasattr(nn.utils.parametrizations, "weight_norm"):440 weight_norm = nn.utils.parametrizations.weight_norm441 442 for layer in self.convs1:443 weight_norm(layer)444 for layer in self.convs2:445 weight_norm(layer)446 447 def remove_weight_norm(self):448 for layer in self.convs1:449 nn.utils.remove_weight_norm(layer)450 for layer in self.convs2:451 nn.utils.remove_weight_norm(layer)452 453 def forward(self, hidden_states):454 for conv1, conv2 in zip(self.convs1, self.convs2):455 residual = hidden_states456 hidden_states = nn.functional.leaky_relu(hidden_states, self.leaky_relu_slope)457 hidden_states = conv1(hidden_states)458 hidden_states = nn.functional.leaky_relu(hidden_states, self.leaky_relu_slope)459 hidden_states = conv2(hidden_states)460 hidden_states = hidden_states + residual461 return hidden_states462 463 464class VitsHifiGan(nn.Module):465 def __init__(self, config: VitsConfig):466 super().__init__()467 self.config = config468 self.num_kernels = len(config.resblock_kernel_sizes)469 self.num_upsamples = len(config.upsample_rates)470 self.conv_pre = nn.Conv1d(471 config.flow_size,472 config.upsample_initial_channel,473 kernel_size=7,474 stride=1,475 padding=3,476 )477 478 self.upsampler = nn.ModuleList()479 for i, (upsample_rate, kernel_size) in enumerate(zip(config.upsample_rates, config.upsample_kernel_sizes)):480 self.upsampler.append(481 nn.ConvTranspose1d(482 config.upsample_initial_channel // (2**i),483 config.upsample_initial_channel // (2 ** (i + 1)),484 kernel_size=kernel_size,485 stride=upsample_rate,486 padding=(kernel_size - upsample_rate) // 2,487 )488 )489 490 self.resblocks = nn.ModuleList()491 for i in range(len(self.upsampler)):492 channels = config.upsample_initial_channel // (2 ** (i + 1))493 for kernel_size, dilation in zip(config.resblock_kernel_sizes, config.resblock_dilation_sizes):494 self.resblocks.append(HifiGanResidualBlock(channels, kernel_size, dilation, config.leaky_relu_slope))495 496 self.conv_post = nn.Conv1d(channels, 1, kernel_size=7, stride=1, padding=3, bias=False)497 498 if config.speaker_embedding_size != 0:499 self.cond = nn.Conv1d(config.speaker_embedding_size, config.upsample_initial_channel, 1)500 501 def apply_weight_norm(self):502 weight_norm = nn.utils.weight_norm503 if hasattr(nn.utils.parametrizations, "weight_norm"):504 weight_norm = nn.utils.parametrizations.weight_norm505 506 for layer in self.upsampler:507 weight_norm(layer)508 for layer in self.resblocks:509 layer.apply_weight_norm()510 511 def remove_weight_norm(self):512 for layer in self.upsampler:513 nn.utils.remove_weight_norm(layer)514 for layer in self.resblocks:515 layer.remove_weight_norm()516 517 def forward(518 self, spectrogram: torch.FloatTensor, global_conditioning: Optional[torch.FloatTensor] = None519 ) -> torch.FloatTensor:520 r"""521 Converts a spectrogram into a speech waveform.522 523 Args:524 spectrogram (`torch.FloatTensor` of shape `(batch_size, config.spectrogram_bins, sequence_length)`):525 Tensor containing the spectrograms.526 global_conditioning (`torch.FloatTensor` of shape `(batch_size, config.speaker_embedding_size, 1)`, *optional*):527 Tensor containing speaker embeddings, for multispeaker models.528 529 Returns:530 `torch.FloatTensor`: Tensor of shape shape `(batch_size, 1, num_frames)` containing the speech waveform.531 """532 hidden_states = self.conv_pre(spectrogram)533 534 if global_conditioning is not None:535 hidden_states = hidden_states + self.cond(global_conditioning)536 537 for i in range(self.num_upsamples):538 hidden_states = nn.functional.leaky_relu(hidden_states, self.config.leaky_relu_slope)539 hidden_states = self.upsampler[i](hidden_states)540 541 res_state = self.resblocks[i * self.num_kernels](hidden_states)542 for j in range(1, self.num_kernels):543 res_state += self.resblocks[i * self.num_kernels + j](hidden_states)544 hidden_states = res_state / self.num_kernels545 546 hidden_states = nn.functional.leaky_relu(hidden_states)547 hidden_states = self.conv_post(hidden_states)548 waveform = torch.tanh(hidden_states)549 return waveform550 551 552class VitsResidualCouplingLayer(nn.Module):553 def __init__(self, config: VitsConfig):554 super().__init__()555 self.half_channels = config.flow_size // 2556 557 self.conv_pre = nn.Conv1d(self.half_channels, config.hidden_size, 1)558 self.wavenet = VitsWaveNet(config, num_layers=config.prior_encoder_num_wavenet_layers)559 self.conv_post = nn.Conv1d(config.hidden_size, self.half_channels, 1)560 561 def forward(self, inputs, padding_mask, global_conditioning=None, reverse=False):562 first_half, second_half = torch.split(inputs, [self.half_channels] * 2, dim=1)563 hidden_states = self.conv_pre(first_half) * padding_mask564 hidden_states = self.wavenet(hidden_states, padding_mask, global_conditioning)565 mean = self.conv_post(hidden_states) * padding_mask566 log_stddev = torch.zeros_like(mean)567 568 if not reverse:569 second_half = mean + second_half * torch.exp(log_stddev) * padding_mask570 outputs = torch.cat([first_half, second_half], dim=1)571 log_determinant = torch.sum(log_stddev, [1, 2])572 return outputs, log_determinant573 else:574 second_half = (second_half - mean) * torch.exp(-log_stddev) * padding_mask575 outputs = torch.cat([first_half, second_half], dim=1)576 return outputs, None577 578 579class VitsResidualCouplingBlock(nn.Module):580 def __init__(self, config: VitsConfig):581 super().__init__()582 self.flows = nn.ModuleList()583 for _ in range(config.prior_encoder_num_flows):584 self.flows.append(VitsResidualCouplingLayer(config))585 586 def forward(self, inputs, padding_mask, global_conditioning=None, reverse=False):587 if not reverse:588 for flow in self.flows:589 inputs, _ = flow(inputs, padding_mask, global_conditioning)590 inputs = torch.flip(inputs, [1])591 else:592 for flow in reversed(self.flows):593 inputs = torch.flip(inputs, [1])594 inputs, _ = flow(inputs, padding_mask, global_conditioning, reverse=True)595 return inputs596 597 598class VitsDilatedDepthSeparableConv(nn.Module):599 def __init__(self, config: VitsConfig, dropout_rate=0.0):600 super().__init__()601 kernel_size = config.duration_predictor_kernel_size602 channels = config.hidden_size603 self.num_layers = config.depth_separable_num_layers604 605 self.dropout = nn.Dropout(dropout_rate)606 self.convs_dilated = nn.ModuleList()607 self.convs_pointwise = nn.ModuleList()608 self.norms_1 = nn.ModuleList()609 self.norms_2 = nn.ModuleList()610 for i in range(self.num_layers):611 dilation = kernel_size**i612 padding = (kernel_size * dilation - dilation) // 2613 self.convs_dilated.append(614 nn.Conv1d(615 in_channels=channels,616 out_channels=channels,617 kernel_size=kernel_size,618 groups=channels,619 dilation=dilation,620 padding=padding,621 )622 )623 self.convs_pointwise.append(nn.Conv1d(channels, channels, 1))624 self.norms_1.append(nn.LayerNorm(channels))625 self.norms_2.append(nn.LayerNorm(channels))626 627 def forward(self, inputs, padding_mask, global_conditioning=None):628 if global_conditioning is not None:629 inputs = inputs + global_conditioning630 631 for i in range(self.num_layers):632 hidden_states = self.convs_dilated[i](inputs * padding_mask)633 hidden_states = self.norms_1[i](hidden_states.transpose(1, -1)).transpose(1, -1)634 hidden_states = nn.functional.gelu(hidden_states)635 hidden_states = self.convs_pointwise[i](hidden_states)636 hidden_states = self.norms_2[i](hidden_states.transpose(1, -1)).transpose(1, -1)637 hidden_states = nn.functional.gelu(hidden_states)638 hidden_states = self.dropout(hidden_states)639 inputs = inputs + hidden_states640 641 return inputs * padding_mask642 643 644class VitsConvFlow(nn.Module):645 def __init__(self, config: VitsConfig):646 super().__init__()647 self.filter_channels = config.hidden_size648 self.half_channels = config.depth_separable_channels // 2649 self.num_bins = config.duration_predictor_flow_bins650 self.tail_bound = config.duration_predictor_tail_bound651 652 self.conv_pre = nn.Conv1d(self.half_channels, self.filter_channels, 1)653 self.conv_dds = VitsDilatedDepthSeparableConv(config)654 self.conv_proj = nn.Conv1d(self.filter_channels, self.half_channels * (self.num_bins * 3 - 1), 1)655 656 def forward(self, inputs, padding_mask, global_conditioning=None, reverse=False):657 first_half, second_half = torch.split(inputs, [self.half_channels] * 2, dim=1)658 659 hidden_states = self.conv_pre(first_half)660 hidden_states = self.conv_dds(hidden_states, padding_mask, global_conditioning)661 hidden_states = self.conv_proj(hidden_states) * padding_mask662 663 batch_size, channels, length = first_half.shape664 hidden_states = hidden_states.reshape(batch_size, channels, -1, length).permute(0, 1, 3, 2)665 666 unnormalized_widths = hidden_states[..., : self.num_bins] / math.sqrt(self.filter_channels)667 unnormalized_heights = hidden_states[..., self.num_bins : 2 * self.num_bins] / math.sqrt(self.filter_channels)668 unnormalized_derivatives = hidden_states[..., 2 * self.num_bins :]669 670 second_half, log_abs_det = _unconstrained_rational_quadratic_spline(671 second_half,672 unnormalized_widths,673 unnormalized_heights,674 unnormalized_derivatives,675 reverse=reverse,676 tail_bound=self.tail_bound,677 )678 679 outputs = torch.cat([first_half, second_half], dim=1) * padding_mask680 if not reverse:681 log_determinant = torch.sum(log_abs_det * padding_mask, [1, 2])682 return outputs, log_determinant683 else:684 return outputs, None685 686 687class VitsElementwiseAffine(nn.Module):688 def __init__(self, config: VitsConfig):689 super().__init__()690 self.channels = config.depth_separable_channels691 self.translate = nn.Parameter(torch.zeros(self.channels, 1))692 self.log_scale = nn.Parameter(torch.zeros(self.channels, 1))693 694 def forward(self, inputs, padding_mask, global_conditioning=None, reverse=False):695 if not reverse:696 outputs = self.translate + torch.exp(self.log_scale) * inputs697 outputs = outputs * padding_mask698 log_determinant = torch.sum(self.log_scale * padding_mask, [1, 2])699 return outputs, log_determinant700 else:701 outputs = (inputs - self.translate) * torch.exp(-self.log_scale) * padding_mask702 return outputs, None703 704 705class VitsStochasticDurationPredictor(nn.Module):706 def __init__(self, config):707 super().__init__()708 embed_dim = config.speaker_embedding_size709 filter_channels = config.hidden_size710 711 self.conv_pre = nn.Conv1d(filter_channels, filter_channels, 1)712 self.conv_proj = nn.Conv1d(filter_channels, filter_channels, 1)713 self.conv_dds = VitsDilatedDepthSeparableConv(714 config,715 dropout_rate=config.duration_predictor_dropout,716 )717 718 if embed_dim != 0:719 self.cond = nn.Conv1d(embed_dim, filter_channels, 1)720 721 self.flows = nn.ModuleList()722 self.flows.append(VitsElementwiseAffine(config))723 for _ in range(config.duration_predictor_num_flows):724 self.flows.append(VitsConvFlow(config))725 726 self.post_conv_pre = nn.Conv1d(1, filter_channels, 1)727 self.post_conv_proj = nn.Conv1d(filter_channels, filter_channels, 1)728 self.post_conv_dds = VitsDilatedDepthSeparableConv(729 config,730 dropout_rate=config.duration_predictor_dropout,731 )732 733 self.post_flows = nn.ModuleList()734 self.post_flows.append(VitsElementwiseAffine(config))735 for _ in range(config.duration_predictor_num_flows):736 self.post_flows.append(VitsConvFlow(config))737 738 def forward(self, inputs, padding_mask, global_conditioning=None, durations=None, reverse=False, noise_scale=1.0):739 inputs = torch.detach(inputs)740 inputs = self.conv_pre(inputs)741 742 if global_conditioning is not None:743 global_conditioning = torch.detach(global_conditioning)744 inputs = inputs + self.cond(global_conditioning)745 746 inputs = self.conv_dds(inputs, padding_mask)747 inputs = self.conv_proj(inputs) * padding_mask748 749 if not reverse:750 hidden_states = self.post_conv_pre(durations)751 hidden_states = self.post_conv_dds(hidden_states, padding_mask)752 hidden_states = self.post_conv_proj(hidden_states) * padding_mask753 754 random_posterior = (755 torch.randn(durations.size(0), 2, durations.size(2)).to(device=inputs.device, dtype=inputs.dtype)756 * padding_mask757 )758 log_determinant_posterior_sum = 0759 latents_posterior = random_posterior760 for flow in self.post_flows:761 latents_posterior, log_determinant = flow(762 latents_posterior, padding_mask, global_conditioning=inputs + hidden_states763 )764 latents_posterior = torch.flip(latents_posterior, [1])765 log_determinant_posterior_sum += log_determinant766 767 first_half, second_half = torch.split(latents_posterior, [1, 1], dim=1)768 769 log_determinant_posterior_sum += torch.sum(770 (nn.functional.logsigmoid(first_half) + nn.functional.logsigmoid(-first_half)) * padding_mask, [1, 2]771 )772 logq = (773 torch.sum(-0.5 * (math.log(2 * math.pi) + (random_posterior**2)) * padding_mask, [1, 2])774 - log_determinant_posterior_sum775 )776 777 first_half = (durations - torch.sigmoid(first_half)) * padding_mask778 first_half = torch.log(torch.clamp_min(first_half, 1e-5)) * padding_mask779 log_determinant_sum = torch.sum(-first_half, [1, 2])780 781 latents = torch.cat([first_half, second_half], dim=1)782 for flow in self.flows:783 latents, log_determinant = flow(latents, padding_mask, global_conditioning=inputs)784 latents = torch.flip(latents, [1])785 log_determinant_sum += log_determinant786 787 nll = torch.sum(0.5 * (math.log(2 * math.pi) + (latents**2)) * padding_mask, [1, 2]) - log_determinant_sum788 return nll + logq789 else:790 flows = list(reversed(self.flows))791 flows = flows[:-2] + [flows[-1]] # remove a useless vflow792 793 latents = (794 torch.randn(inputs.size(0), 2, inputs.size(2)).to(device=inputs.device, dtype=inputs.dtype)795 * noise_scale796 )797 for flow in flows:798 latents = torch.flip(latents, [1])799 latents, _ = flow(latents, padding_mask, global_conditioning=inputs, reverse=True)800 801 log_duration, _ = torch.split(latents, [1, 1], dim=1)802 return log_duration803 804 805class VitsDurationPredictor(nn.Module):806 def __init__(self, config):807 super().__init__()808 kernel_size = config.duration_predictor_kernel_size809 filter_channels = config.duration_predictor_filter_channels810 811 self.dropout = nn.Dropout(config.duration_predictor_dropout)812 self.conv_1 = nn.Conv1d(config.hidden_size, filter_channels, kernel_size, padding=kernel_size // 2)813 self.norm_1 = nn.LayerNorm(filter_channels, eps=config.layer_norm_eps)814 self.conv_2 = nn.Conv1d(filter_channels, filter_channels, kernel_size, padding=kernel_size // 2)815 self.norm_2 = nn.LayerNorm(filter_channels, eps=config.layer_norm_eps)816 self.proj = nn.Conv1d(filter_channels, 1, 1)817 818 if config.speaker_embedding_size != 0:819 self.cond = nn.Conv1d(config.speaker_embedding_size, config.hidden_size, 1)820 821 def forward(self, inputs, padding_mask, global_conditioning=None):822 inputs = torch.detach(inputs)823 824 if global_conditioning is not None:825 global_conditioning = torch.detach(global_conditioning)826 inputs = inputs + self.cond(global_conditioning)827 828 inputs = self.conv_1(inputs * padding_mask)829 inputs = torch.relu(inputs)830 inputs = self.norm_1(inputs.transpose(1, -1)).transpose(1, -1)831 inputs = self.dropout(inputs)832 833 inputs = self.conv_2(inputs * padding_mask)834 inputs = torch.relu(inputs)835 inputs = self.norm_2(inputs.transpose(1, -1)).transpose(1, -1)836 inputs = self.dropout(inputs)837 838 inputs = self.proj(inputs * padding_mask)839 return inputs * padding_mask840 841 842class VitsAttention(nn.Module):843 """Multi-headed attention with relative positional representation."""844 845 def __init__(self, config: VitsConfig):846 super().__init__()847 self.embed_dim = config.hidden_size848 self.num_heads = config.num_attention_heads849 self.dropout = config.attention_dropout850 self.window_size = config.window_size851 852 self.head_dim = self.embed_dim // self.num_heads853 self.scaling = self.head_dim**-0.5854 855 if (self.head_dim * self.num_heads) != self.embed_dim:856 raise ValueError(857 f"hidden_size must be divisible by num_attention_heads (got `hidden_size`: {self.embed_dim}"858 f" and `num_attention_heads`: {self.num_heads})."859 )860 861 self.k_proj = nn.Linear(self.embed_dim, self.embed_dim, bias=config.use_bias)862 self.v_proj = nn.Linear(self.embed_dim, self.embed_dim, bias=config.use_bias)863 self.q_proj = nn.Linear(self.embed_dim, self.embed_dim, bias=config.use_bias)864 self.out_proj = nn.Linear(self.embed_dim, self.embed_dim, bias=config.use_bias)865 866 if self.window_size:867 self.emb_rel_k = nn.Parameter(torch.randn(1, self.window_size * 2 + 1, self.head_dim) * self.scaling)868 self.emb_rel_v = nn.Parameter(torch.randn(1, self.window_size * 2 + 1, self.head_dim) * self.scaling)869 870 def _shape(self, tensor: torch.Tensor, seq_len: int, bsz: int):871 return tensor.view(bsz, seq_len, self.num_heads, self.head_dim).transpose(1, 2).contiguous()872 873 def forward(874 self,875 hidden_states: torch.Tensor,876 key_value_states: Optional[torch.Tensor] = None,877 attention_mask: Optional[torch.Tensor] = None,878 layer_head_mask: Optional[torch.Tensor] = None,879 output_attentions: bool = False,880 ) -> tuple[torch.Tensor, Optional[torch.Tensor]]:881 """Input shape: Batch x Time x Channel"""882 883 # if key_value_states are provided this layer is used as a cross-attention layer884 # for the decoder885 886 bsz, tgt_len, _ = hidden_states.size()887 888 # get query proj889 query_states = self.q_proj(hidden_states) * self.scaling890 891 # self_attention892 key_states = self._shape(self.k_proj(hidden_states), -1, bsz)893 value_states = self._shape(self.v_proj(hidden_states), -1, bsz)894 895 proj_shape = (bsz * self.num_heads, -1, self.head_dim)896 query_states = self._shape(query_states, tgt_len, bsz).view(*proj_shape)897 key_states = key_states.view(*proj_shape)898 value_states = value_states.view(*proj_shape)899 900 src_len = key_states.size(1)901 attn_weights = torch.bmm(query_states, key_states.transpose(1, 2))902 903 if attn_weights.size() != (bsz * self.num_heads, tgt_len, src_len):904 raise ValueError(905 f"Attention weights should be of size {(bsz * self.num_heads, tgt_len, src_len)}, but is"906 f" {attn_weights.size()}"907 )908 909 if self.window_size is not None:910 key_relative_embeddings = self._get_relative_embeddings(self.emb_rel_k, src_len)911 relative_logits = torch.matmul(query_states, key_relative_embeddings.transpose(-2, -1))912 rel_pos_bias = self._relative_position_to_absolute_position(relative_logits)913 attn_weights += rel_pos_bias914 915 if attention_mask is not None:916 if attention_mask.size() != (bsz, 1, tgt_len, src_len):917 raise ValueError(918 f"Attention mask should be of size {(bsz, 1, tgt_len, src_len)}, but is {attention_mask.size()}"919 )920 attn_weights = attn_weights.view(bsz, self.num_heads, tgt_len, src_len) + attention_mask921 attn_weights = attn_weights.view(bsz * self.num_heads, tgt_len, src_len)922 923 attn_weights = nn.functional.softmax(attn_weights, dim=-1)924 925 if layer_head_mask is not None:926 if layer_head_mask.size() != (self.num_heads,):927 raise ValueError(928 f"Head mask for a single layer should be of size {(self.num_heads,)}, but is"929 f" {layer_head_mask.size()}"930 )931 attn_weights = layer_head_mask.view(1, -1, 1, 1) * attn_weights.view(bsz, self.num_heads, tgt_len, src_len)932 attn_weights = attn_weights.view(bsz * self.num_heads, tgt_len, src_len)933 934 if output_attentions:935 # this operation is a bit awkward, but it's required to936 # make sure that attn_weights keeps its gradient.937 # In order to do so, attn_weights have to be reshaped938 # twice and have to be reused in the following939 attn_weights_reshaped = attn_weights.view(bsz, self.num_heads, tgt_len, src_len)940 attn_weights = attn_weights_reshaped.view(bsz * self.num_heads, tgt_len, src_len)941 else:942 attn_weights_reshaped = None943 944 attn_probs = nn.functional.dropout(attn_weights, p=self.dropout, training=self.training)945 946 attn_output = torch.bmm(attn_probs, value_states)947 948 if attn_output.size() != (bsz * self.num_heads, tgt_len, self.head_dim):949 raise ValueError(950 f"`attn_output` should be of size {(bsz, self.num_heads, tgt_len, self.head_dim)}, but is"951 f" {attn_output.size()}"952 )953 954 if self.window_size is not None:955 value_relative_embeddings = self._get_relative_embeddings(self.emb_rel_v, src_len)956 relative_weights = self._absolute_position_to_relative_position(attn_probs)957 rel_pos_bias = torch.matmul(relative_weights, value_relative_embeddings)958 attn_output += rel_pos_bias959 960 attn_output = attn_output.view(bsz, self.num_heads, tgt_len, self.head_dim)961 attn_output = attn_output.transpose(1, 2)962 963 # Use the `embed_dim` from the config (stored in the class) rather than `hidden_state` because `attn_output` can be964 # partitioned across GPUs when using tensor-parallelism.965 attn_output = attn_output.reshape(bsz, tgt_len, self.embed_dim)966 967 attn_output = self.out_proj(attn_output)968 969 return attn_output, attn_weights_reshaped970 971 def _get_relative_embeddings(self, relative_embeddings, length):972 pad_length = max(length - (self.window_size + 1), 0)973 if pad_length > 0:974 relative_embeddings = nn.functional.pad(relative_embeddings, [0, 0, pad_length, pad_length, 0, 0])975 976 slice_start_position = max((self.window_size + 1) - length, 0)977 slice_end_position = slice_start_position + 2 * length - 1978 return relative_embeddings[:, slice_start_position:slice_end_position]979 980 def _relative_position_to_absolute_position(self, x):981 batch_heads, length, _ = x.size()982 983 # Concat columns of pad to shift from relative to absolute indexing.984 x = nn.functional.pad(x, [0, 1, 0, 0, 0, 0])985 986 # Concat extra elements so to add up to shape (len+1, 2*len-1).987 x_flat = x.view([batch_heads, length * 2 * length])988 x_flat = nn.functional.pad(x_flat, [0, length - 1, 0, 0])989 990 # Reshape and slice out the padded elements.991 x_final = x_flat.view([batch_heads, length + 1, 2 * length - 1])992 x_final = x_final[:, :length, length - 1 :]993 return x_final994 995 def _absolute_position_to_relative_position(self, x):996 batch_heads, length, _ = x.size()997 998 # Pad along column999 x = nn.functional.pad(x, [0, length - 1, 0, 0, 0, 0])1000 x_flat = x.view([batch_heads, length * (2 * length - 1)])1001 1002 # Add 0's in the beginning that will skew the elements after reshape1003 x_flat = nn.functional.pad(x_flat, [length, 0, 0, 0])1004 x_final = x_flat.view([batch_heads, length, 2 * length])[:, :, 1:]1005 return x_final1006 1007 1008class VitsFeedForward(nn.Module):1009 def __init__(self, config):1010 super().__init__()1011 self.conv_1 = nn.Conv1d(config.hidden_size, config.ffn_dim, config.ffn_kernel_size)1012 self.conv_2 = nn.Conv1d(config.ffn_dim, config.hidden_size, config.ffn_kernel_size)1013 self.dropout = nn.Dropout(config.activation_dropout)1014 1015 if isinstance(config.hidden_act, str):1016 self.act_fn = ACT2FN[config.hidden_act]1017 else:1018 self.act_fn = config.hidden_act1019 1020 if config.ffn_kernel_size > 1:1021 pad_left = (config.ffn_kernel_size - 1) // 21022 pad_right = config.ffn_kernel_size // 21023 self.padding = [pad_left, pad_right, 0, 0, 0, 0]1024 else:1025 self.padding = None1026 1027 def forward(self, hidden_states, padding_mask):1028 hidden_states = hidden_states.permute(0, 2, 1)1029 padding_mask = padding_mask.permute(0, 2, 1)1030 1031 hidden_states = hidden_states * padding_mask1032 if self.padding is not None:1033 hidden_states = nn.functional.pad(hidden_states, self.padding)1034 1035 hidden_states = self.conv_1(hidden_states)1036 hidden_states = self.act_fn(hidden_states)1037 hidden_states = self.dropout(hidden_states)1038 1039 hidden_states = hidden_states * padding_mask1040 if self.padding is not None:1041 hidden_states = nn.functional.pad(hidden_states, self.padding)1042 1043 hidden_states = self.conv_2(hidden_states)1044 hidden_states = hidden_states * padding_mask1045 1046 hidden_states = hidden_states.permute(0, 2, 1)1047 return hidden_states1048 1049 1050class VitsEncoderLayer(GradientCheckpointingLayer):1051 def __init__(self, config: VitsConfig):1052 super().__init__()1053 self.attention = VitsAttention(config)1054 self.dropout = nn.Dropout(config.hidden_dropout)1055 self.layer_norm = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps)1056 self.feed_forward = VitsFeedForward(config)1057 self.final_layer_norm = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps)1058 1059 def forward(1060 self,1061 hidden_states: torch.Tensor,1062 padding_mask: torch.FloatTensor,1063 attention_mask: Optional[torch.Tensor] = None,1064 output_attentions: bool = False,1065 ):1066 residual = hidden_states1067 hidden_states, attn_weights = self.attention(1068 hidden_states=hidden_states,1069 attention_mask=attention_mask,1070 output_attentions=output_attentions,1071 )1072 1073 hidden_states = self.dropout(hidden_states)1074 hidden_states = self.layer_norm(residual + hidden_states)1075 1076 residual = hidden_states1077 hidden_states = self.feed_forward(hidden_states, padding_mask)1078 hidden_states = self.dropout(hidden_states)1079 hidden_states = self.final_layer_norm(residual + hidden_states)1080 1081 outputs = (hidden_states,)1082 1083 if output_attentions:1084 outputs += (attn_weights,)1085 1086 return outputs1087 1088 1089class VitsEncoder(nn.Module):1090 def __init__(self, config: VitsConfig):1091 super().__init__()1092 self.config = config1093 self.layers = nn.ModuleList([VitsEncoderLayer(config) for _ in range(config.num_hidden_layers)])1094 self.gradient_checkpointing = False1095 self.layerdrop = config.layerdrop1096 1097 def forward(1098 self,1099 hidden_states: torch.FloatTensor,1100 padding_mask: torch.FloatTensor,1101 attention_mask: Optional[torch.Tensor] = None,1102 output_attentions: Optional[bool] = None,1103 output_hidden_states: Optional[bool] = None,1104 return_dict: Optional[bool] = None,1105 ) -> Union[tuple, BaseModelOutput]:1106 all_hidden_states = () if output_hidden_states else None1107 all_self_attentions = () if output_attentions else None1108 1109 # expand attention_mask1110 if attention_mask is not None:1111 # [bsz, seq_len] -> [bsz, 1, tgt_seq_len, src_seq_len]1112 attention_mask = _prepare_4d_attention_mask(attention_mask, hidden_states.dtype)1113 1114 hidden_states = hidden_states * padding_mask1115 1116 synced_gpus = is_deepspeed_zero3_enabled() or is_fsdp_managed_module(self)1117 1118 for encoder_layer in self.layers:1119 if output_hidden_states:1120 all_hidden_states = all_hidden_states + (hidden_states,)1121 1122 # add LayerDrop (see https://huggingface.co/papers/1909.11556 for description)1123 dropout_probability = np.random.uniform(0, 1)1124 1125 skip_the_layer = self.training and (dropout_probability < self.layerdrop)1126 if not skip_the_layer or synced_gpus:1127 # under fsdp or deepspeed zero3 all gpus must run in sync1128 layer_outputs = encoder_layer(1129 hidden_states,1130 attention_mask=attention_mask,1131 padding_mask=padding_mask,1132 output_attentions=output_attentions,1133 )1134 hidden_states = layer_outputs[0]1135 1136 if skip_the_layer:1137 layer_outputs = (None, None)1138 1139 if output_attentions:1140 all_self_attentions = all_self_attentions + (layer_outputs[1],)1141 1142 hidden_states = hidden_states * padding_mask1143 1144 if output_hidden_states:1145 all_hidden_states = all_hidden_states + (hidden_states,)1146 1147 if not return_dict:1148 return tuple(v for v in [hidden_states, all_hidden_states, all_self_attentions] if v is not None)1149 1150 return BaseModelOutput(1151 last_hidden_state=hidden_states,1152 hidden_states=all_hidden_states,1153 attentions=all_self_attentions,1154 )1155 1156 1157class VitsTextEncoder(nn.Module):1158 """1159 Transformer encoder that uses relative positional representation instead of absolute positional encoding.1160 """1161 1162 def __init__(self, config: VitsConfig):1163 super().__init__()1164 self.config = config1165 self.embed_tokens = nn.Embedding(config.vocab_size, config.hidden_size, config.pad_token_id)1166 self.encoder = VitsEncoder(config)1167 self.project = nn.Conv1d(config.hidden_size, config.flow_size * 2, kernel_size=1)1168 1169 def forward(1170 self,1171 input_ids: torch.Tensor,1172 padding_mask: torch.FloatTensor,1173 attention_mask: Optional[torch.Tensor] = None,1174 output_attentions: Optional[bool] = None,1175 output_hidden_states: Optional[bool] = None,1176 return_dict: Optional[bool] = True,1177 ) -> Union[tuple[torch.Tensor], VitsTextEncoderOutput]:1178 hidden_states = self.embed_tokens(input_ids) * math.sqrt(self.config.hidden_size)1179 1180 encoder_outputs = self.encoder(1181 hidden_states=hidden_states,1182 padding_mask=padding_mask,1183 attention_mask=attention_mask,1184 output_attentions=output_attentions,1185 output_hidden_states=output_hidden_states,1186 return_dict=return_dict,1187 )1188 1189 last_hidden_state = encoder_outputs[0] if not return_dict else encoder_outputs.last_hidden_state1190 1191 stats = self.project(last_hidden_state.transpose(1, 2)).transpose(1, 2) * padding_mask1192 prior_means, prior_log_variances = torch.split(stats, self.config.flow_size, dim=2)1193 1194 if not return_dict:1195 outputs = (last_hidden_state, prior_means, prior_log_variances) + encoder_outputs[1:]1196 return outputs1197 1198 return VitsTextEncoderOutput(1199 last_hidden_state=last_hidden_state,1200 prior_means=prior_means,