Skip to content

FastSpeech 2 Loss

FastSpeech2LossGen

Bases: Module

Source code in training/loss/fast_speech_2_loss_gen.py 
 13
 14
 15
 16
 17
 18
 19
 20
 21
 22
 23
 24
 25
 26
 27
 28
 29
 30
 31
 32
 33
 34
 35
 36
 37
 38
 39
 40
 41
 42
 43
 44
 45
 46
 47
 48
 49
 50
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
class FastSpeech2LossGen(Module):
    def __init__(
        self,
        bin_warmup: bool = True,
        binarization_loss_enable_steps: int = 1260,
        binarization_loss_warmup_steps: int = 700,
    ):
        r"""Initializes the FastSpeech2LossGen module.

        Args:
            bin_warmup (bool, optional): Whether to use binarization warmup. Defaults to True. NOTE: Switch this off if you preload the model with a checkpoint that has already passed the warmup phase.
            binarization_loss_enable_steps (int, optional): Number of steps to enable the binarization loss. Defaults to 1260.
            binarization_loss_warmup_steps (int, optional): Number of warmup steps for the binarization loss. Defaults to 700.
        """
        super().__init__()

        self.mse_loss = nn.MSELoss()
        self.mae_loss = nn.L1Loss()
        self.ssim_loss = SSIMLoss()
        self.sum_loss = ForwardSumLoss()
        self.bin_loss = BinLoss()

        self.bin_warmup = bin_warmup
        self.binarization_loss_enable_steps = binarization_loss_enable_steps
        self.binarization_loss_warmup_steps = binarization_loss_warmup_steps

    def forward(
        self,
        src_masks: torch.Tensor,
        mel_masks: torch.Tensor,
        mel_targets: torch.Tensor,
        mel_predictions: torch.Tensor,
        log_duration_predictions: torch.Tensor,
        u_prosody_ref: torch.Tensor,
        u_prosody_pred: torch.Tensor,
        p_prosody_ref: torch.Tensor,
        p_prosody_pred: torch.Tensor,
        durations: torch.Tensor,
        pitch_predictions: torch.Tensor,
        p_targets: torch.Tensor,
        attn_logprob: torch.Tensor,
        attn_soft: torch.Tensor,
        attn_hard: torch.Tensor,
        step: int,
        src_lens: torch.Tensor,
        mel_lens: torch.Tensor,
        energy_pred: torch.Tensor,
        energy_target: torch.Tensor,
    ) -> Tuple[
        torch.Tensor,
        torch.Tensor,
        torch.Tensor,
        torch.Tensor,
        torch.Tensor,
        torch.Tensor,
        torch.Tensor,
        torch.Tensor,
        torch.Tensor,
        torch.Tensor,
    ]:
        r"""Computes the loss for the FastSpeech2 model.

        Args:
            src_masks (torch.Tensor): Mask for the source sequence.
            mel_masks (torch.Tensor): Mask for the mel-spectrogram.
            mel_targets (torch.Tensor): Target mel-spectrogram.
            mel_predictions (torch.Tensor): Predicted mel-spectrogram.
            log_duration_predictions (torch.Tensor): Predicted log-duration.
            u_prosody_ref (torch.Tensor): Reference unvoiced prosody.
            u_prosody_pred (torch.Tensor): Predicted unvoiced prosody.
            p_prosody_ref (torch.Tensor): Reference voiced prosody.
            p_prosody_pred (torch.Tensor): Predicted voiced prosody.
            durations (torch.Tensor): Ground-truth durations.
            pitch_predictions (torch.Tensor): Predicted pitch.
            p_targets (torch.Tensor): Ground-truth pitch.
            attn_logprob (torch.Tensor): Log-probability of attention.
            attn_soft (torch.Tensor): Soft attention.
            attn_hard (torch.Tensor): Hard attention.
            step (int): Current training step.
            src_lens (torch.Tensor): Lengths of the source sequences.
            mel_lens (torch.Tensor): Lengths of the mel-spectrograms.
            energy_pred (torch.Tensor): Predicted energy.
            energy_target (torch.Tensor): Ground-truth energy.

        Returns:
            Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: The total loss and its components.

        Note:
            Here is the description of the returned loss components:
            `total_loss`: This is the total loss computed as the sum of all the other losses.
            `mel_loss`: This is the mean absolute error (MAE) loss between the predicted and target mel-spectrograms. It measures how well the model predicts the mel-spectrograms.
            `sc_mag_loss`: This is the spectral convergence loss between the predicted and target mel-spectrograms. It measures how well the model predicts the mel-spectrograms in terms of their spectral structure.
            `log_mag_loss`: This is the log STFT magnitude loss between the predicted and target mel-spectrograms. It measures how well the model predicts the mel-spectrograms in terms of their spectral structure.
            `ssim_loss`: This is the Structural Similarity Index (SSIM) loss between the predicted and target mel-spectrograms. It measures the similarity between the two mel-spectrograms in terms of their structure, contrast, and luminance.
            `duration_loss`: This is the mean squared error (MSE) loss between the predicted and target log-durations. It measures how well the model predicts the durations of the phonemes.
            `u_prosody_loss`: This is the MAE loss between the predicted and reference unvoiced prosody. It measures how well the model predicts the prosody (rhythm, stress, and intonation) of the unvoiced parts of the speech.
            `p_prosody_loss`: This is the MAE loss between the predicted and reference voiced prosody. It measures how well the model predicts the prosody of the voiced parts of the speech.
            `pitch_loss`: This is the MSE loss between the predicted and target pitch. It measures how well the model predicts the pitch of the speech.
            `ctc_loss`: This is the Connectionist Temporal Classification (CTC) loss computed from the log-probability of attention and the lengths of the source sequences and mel-spectrograms. It measures how well the model aligns the input and output sequences.
            `bin_loss`: This is the binarization loss computed from the hard and soft attention. It measures how well the model learns to attend to the correct parts of the input sequence.
            `energy_loss`: This is the MSE loss between the predicted and target energy. It measures how well the model predicts the energy of the speech.
        """
        log_duration_targets = torch.log(durations.float() + 1).to(src_masks.device)

        log_duration_targets.requires_grad = False
        mel_targets.requires_grad = False
        p_targets.requires_grad = False
        energy_target.requires_grad = False

        log_duration_predictions = log_duration_predictions.masked_select(~src_masks)
        log_duration_targets = log_duration_targets.masked_select(~src_masks)

        mel_masks_expanded = mel_masks.unsqueeze(1)

        mel_predictions_normalized = (
            sample_wise_min_max(mel_predictions).float().to(mel_predictions.device)
        )
        mel_targets_normalized = (
            sample_wise_min_max(mel_targets).float().to(mel_predictions.device)
        )

        ssim_loss: torch.Tensor = self.ssim_loss(
            mel_predictions_normalized.unsqueeze(1),
            mel_targets_normalized.unsqueeze(1),
        )

        if ssim_loss.item() > 1.0 or ssim_loss.item() < 0.0:
            ssim_loss = torch.tensor([1.0], device=mel_predictions.device)

        masked_mel_predictions = mel_predictions.masked_select(~mel_masks_expanded)

        masked_mel_targets = mel_targets.masked_select(~mel_masks_expanded)

        mel_loss: torch.Tensor = self.mae_loss(
            masked_mel_predictions,
            masked_mel_targets,
        )

        p_prosody_ref = p_prosody_ref.permute((0, 2, 1))
        p_prosody_pred = p_prosody_pred.permute((0, 2, 1))

        p_prosody_ref = p_prosody_ref.masked_fill(src_masks.unsqueeze(1), 0.0)
        p_prosody_pred = p_prosody_pred.masked_fill(src_masks.unsqueeze(1), 0.0)

        p_prosody_ref = p_prosody_ref.detach()

        p_prosody_loss: torch.Tensor = 0.5 * self.mae_loss(
            p_prosody_ref.masked_select(~src_masks.unsqueeze(1)),
            p_prosody_pred.masked_select(~src_masks.unsqueeze(1)),
        )

        u_prosody_ref = u_prosody_ref.detach()
        u_prosody_loss: torch.Tensor = 0.5 * self.mae_loss(
            u_prosody_ref,
            u_prosody_pred,
        )

        duration_loss: torch.Tensor = self.mse_loss(
            log_duration_predictions,
            log_duration_targets,
        )

        pitch_predictions = pitch_predictions.masked_select(~src_masks)
        p_targets = p_targets.masked_select(~src_masks)

        pitch_loss: torch.Tensor = self.mse_loss(pitch_predictions, p_targets)

        ctc_loss: torch.Tensor = self.sum_loss(
            attn_logprob=attn_logprob,
            in_lens=src_lens,
            out_lens=mel_lens,
        )

        if self.bin_warmup:
            if step < self.binarization_loss_enable_steps:
                bin_loss_weight = 0.0
            else:
                bin_loss_weight = (
                    min(
                        (step - self.binarization_loss_enable_steps)
                        / self.binarization_loss_warmup_steps,
                        1.0,
                    )
                    * 1.0
                )

            bin_loss: torch.Tensor = (
                self.bin_loss(hard_attention=attn_hard, soft_attention=attn_soft)
                * bin_loss_weight
            )
        else:
            bin_loss: torch.Tensor = self.bin_loss(
                hard_attention=attn_hard,
                soft_attention=attn_soft,
            )

        energy_loss: torch.Tensor = self.mse_loss(energy_pred, energy_target)

        total_loss = (
            mel_loss
            + duration_loss
            + u_prosody_loss
            + p_prosody_loss
            + ssim_loss
            + pitch_loss
            + ctc_loss
            + bin_loss
            + energy_loss
        )

        return (
            total_loss,
            mel_loss,
            ssim_loss,
            duration_loss,
            u_prosody_loss,
            p_prosody_loss,
            pitch_loss,
            ctc_loss,
            bin_loss,
            energy_loss,
        )

__init__(bin_warmup=True, binarization_loss_enable_steps=1260, binarization_loss_warmup_steps=700)

Initializes the FastSpeech2LossGen module.

Parameters:

Name Type Description Default
bin_warmup bool

Whether to use binarization warmup. Defaults to True. NOTE: Switch this off if you preload the model with a checkpoint that has already passed the warmup phase.

 True
binarization_loss_enable_steps int

Number of steps to enable the binarization loss. Defaults to 1260.

 1260
binarization_loss_warmup_steps int

Number of warmup steps for the binarization loss. Defaults to 700.

 700
Source code in training/loss/fast_speech_2_loss_gen.py 
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
def __init__(
    self,
    bin_warmup: bool = True,
    binarization_loss_enable_steps: int = 1260,
    binarization_loss_warmup_steps: int = 700,
):
    r"""Initializes the FastSpeech2LossGen module.

    Args:
        bin_warmup (bool, optional): Whether to use binarization warmup. Defaults to True. NOTE: Switch this off if you preload the model with a checkpoint that has already passed the warmup phase.
        binarization_loss_enable_steps (int, optional): Number of steps to enable the binarization loss. Defaults to 1260.
        binarization_loss_warmup_steps (int, optional): Number of warmup steps for the binarization loss. Defaults to 700.
    """
    super().__init__()

    self.mse_loss = nn.MSELoss()
    self.mae_loss = nn.L1Loss()
    self.ssim_loss = SSIMLoss()
    self.sum_loss = ForwardSumLoss()
    self.bin_loss = BinLoss()

    self.bin_warmup = bin_warmup
    self.binarization_loss_enable_steps = binarization_loss_enable_steps
    self.binarization_loss_warmup_steps = binarization_loss_warmup_steps

forward(src_masks, mel_masks, mel_targets, mel_predictions, log_duration_predictions, u_prosody_ref, u_prosody_pred, p_prosody_ref, p_prosody_pred, durations, pitch_predictions, p_targets, attn_logprob, attn_soft, attn_hard, step, src_lens, mel_lens, energy_pred, energy_target)

Computes the loss for the FastSpeech2 model.

Parameters:

Name Type Description Default
src_masks Tensor

Mask for the source sequence.

 required
mel_masks Tensor

Mask for the mel-spectrogram.

 required
mel_targets Tensor

Target mel-spectrogram.

 required
mel_predictions Tensor

Predicted mel-spectrogram.

 required
log_duration_predictions Tensor

Predicted log-duration.

 required
u_prosody_ref Tensor

Reference unvoiced prosody.

 required
u_prosody_pred Tensor

Predicted unvoiced prosody.

 required
p_prosody_ref Tensor

Reference voiced prosody.

 required
p_prosody_pred Tensor

Predicted voiced prosody.

 required
durations Tensor

Ground-truth durations.

 required
pitch_predictions Tensor

Predicted pitch.

 required
p_targets Tensor

Ground-truth pitch.

 required
attn_logprob Tensor

Log-probability of attention.

 required
attn_soft Tensor

Soft attention.

 required
attn_hard Tensor

Hard attention.

 required
step int

Current training step.

 required
src_lens Tensor

Lengths of the source sequences.

 required
mel_lens Tensor

Lengths of the mel-spectrograms.

 required
energy_pred Tensor

Predicted energy.

 required
energy_target Tensor

Ground-truth energy.

 required

Returns:

Type Description
Tuple[Tensor, Tensor, Tensor, Tensor, Tensor, Tensor, Tensor, Tensor, Tensor, Tensor]

Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: The total loss and its components.
Note

Here is the description of the returned loss components: total_loss: This is the total loss computed as the sum of all the other losses. mel_loss: This is the mean absolute error (MAE) loss between the predicted and target mel-spectrograms. It measures how well the model predicts the mel-spectrograms. sc_mag_loss: This is the spectral convergence loss between the predicted and target mel-spectrograms. It measures how well the model predicts the mel-spectrograms in terms of their spectral structure. log_mag_loss: This is the log STFT magnitude loss between the predicted and target mel-spectrograms. It measures how well the model predicts the mel-spectrograms in terms of their spectral structure. ssim_loss: This is the Structural Similarity Index (SSIM) loss between the predicted and target mel-spectrograms. It measures the similarity between the two mel-spectrograms in terms of their structure, contrast, and luminance. duration_loss: This is the mean squared error (MSE) loss between the predicted and target log-durations. It measures how well the model predicts the durations of the phonemes. u_prosody_loss: This is the MAE loss between the predicted and reference unvoiced prosody. It measures how well the model predicts the prosody (rhythm, stress, and intonation) of the unvoiced parts of the speech. p_prosody_loss: This is the MAE loss between the predicted and reference voiced prosody. It measures how well the model predicts the prosody of the voiced parts of the speech. pitch_loss: This is the MSE loss between the predicted and target pitch. It measures how well the model predicts the pitch of the speech. ctc_loss: This is the Connectionist Temporal Classification (CTC) loss computed from the log-probability of attention and the lengths of the source sequences and mel-spectrograms. It measures how well the model aligns the input and output sequences. bin_loss: This is the binarization loss computed from the hard and soft attention. It measures how well the model learns to attend to the correct parts of the input sequence. energy_loss: This is the MSE loss between the predicted and target energy. It measures how well the model predicts the energy of the speech.

Source code in training/loss/fast_speech_2_loss_gen.py 
 39
 40
 41
 42
 43
 44
 45
 46
 47
 48
 49
 50
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
def forward(
    self,
    src_masks: torch.Tensor,
    mel_masks: torch.Tensor,
    mel_targets: torch.Tensor,
    mel_predictions: torch.Tensor,
    log_duration_predictions: torch.Tensor,
    u_prosody_ref: torch.Tensor,
    u_prosody_pred: torch.Tensor,
    p_prosody_ref: torch.Tensor,
    p_prosody_pred: torch.Tensor,
    durations: torch.Tensor,
    pitch_predictions: torch.Tensor,
    p_targets: torch.Tensor,
    attn_logprob: torch.Tensor,
    attn_soft: torch.Tensor,
    attn_hard: torch.Tensor,
    step: int,
    src_lens: torch.Tensor,
    mel_lens: torch.Tensor,
    energy_pred: torch.Tensor,
    energy_target: torch.Tensor,
) -> Tuple[
    torch.Tensor,
    torch.Tensor,
    torch.Tensor,
    torch.Tensor,
    torch.Tensor,
    torch.Tensor,
    torch.Tensor,
    torch.Tensor,
    torch.Tensor,
    torch.Tensor,
]:
    r"""Computes the loss for the FastSpeech2 model.

    Args:
        src_masks (torch.Tensor): Mask for the source sequence.
        mel_masks (torch.Tensor): Mask for the mel-spectrogram.
        mel_targets (torch.Tensor): Target mel-spectrogram.
        mel_predictions (torch.Tensor): Predicted mel-spectrogram.
        log_duration_predictions (torch.Tensor): Predicted log-duration.
        u_prosody_ref (torch.Tensor): Reference unvoiced prosody.
        u_prosody_pred (torch.Tensor): Predicted unvoiced prosody.
        p_prosody_ref (torch.Tensor): Reference voiced prosody.
        p_prosody_pred (torch.Tensor): Predicted voiced prosody.
        durations (torch.Tensor): Ground-truth durations.
        pitch_predictions (torch.Tensor): Predicted pitch.
        p_targets (torch.Tensor): Ground-truth pitch.
        attn_logprob (torch.Tensor): Log-probability of attention.
        attn_soft (torch.Tensor): Soft attention.
        attn_hard (torch.Tensor): Hard attention.
        step (int): Current training step.
        src_lens (torch.Tensor): Lengths of the source sequences.
        mel_lens (torch.Tensor): Lengths of the mel-spectrograms.
        energy_pred (torch.Tensor): Predicted energy.
        energy_target (torch.Tensor): Ground-truth energy.

    Returns:
        Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: The total loss and its components.

    Note:
        Here is the description of the returned loss components:
        `total_loss`: This is the total loss computed as the sum of all the other losses.
        `mel_loss`: This is the mean absolute error (MAE) loss between the predicted and target mel-spectrograms. It measures how well the model predicts the mel-spectrograms.
        `sc_mag_loss`: This is the spectral convergence loss between the predicted and target mel-spectrograms. It measures how well the model predicts the mel-spectrograms in terms of their spectral structure.
        `log_mag_loss`: This is the log STFT magnitude loss between the predicted and target mel-spectrograms. It measures how well the model predicts the mel-spectrograms in terms of their spectral structure.
        `ssim_loss`: This is the Structural Similarity Index (SSIM) loss between the predicted and target mel-spectrograms. It measures the similarity between the two mel-spectrograms in terms of their structure, contrast, and luminance.
        `duration_loss`: This is the mean squared error (MSE) loss between the predicted and target log-durations. It measures how well the model predicts the durations of the phonemes.
        `u_prosody_loss`: This is the MAE loss between the predicted and reference unvoiced prosody. It measures how well the model predicts the prosody (rhythm, stress, and intonation) of the unvoiced parts of the speech.
        `p_prosody_loss`: This is the MAE loss between the predicted and reference voiced prosody. It measures how well the model predicts the prosody of the voiced parts of the speech.
        `pitch_loss`: This is the MSE loss between the predicted and target pitch. It measures how well the model predicts the pitch of the speech.
        `ctc_loss`: This is the Connectionist Temporal Classification (CTC) loss computed from the log-probability of attention and the lengths of the source sequences and mel-spectrograms. It measures how well the model aligns the input and output sequences.
        `bin_loss`: This is the binarization loss computed from the hard and soft attention. It measures how well the model learns to attend to the correct parts of the input sequence.
        `energy_loss`: This is the MSE loss between the predicted and target energy. It measures how well the model predicts the energy of the speech.
    """
    log_duration_targets = torch.log(durations.float() + 1).to(src_masks.device)

    log_duration_targets.requires_grad = False
    mel_targets.requires_grad = False
    p_targets.requires_grad = False
    energy_target.requires_grad = False

    log_duration_predictions = log_duration_predictions.masked_select(~src_masks)
    log_duration_targets = log_duration_targets.masked_select(~src_masks)

    mel_masks_expanded = mel_masks.unsqueeze(1)

    mel_predictions_normalized = (
        sample_wise_min_max(mel_predictions).float().to(mel_predictions.device)
    )
    mel_targets_normalized = (
        sample_wise_min_max(mel_targets).float().to(mel_predictions.device)
    )

    ssim_loss: torch.Tensor = self.ssim_loss(
        mel_predictions_normalized.unsqueeze(1),
        mel_targets_normalized.unsqueeze(1),
    )

    if ssim_loss.item() > 1.0 or ssim_loss.item() < 0.0:
        ssim_loss = torch.tensor([1.0], device=mel_predictions.device)

    masked_mel_predictions = mel_predictions.masked_select(~mel_masks_expanded)

    masked_mel_targets = mel_targets.masked_select(~mel_masks_expanded)

    mel_loss: torch.Tensor = self.mae_loss(
        masked_mel_predictions,
        masked_mel_targets,
    )

    p_prosody_ref = p_prosody_ref.permute((0, 2, 1))
    p_prosody_pred = p_prosody_pred.permute((0, 2, 1))

    p_prosody_ref = p_prosody_ref.masked_fill(src_masks.unsqueeze(1), 0.0)
    p_prosody_pred = p_prosody_pred.masked_fill(src_masks.unsqueeze(1), 0.0)

    p_prosody_ref = p_prosody_ref.detach()

    p_prosody_loss: torch.Tensor = 0.5 * self.mae_loss(
        p_prosody_ref.masked_select(~src_masks.unsqueeze(1)),
        p_prosody_pred.masked_select(~src_masks.unsqueeze(1)),
    )

    u_prosody_ref = u_prosody_ref.detach()
    u_prosody_loss: torch.Tensor = 0.5 * self.mae_loss(
        u_prosody_ref,
        u_prosody_pred,
    )

    duration_loss: torch.Tensor = self.mse_loss(
        log_duration_predictions,
        log_duration_targets,
    )

    pitch_predictions = pitch_predictions.masked_select(~src_masks)
    p_targets = p_targets.masked_select(~src_masks)

    pitch_loss: torch.Tensor = self.mse_loss(pitch_predictions, p_targets)

    ctc_loss: torch.Tensor = self.sum_loss(
        attn_logprob=attn_logprob,
        in_lens=src_lens,
        out_lens=mel_lens,
    )

    if self.bin_warmup:
        if step < self.binarization_loss_enable_steps:
            bin_loss_weight = 0.0
        else:
            bin_loss_weight = (
                min(
                    (step - self.binarization_loss_enable_steps)
                    / self.binarization_loss_warmup_steps,
                    1.0,
                )
                * 1.0
            )

        bin_loss: torch.Tensor = (
            self.bin_loss(hard_attention=attn_hard, soft_attention=attn_soft)
            * bin_loss_weight
        )
    else:
        bin_loss: torch.Tensor = self.bin_loss(
            hard_attention=attn_hard,
            soft_attention=attn_soft,
        )

    energy_loss: torch.Tensor = self.mse_loss(energy_pred, energy_target)

    total_loss = (
        mel_loss
        + duration_loss
        + u_prosody_loss
        + p_prosody_loss
        + ssim_loss
        + pitch_loss
        + ctc_loss
        + bin_loss
        + energy_loss
    )

    return (
        total_loss,
        mel_loss,
        ssim_loss,
        duration_loss,
        u_prosody_loss,
        p_prosody_loss,
        pitch_loss,
        ctc_loss,
        bin_loss,
        energy_loss,
    )