import os import torch import torch.nn as nn from torch.utils.data import Dataset import soundfile as sf from huggingface_hub import PyTorchModelHubMixin # ===================================================================== # 1. ARCHITECTURE: Multi-Masking U-Net # ===================================================================== class UNetBlock(nn.Module): def __init__(self, in_channels, out_channels): super().__init__() self.conv = nn.Sequential( nn.Conv2d(in_channels, out_channels, kernel_size=3, padding=1, bias=False), nn.GroupNorm(num_groups=4, num_channels=out_channels), nn.LeakyReLU(0.2, inplace=True), nn.Conv2d(out_channels, out_channels, kernel_size=3, padding=1, bias=False), nn.GroupNorm(num_groups=4, num_channels=out_channels), nn.LeakyReLU(0.2, inplace=True) ) def forward(self, x): return self.conv(x) class MultiStemSeparationModel(nn.Module, PyTorchModelHubMixin): def __init__(self, in_channels: int = 1, num_stems: int = 6, base_features: int = 16): super().__init__() self.in_channels = in_channels self.num_stems = num_stems self.base_features = base_features # Downsampling Encoder self.enc1 = UNetBlock(in_channels, base_features) self.pool1 = nn.MaxPool2d(kernel_size=2, stride=2) self.enc2 = UNetBlock(base_features, base_features * 2) self.pool2 = nn.MaxPool2d(kernel_size=2, stride=2) self.enc3 = UNetBlock(base_features * 2, base_features * 4) self.pool3 = nn.MaxPool2d(kernel_size=2, stride=2) # Bottleneck self.bottleneck = UNetBlock(base_features * 4, base_features * 8) # Upsampling Decoder self.up3 = nn.ConvTranspose2d(base_features * 8, base_features * 4, kernel_size=2, stride=2) self.dec3 = UNetBlock(base_features * 8, base_features * 4) self.up2 = nn.ConvTranspose2d(base_features * 4, base_features * 2, kernel_size=2, stride=2) self.dec2 = UNetBlock(base_features * 4, base_features * 2) self.up1 = nn.ConvTranspose2d(base_features * 2, base_features, kernel_size=2, stride=2) self.dec1 = UNetBlock(base_features * 2, base_features) # Output projection self.final_conv = nn.Conv2d(base_features, num_stems, kernel_size=1) self.activation = nn.Softmax(dim=1) def forward(self, x): # Encoder x1 = self.enc1(x) p1 = self.pool1(x1) x2 = self.enc2(p1) p2 = self.pool2(x2) x3 = self.enc3(p2) p3 = self.pool3(x3) # Bottleneck b = self.bottleneck(p3) # Decoder u3 = self.up3(b) if u3.shape != x3.shape: u3 = nn.functional.interpolate(u3, size=x3.shape[2:]) m3 = torch.cat([u3, x3], dim=1) d3 = self.dec3(m3) u2 = self.up2(d3) if u2.shape != x2.shape: u2 = nn.functional.interpolate(u2, size=x2.shape[2:]) m2 = torch.cat([u2, x2], dim=1) d2 = self.dec2(m2) u1 = self.up1(d2) if u1.shape != x1.shape: u1 = nn.functional.interpolate(u1, size=x1.shape[2:]) m1 = torch.cat([u1, x1], dim=1) d1 = self.dec1(m1) masks = self.activation(self.final_conv(d1)) separated_stems = x * masks return separated_stems, masks # ===================================================================== # 2. DATASET: Torchaudio-Free Loader # ===================================================================== class MultiStemDataset(Dataset): def __init__(self, dataset_dir, stems_list, sample_rate=16000, duration_seconds=4.0): super().__init__() self.dataset_dir = dataset_dir self.stems_list = stems_list self.sample_rate = sample_rate self.num_samples = int(sample_rate * duration_seconds) self.song_folders = sorted([ os.path.join(dataset_dir, d) for d in os.listdir(dataset_dir) if os.path.isdir(os.path.join(dataset_dir, d)) ]) def __len__(self): return len(self.song_folders) def _load_and_pad_crop(self, filepath): if not os.path.exists(filepath): return torch.zeros((1, self.num_samples)) try: # Read with soundfile (bypassing torchaudio) data, sr = sf.read(filepath) waveform = torch.tensor(data, dtype=torch.float32) # Mono conversion if len(waveform.shape) > 1: waveform = torch.mean(waveform, dim=1, keepdim=True) waveform = waveform.t() else: waveform = waveform.unsqueeze(0) # Mathematical resampling using PyTorch 1D interpolation if sr != self.sample_rate: target_len = int(waveform.shape[1] * self.sample_rate / sr) # Reshape to [1, 1, samples] for interpolation, then back waveform = nn.functional.interpolate( waveform.unsqueeze(0), size=target_len, mode='linear', align_corners=False ).squeeze(0) except Exception as e: print(f"Failed to load audio {filepath}: {e}") return torch.zeros((1, self.num_samples)) length = waveform.shape[1] if length > self.num_samples: waveform = waveform[:, :self.num_samples] elif length < self.num_samples: padding = self.num_samples - length waveform = torch.nn.functional.pad(waveform, (0, padding)) return waveform def __getitem__(self, idx): song_path = self.song_folders[idx] stem_tensors = [] for stem_name in self.stems_list: stem_file = os.path.join(song_path, f"{stem_name}.wav") stem_wav = self._load_and_pad_crop(stem_file) stem_tensors.append(stem_wav) mix_wav = torch.stack(stem_tensors, dim=0).sum(dim=0) max_val = torch.max(torch.abs(mix_wav)) if max_val > 0: mix_wav = mix_wav / max_val for i in range(len(stem_tensors)): stem_tensors[i] = stem_tensors[i] / max_val # Core PyTorch native spectrogram extraction (no torchaudio transforms) window = torch.hann_window(512) mix_spec = torch.stft( mix_wav, n_fft=512, win_length=512, hop_length=128, window=window, normalized=True, return_complex=True ) mix_spec = torch.abs(mix_spec) target_specs = [] for stem_wav in stem_tensors: spec = torch.stft( stem_wav, n_fft=512, win_length=512, hop_length=128, window=window, normalized=True, return_complex=True ) target_specs.append(torch.abs(spec)) target_specs_tensor = torch.cat(target_specs, dim=0) return mix_spec, target_specs_tensor