ae_model.py

import torch
import torch.nn as nn
from huggingface_hub import PyTorchModelHubMixin
import torch.nn.functional as F

class ResidualBlock(nn.Module):
    def __init__(self, in_channels, out_channels, downsample=False):
        super(ResidualBlock, self).__init__()

        stride = 2 if downsample else 1

        self.conv1 = nn.Conv2d(in_channels, out_channels, kernel_size=3, stride=stride, padding=1)
        self.bn1 = nn.BatchNorm2d(out_channels)
        self.relu = nn.ReLU(inplace=True)
        self.conv2 = nn.Conv2d(out_channels, out_channels, kernel_size=3, padding=1)
        self.bn2 = nn.BatchNorm2d(out_channels)

        if in_channels != out_channels or downsample:
            self.shortcut = nn.Sequential(
                nn.Conv2d(in_channels, out_channels, kernel_size=1, stride=stride),
                nn.BatchNorm2d(out_channels)
            )
        else:
            self.shortcut = nn.Identity()

    def forward(self, x):
        identity = self.shortcut(x)

        out = self.conv1(x)
        out = self.bn1(out)
        out = self.relu(out)

        out = self.conv2(out)
        out = self.bn2(out)

        out += identity
        out = self.relu(out)

        return out


class ResidualUpBlock(nn.Module):
    def __init__(self, in_channels, out_channels):
        super(ResidualUpBlock, self).__init__()

        self.upsample = nn.ConvTranspose2d(in_channels, out_channels, kernel_size=4, stride=2, padding=1)
        self.bn1 = nn.BatchNorm2d(out_channels)
        self.relu = nn.ReLU(inplace=True)
        self.conv = nn.Conv2d(out_channels, out_channels, kernel_size=3, padding=1)
        self.bn2 = nn.BatchNorm2d(out_channels)

        self.shortcut = nn.ConvTranspose2d(in_channels, out_channels, kernel_size=4, stride=2, padding=1)

    def forward(self, x):
        identity = self.shortcut(x)

        out = self.upsample(x)
        out = self.bn1(out)
        out = self.relu(out)

        out = self.conv(out)
        out = self.bn2(out)

        out += identity
        out = self.relu(out)

        return out


class Encoder(nn.Module):

    def __init__(self, input_channels=1, num_layers=2, latent_size=4096, hidden_channels=64, use_lstm=True):
        super(Encoder, self).__init__()

        self.latent_size = latent_size
        self.hidden_channels = hidden_channels
        self.spatial_cnn = nn.Sequential(
            # 256 -> 128
            #ResidualBlock(input_channels, self.hidden_channels, downsample=True),

            # 128 -> 64 
            #ResidualBlock(self.hidden_channels, self.hidden_channels, downsample=True),
            ResidualBlock(input_channels, self.hidden_channels, downsample=True),

            # 64 -> 32 
            ResidualBlock(self.hidden_channels, self.hidden_channels, downsample=True),

            # 32 -> 16 
            ResidualBlock(self.hidden_channels, self.hidden_channels, downsample=True),

            # 16 -> 8
            #ResidualBlock(self.hidden_channels, self.hidden_channels, downsample=True),

        )
        self.final_resolution = 16 #2 ** (8 - len([module for module in self.spatial_cnn.modules() if not isinstance(self.spatial_cnn, nn.Sequential)]))
        self.use_lstm = use_lstm
        """if not self.use_convlstm:
            self.convlstm = None 
        else:
            self.convlstm = ConvLSTM(
                input_dim=256,
                hidden_dim=hidden_dim,
                kernel_size=(3, 3),
                num_layers=num_layers,
                batch_first=True,
                return_all_layers=False
                )"""
        
        self.dropout = nn.Dropout(0.1)

        self.latent_compress = nn.Linear(self.hidden_channels * self.final_resolution * self.final_resolution, latent_size)
        if self.use_lstm:
            self.lstm_enc = nn.LSTM(latent_size, latent_size, batch_first=True)
        else:
            self.lstm_enc = None

        self.lin1 = nn.Linear(latent_size,latent_size)



    def forward(self, x):
        B, T, C, H, W = x.shape

        x = x.view(B * T, C, H, W)  # (B*T, 1, H, W)

        x = self.spatial_cnn(x)      
        _, C2, H2, W2 = x.shape
        x = x.view(B, T, C2, H2, W2)  

        """# ConvLSTM processes temporal sequence
        if(self.use_convlstm):
            lstm_out, _ = self.convlstm(x)  # list of (B, T, hidden_dim, 16, 16)
            h_seq = lstm_out[0]             # (B, T, hidden_dim, 16, 16)
        else:
            h_seq = x # just pass it forward if not"""
        h_seq = x

        # Flatten and compress spatial dimensions with linear layer
        B, T, C, H, W = h_seq.shape
        h_flat = h_seq.view(B, T, C * H * W)  # Linear just works on bottom most dim
        z_compressed = F.relu(self.latent_compress(h_flat))
        if self.use_lstm:
            z_compressed, _ = self.lstm_enc(z_compressed)
        z_compressed = self.dropout(z_compressed)
        z_compressed = self.lin1(z_compressed)
        z_seq = z_compressed.view(B, T, self.latent_size)  


        return z_seq


class Decoder(nn.Module):

    def __init__(self, latent_size=4096, num_layers=2, hidden_channels=64, initial_resolution=16,final_size=128, use_lstm=True):
        super(Decoder, self).__init__()
        self.latent_size = latent_size
        self.hidden_channels = hidden_channels

        #self.latent_expand = nn.Linear(latent_size, self.hidden_channels * 16 * 16)

        self.use_lstm = use_lstm
        """if not self.use_convlstm:
            self.convlstm = None 
        else:
            self.convlstm = ConvLSTM(
                input_dim=latent_dim,
                hidden_dim=hidden_dim,
                kernel_size=(3, 3),
                num_layers=num_layers,
                batch_first=True,
                return_all_layers=False
            )"""
        if self.use_lstm:
            self.lstm_dec = nn.LSTM(latent_size, latent_size, batch_first=True)
        else:
            self.lstm_dec = None
    
        self.dropout = nn.Dropout(0.1)
        self.spatial_decoder = nn.Sequential(
            # 8 -> 16
            #ResidualUpBlock(self.hidden_channels, self.hidden_channels),

            # 16 -> 32 
            ResidualUpBlock(self.hidden_channels, self.hidden_channels),

            # 32 -> 64 
            ResidualUpBlock(self.hidden_channels, self.hidden_channels),

            # 64 -> 128 
            ResidualUpBlock(self.hidden_channels, self.hidden_channels),

            # 128 -> 256
            #ResidualUpBlock(self.hidden_channels, self.hidden_channels),


            nn.Conv2d(self.hidden_channels, 1, kernel_size=3, padding=1),
            nn.Sigmoid()
        )
        
        self.initial_resolution = initial_resolution
        self.lin1 = nn.Linear(latent_size,latent_size)

        self.latent_expand = nn.Linear(latent_size, self.hidden_channels * self.initial_resolution * self.initial_resolution)
        self.final_size = final_size
    def forward(self, z_seq):
        B, T, L = z_seq.shape

        z_flat = z_seq # linear works on bottom most dim
        z_flat = F.relu(self.lin1(z_flat))
        if self.use_lstm:
            z_flat, _ = self.lstm_dec(z_flat)
        z_flat = self.dropout(z_flat) 
        z_expanded = F.relu(self.latent_expand(z_flat)) 
        assert z_expanded.shape == (B, T, self.hidden_channels * (self.initial_resolution ** 2)), f"BAD z_expanded shape: {z_expanded.shape}"
        z_spatial = z_expanded.view(B, T, self.hidden_channels, self.initial_resolution, self.initial_resolution)  

        """# ConvLSTM decodes temporal dimension
        if(self.use_convlstm):
            lstm_out, _ = self.convlstm(z_spatial)  # list of (B, T, hidden_dim, 16, 16)
            h_seq = lstm_out[0]                 # (B, T, hidden_dim, 16, 16)
        else:
            h_seq = z_spatial # just pass it forward"""
        h_seq = z_spatial
        # Spatial decoding: process each timestep separately
        B, T, C, H, W = h_seq.shape
        h_seq = h_seq.view(B * T, C, H, W)  # (B*T, hidden_dim, 16, 16)
        x_rec = self.spatial_decoder(h_seq)  # (B*T, 1, 128, 128)
        x_rec = x_rec.view(B, T, 1, self.final_size, self.final_size)  # (B, T, 1, 128, 128)

        return x_rec



class ConvLSTMAutoencoder(nn.Module, PyTorchModelHubMixin):
    def __init__(
        self,
        config=None,
        input_channels=1,
        encoder_layers=2,
        decoder_layers=2,
        latent_size=4096,
        use_classifier=True,
        num_classes=2,
        use_latent_split=False,
        # Ablation parameters
        dropout_rate=0.1,
        use_lstm=True,
        use_residual=True,
        use_batchnorm=True,
        hidden_channels=64
        ):
        super(ConvLSTMAutoencoder, self).__init__()
        self.use_classifier = use_classifier
        self.latent_size = latent_size
        self.use_latent_split = use_latent_split
        # Store ablation settings for reproducibility
        self.dropout_rate = dropout_rate
        self.use_lstm = use_lstm
        self.use_residual = use_residual
        self.use_batchnorm = use_batchnorm
        self.hidden_channels = hidden_channels
        if(config != None):
            if isinstance(config, dict):
                self.use_classifier = config.get('use_classifier', use_classifier)
                self.latent_size = config.get('latent_size', latent_size)
                self.use_latent_split = config.get('use_latent_split', use_latent_split)
                self.dropout_rate = config.get('dropout_rate', dropout_rate)
                self.use_lstm = config.get('use_lstm', use_lstm)
                self.use_residual = config.get('use_residual', use_residual)
                self.use_batchnorm = config.get('use_batchnorm', use_batchnorm)
                self.hidden_channels = config.get('hidden_channels', hidden_channels)
            else:
                self.use_classifier = config.use_classifier
                self.latent_size = config.latent_size
                self.use_latent_split = config.use_latent_split
                self.dropout_rate = config.dropout_rate
                self.use_lstm = config.use_lstm
                self.use_residual = config.use_residual
                self.use_batchnorm = config.use_batchnorm
                self.hidden_channels = config.hidden_channels

        self.encoder = Encoder(
            latent_size=self.latent_size,
            use_lstm=self.use_lstm,
            hidden_channels = self.hidden_channels
        )

        self.decoder = Decoder(
            latent_size=self.latent_size,
            use_lstm=self.use_lstm,
            hidden_channels = self.hidden_channels
        )


    def forward(self, x, return_all=False, hidden=None):

        z_seq = self.encoder(x)

        x_rec = self.decoder(z_seq)


        return x_rec, z_seq

    def encode(self, x):
        z_seq, z_last = self.encoder(x)
        return z_seq, z_last

    def decode(self, z_seq):
        return self.decoder(z_seq)