import torch
import torch.nn as nn
import torch.nn.functional as F
import math

class SEBlock(nn.Module):
    def __init__(self, channels, reduction=16):
        super().__init__()
        self.se = nn.Sequential(
            nn.AdaptiveAvgPool2d(1),
            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()
        )
    
    def forward(self, x):
        se_weights = self.se(x)
        return x * se_weights

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),
            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)  # (batch, dim, 1, 1)
        if cond is not None:
            cond = self.cond_proj(cond).unsqueeze(2).unsqueeze(3)  # (batch, dim, 1, 1)
            t = t + cond
        return self.layers(x) + t

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 = skip_connections[0]
        for skip in skip_connections[1:]:
            out = out + F.interpolate(skip, size=out.shape[2:], mode='nearest')
        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.Sequential(
            nn.Conv2d(ch_in, dim, kernel_size=3, padding=1),
            nn.BatchNorm2d(dim),
            nn.ReLU(inplace=True)
        )
        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.register_buffer('inv_freq', 1.0 / (10000 ** (torch.arange(0, time_dim, 2).float() / time_dim)))
    
    def pos_encoding(self, t, channels):
        device = t.device
        half_dim = channels // 2
        emb = math.log(10000) / (half_dim - 1)
        emb = torch.exp(torch.arange(half_dim, device=device) * -emb)
        emb = t[:, None] * emb[None, :]  # Shape: (batch, half_dim)
        pos_enc = torch.cat([torch.sin(emb), torch.cos(emb)], dim=1)  # Shape: (batch, channels)
        return pos_enc
    
    def forward(self, x, t, label=None):
        # Positional Encoding
        t = t.unsqueeze(-1).type(torch.float)  # Shape: (batch, time_dim, 1)
        t = self.pos_encoding(t.squeeze(-1), channels=self.time_dim)  # Shape: (batch, time_dim)
        
        # Initial Convolution
        x = self.init_conv(x)
        
        # Downsampling Path
        skip_connections = []
        for down in self.down:
            x = down(x, t, label)
            skip_connections.append(x)
            x = F.interpolate(x, scale_factor=0.5, mode='nearest')
        
        # Aggregating Skip Connections
        out = skip_connections[0]
        for skip in skip_connections[1:]:
            out = out + F.interpolate(skip, size=out.shape[2:], mode='nearest')
        
        # Final Convolution
        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 single batch but it can't do this even after 2000 epochs, so I think model has 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 module) or better "macro architecture" (now I'm using 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.