ASesYusuf1/SESA_Audio_Separation

1import torch2import torch.nn as nn3import torch.nn.functional as F4from typing import Optional5from conformer import Conformer6 7class NeuralModel(nn.Module):8    """9    Принимает |X| STFT: (B, C, F, T_spec) и предсказывает комплексные маски10    в свернутом виде: (B, 2 * (sources*channels), F, T_spec)11    где 2 — это [real, imag].12    """13    def __init__(14        self,15        in_channels: int = 2,16        sources: int = 2,17        freq_bins: int = 2049,18        embed_dim: int = 512,19        depth: int = 8,20        dim_head: int = 64,21        heads: int = 8,22        ff_mult: int = 4,23        conv_expansion_factor: int = 2,24        conv_kernel_size: int = 31,25        attn_dropout: float = 0.1,26        ff_dropout: float = 0.1,27        conv_dropout: float = 0.1,28    ):29        super().__init__()30        self.freq_bins = freq_bins31        self.in_channels = in_channels32        self.sources = sources33        self.out_masks = sources * in_channels34        self.embed_dim = embed_dim35 36        self.input_proj_stft = nn.Linear(freq_bins * in_channels, embed_dim)37        self.model = Conformer(38            dim=embed_dim,39            depth=depth,40            dim_head=dim_head,41            heads=heads,42            ff_mult=ff_mult,43            conv_expansion_factor=conv_expansion_factor,44            conv_kernel_size=conv_kernel_size,45            attn_dropout=attn_dropout,46            ff_dropout=ff_dropout,47            conv_dropout=conv_dropout,48        )49        # 2 = [real, imag]50        self.output_proj = nn.Linear(embed_dim, freq_bins * self.out_masks * 2)51 52    def forward(self, x_stft_mag: torch.Tensor) -> torch.Tensor:53        """54        x_stft_mag: (B, C, F, T_spec)55        returns: (B, 2 * (sources*channels), F, T_spec)  — real/imag масок56        """57        assert x_stft_mag.dim() == 4, f"Expected (B,C,F,T), got {tuple(x_stft_mag.shape)}"58        B, C, F, T_spec = x_stft_mag.shape59        # (B, T_spec, C*F)60        x_stft_mag = x_stft_mag.permute(0, 3, 1, 2).contiguous().view(B, T_spec, C * F)61 62        x = self.input_proj_stft(x_stft_mag)     # (B, T_spec, E)63        x = self.model(x)                        # (B, T_spec, E)64        x = torch.tanh(x)                        # стабилизируем65        x = self.output_proj(x)                  # (B, T_spec, F * out_masks * 2)66 67        # back to (B, 2*out_masks, F, T_spec)68        x = x.reshape(B, T_spec, self.out_masks * 2, F).permute(0, 2, 3, 1).contiguous()69        return x70 71 72class ConformerMSS(nn.Module):73    """74    Совместимо с твоим train:75      forward(x: (B, C, T)) -> y_hat: (B, S, C, T)76    где S = число источников (sources).77    Внутри: STFT -> NeuralModel -> комплексные маски -> iSTFT.78    """79    def __init__(80        self,81        core: NeuralModel,82        n_fft: int = 4096,83        hop_length: int = 1024,84        win_length: Optional[int] = None,85        center: bool = True,86    ):87        super().__init__()88        self.core = core89        self.n_fft = n_fft90        self.hop_length = hop_length91        self.win_length = win_length if win_length is not None else n_fft92        self.center = center93 94        window = torch.hann_window(self.win_length)95        # окно — буфер, чтобы таскалось на .to(device)96        self.register_buffer("window", window, persistent=False)97 98        # sanity-check: freq_bins у core должен совпадать с n_fft//2 + 199        expected_bins = n_fft // 2 + 1100        assert core.freq_bins == expected_bins, (101            f"NeuralModel.freq_bins={core.freq_bins} != n_fft//2+1={expected_bins}. "102            f"Поставь freq_bins={expected_bins} при создании core."103        )104 105    def _stft(self, x: torch.Tensor) -> torch.Tensor:106        """107        x: (B, C, T) -> spec: complex (B, C, F, TT)108        """109        assert x.dim() == 3, f"Expected (B,C,T), got {tuple(x.shape)}"110        B, C, T = x.shape111        x_bc_t = x.reshape(B * C, T)112        spec = torch.stft(113            x_bc_t,114            n_fft=self.n_fft,115            hop_length=self.hop_length,116            win_length=self.win_length,117            window=self.window.to(x.device),118            center=self.center,119            return_complex=True,120        )  # (B*C, F, TT)121        F, TT = spec.shape[-2], spec.shape[-1]122        spec = spec.reshape(B, C, F, TT)123        return spec124 125    def _istft(self, spec: torch.Tensor, length: int) -> torch.Tensor:126        """127        spec: complex (B, C, F, TT) -> audio: (B, C, T)128        """129        B, C, F, TT = spec.shape130        spec_bc = spec.reshape(B * C, F, TT)131        y_bc_t = torch.istft(132            spec_bc,133            n_fft=self.n_fft,134            hop_length=self.hop_length,135            win_length=self.win_length,136            window=self.window.to(spec.device),137            center=self.center,138            length=length,139        )140        return y_bc_t.reshape(B, C, -1)141 142    def forward(self, x: torch.Tensor) -> torch.Tensor:143        """144        x: (B, C, T)  (микс в волне)145        returns y_hat: (B, S, C, T) — предсказанные источники в волне146        """147        B, C, T = x.shape148        # 1) STFT149        mix_spec = self._stft(x)                     # (B, C, F, TT)150        mix_mag = mix_spec.abs()                     # (B, C, F, TT)151 152        # 2) Прогон через core -> real/imag масок153        mask_ri = self.core(mix_mag)                 # (B, 2*(S*C), F, TT2)154        _, two_sc, F, TT2 = mask_ri.shape155 156        S = self.core.sources157        assert two_sc == 2 * (S * C), (158            f"core вернул {two_sc} каналов масок, ожидалось {2*(S*C)} "159            f"(2*[real/imag]*[sources*channels]). Проверь in_channels/sources."160        )161 162        # 3) Синхронизация по времени (если вдруг TT != TT2)163        TT = mix_spec.shape[-1]164        TT_min = min(TT, TT2)165        if TT != TT_min:166            mix_spec = mix_spec[..., :TT_min]167        if TT2 != TT_min:168            mask_ri = mask_ri[..., :TT_min]169        TT = TT_min170        # теперь у обоих время = TT171 172        # 4) Преобразуем к (B, 2, S, C, F, TT)173        mask_ri = mask_ri.view(B, 2, S, C, F, TT).contiguous()174        mask_real = mask_ri[:, 0]                    # (B, S, C, F, TT)175        mask_imag = mask_ri[:, 1]                    # (B, S, C, F, TT)176        masks_c = torch.complex(mask_real, mask_imag)177 178        # 5) Применяем маски к комплексному спектру микса179        mix_spec_bc = mix_spec.unsqueeze(1)          # (B, 1, C, F, TT)180        est_specs = masks_c * mix_spec_bc            # (B, S, C, F, TT)181 182        # 6) iSTFT по каждому источнику183        outs = []184        for s in range(S):185            y_s = self._istft(est_specs[:, s], length=T)  # (B, C, T)186            outs.append(y_s)187        y_hat = torch.stack(outs, dim=1)             # (B, S, C, T)188        return y_hat