class SEBlock(nn.Module):
    def __init__(self, channels, reduction=16):
        """
        Squeeze-and-Excitation Block.

        Args:
            channels (int): Number of input channels.
            reduction (int): Reduction ratio for the intermediate dense layer.
        """
        super().__init__()
        self.se = nn.Sequential(
            nn.AdaptiveAvgPool2d(1),  # Squeeze: Global average pooling
            nn.Conv2d(channels, channels // reduction, kernel_size=1, bias=True),
            nn.ReLU(inplace=True),
            nn.Conv2d(channels // reduction, channels, kernel_size=1, bias=True),
            nn.Sigmoid()  # Excitation: Channel-wise weights
        )

    def forward(self, x):
        se_weights = self.se(x)
        return x * se_weights
    
# class MobileNetModule(nn.Module):    
#     def __init__(self, dim, time_dim, cond_dim, expand_ratio=6):
#         super().__init__()
#         hidden_dim = dim*expand_ratio
#         self.expand_conv = nn.Sequential(
#             nn.Conv2d(dim, hidden_dim, kernel_size=1, bias=False),
#             nn.BatchNorm2d(hidden_dim),
#             nn.ReLU6()
#         )
#         self.dwise_conv = nn.Sequential(
#             nn.Conv2d(hidden_dim, hidden_dim, kernel_size=3, groups=hidden_dim, padding=1, bias=False),
#             nn.BatchNorm2d(hidden_dim),
#             nn.ReLU6()
#         )
#         self.proj_conv = nn.Sequential(
#             nn.Conv2d(hidden_dim, dim, kernel_size=1, bias=False),
#             nn.BatchNorm2d(dim),
#         )
#         self.t_emb = nn.Sequential(
#             nn.Linear(time_dim, dim),
#             nn.ReLU()
#         )
#         self.cond_proj = nn.Linear(cond_dim, dim)
    
#     def forward(self, x, t, cond):
#         t = self.t_emb(t).unsqueeze(2).unsqueeze(3)
#         if cond is not None:
#             cond = self.cond_proj(cond).unsqueeze(2).unsqueeze(3)
#             t = t + cond
#         out = self.expand_conv(x)
#         out = self.dwise_conv(out)
#         out = self.proj_conv(out)
#         return out + t
    

class ConvNextModule(nn.Module):
    def __init__(self, dim, time_dim, cond_dim):
        super().__init__()
        self.layers = nn.Sequential(
            nn.Conv2d(dim, dim, kernel_size=7, padding=3, groups=dim, bias=False),
            nn.GroupNorm(1, dim),  # Use GroupNorm instead of LayerNorm
            nn.Conv2d(dim, dim*4, kernel_size=1, bias=False),
            nn.GELU(),
            nn.Conv2d(dim*4, dim, kernel_size=1, bias=False),
        )
        self.t_emb = nn.Sequential(
            nn.Linear(time_dim, dim),
            nn.ReLU()
        )
        self.cond_proj = nn.Linear(cond_dim, dim)
    
    def forward(self, x, t, cond):
        t = self.t_emb(t).unsqueeze(2).unsqueeze(3)
        if cond is not None:
            cond = self.cond_proj(cond).unsqueeze(2).unsqueeze(3)
            t = t + cond
            
        return self.layers(x) + t


# class GhostModule(nn.Module):
#     def __init__(self, dim, kernel_size=1, ratio=2, dw_size=3, stride=1, relu=True, cond_dim=128):
#         super(GhostModule, self).__init__()
#         init_channels = math.ceil(dim / ratio)
#         new_channels = init_channels * (ratio - 1)

#         self.primary_conv = nn.Sequential(
#             nn.Conv2d(dim, init_channels, kernel_size, stride, kernel_size // 2, bias=False),
#             nn.BatchNorm2d(init_channels),
#             nn.ReLU(inplace=True) if relu else nn.Identity(),
#         )

#         self.cheap_operation = nn.Sequential(
#             nn.Conv2d(init_channels, new_channels, dw_size, 1, dw_size // 2, groups=init_channels, bias=False),
#             nn.BatchNorm2d(new_channels),
#             nn.ReLU(inplace=True) if relu else nn.Identity(),
#         )

#         self.t_emb = nn.Sequential(
#             nn.Linear(cond_dim, dim),
#             nn.ReLU(inplace=True)
#         )

#     def forward(self, x, t):
#         t = self.t_emb(t).unsqueeze(2).unsqueeze(3).expand(x.size(0), -1, x.size(2), x.size(3))
#         x1 = self.primary_conv(x)
#         x2 = self.cheap_operation(x1)
#         out = torch.cat([x1, x2], dim=1)
#         out = out[:, :self.oup, :, :] + t
#         return out


class HalfUNetBlock(nn.Module):
    def __init__(self, block: nn.Module, dim, time_dim, cond_dim, reduction=16):
        super().__init__()
        self.block1 = block(dim, time_dim=time_dim, cond_dim=cond_dim)
        self.block2 = block(dim, time_dim=time_dim, cond_dim=cond_dim)
        self.pool = nn.MaxPool2d(kernel_size=2, stride=2)
        
        self.se = SEBlock(dim, reduction)

    def forward(self, x, t, cond):
        x = self.block1(x, t, cond)
        x = self.block2(x, t, cond)
        x = self.se(x)
        p = self.pool(x)
        return x, p

class HalfUNetModel(nn.Module):
    def __init__(self, block: nn.Module, dim, depth, time_dim=128, cond_dim=1000):
        super().__init__()
        self.depth = depth
        self.blocks = nn.ModuleList()
        for _ in range(depth):
            self.blocks.append(HalfUNetBlock(block, dim=dim, time_dim=time_dim, cond_dim=cond_dim))
    
    def forward(self, x, t, cond):
        skip_connections = []
        for block in self.blocks:
            x, p = block(x, t, cond)
            skip_connections.append(x)
            x = p
        out = torch.zeros_like(skip_connections[0])
        for i, skip in enumerate(skip_connections):
            out += F.interpolate(skip, size=out.shape[2:])
        return out

class NestedHalfUNet(nn.Module):
    def __init__(self, block: nn.Module, ch_in, ch_out, depth, dim, time_dim, num_classes=1000):
        super().__init__()
        self.time_dim = time_dim
        self.init_conv = nn.Conv2d(ch_in, dim, kernel_size=1)
        self.down = nn.ModuleList()
        for i in range(depth):
            self.down.append(HalfUNetModel(block, depth=depth-i, dim=dim, time_dim=time_dim, cond_dim=num_classes))
        self.final_conv = nn.Conv2d(dim, ch_out, kernel_size=1)
        self.buff1 = self.register_buffer('inv_freq', 1.0 / (
            10000 ** (torch.arange(0, time_dim, 2).float() / time_dim)
        ))
        # self.buff2 = self.register_buffer('t_repeat', torch.ones((1, self.time_dim//2)))
    
    def pos_encoding(self, t, channels):
        pos_enc_a = torch.sin(t.repeat(1, channels // 2) * self.inv_freq)
        pos_enc_b = torch.cos(t.repeat(1, channels // 2) * self.inv_freq)
        pos_enc = torch.cat([pos_enc_a, pos_enc_b], dim=-1)
        return pos_enc
    
    def forward(self, x, t, label=None):
        t = t.unsqueeze(-1).type(torch.float)
        t = self.pos_encoding(t, channels=self.time_dim)
        x = self.init_conv(x)
        skip_connections = []
        for down in self.down:
            x = down(x, t, label)
            skip_connections.append(x)
            x = F.interpolate(x, scale_factor=0.5)
        out = torch.zeros_like(skip_connections[0])
        for i, skip in enumerate(skip_connections):
            out += F.interpolate(skip, size=out.shape[2:])
        return self.final_conv(out)

model = NestedHalfUNet(block=ConvNextModule, ch_in=4, ch_out=4, depth=4, dim=512, time_dim=512) # ~60M params

how to improve this model? I tried to overfit it on a few images but it can't do even this after 2000 epochs, so I think it has some bug or too small capacity to learn ImageNet dataset. My goal is to train the smallest latent diffusion model on ImageNet that can be trained on single GPU (I'm using precomputed vae latents from SDXL). Look for potential bugs in code or come up with better building block (I have tried already convnext, mobilenetv2 and ghost modules) or better "macro architecture" (now Half Unet because I read in the paper that upsampling part of unet is redundant for segmentation task and this version is more efficient). If you will be thinking about macro architecture then keep in mind to avoid memory usage spikes (in unet for example feature maps from upsampling path needs to be temp. stored for downsampling and then concatenated [avoid concat feature maps - expensive for mem.]). Feature maps should be small because latents are size 32x32x4 so it shouldn't be too much problem.