188 lines
7.0 KiB
Python
188 lines
7.0 KiB
Python
# Copyright (c) 2025 Salvador E. Tropea
|
|
# Copyright (c) 2025 Instituto Nacional de Tecnología Industrial
|
|
# License: GPLv3
|
|
# Project: ComfyUI-AudioSeparation
|
|
# Developed using:
|
|
# - Netron to inspect the topology
|
|
# - onnx2pytorch.ConvertModel to find mapping details
|
|
# - Gemini 2.5 Pro to analyze the network, write the code and debug it
|
|
# I never saw the original implementation, this is a reconstruction from Kim_Vocals_2.onnx
|
|
#
|
|
# Found geometries:
|
|
#
|
|
# dim_f | channels | stages | ~params
|
|
# -------|----------|--------|------------
|
|
# 3072 | 48 | 5 | 16 684 228
|
|
# 2560 | 48 | 5 | 14 763 012
|
|
# 2048 | 48 | 5 | 13 191 108
|
|
# 2048 | 32 | 5 | 7 420 548
|
|
# 2048 | 32 | 4 | 5 478 276
|
|
from torch import nn
|
|
|
|
|
|
class FrequencyBranch(nn.Module):
|
|
""" Frequency-domain branch with Linear -> BatchNorm2d -> ReLU sequences. """
|
|
def __init__(self, channels, freq_dim, hidden_dim, bn_eps):
|
|
super().__init__()
|
|
self.sequence = nn.Sequential(
|
|
nn.Linear(freq_dim, hidden_dim, bias=False),
|
|
# Using the standard, verified nn.BatchNorm2d
|
|
nn.BatchNorm2d(num_features=channels, eps=bn_eps),
|
|
nn.ReLU(True),
|
|
nn.Linear(hidden_dim, freq_dim, bias=False),
|
|
nn.BatchNorm2d(num_features=channels, eps=bn_eps),
|
|
nn.ReLU(True)
|
|
)
|
|
|
|
def forward(self, x):
|
|
# Relies on PyTorch's nn.Linear broadcasting over the first 3 dims of (B,C,T,F)
|
|
# And nn.BatchNorm2d operating on the C dimension of the 4D tensor.
|
|
return self.sequence(x)
|
|
|
|
|
|
class TimeBranch(nn.Module):
|
|
"""
|
|
Time-domain branch using 3x3 convolutions.
|
|
"""
|
|
def __init__(self, channels):
|
|
super().__init__()
|
|
self.sequence = nn.Sequential(
|
|
nn.Conv2d(channels, channels, kernel_size=3, padding='same', bias=True),
|
|
nn.ReLU(True),
|
|
nn.Conv2d(channels, channels, kernel_size=3, padding='same', bias=True),
|
|
nn.ReLU(True),
|
|
nn.Conv2d(channels, channels, kernel_size=3, padding='same', bias=True),
|
|
nn.ReLU(True)
|
|
)
|
|
|
|
def forward(self, x):
|
|
return self.sequence(x)
|
|
|
|
|
|
class TDF_Block(nn.Module):
|
|
""" The main processing block, combining time and frequency branches.
|
|
This is a sequential-residual block, as shown in the ONNX graph. """
|
|
def __init__(self, channels, freq_dim, hidden_dim, bn_eps):
|
|
super().__init__()
|
|
self.time_branch = TimeBranch(channels)
|
|
self.freq_branch = FrequencyBranch(channels, freq_dim, hidden_dim, bn_eps)
|
|
|
|
def forward(self, x):
|
|
# 1. The input 'x' goes through the time branch first.
|
|
time_out = self.time_branch(x)
|
|
|
|
# 2. The output of the time branch is then fed into the frequency branch.
|
|
freq_out = self.freq_branch(time_out)
|
|
|
|
# 3. The final result is a residual connection:
|
|
# Output of Time Branch + Output of Frequency Branch
|
|
return time_out + freq_out
|
|
|
|
|
|
class Transpose(nn.Module):
|
|
""" A simple nn.Module wrapper for the permute operation """
|
|
def __init__(self, dims):
|
|
super().__init__()
|
|
self.dims = dims
|
|
|
|
def forward(self, x):
|
|
return x.permute(self.dims)
|
|
|
|
|
|
class MDX_Net(nn.Module):
|
|
"""
|
|
The complete U-Net architecture.
|
|
Fully parametric for frequency bins, channels, and number of stages.
|
|
This version uses your elegant interlaced ModuleList design for a clean,
|
|
dynamic structure that correctly matches the ONNX graph order.
|
|
"""
|
|
def __init__(self, dim_f=3072, ch=48, num_stages=5):
|
|
super().__init__()
|
|
# Validate input
|
|
if num_stages < 1 or num_stages > 12:
|
|
raise ValueError(f"num_stages must be between 1 and 12, but got {num_stages}")
|
|
|
|
self.num_stages = num_stages
|
|
|
|
# Define shared BatchNorm parameters
|
|
BN_EPS = 9.999999747378752e-06
|
|
freq_hidden_dim = dim_f // 8
|
|
# Allow others to know our creation parameters
|
|
self.dim_f = dim_f
|
|
self.ch = ch
|
|
self.num_stages = num_stages
|
|
|
|
# --- Initial Chain (always exists) ---
|
|
self.initial_conv = nn.Conv2d(4, ch, 1, bias=True)
|
|
self.initial_relu = nn.ReLU(True)
|
|
self.initial_transpose = Transpose(dims=(0, 1, 3, 2))
|
|
|
|
# --- Encoder Path with Interlaced Layers ---
|
|
# We create all stages and downsamplers as ModuleLists
|
|
self.enc_stages = nn.ModuleList()
|
|
|
|
# This list defines the channel progression
|
|
# e.g., for ch=48: (48, 96, 144, 192, 240, 288)
|
|
channels = [ch * (i + 1) for i in range(num_stages + 1)]
|
|
|
|
for i in range(num_stages):
|
|
# Append the TDF_Block stage
|
|
self.enc_stages.append(TDF_Block(channels[i], dim_f // (2**i), freq_hidden_dim // (2**i), BN_EPS))
|
|
# Append the downsampling block immediately after
|
|
self.enc_stages.append(nn.Sequential(nn.Conv2d(channels[i], channels[i+1], 2, 2, bias=True), nn.ReLU(True)))
|
|
|
|
# --- Bottleneck ---
|
|
bottleneck_in_ch = channels[num_stages]
|
|
self.bottleneck = TDF_Block(bottleneck_in_ch, dim_f // (2**num_stages), freq_hidden_dim // (2**num_stages), BN_EPS)
|
|
|
|
# --- Decoder Path with Interlaced Layers ---
|
|
self.dec_stages = nn.ModuleList()
|
|
for i in range(num_stages):
|
|
dec_idx = num_stages - 1 - i
|
|
# Upsampler takes bottleneck/previous stage channels and outputs encoder stage channels
|
|
in_ch = channels[dec_idx + 1]
|
|
out_ch = channels[dec_idx]
|
|
|
|
# Append the upsampling block
|
|
seq = nn.Sequential(nn.ConvTranspose2d(in_ch, out_ch, 2, 2, bias=True),
|
|
nn.BatchNorm2d(out_ch, eps=BN_EPS),
|
|
nn.ReLU(True))
|
|
self.dec_stages.append(seq)
|
|
# Append the TDF_Block stage immediately after
|
|
self.dec_stages.append(TDF_Block(out_ch, dim_f // (2**dec_idx), freq_hidden_dim // (2**dec_idx), BN_EPS))
|
|
|
|
# --- Final Chain (always exists) ---
|
|
self.final_transpose = Transpose(dims=(0, 1, 3, 2))
|
|
self.final_conv = nn.Conv2d(ch, 4, 1, bias=True)
|
|
|
|
def forward(self, x):
|
|
# Initial processing
|
|
x = self.initial_conv(x)
|
|
x = self.initial_relu(x)
|
|
x = self.initial_transpose(x)
|
|
|
|
# --- Dynamic Encoder Path ---
|
|
skip_connections = []
|
|
# Encoder runs through the interlaced list
|
|
for i in range(0, self.num_stages*2, 2):
|
|
s = self.enc_stages[i](x)
|
|
skip_connections.append(s)
|
|
x = self.enc_stages[i+1](s)
|
|
|
|
# --- Bottleneck ---
|
|
x = self.bottleneck(x)
|
|
|
|
# --- Dynamic Decoder Path ---
|
|
skip_connections.reverse() # Reverse for easy lookup
|
|
# Decoder also runs through its interlaced list
|
|
for i in range(0, self.num_stages * 2, 2):
|
|
x = self.dec_stages[i](x)
|
|
x = x * skip_connections[i//2]
|
|
x = self.dec_stages[i+1](x)
|
|
|
|
# Final processing
|
|
output = self.final_transpose(x)
|
|
output = self.final_conv(output)
|
|
|
|
return output
|