train_aio.py

import glob
import hashlib
import os
import random
import subprocess

import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.utils.data
import torchaudio
from torch.nn.utils import weight_norm, spectral_norm, remove_weight_norm
from tqdm import tqdm

torch_eps = torch.finfo(torch.float).eps


# ==============================================================================
# 0. GAN 相关模块
# ==============================================================================

def apply_parametrization_norm(module: nn.Module, norm: str = 'none') -> nn.Module:
    assert norm in frozenset(['none', 'weight_norm', 'spectral_norm'])
    if norm == 'weight_norm':
        return weight_norm(module)
    elif norm == 'spectral_norm':
        return spectral_norm(module)
    else:
        return module


class NormConv2d(nn.Module):
    """ 2D 卷积的归一化包装器 """

    def __init__(self, *args, norm: str = 'none', **kwargs):
        super().__init__()
        self.conv = apply_parametrization_norm(nn.Conv2d(*args, **kwargs), norm)
        self.norm_type = norm

    def forward(self, x):
        return self.conv(x)

    def remove_weight_norm(self):
        if self.norm_type == 'weight_norm':
            remove_weight_norm(self.conv)


class ResnetBlock(nn.Module):
    """ STFT 判别器中使用的 ResNet 块 (参考实现) """

    def __init__(self, in_channels, N, m, st, sf, norm='spectral_norm'):
        super(ResnetBlock, self).__init__()
        self.conv1 = NormConv2d(in_channels, N, kernel_size=3, padding=1, norm=norm)
        self.act1 = nn.LeakyReLU(0.2, True)
        self.shortcut = NormConv2d(in_channels, m * N, kernel_size=1, stride=(st, sf), padding=0, norm=norm)
        # 卷积核和填充根据 stride 动态调整，以匹配参考实现
        self.conv2 = NormConv2d(N, m * N, kernel_size=(st + 2, sf + 2), stride=(st, sf), padding=(2, 2), norm=norm)

    def forward(self, x):
        y = self.act1(self.conv1(x))
        y = self.conv2(y)
        shortcut = self.shortcut(x)
        # 裁剪输出以匹配 shortcut 的维度
        result = shortcut + y[:, :, :shortcut.shape[2], :shortcut.shape[3]]
        return result


class STFTDiscriminator(nn.Module):
    """ 单个 STFT 分辨率的判别器 (参考实现) """

    def __init__(self, C_out=8, C_in=1, norm='spectral_norm'):
        super().__init__()
        self.model = nn.ModuleList([
            nn.Sequential(
                NormConv2d(C_in, C_out, kernel_size=7, padding=3, norm=norm),
                nn.LeakyReLU(0.2, True),
            ),
            ResnetBlock(C_out, C_out, 2, 1, 2, norm),  # (B, 16, F, T/2)
            ResnetBlock(2 * C_out, 2 * C_out, 2, 2, 1, norm),  # (B, 32, F/2, T/2)
            ResnetBlock(4 * C_out, 4 * C_out, 1, 1, 2, norm),  # (B, 32, F/2, T/4)
            ResnetBlock(4 * C_out, 4 * C_out, 2, 2, 1, norm),  # (B, 64, F/4, T/4)
            ResnetBlock(8 * C_out, 8 * C_out, 1, 1, 2, norm),  # (B, 64, F/4, T/8)
            ResnetBlock(8 * C_out, 8 * C_out, 2, 2, 1, norm),  # (B, 128, F/8, T/8)
            NormConv2d(16 * C_out, 1, kernel_size=(7, 7), padding=3, norm=norm)
        ])

    def forward(self, x):
        feature_maps = []
        for layer in self.model:
            x = layer(x)
            feature_maps.append(x)
        # 最后一个特征图是判别分数，其余是中间特征
        return x, feature_maps[:-1]


class MultiResolutionSTFTDiscriminator(nn.Module):
    """ 多分辨率 STFT 判别器 """

    def __init__(self, stft_params, norm='spectral_norm'):
        super().__init__()
        self.discriminators = nn.ModuleList()
        self.stft_params = stft_params
        for params in stft_params:
            self.discriminators.append(STFTDiscriminator(norm=norm))
            self.register_buffer(f"window_{params['n_fft']}", torch.hann_window(params['win_length']))

    def forward(self, x: torch.Tensor):
        x = x.squeeze(1)  # (B, T)
        scores = []
        feature_maps = []
        for i, params in enumerate(self.stft_params):
            window = getattr(self, f"window_{params['n_fft']}")
            stft_out = torch.stft(
                x,
                n_fft=params['n_fft'],
                hop_length=params['hop_length'],
                win_length=params['win_length'],
                window=window,
                return_complex=True,
                pad_mode='reflect',
                center=True,
            )
            # 使用幅度谱作为输入, (B, F, T)
            mag = torch.abs(stft_out)
            # 增加通道维度并调整为 (B, C, F, T) -> Conv2d期望的格式
            mag = mag.unsqueeze(1)

            # 注意：参考代码的 ResNetBlock 期望 (B, C, T, F) 格式，这里进行适配
            # 如果 ResNetBlock 设计为处理 (F, T) 格式，则不需要 permute
            # mag = mag.permute(0, 1, 3, 2) # (B, C, T, F)

            score, fmap = self.discriminators[i](mag)
            scores.append(score)
            feature_maps.append(fmap)
        return scores, feature_maps


def feature_matching_loss(fmap_r, fmap_g):
    """ 特征匹配损失 """
    loss = 0
    for dr, dg in zip(fmap_r, fmap_g):
        for rl, gl in zip(dr, dg):
            loss += torch.mean(torch.abs(rl - gl))
    return loss


def discriminator_loss(disc_real_outputs, disc_fake_outputs):
    """ 判别器损失 (LSGAN) """
    loss = 0
    for dr, dg in zip(disc_real_outputs, disc_fake_outputs):
        r_loss = torch.mean((dr - 1) ** 2)
        f_loss = torch.mean(dg ** 2)
        loss += (r_loss + f_loss)
    return loss


def generator_adversarial_loss(disc_fake_outputs):
    """ 生成器对抗损失 (LSGAN) """
    loss = 0
    for dg in disc_fake_outputs:
        loss += torch.mean((dg - 1) ** 2)
    return loss


# ==============================================================================
# 1. 网络定义代码
# ==============================================================================

CONV_NORMALIZATIONS = frozenset(['none', 'weight_norm', 'spectral_norm',
                                 'time_layer_norm', 'layer_norm', 'time_group_norm'])


class NormSwitch(nn.Module):
    def __init__(self, norm_type, format, dim):
        super(NormSwitch, self).__init__()
        self.norm_type = norm_type
        self.format = format
        self.dim = dim
        if self.norm_type == "BN":
            if self.format == "1D":
                self.norm = nn.BatchNorm1d(self.dim)
            elif self.format == "2D":
                self.norm = nn.BatchNorm2d(self.dim)
        elif self.norm_type == "IN":
            if self.format == "1D":
                self.norm = nn.InstanceNorm1d(self.dim, affine=True)
            elif self.format == "2D":
                self.norm = nn.InstanceNorm2d(self.dim, affine=True)
        elif self.norm_type == "LN":
            # LayerNorm 通常在最后一个维度上操作
            self.norm = nn.LayerNorm(self.dim)

    def forward(self, x):
        return self.norm(x)


class NormConv1d(nn.Module):
    def __init__(self, *args, norm: str = 'none', **kwargs):
        super().__init__()
        self.conv = apply_parametrization_norm(nn.Conv1d(*args, **kwargs), norm)
        self.norm_type = norm

    def forward(self, x):
        x = self.conv(x)
        return x

    def remove_weight_norm(self):
        remove_weight_norm(self.conv)


class interction(nn.Module):
    def __init__(self, input_size):
        super(interction, self).__init__()
        self.inter = nn.Sequential(
            NormConv1d(input_size, input_size, kernel_size=1),
            nn.Sigmoid()
        )

    def forward(self, input1, input2):
        input_merge = torch.add(input1, input2)
        output_mask = self.inter(input_merge)
        output = input1 + input2 * output_mask
        return output


class Phase_Encoder(nn.Module):
    def __init__(self, input_size, kernel_size, norm: str = 'weight_norm'):
        super().__init__()
        kernel_size = kernel_size
        # 增加了通道数以提升模型容量
        enc_c_in = [512, 256, 256, 128, 128]
        self.enc_c_in = enc_c_in

        enc_1 = nn.Sequential(
            nn.ConstantPad1d((kernel_size[0] - 1, 0), 0),
            NormConv1d(input_size, enc_c_in[0], kernel_size=kernel_size[0], groups=2, stride=1, norm=norm),
            nn.LeakyReLU(0.2),
        )
        enc_2 = nn.Sequential(
            nn.ConstantPad1d((kernel_size[1] - 1, 0), 0),
            NormConv1d(enc_c_in[0], enc_c_in[1], kernel_size=kernel_size[1], groups=2, stride=1, norm=norm),
            nn.LeakyReLU(0.2),
        )
        enc_3 = nn.Sequential(
            nn.ConstantPad1d((kernel_size[2] - 1, 0), 0),
            NormConv1d(enc_c_in[1], enc_c_in[2], kernel_size=kernel_size[2], groups=2, stride=1, norm=norm),
            nn.LeakyReLU(0.2),
        )
        enc_4 = nn.Sequential(
            nn.ConstantPad1d((kernel_size[3] - 1, 0), 0),
            NormConv1d(enc_c_in[2], enc_c_in[3], kernel_size=kernel_size[3], groups=1, stride=1, norm=norm),
            nn.LeakyReLU(0.2),
        )
        enc_5 = nn.Sequential(
            nn.ConstantPad1d((kernel_size[4] - 1, 0), 0),
            NormConv1d(enc_c_in[3], enc_c_in[4], kernel_size=kernel_size[4], groups=1, stride=1, norm=norm),
            nn.LeakyReLU(0.2),
        )
        self.en = nn.ModuleList([enc_1, enc_2, enc_3, enc_4, enc_5])

        # 交互模块的通道数现在与新的 enc_c_in 匹配
        inter_com_1 = interction(enc_c_in[1])
        inter_com_2 = interction(enc_c_in[2])
        inter_com_3 = interction(enc_c_in[3])
        inter_com_4 = interction(enc_c_in[4])
        self.inter_com = nn.ModuleList([inter_com_1, inter_com_2, inter_com_3, inter_com_4])

    def forward(self, x_ri, mag_list):
        phase_list = []

        enc_1 = self.en[0](x_ri)
        enc_2 = self.en[1](enc_1)

        # 交互1: enc_2 (256ch) 与 mag_list[0] (enc_1_mag, 256ch)
        inter_1 = self.inter_com[0](enc_2, mag_list[0])
        enc_3 = self.en[2](inter_1)

        # 交互2: enc_3 (256ch) 与 mag_list[1] (enc_2_mag, 256ch)
        inter_2 = self.inter_com[1](enc_3, mag_list[1])
        enc_4 = self.en[3](inter_2)

        # 交互3: enc_4 (128ch) 与 mag_list[2] (enc_3_mag, 128ch)
        inter_3 = self.inter_com[2](enc_4, mag_list[2])
        enc_5 = self.en[4](inter_3)

        # 交互4: enc_5 (128ch) 与 mag_list[3] (enc_4_mag, 128ch)
        inter_4 = self.inter_com[3](enc_5, mag_list[3])

        phase_list.append(enc_1)
        phase_list.append(inter_1)
        phase_list.append(inter_2)
        phase_list.append(inter_3)

        return inter_4, phase_list

    def remove_weight_norm(self):
        for module in self.en:
            for layer in module:
                if isinstance(layer, NormConv1d):
                    layer.remove_weight_norm()


class Phase_Decoder(nn.Module):
    def __init__(self, input_size, output_freq_bins, kernel_size, norm: str = 'weight_norm'):
        super().__init__()
        kernel_size = kernel_size
        # 增加了通道数, output_freq_bins 动态传入
        de_c_in = [256, 256, 512]

        dec_1 = nn.Sequential(
            nn.ConstantPad1d((kernel_size[0] - 1, 0), 0),
            NormConv1d(input_size, de_c_in[0], kernel_size=kernel_size[0], groups=1, stride=1, norm=norm),
            nn.LeakyReLU(0.2)
        )
        dec_2 = nn.Sequential(
            nn.ConstantPad1d((kernel_size[1] - 1, 0), 0),
            NormConv1d(de_c_in[0], de_c_in[1], kernel_size=kernel_size[1], groups=2, stride=1, norm=norm),
            nn.LeakyReLU(0.2)
        )
        dec_3 = nn.Sequential(
            nn.ConstantPad1d((kernel_size[2] - 1, 0), 0),
            NormConv1d(de_c_in[1], de_c_in[2], kernel_size=kernel_size[2], groups=2, stride=1, norm=norm),
            nn.LeakyReLU(0.2)
        )
        self.de = nn.ModuleList([dec_1, dec_2, dec_3])

        # 交互模块的通道数现在与新的解码器和编码器通道数匹配
        inter_channel = [128, 256, 256, 512]
        inter_com_0 = interction(inter_channel[0])
        inter_com_1 = interction(inter_channel[1])
        inter_com_2 = interction(inter_channel[2])
        inter_com_3 = interction(inter_channel[3])
        self.inter_com = nn.ModuleList([inter_com_0, inter_com_1, inter_com_2, inter_com_3])

        # 输出层现在动态地使用 output_freq_bins
        self.de_real_out = nn.Linear(de_c_in[2], output_freq_bins)
        self.de_imag_out = nn.Linear(de_c_in[2], output_freq_bins)

    def forward(self, x, x_list, out_mag_list):
        # x_list is [enc_1, inter_1, inter_2, inter_3] from PhaseEncoder
        x = torch.add(x, x_list[-1])
        x = self.inter_com[0](x, out_mag_list[0])  # inter(128, 128)
        x = self.de[0](x)  # dec(128->256)

        x = torch.add(x, x_list[-2])
        x = self.inter_com[1](x, out_mag_list[1])  # inter(256, 256)
        x = self.de[1](x)  # dec(256->256)

        x = torch.add(x, x_list[-3])
        x = self.inter_com[2](x, out_mag_list[2])  # inter(256, 256)
        x = self.de[2](x)  # dec(256->512)

        x = torch.add(x, x_list[0])  # add enc_1 (512)
        # x = self.inter_com[3](x, out_mag_list[3]) # This interaction might be misaligned, skipping for now

        x_fc_in = x.permute(0, 2, 1).contiguous()
        x_r = self.de_real_out(x_fc_in).permute(0, 2, 1).contiguous()
        x_i = self.de_imag_out(x_fc_in).permute(0, 2, 1).contiguous()
        return x_r, x_i

    def remove_weight_norm(self):
        for module in self.de:
            for layer in module:
                if isinstance(layer, NormConv1d):
                    layer.remove_weight_norm()


class Mag_Encoder(nn.Module):
    def __init__(self, input_size, kernel_size, norm: str = 'weight_norm'):
        super().__init__()
        # input_size 是 STFT 的频点数, e.g., 1537
        linear_freq_bins = input_size
        kernel_size = kernel_size

        # 增加了通道数以提升模型容量
        enc_c_in = [256, 256, 128, 128]
        self.enc_c_in = enc_c_in

        # 第一层编码器接收 STFT 频点作为输入通道
        enc_1 = nn.Sequential(
            nn.ConstantPad1d((kernel_size[0] - 1, 0), 0),
            NormConv1d(linear_freq_bins, enc_c_in[0], kernel_size=kernel_size[0], groups=1, stride=1, norm=norm),
            nn.LeakyReLU(0.2),
        )
        enc_2 = nn.Sequential(
            nn.ConstantPad1d((kernel_size[1] - 1, 0), 0),
            NormConv1d(enc_c_in[0], enc_c_in[1], kernel_size=kernel_size[1], groups=2, stride=1, norm=norm),
            nn.LeakyReLU(0.2),
        )
        enc_3 = nn.Sequential(
            nn.ConstantPad1d((kernel_size[2] - 1, 0), 0),
            NormConv1d(enc_c_in[1], enc_c_in[2], kernel_size=kernel_size[2], groups=2, stride=1, norm=norm),
            nn.LeakyReLU(0.2),
        )
        enc_4 = nn.Sequential(
            nn.ConstantPad1d((kernel_size[3] - 1, 0), 0),
            NormConv1d(enc_c_in[2], enc_c_in[3], kernel_size=kernel_size[3], groups=1, stride=1, norm=norm),
            nn.LeakyReLU(0.2),
        )
        self.en = nn.ModuleList([enc_1, enc_2, enc_3, enc_4])

    def forward(self, x):
        x_list = []
        x_list.append(x)
        for i in range(len(self.en)):
            x = self.en[i](x)
            x_list.append(x)
        return x, x_list

    def remove_weight_norm(self):
        for module in self.en:
            for layer in module:
                if isinstance(layer, NormConv1d):
                    layer.remove_weight_norm()


class Mag_Decoder(nn.Module):
    def __init__(self, input_size, output_freq_bins, kernel_size, norm: str = 'weight_norm'):
        super().__init__()
        kernel_size = kernel_size
        # 增加了通道数, output_freq_bins 动态传入
        de_c_in = [128, 256, 256]

        dec_1 = nn.Sequential(
            nn.ConstantPad1d((kernel_size[0] - 1, 0), 0),
            NormConv1d(input_size, de_c_in[0], kernel_size=kernel_size[0], groups=1, stride=1, norm=norm),
            nn.LeakyReLU(0.2)
        )
        dec_2 = nn.Sequential(
            nn.ConstantPad1d((kernel_size[1] - 1, 0), 0),
            NormConv1d(de_c_in[0], de_c_in[1], kernel_size=kernel_size[1], groups=2, stride=1, norm=norm),
            nn.LeakyReLU(0.2)
        )
        dec_3 = nn.Sequential(
            nn.ConstantPad1d((kernel_size[2] - 1, 0), 0),
            NormConv1d(de_c_in[1], de_c_in[2], kernel_size=kernel_size[2], groups=2, stride=1, norm=norm),
            nn.LeakyReLU(0.2)
        )
        self.de = nn.ModuleList([dec_1, dec_2, dec_3])

        # 输出层现在动态地使用 output_freq_bins
        self.de_out = nn.Sequential(
            nn.ConstantPad1d((kernel_size[3] - 1, 0), 0),
            NormConv1d(de_c_in[2], output_freq_bins, kernel_size=kernel_size[3], groups=1, stride=1, norm=norm)
        )

        self.mask1 = nn.Sequential(nn.Conv2d(in_channels=1, out_channels=1, kernel_size=1), nn.Sigmoid())
        self.mask2 = nn.Sequential(nn.Conv2d(in_channels=1, out_channels=1, kernel_size=1), nn.Tanh())
        self.maskconv = nn.Conv2d(in_channels=1, out_channels=1, kernel_size=1)
        self.mask_gate = nn.Sigmoid()

    def forward(self, x, x_list):
        out_mag_list = []
        # x_list 包含 [orig_mag, enc1, enc2, enc3, enc4]
        x = torch.add(x, x_list[-1])  # Add enc4
        x = self.de[0](x)  # dec(128->128)
        out_mag_list.append(x)

        x = torch.add(x, x_list[-2])  # Add enc3
        x = self.de[1](x)  # dec(128->256)
        out_mag_list.append(x)

        x = torch.add(x, x_list[-3])  # Add enc2
        x = self.de[2](x)  # dec(256->256)
        out_mag_list.append(x)

        x = torch.add(x, x_list[-4])  # Add enc1
        x = self.de_out(x)

        x_out = x.unsqueeze(1)
        x_in = x_list[0].unsqueeze(1)
        x_mask_in = torch.add(x_out, x_in)
        x_dual_mask = self.mask1(x_mask_in) * self.mask2(x_mask_in)
        out_mask = self.mask_gate(self.maskconv(x_dual_mask))
        out_full = x_out * out_mask
        out_full = out_full.squeeze(dim=1)
        out_full = torch.add(out_full, x_list[0])
        return out_full, out_mag_list

    def remove_weight_norm(self):
        for module in self.de:
            for layer in module:
                if isinstance(layer, NormConv1d):
                    layer.remove_weight_norm()
        for layer in self.de_out:
            if isinstance(layer, NormConv1d):
                layer.remove_weight_norm()


class GroupRNN(nn.Module):
    def __init__(self,
                 input_size: int,
                 hidden_size: int,
                 split_group: int,
                 rnn_type: str,
                 is_causal: bool,
                 ):
        super(GroupRNN, self).__init__()
        self.input_size = input_size
        self.hidden_size = hidden_size
        self.split_group = split_group
        self.rnn_type = rnn_type
        self.is_causal = is_causal
        input_size_t = input_size // split_group
        hidden_size_t = hidden_size // split_group
        causal_flag = 1 if is_causal else 2
        self.rnn_list1 = nn.ModuleList(
            [getattr(nn, rnn_type)(input_size=input_size_t,
                                   hidden_size=hidden_size_t // causal_flag,
                                   num_layers=1,
                                   batch_first=True,
                                   bidirectional=not is_causal) for _ in range(split_group)])
        self.rnn_list2 = nn.ModuleList(
            [getattr(nn, rnn_type)(input_size=hidden_size_t,
                                   hidden_size=hidden_size_t // causal_flag,
                                   num_layers=1,
                                   batch_first=True,
                                   bidirectional=not is_causal) for _ in range(split_group)])
        self.norm1 = NormSwitch('BN', "1D", hidden_size)
        self.norm2 = NormSwitch('BN', "1D", hidden_size)

    def forward(self, inpt):
        # inpt: (B, T, F) -> return: (B, T, F)
        x_list = torch.chunk(inpt, self.split_group, dim=-1)
        x = torch.stack([self.rnn_list1[i](x_list[i])[0] for i in range(self.split_group)], dim=-1)
        x = torch.flatten(x, start_dim=-2, end_dim=-1)

        x = x.permute(0, 2, 1).contiguous()
        x = self.norm1(x)
        x = x.permute(0, 2, 1).contiguous()

        x_list = torch.chunk(x, self.split_group, dim=-1)
        x = torch.cat([self.rnn_list2[i](x_list[i])[0] for i in range(self.split_group)], dim=-1)

        x = x.permute(0, 2, 1).contiguous()
        x = self.norm2(x)
        x = x.permute(0, 2, 1).contiguous()
        return x


class Generator(nn.Module):
    def __init__(self, input_size, dec_dim, norm='weight_norm'):
        super().__init__()
        self.input_size = input_size
        self.dec_dim = dec_dim

        self.enc_kernel_size_mag = [3, 3, 3, 3]
        self.enc_kernel_size_ri = [3, 3, 3, 3, 3]
        self.dec_kernel_size = [3, 3, 3, 3]

        # dec_dim 增加，并传递给相应模块
        self.enc_mag = Mag_Encoder(input_size=self.input_size, kernel_size=self.enc_kernel_size_mag, norm=norm)
        # Mag_Decoder 的输入维度是 GRU 的输出维度 (dec_dim)
        self.dec_mag = Mag_Decoder(input_size=self.dec_dim, output_freq_bins=self.input_size,
                                   kernel_size=self.dec_kernel_size, norm=norm)
        self.gru_mag = GroupRNN(self.enc_mag.enc_c_in[-1], self.dec_dim, split_group=4, rnn_type="GRU", is_causal=True)

        # Phase_Encoder 的输入是实部+虚部，所以是 input_size * 2
        self.enc_ri = Phase_Encoder(input_size=self.input_size * 2, kernel_size=self.enc_kernel_size_ri, norm=norm)
        # Phase_Decoder 的输入维度是 GRU 的输出维度 (dec_dim)
        self.dec_ri = Phase_Decoder(input_size=self.dec_dim, output_freq_bins=self.input_size,
                                    kernel_size=self.dec_kernel_size, norm=norm)
        self.gru_ri = GroupRNN(self.enc_ri.enc_c_in[-1], self.dec_dim, split_group=4, rnn_type="GRU", is_causal=True)

    def forward(self, input_ri):
        # input_ri shape: (B, 2, F, T)
        mag_input = (torch.norm(input_ri, dim=1))  # (B, F, T)
        phase_ori = torch.atan2(input_ri[:, 1, :, :], input_ri[:, 0, :, :])  # (B, F, T)

        # 幅度分支
        enc_out_mag, enc_list_mag = self.enc_mag(mag_input)
        enc_out_mag = enc_out_mag.permute(0, 2, 1).contiguous()
        enc_out_mag = self.gru_mag(enc_out_mag)
        enc_out_mag = enc_out_mag.permute(0, 2, 1).contiguous()
        dec_out_mag, dec_out_mag_list = self.dec_mag(enc_out_mag, enc_list_mag)

        # 相位分支
        ri_input = torch.cat((input_ri[:, 0, :, :], input_ri[:, 1, :, :]), dim=1)  # (B, 2*F, T)

        # 用于相位交互的幅度特征列表 [enc1_mag, enc2_mag, enc3_mag, enc4_mag]
        phase_interaction_mag_list = enc_list_mag[1:]
        enc_out_ri, enc_list_ri = self.enc_ri(ri_input, phase_interaction_mag_list)

        enc_out_ri = enc_out_ri.permute(0, 2, 1).contiguous()
        enc_out_ri = self.gru_ri(enc_out_ri)
        enc_out_ri = enc_out_ri.permute(0, 2, 1).contiguous()

        dec_out_r, dec_out_i = self.dec_ri(enc_out_ri, enc_list_ri, dec_out_mag_list)

        # 结合输出
        dec_out_final_r = dec_out_mag * torch.cos(phase_ori) + dec_out_r
        dec_out_final_i = dec_out_mag * torch.sin(phase_ori) + dec_out_i
        x_com_out = torch.stack((dec_out_final_r, dec_out_final_i), dim=1)
        return x_com_out

    def eval(self, inference=False):
        super(Generator, self).eval()
        if inference:
            self.remove_weight_norm()

    def remove_weight_norm(self):
        self.enc_mag.remove_weight_norm()
        self.dec_mag.remove_weight_norm()
        self.enc_ri.remove_weight_norm()
        self.dec_ri.remove_weight_norm()


# ==============================================================================
# 2. 新实现的 2x 音频上采样包装模型
# ==============================================================================
class AudioUpsampler(nn.Module):
    """
    一个包装模型，用于将谱域增强模型 (Generator) 应用于时域音频上采样任务。
    修改了 STFT 参数和模型维度以适应一般音频。
    """

    def __init__(self, upscale_factor: int = 2):
        super().__init__()

        if upscale_factor != 2:
            raise ValueError("This model is specifically designed for 2x upsampling.")
        self.upscale_factor = upscale_factor

        # STFT 参数更新
        # 增加 n_fft 以获得更好的频率分辨率，适用于一般音频
        self.n_fft = 3072
        self.hop_length = self.n_fft // 4
        self.win_length = self.n_fft

        # 计算 Generator 的输入频点数
        input_freq_bins = self.n_fft // 2 + 1  # 3072 / 2 + 1 = 1537

        self.register_buffer('window', torch.hann_window(self.win_length))

        # 增加解码器和 GRU 的维度 (dec_dim) 以提升模型容量
        generator_dec_dim = 128

        # 实例化谱域增强模型
        self.generator = Generator(
            input_size=input_freq_bins,
            dec_dim=generator_dec_dim,
            norm='weight_norm'
        )

    def forward(self, x: torch.Tensor):
        # 1. STFT
        spec = torch.stft(
            x,
            n_fft=self.n_fft,
            hop_length=self.hop_length,
            win_length=self.win_length,
            window=self.window,
            return_complex=True,
            pad_mode='reflect',
            center=True,
        )

        # 2. 分离实部虚部
        spec_ri = torch.stack([spec.real, spec.imag], dim=1)

        # 3. Generator 谱域增强
        enhanced_spec_ri = self.generator(spec_ri)

        # 4. 时间维度上采样
        upsampled_spec_ri = F.interpolate(
            enhanced_spec_ri,
            scale_factor=(1, self.upscale_factor),
            mode='bilinear',
            align_corners=False
        )

        # 5. 合并回复数谱图
        upsampled_spec = torch.complex(upsampled_spec_ri[:, 0], upsampled_spec_ri[:, 1])

        # 6. iSTFT
        expected_output_len = x.shape[-1] * self.upscale_factor
        y_residual = torch.istft(
            upsampled_spec,
            n_fft=self.n_fft,
            hop_length=self.hop_length,
            win_length=self.win_length,
            window=self.window,
            length=expected_output_len,
            center=True,
        )

        # 7. 基线上采样
        x_baseline = F.interpolate(
            x.unsqueeze(1),
            scale_factor=self.upscale_factor,
            mode='linear',
            align_corners=False
        ).squeeze(1)

        # 8. 最终输出
        y_hat = x_baseline + y_residual
        return y_hat


# ==============================================================================
# 3. 数据处理/加载和训练代码
# ==============================================================================

def preprocess_audio_dataset(
        input_dir: str,
        output_x_dir: str,
        output_y_dir: str,
        target_sr_y: int,
        upscale_factor: int,
        bitrate_x: str = '64k',
        overwrite: bool = False
):
    os.makedirs(output_x_dir, exist_ok=True)
    os.makedirs(output_y_dir, exist_ok=True)

    audio_files = []
    for ext in ['*.wav', '*.flac', '*.aiff', '*.alac', '*.m4a']:
        audio_files.extend(glob.glob(os.path.join(input_dir, '**', ext), recursive=True))

    if not audio_files:
        print(f"No audio files found in {input_dir}. Please check the directory and file extensions.")
        return

    print(f"Found {len(audio_files)} audio files for preprocessing. Target SR {target_sr_y}.")

    encoders_x = ['libmp3lame', 'aac', 'libfdk_aac', 'libopus', 'flac']
    target_sr_x = target_sr_y // upscale_factor

    for i, input_filepath in enumerate(tqdm(audio_files, desc="Preprocessing Audio")):
        original_file_hash_base = hashlib.blake2s(input_filepath.encode('utf-8')).hexdigest()

        for ch_idx in ['FL', 'FR']:
            y_filename_unique = f"Y_{original_file_hash_base}_ch{ch_idx}_{target_sr_y}.wav"
            output_y_filepath = os.path.join(output_y_dir, y_filename_unique)

            if overwrite or not os.path.exists(output_y_filepath):
                cmd_y = [
                    './ffmpeg', '-y' if overwrite else '-n',
                    '-i', input_filepath,
                    '-af', f'pan=mono|c0={ch_idx}',
                    '-ar', str(target_sr_y),
                    '-loglevel', 'error',
                    output_y_filepath
                ]
                try:
                    subprocess.run(cmd_y, check=True, stdout=subprocess.PIPE, stderr=subprocess.PIPE)
                except subprocess.CalledProcessError as e:
                    print(f"Error generating Y for {input_filepath} channel {ch_idx}: {e.stderr.decode()}")
                    continue
                except FileNotFoundError:
                    print("FFmpeg command not found. Please ensure FFmpeg is installed and in your PATH.")
                    return

            for encoding_x in encoders_x:
                output_x_ext = '.wav'
                if 'mp3' in encoding_x:
                    output_x_ext = '.mp3'
                    if target_sr_x not in [8000, 11025, 12000, 16000, 22050, 24000, 32000, 44100, 48000]: continue
                elif 'flac' in encoding_x:
                    output_x_ext = '.flac'
                elif 'opus' in encoding_x:
                    output_x_ext = '.opus'
                    if target_sr_x not in [8000, 12000, 16000, 24000, 48000]: continue
                elif 'aac' in encoding_x:
                    output_x_ext = '.aac'
                    if target_sr_x not in [8000, 11025, 12000, 16000, 22050, 24000, 32000, 44100, 48000, 64000, 88200,
                                           96000]: continue

                x_filename_unique = f"{original_file_hash_base}_ch{ch_idx}_{target_sr_x}_{encoding_x}_{bitrate_x}{output_x_ext}"
                output_x_filepath = os.path.join(output_x_dir, x_filename_unique)

                if not overwrite and os.path.exists(output_x_filepath): continue

                cmd_x = [
                    './ffmpeg', '-y' if overwrite else '-n',
                    '-i', output_y_filepath,
                    '-ar', str(target_sr_x),
                    '-b:a', bitrate_x,
                    '-c:a', encoding_x,
                    '-loglevel', 'error',
                    output_x_filepath
                ]
                try:
                    subprocess.run(cmd_x, check=True, stdout=subprocess.PIPE, stderr=subprocess.PIPE)
                except subprocess.CalledProcessError as e:
                    print(f"Error generating X for {input_filepath} channel {ch_idx}: {e.stderr.decode()}")
                    continue
                except FileNotFoundError:
                    print("FFmpeg command not found. Please ensure FFmpeg is installed and in your PATH.")
                    return
    print("Preprocessing complete!")


class AudioDataset(torch.utils.data.Dataset):
    def __init__(self, x_dir, y_dir, segment_seconds, upscale_factor, hop_length):
        self.x_filepaths = []
        self.y_filepaths = []
        y_files_map = {os.path.basename(f): f for f in glob.glob(os.path.join(y_dir, '*.wav'))}
        for i in sorted(glob.glob(os.path.join(x_dir, '*.*'))):
            meta_x = os.path.basename(i).split('_')
            # Handle potential filename parsing issues
            if len(meta_x) < 3: continue
            try:
                y_target_sr = upscale_factor * int(meta_x[2])
            except ValueError:
                continue  # Skip if sample rate part is not an integer
            y_name = f'Y_{meta_x[0]}_{meta_x[1]}_{y_target_sr}.wav'
            if y_name in y_files_map:
                self.x_filepaths.append(i)
                self.y_filepaths.append(y_files_map[y_name])

        if not self.x_filepaths or not y_files_map:
            raise RuntimeError(
                f"No matching X/Y files found in {x_dir}, {y_dir}. Please ensure preprocessing was successful.")

        self.loaded_x_data = {}
        self.loaded_y_data = {}
        self.y_norm_sr = 64000
        self.x_norm_sr = self.y_norm_sr // upscale_factor

        for i in tqdm(self.y_filepaths, desc="Loading Y audio to RAM"):
            if i in self.loaded_y_data: continue
            y_tensor, y_sr = torchaudio.load(i)
            if y_sr != self.y_norm_sr:
                y_tensor = torchaudio.functional.resample(y_tensor, y_sr, self.y_norm_sr)
            self.loaded_y_data[i] = y_tensor.squeeze(0)
        for i in tqdm(self.x_filepaths, desc="Loading X audio to RAM"):
            x_tensor, x_sr = torchaudio.load(i)
            if x_sr != self.x_norm_sr:
                x__mid_sr = x_sr
                # if random.random() > 0.2:
                #     x__mid_sr = self.x_norm_sr // random.choice([1, 3, 5, 7, 11, 13])
                #     x_tensor = torchaudio.functional.resample(x_tensor, x_sr, x__mid_sr)
                x_tensor = torchaudio.functional.resample(x_tensor, x__mid_sr, self.x_norm_sr)
            self.loaded_x_data[i] = x_tensor.squeeze(0)

        if not self.loaded_x_data:
            raise RuntimeError("No audio data could be loaded. Check file integrity.")

        self.upscale_factor = upscale_factor
        self.segment_seconds = segment_seconds
        self.hop_length = hop_length

    def __len__(self):
        return len(self.x_filepaths) * int(60 * 4.0 / self.segment_seconds)

    def __getitem__(self, idx):
        mapping_idx = idx % len(self.x_filepaths)
        x_audio = self.loaded_x_data[self.x_filepaths[mapping_idx]]
        y_audio = self.loaded_y_data[self.y_filepaths[mapping_idx]]

        segment_samples_x = (int(self.segment_seconds * self.x_norm_sr) // self.hop_length) * self.hop_length
        segment_samples_y = segment_samples_x * self.upscale_factor

        if x_audio.shape[0] < segment_samples_x:
            x_segment = F.pad(x_audio, (0, segment_samples_x - x_audio.shape[0]))
            start_idx_x = 0
        else:
            start_idx_x = random.randint(0, x_audio.shape[0] - segment_samples_x)
            x_segment = x_audio[start_idx_x: start_idx_x + segment_samples_x]

        start_idx_y = start_idx_x * self.upscale_factor
        end_idx_y = start_idx_y + segment_samples_y

        if y_audio.shape[0] < end_idx_y:
            y_segment = F.pad(y_audio[start_idx_y:], (0, end_idx_y - y_audio.shape[0]))
        else:
            y_segment = y_audio[start_idx_y: end_idx_y]

        if x_segment.shape[0] != segment_samples_x: x_segment = F.pad(x_segment,
                                                                      (0, segment_samples_x - x_segment.shape[0]))
        if y_segment.shape[0] != segment_samples_y: y_segment = F.pad(y_segment,
                                                                      (0, segment_samples_y - y_segment.shape[0]))

        return x_segment, y_segment


class MultiResolutionSTFTLoss(nn.Module):
    def __init__(self,
                 fft_sizes=(1024, 2048, 512),
                 hop_sizes=(256, 512, 128),
                 win_lengths=(1024, 2048, 512),
                 window=torch.hann_window,
                 comp_weight=1.0,
                 log_mag_weight=1.0,
                 derivative_loss_weight=1.0,
                 epsilon=1e-8
                 ):
        super().__init__()
        assert len(fft_sizes) == len(hop_sizes) == len(win_lengths)
        self.fft_sizes = fft_sizes
        self.hop_sizes = hop_sizes
        self.win_lengths = win_lengths
        self.window_fn = window
        self.comp_weight = comp_weight
        self.log_mag_weight = log_mag_weight
        self.derivative_loss_weight = derivative_loss_weight
        self.epsilon = epsilon
        for i in range(len(fft_sizes)):
            self.register_buffer(f'window_{i}', self.window_fn(win_lengths[i]))

    def forward(self, x, y):
        total_loss = 0.0
        for i in range(len(self.fft_sizes)):
            n_fft, hop_length, win_length = self.fft_sizes[i], self.hop_sizes[i], self.win_lengths[i]
            window = getattr(self, f'window_{i}')
            x_stft = torch.stft(x, n_fft, hop_length, win_length, window=window, return_complex=True,
                                pad_mode='reflect', center=True)
            y_stft = torch.stft(y, n_fft, hop_length, win_length, window=window, return_complex=True,
                                pad_mode='reflect', center=True)

            min_dim_t = min(x_stft.shape[-1], y_stft.shape[-1])
            x_stft, y_stft = x_stft[..., :min_dim_t], y_stft[..., :min_dim_t]

            x_mag_deriv = torch.diff(torch.abs(x_stft), dim=1)
            y_mag_deriv = torch.diff(torch.abs(y_stft), dim=1)
            deriv_loss = F.l1_loss(x_mag_deriv, y_mag_deriv)
            total_loss += deriv_loss * self.derivative_loss_weight

            x_mag, y_mag = torch.abs(x_stft), torch.abs(y_stft)
            log_mag_loss = F.l1_loss(torch.log(x_mag + self.epsilon),
                                     torch.log(y_mag + self.epsilon)) * self.log_mag_weight
            total_loss += log_mag_loss

            x_comp = torch.cat([x_stft.real, x_stft.imag], dim=1)
            y_comp = torch.cat([y_stft.real, y_stft.imag], dim=1)
            comp_loss = F.l1_loss(x_comp, y_comp) * self.comp_weight
            total_loss += comp_loss
        return total_loss


def train_model(
        generator: nn.Module,
        discriminator: nn.Module,
        dataset: torch.utils.data.Dataset,
        output_model_path: str,
        batch_size: int = 4,
        epochs: int = 100,
        learning_rate: float = 1e-4,
        warmup_epochs: int = 5,
        device: str = 'cpu',
        save_interval: int = 10,
        max_grad_norm: float = 1.0,
        time_loss_weight: float = 1.0,
        stft_loss_params: dict = None,
        adv_loss_weight: float = 1.0,
        feat_match_weight: float = 10.0,
        gan_start_epoch: int = 10
):
    num_workers = min(os.cpu_count(), 8)
    print(f"Using {num_workers} data loader workers.")

    dataloader = torch.utils.data.DataLoader(
        dataset, batch_size=batch_size, shuffle=True,
        num_workers=num_workers, pin_memory=True, drop_last=True
    )
    optimizer_g = torch.optim.AdamW(generator.parameters(), lr=learning_rate, betas=(0.8, 0.99))
    optimizer_d = torch.optim.AdamW(discriminator.parameters(), lr=learning_rate / 2, betas=(0.8, 0.99))
    scheduler_g = torch.optim.lr_scheduler.CosineAnnealingLR(optimizer_g, T_max=epochs - warmup_epochs, eta_min=1e-6)
    scheduler_d = torch.optim.lr_scheduler.CosineAnnealingLR(optimizer_d, T_max=epochs - warmup_epochs, eta_min=1e-6)

    criterion_time = nn.L1Loss()
    if stft_loss_params is None:
        stft_loss_params = {
            'fft_sizes': (512, 1024, 2048), 'hop_sizes': (128, 256, 512), 'win_lengths': (512, 1024, 2048),
            'log_mag_weight': 1.0, 'comp_weight': 1.0, 'derivative_loss_weight': 1.0,
        }
    criterion_stft = MultiResolutionSTFTLoss(**stft_loss_params).to(device)

    generator.to(device)
    discriminator.to(device)
    generator.train()
    discriminator.train()
    scaler = torch.amp.GradScaler(enabled=(device == 'cuda'))

    print(f"Training started on {device} with AMP {'enabled' if device == 'cuda' else 'disabled'}...")
    print(f"GAN training will start at epoch {gan_start_epoch}.")

    for epoch in range(epochs):
        # 手动学习率预热
        if epoch < warmup_epochs:
            current_g_lr = learning_rate * (epoch + 1) / warmup_epochs
            current_d_lr = (learning_rate / 2) * (epoch + 1) / warmup_epochs
            for param_group in optimizer_g.param_groups: param_group['lr'] = current_g_lr
            for param_group in optimizer_d.param_groups: param_group['lr'] = current_d_lr

        total_d_loss_epoch = 0
        total_g_loss_epoch = 0
        pbar = tqdm(dataloader, desc=f"Epoch {epoch + 1}/{epochs}", ncols=150)
        for batch_idx, (x_batch, y_batch) in enumerate(pbar):
            x_batch, y_batch = x_batch.to(device), y_batch.to(device)

            with torch.amp.autocast(device_type=device, enabled=(device == 'cuda')):
                output_audio = generator(x_batch)
                min_len = min(output_audio.shape[-1], y_batch.shape[-1])
                output_audio_synced = output_audio[..., :min_len]
                y_batch_synced = y_batch[..., :min_len]

            # ---- 训练判别器 ----
            optimizer_d.zero_grad()
            with torch.amp.autocast(device_type=device, enabled=(device == 'cuda')):
                fake_audio_detached = output_audio_synced.detach()
                real_scores, _ = discriminator(y_batch_synced.unsqueeze(1))
                fake_scores, _ = discriminator(fake_audio_detached.unsqueeze(1))

                loss_d = torch.tensor(0.0, device=device)
                if epoch >= gan_start_epoch:
                    loss_d = discriminator_loss(real_scores, fake_scores)

            if epoch >= gan_start_epoch:
                scaler.scale(loss_d).backward()
                scaler.unscale_(optimizer_d)
                torch.nn.utils.clip_grad_norm_(discriminator.parameters(), max_norm=max_grad_norm)
                scaler.step(optimizer_d)

            # ---- 训练生成器 ----
            optimizer_g.zero_grad()
            with torch.amp.autocast(device_type=device, enabled=(device == 'cuda')):
                loss_time = criterion_time(output_audio_synced, y_batch_synced) * time_loss_weight
                loss_stft = criterion_stft(output_audio_synced, y_batch_synced)
                loss_recon = loss_time + loss_stft

                loss_g_total = loss_recon
                loss_g_adv = torch.tensor(0.0, device=device)
                loss_g_fm = torch.tensor(0.0, device=device)

                if epoch >= gan_start_epoch:
                    real_scores_g, real_fmaps_g = discriminator(y_batch_synced.unsqueeze(1))
                    fake_scores_g, fake_fmaps_g = discriminator(output_audio_synced.unsqueeze(1))

                    loss_g_adv = generator_adversarial_loss(fake_scores_g) * adv_loss_weight
                    loss_g_fm = feature_matching_loss(real_fmaps_g, fake_fmaps_g) * feat_match_weight
                    loss_g_total += loss_g_adv + loss_g_fm

            scaler.scale(loss_g_total).backward()
            scaler.unscale_(optimizer_g)
            torch.nn.utils.clip_grad_norm_(generator.parameters(), max_norm=max_grad_norm)
            scaler.step(optimizer_g)
            scaler.update()

            total_d_loss_epoch += loss_d.item()
            total_g_loss_epoch += loss_g_total.item()
            pbar.set_postfix(
                _D=f"{loss_d.item():.3f}",
                _G=f"{loss_g_total.item():.3f}",
                T=f"{loss_time.item():.3f}",
                S=f"{loss_stft.item():.3f}",
                Adv=f"{loss_g_adv.item():.3f}",
                FM=f"{loss_g_fm.item():.3f}"
            )

        avg_d_loss = total_d_loss_epoch / len(dataloader)
        avg_g_loss = total_g_loss_epoch / len(dataloader)
        current_lr_g = optimizer_g.param_groups[0]['lr']
        current_lr_d = optimizer_d.param_groups[0]['lr']
        print(
            f"Epoch {epoch + 1} finished. Avg _D: {avg_d_loss}, Avg _G: {avg_g_loss}, "
            f"LR_G: {current_lr_g:.6f}, LR_D: {current_lr_d:.6f}"
        )

        # 在预热期之后，开始使用调度器
        if epoch >= warmup_epochs:
            scheduler_g.step()
            scheduler_d.step()

        # 保存模型
        if (epoch + 1) % save_interval == 0:
            torch.save(generator.state_dict(), f"{output_model_path}_G_epoch_{epoch + 1}.pth")
            torch.save(discriminator.state_dict(), f"{output_model_path}_D_epoch_{epoch + 1}.pth")
            print(f"Models saved to {output_model_path}_(G/D)_epoch_{epoch + 1}.pth")

    torch.save(generator.state_dict(), f"{output_model_path}_G_final.pth")
    torch.save(discriminator.state_dict(), f"{output_model_path}_D_final.pth")
    print(f"Training complete. Final models saved to {output_model_path}_(G/D)_final.pth")


# ==============================================================================
# 4. 主执行代码
# ==============================================================================

if __name__ == "__main__":
    torch.backends.cudnn.benchmark = True
    print(f"torch.backends.cudnn.benchmark set to {torch.backends.cudnn.benchmark}")

    # --- 配置 ---
    RAW_DATA_DIR = 'dataset'
    PREPROCESSED_X_DIR = f'preprocessed_{RAW_DATA_DIR}/low_res_X'
    PREPROCESSED_Y_DIR = f'preprocessed_{RAW_DATA_DIR}/high_res_Y'

    UPSCALE_FACTOR = 2
    TARGET_SR_Y = [96000]  # 仅使用一个目标采样率以简化
    BITRATE_X = '128k'

    # 训练参数
    TARGET_SEGMENT_SECONDS_X = 1.0
    BATCH_SIZE = 256
    EPOCHS = 1000
    LEARNING_RATE = 1e-4
    WARMUP_EPOCHS = 5
    MODEL_OUTPUT_DIR = 'checkpoints'
    MODEL_NAME = f'audio_upsampler_{UPSCALE_FACTOR}x_bae_ex_gan'
    SAVE_INTERVAL = 10
    MAX_GRAD_NORM = 1.0

    # GAN 相关参数
    GAN_TRAINING_START_EPOCH = 10  # 在第10个 epoch 后开始对抗训练

    # 损失权重
    TIME_LOSS_WEIGHT = 100.0
    STFT_LOG_MAG_WEIGHT = 1.0
    STFT_COMP_WEIGHT = 1.0
    STFT_DERIV_WEIGHT = 0.5
    ADV_LOSS_WEIGHT = 1.0  # 对抗损失权重
    FEAT_MATCH_WEIGHT = 10.0  # 特征匹配损失权重

    DEVICE = 'cuda' if torch.cuda.is_available() else 'cpu'
    print(f"Using device: {DEVICE}")
    os.makedirs(MODEL_OUTPUT_DIR, exist_ok=True)

    # --- 步骤 1: 预处理数据集 (如果需要) ---
    print("\n--- Starting Data Preprocessing ---")
    for i_target_sr_y in TARGET_SR_Y:
        preprocess_audio_dataset(
            input_dir=RAW_DATA_DIR,
            output_x_dir=PREPROCESSED_X_DIR,
            output_y_dir=PREPROCESSED_Y_DIR,
            target_sr_y=i_target_sr_y,
            upscale_factor=UPSCALE_FACTOR,
            bitrate_x=BITRATE_X,
            overwrite=False
        )
    print("--- Data Preprocessing Skipped (manual run required) ---\n")

    # --- 步骤 2: 初始化模型 ---
    generator = AudioUpsampler(upscale_factor=UPSCALE_FACTOR)
    print(f"Model {type(generator).__name__} initialized for {UPSCALE_FACTOR}x upsampling.")

    # 判别器的 STFT 参数
    discriminator_stft_params = [
        {'n_fft': 1024, 'hop_length': 256, 'win_length': 1024},
        {'n_fft': 1536, 'hop_length': 480, 'win_length': 1536},
    ]
    discriminator = MultiResolutionSTFTDiscriminator(stft_params=discriminator_stft_params)
    print(f"Discriminator {type(discriminator).__name__} initialized.")

    # 计算模型参数量
    total_params_g = sum(p.numel() for p in generator.parameters() if p.requires_grad)
    total_params_d = sum(p.numel() for p in discriminator.parameters() if p.requires_grad)
    print(f"Total trainable parameters (Generator): {total_params_g / 1e6:.2f}M")
    print(f"Total trainable parameters (Discriminator): {total_params_d / 1e6:.2f}M")

    # 模拟一次前向传播，检查维度是否匹配
    try:
        print("Performing a test forward pass...")
        test_input = torch.randn(BATCH_SIZE, int(47616 * TARGET_SEGMENT_SECONDS_X))  # 1秒 48kHz 的音频
        _ = generator(test_input)
        print("Test forward pass successful. Dimensions seem correct.")
    except Exception as e:
        print(f"Error during test forward pass: {e}")

    # --- 步骤 3: 准备训练数据集 ---
    # hop_length 需要与模型内部的 STFT 参数一致，以确保数据分段正确
    dataset_hop_length = generator.hop_length

    train_dataset = AudioDataset(
        x_dir=PREPROCESSED_X_DIR,
        y_dir=PREPROCESSED_Y_DIR,
        segment_seconds=TARGET_SEGMENT_SECONDS_X,
        upscale_factor=UPSCALE_FACTOR,
        hop_length=dataset_hop_length
    )
    print(f"Dataset loaded. Number of training steps per epoch (heuristic): {len(train_dataset)}")

    # --- 步骤 4: 训练模型 ---
    train_model(
        generator=generator,
        discriminator=discriminator,
        dataset=train_dataset,
        output_model_path=os.path.join(MODEL_OUTPUT_DIR, MODEL_NAME),
        batch_size=BATCH_SIZE,
        epochs=EPOCHS,
        learning_rate=LEARNING_RATE,
        warmup_epochs=WARMUP_EPOCHS,
        device=DEVICE,
        save_interval=SAVE_INTERVAL,
        max_grad_norm=MAX_GRAD_NORM,
        time_loss_weight=TIME_LOSS_WEIGHT,
        stft_loss_params={
            "fft_sizes": (512, 1024, 2048, 4096),
            "hop_sizes": (50, 120, 240, 480),
            "win_lengths": (240, 600, 1200, 2400),
            "log_mag_weight": STFT_LOG_MAG_WEIGHT,
            "comp_weight": STFT_COMP_WEIGHT,
            "derivative_loss_weight": STFT_DERIV_WEIGHT,
        },
        adv_loss_weight=ADV_LOSS_WEIGHT,
        feat_match_weight=FEAT_MATCH_WEIGHT,
        gan_start_epoch=GAN_TRAINING_START_EPOCH
    )

    print("\nTraining script finished.")