Spaces:
Runtime error
Runtime error
| # Modelo Attention Res-UNet 2D para Audio Super Resolution | |
| import torch | |
| import torch.nn as nn | |
| import torch.nn.functional as F | |
| class AttentionGate(nn.Module): | |
| """Módulo de atención para ponderar skip connections.""" | |
| def __init__(self, F_g, F_l, F_int): | |
| """ | |
| Inicializa el módulo de atención. | |
| Args: | |
| F_g (int): Número de canales del gating signal. | |
| F_l (int): Número de canales de la skip connection. | |
| F_int (int): Número de canales intermedios. | |
| """ | |
| super().__init__() | |
| # Uso de Attention Gates basado en el paper "https://arxiv.org/abs/1804.03999" | |
| self.W_g = nn.Conv2d(F_g, F_int, kernel_size=1) | |
| self.W_x = nn.Conv2d(F_l, F_int, kernel_size=1) | |
| self.psi = nn.Conv2d(F_int, 1, kernel_size=1) | |
| self.relu = nn.ReLU(inplace=True) | |
| self.sigmoid = nn.Sigmoid() | |
| def forward(self, g, x): | |
| # Interpolar g para que tenga el mismo tamaño que x | |
| if g.shape[-2:] != x.shape[-2:]: | |
| g = F.interpolate(g, size=x.shape[-2:], mode="bilinear", align_corners=False) | |
| # Calcular atención | |
| att = self.sigmoid(self.psi(self.relu(self.W_g(g) + self.W_x(x)))) | |
| return x * att | |
| class DilatedBlock(nn.Module): | |
| """Bloque con capas convolucionales dilatadas para capturar contexto de largo alcance.""" | |
| def __init__(self, channels): | |
| """ | |
| Inicializa el bloque dilatado. | |
| Args: | |
| channels (int): Número de canales de entrada y salida. | |
| """ | |
| super().__init__() | |
| self.net = nn.Sequential( | |
| nn.Conv2d(channels, channels, kernel_size=(7,3), padding=(3,1), dilation=(1,1)), | |
| nn.LeakyReLU(0.2, inplace=True), | |
| nn.Conv2d(channels, channels, kernel_size=(7,3), padding=(6,2), dilation=(2,2)), | |
| nn.LeakyReLU(0.2, inplace=True), | |
| nn.Conv2d(channels, channels, kernel_size=(7,3), padding=(12,4), dilation=(4,4)), | |
| nn.LeakyReLU(0.2, inplace=True), | |
| ) | |
| def forward(self, x): | |
| return self.net(x) + x | |
| class ResBlock(nn.Module): | |
| """Bloque Residual para Attention Res-UNet.""" | |
| def __init__(self, in_ch, out_ch): | |
| """ | |
| Inicializa el bloque residual. | |
| Args: | |
| in_ch (int): Número de canales de entrada. | |
| out_ch (int): Número de canales de salida. | |
| """ | |
| super().__init__() | |
| self.conv1 = nn.Conv2d(in_ch, out_ch, kernel_size=(7,3), padding=(3,1)) | |
| self.norm1 = nn.GroupNorm(out_ch//4, out_ch) # GroupNorm para batchs pequeños | |
| self.relu1 = nn.LeakyReLU(0.2, inplace=True) | |
| self.conv2 = nn.Conv2d(out_ch, out_ch, kernel_size=(7,3), padding=(3,1)) | |
| self.norm2 = nn.GroupNorm(out_ch//4, out_ch) | |
| self.relu2 = nn.LeakyReLU(0.2, inplace=True) | |
| # Skip connection | |
| if in_ch != out_ch: | |
| self.skip = nn.Sequential( | |
| nn.Conv2d(in_ch, out_ch, kernel_size=1, bias=False), | |
| nn.GroupNorm(out_ch//4, out_ch) | |
| ) | |
| else: | |
| self.skip = nn.Identity() | |
| def forward(self, x): | |
| identity = self.skip(x) | |
| out = self.conv1(x) | |
| out = self.norm1(out) | |
| out = self.relu1(out) | |
| out = self.conv2(out) | |
| out = self.norm2(out) | |
| out += identity | |
| out = self.relu2(out) | |
| return out | |
| class UNetAudio2D(nn.Module): | |
| """ | |
| Modelo Attention Res-UNet 2D para super resolución de audio. | |
| Entrada y salida con shape (B, 2, F, T). | |
| """ | |
| def __init__(self): | |
| """Inicializa la arquitectura UNet con encoder, bottleneck y decoder.""" | |
| super().__init__() | |
| # Encoder | |
| # Entrada: (B, 2, F, T) | |
| self.enc1 = ResBlock(2, 32) | |
| self.down1 = nn.Conv2d(32, 32, kernel_size=(4,4), stride=(2,2), padding=(1,1)) # Strided conv | |
| self.enc2 = ResBlock(32, 64) | |
| self.down2 = nn.Conv2d(64, 64, kernel_size=(4,4), stride=(2,2), padding=(1,1)) | |
| self.enc3 = ResBlock(64, 128) | |
| self.down3 = nn.Conv2d(128, 128, kernel_size=(4,4), stride=(2,2), padding=(1,1)) | |
| self.enc4 = ResBlock(128, 256) | |
| self.down4 = nn.Conv2d(256, 256, kernel_size=(4,4), stride=(2,2), padding=(1,1)) | |
| # Bottleneck | |
| self.bottleneck_conv = ResBlock(256, 512) | |
| self.bottleneck_dilated = DilatedBlock(512) | |
| # Decoder | |
| self.up4 = self.up_block(512,256) | |
| self.dec4 = ResBlock(512,256) | |
| self.up3 = self.up_block(256,128) | |
| self.dec3 = ResBlock(256,128) | |
| self.up2 = self.up_block(128,64) | |
| self.dec2 = ResBlock(128,64) | |
| self.up1 = self.up_block(64,32) | |
| self.dec1 = ResBlock(64,32) | |
| # Attention gates | |
| self.att4 = AttentionGate(256,256,128) | |
| self.att3 = AttentionGate(128,128,64) | |
| self.att2 = AttentionGate(64,64,32) | |
| self.att1 = AttentionGate(32,32,16) | |
| # Output | |
| self.final = nn.Conv2d(32,2,kernel_size=1) | |
| def up_block(self, in_ch, out_ch): | |
| """ | |
| Crea un bloque de upsampling con ConvTranspose2d. | |
| Args: | |
| in_ch (int): Número de canales de entrada. | |
| out_ch (int): Número de canales de salida. | |
| Returns: | |
| nn.Sequential: Bloque de upsampling en frecuencia y tiempo. | |
| """ | |
| return nn.Sequential( | |
| nn.ConvTranspose2d(in_ch, out_ch, kernel_size=(4,4), stride=(2,2), padding=(1,1)), | |
| nn.LeakyReLU(0.2, inplace=True) | |
| ) | |
| def forward(self, x): | |
| # Encoder | |
| e1 = self.enc1(x) | |
| e2 = self.enc2(self.down1(e1)) | |
| e3 = self.enc3(self.down2(e2)) | |
| e4 = self.enc4(self.down3(e3)) | |
| # Bottleneck | |
| b = self.bottleneck_conv(self.down4(e4)) | |
| b = self.bottleneck_dilated(b) | |
| # Decoder con skip connections y attention gates | |
| up4 = self.up4(b) | |
| d4 = self.dec4(torch.cat([up4, self.att4(up4, e4)], dim=1)) | |
| up3 = self.up3(d4) | |
| d3 = self.dec3(torch.cat([up3, self.att3(up3, e3)], dim=1)) | |
| up2 = self.up2(d3) | |
| d2 = self.dec2(torch.cat([up2, self.att2(up2, e2)], dim=1)) | |
| up1 = self.up1(d2) | |
| d1 = self.dec1(torch.cat([up1, self.att1(up1, e1)], dim=1)) | |
| return self.final(d1) + x |