Update app.py

app.py

@@ -1,5 +1,4 @@
 import torchaudio
-import math
 import torch
 from typing import Dict, Tuple
 from huggingface_hub import hf_hub_download
@@ -9,7 +8,7 @@ from torch import nn, Tensor
 from torch.nn import functional as F
 from torch.utils.data import Dataset , DataLoader
 from torchaudio.transforms import Fade
-from torchaudio.models._hdemucs import HDemucs
+from torchaudio.models import HDemucs
 
 class Crop2d(nn.Module):
     def __init__(self, left, right, top, bottom):
@@ -37,9 +36,9 @@ class DecoderBlock(nn.Module):
     def __init__(self, in_channels: int, out_channels: int) -> None:
         super().__init__()
         self.tconv = nn.ConvTranspose2d(in_channels, out_channels, 5, 2)
+        self.crop = Crop2d(1, 2, 1, 2)   # reverse padding
         self.bn = nn.BatchNorm2d(out_channels,0.001,0.01)
         self.relu = nn.ReLU()
-        self.crop = Crop2d(1, 2, 1, 2)   # reverse padding
     def forward(self, input: Tensor) -> Tensor:
         return self.bn(self.relu(self.crop(self.tconv(input))))
 
@@ -57,6 +56,7 @@ class UNet(nn.Module):
         self.decoder_layers = nn.ModuleList([DecoderBlock(in_ch if i == 0 else in_ch * 2,out_ch) for i, (in_ch, out_ch) in enumerate(zip(up_set[:-1], up_set[1:]))])
         self.up_final = nn.Conv2d(1, in_channels, kernel_size=4, dilation=2, padding=3)
         self.sigmoid = nn.Sigmoid()
+
     def forward(self, input: Tensor) -> Tensor:
         encoder_outputs_pre_act = []
         x = input
@@ -79,78 +79,129 @@ class UNet(nn.Module):
         input = input[..., :min_f, :min_t]
         return mask * input
 
-class ChunkSplitterDataset(Dataset):
-    def __init__(self, wav, win):
-        self.win_length = 4096
-        self.hop_length = 1024
+class STFTChunkDataset(Dataset):
+    def __init__(self, wav, win, win_length=4096, T=512, F=1024):
+        self.win_length = win_length
         self.win = win
-        self.T = 512
-        self.stft_mag = self.batchify(self.compute_stft(wav))
+        self.T = T
+        self.F = F
+
+        wav = wav.view(wav.size(0), -1)
+
+        stft = torch.stft(
+            wav,
+            n_fft=win_length,
+            window=win,
+            return_complex=True,
+            pad_mode="constant"
+        )[:, :F, :]
+
+        self.L = stft.size(-1)
+        self.stft_complex = torch.view_as_real(stft)
+
+        mag = stft.abs().unsqueeze(1)  # (1, 1, F, L)
+
+        # pad time to multiple of T
+        pad_T = (T - self.L % T) % T
+        mag = F.pad(mag, (0, pad_T))
+
+        # split into chunks
+        self.chunks = mag.view(1, 1, F, -1, T)\
+                         .permute(3, 0, 1, 2, 4)\
+                         .squeeze(1)
+        # shape: (num_chunks, 1, F, T)
 
     def __len__(self):
-        return self.stft_mag.size(0)
+        return self.chunks.size(0)
 
     def __getitem__(self, idx):
-        return self.stft_mag[idx]
-
-    def compute_stft(self, wav: Tensor):
-        stft = torch.stft(wav.squeeze(),n_fft=self.win_length,hop_length=self.hop_length,window=self.win,return_complex=False,pad_mode="constant",)
-        stft = stft[:, :1024, :, :]
-        real = stft[:, :, :, 0]
-        imag = stft[:, :, :, 1]
-        self.stft = stft
-        self.L = self.stft.size(2)
-        return torch.sqrt(real**2 + imag**2 + 1e-10).unsqueeze(-1).permute([3, 0, 1, 2])
-
-    def batchify(self, tensor: Tensor) -> Tensor:
-        orig_size = tensor.size(-1)
-        new_size = math.ceil(orig_size / self.T) * self.T
-        tensor = F.pad(tensor, [0, new_size - orig_size])
-        return torch.cat(torch.split(tensor, self.T, dim=-1), dim=0).transpose(2, 3)
-
-    def apply_mask(self,mask,mask_sum):
-        mask = (mask**2 + 1e-10 / 2) / (mask_sum)
-        mask = mask.transpose(2, 3)  # B x 2 X F x T
-        mask = torch.cat(torch.split(mask, 1, dim=0), dim=3)
-        mask = mask.squeeze(0)[:, :, :self.L].unsqueeze(-1)  # 2 x F x L x 1
-        stft = self.stft * mask
-        target_F = self.win_length // 2 + 1
-        if stft.size(1) < target_F:
-            pad = target_F - stft.size(1)
-            stft = F.pad(stft, (0, 0, 0, 0, 0, pad))  # pad along freq dim
-        return torch.istft(torch.view_as_complex(stft),n_fft=self.win_length,hop_length=self.hop_length,win_length=self.win_length,center=True,window=self.win)
-
-    def decoder(self,masks):
-        mask_sum = sum([m**2 for m in masks.values()]) + 1e-10
-        return {name: self.apply_mask(m,mask_sum) for name, m in masks.items()}
+        return self.chunks[idx]
 
 class Splitter(nn.Module):
     CONFIG = {
-    2:['2_other', '2_vocals'],
-    4:['4_bass', '4_drums', '4_other', '4_vocals'],
-    5:['5_piano','5_bass', '5_drums','5_other', '5_vocals']
-}
+        2: ['2_other', '2_vocals'],
+        4: ['4_bass', '4_drums', '4_other', '4_vocals'],
+        5: ['5_piano', '5_bass', '5_drums', '5_other', '5_vocals']
+    }
+
     def __init__(self, stem=2):
         super().__init__()
         self.win_length = 4096
-        self.win = nn.Parameter(torch.hann_window(self.win_length), requires_grad=False)
-        self.stems = nn.ModuleDict({name: UNet() for name in self.CONFIG[stem]})
+        self.T = 512
+        self.F = 1024
+        self.target_F = self.win_length // 2 + 1
+
+        self.win = nn.Parameter(
+            torch.hann_window(self.win_length),
+            requires_grad=False
+        )
+
+        self.stems = nn.ModuleDict({
+            name: UNet() for name in self.CONFIG[stem]
+        })
+
         for name in self.stems:
-            self.stems[name].load_state_dict(load_file(hf_hub_download("shethjenil/spleeter",f"{name}.safetensors")))
+            self.stems[name].load_state_dict(
+                load_file(
+                    hf_hub_download("shethjenil/spleeter", f"{name}.safetensors")
+                )
+            )
+
         self.eval()
 
     @torch.inference_mode()
-    def forward(self, wav: Tensor,sr:int,batch_size) -> Dict[str, Tensor]:
+    def forward(self, wav, sr, batch_size):
         device = next(self.parameters()).device
+
         if sr != 44100:
             wav = torchaudio.functional.resample(wav, sr, 44100)
-        ds = ChunkSplitterDataset(wav.to(device),self.win)
-        masks = {name: [] for name in self.stems}
-        for batch in tqdm(DataLoader(ds,batch_size)):
-            outputs = {name: net(batch) for name, net in self.stems.items()}
-            for name in outputs:
-                masks[name].append(outputs[name])
-        return ds.decoder({k: torch.cat(v, dim=0) for k, v in masks.items()})
+
+        wav = wav.to(device)
+
+        ds = STFTChunkDataset(wav, self.win)
+        loader = DataLoader(
+            ds,
+            batch_size=batch_size,
+            shuffle=False,        # IMPORTANT
+            pin_memory=True
+        )
+
+        masks = {k: [] for k in self.stems}
+
+        for batch in loader:
+            batch = batch.to(device)
+
+            for name, net in self.stems.items():
+                masks[name].append(net(batch))
+
+        masks = {k: torch.cat(v, dim=0) for k, v in masks.items()}
+
+        return self.decode(masks, ds)
+
+    def decode(self, masks, ds):
+        mask_sum = sum(m ** 2 for m in masks.values()) + 1e-10
+        outputs = {}
+
+        for name, m in masks.items():
+            mask = (m ** 2 / mask_sum)
+
+            # (chunks, 1, F, T) → (1, F, time)
+            mask = mask.permute(1, 2, 0, 3).reshape(1, self.F, -1)
+            mask = mask[:, :, :ds.L]
+
+            stft = ds.stft_complex * mask.unsqueeze(-1)
+
+            if stft.size(1) < self.target_F:
+                pad = self.target_F - stft.size(1)
+                stft = F.pad(stft, (0, 0, 0, 0, 0, pad))
+
+            outputs[name] = torch.istft(
+                torch.view_as_complex(stft),
+                n_fft=self.win_length,
+                window=self.win
+            )
+
+        return outputs
 
 def separate_audio_spleeter(audio_path:str,batch_size:int,stem:int,progress=gr.Progress(True)):
     wav, sr = torchaudio.load(audio_path)
@@ -293,4 +344,3 @@ gr.TabbedInterface([
     gr.Interface(separate_audio_spleeter, [gr.Audio(type="filepath"),gr.Number(16),gr.Radio([2,4,5],label="STEM")],gr.Files()),
     gr.Interface(separate_audio_demucs, [gr.Audio(type="filepath"),gr.Number(16),gr.Radio([4],label="STEM")],gr.Files())
 ],['spleeter','demucs']).launch()
-