pcunwa/

BS-Roformer-Inst-FNO

from neuralop.models import FNO1d
class MaskEstimator(Module):
    @beartype
    def __init__(
            self,
            dim,
            dim_inputs: Tuple[int, ...],
            depth,
            mlp_expansion_factor=4
    ):
        super().__init__()
        self.dim_inputs = dim_inputs
        self.to_freqs = ModuleList([])
        dim_hidden = dim * mlp_expansion_factor

        for dim_in in dim_inputs:
            net = []

            mlp = nn.Sequential(
                FNO1d(n_modes_height=64, hidden_channels=dim, in_channels=dim, out_channels=dim_in*2, lifting_channels=dim, projection_channels=dim, n_layers=3, separable=True),
                nn.GLU(dim=-2)
            )

            self.to_freqs.append(mlp)

    def forward(self, x):
        x = x.unbind(dim=-2)

        outs = []

        for band_features, mlp in zip(x, self.to_freqs):
            band_features = rearrange(band_features, 'b t c -> b c t')
            with torch.autocast(device_type='cuda', enabled=False, dtype=torch.float32):
                freq_out = mlp(band_features).float()
            freq_out = rearrange(freq_out, 'b c t -> b t c')
            outs.append(freq_out)

        return torch.cat(outs, dim=-1)

with torch.serialization.safe_globals([torch._C._nn.gelu]):