Upload raffael_model.py with huggingface_hub

raffael_model.py

@@ -0,0 +1,418 @@
+"""
+Complete High-Quality ConvLSTM Autoencoder
+- Uses true ConvLSTM (not regular LSTM)
+- Complete Encoder (2D CNN + ConvLSTM) with flattened latents
+- Complete Decoder (ConvLSTM + ConvTranspose)
+- Optional Empty/Non-empty Classifier
+- Works with 128x128 input images
+- Latent format: (B, T, N) where N is flattened spatial dimensions
+- Includes ResNet-style residual connections in CNN layers
+"""
+import torch
+import torch.nn as nn
+from raffael_conv_lstm import ConvLSTM
+from huggingface_hub import PyTorchModelHubMixin
+
+
+class ResidualBlock(nn.Module):
+    """
+    Residual block for encoder with optional downsampling
+    """
+    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)
+
+        # Projection shortcut if channels change or downsampling
+        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):
+    """
+    Residual block for decoder with upsampling
+    """
+    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)
+
+        # Shortcut with upsampling
+        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):
+    """
+    Encoder: 2D CNN spatial compression + ConvLSTM temporal modeling + flatten to (B, T, N)
+    Output: z_seq (B, T, latent_size) and z_last (B, latent_size)
+    """
+
+    def __init__(self, input_channels=1, hidden_dim=256, num_layers=2, latent_size=4096):
+        super(Encoder, self).__init__()
+
+        self.hidden_dim = hidden_dim
+        self.latent_size = latent_size
+
+        # Spatial convolution with residual connections: process each frame separately
+        # 128x128 -> 64x64 -> 32x32 -> 16x16
+        self.spatial_cnn = nn.Sequential(
+            # Layer 1: 128 -> 64 (with downsampling)
+            ResidualBlock(input_channels, 64, downsample=True),
+
+            # Layer 2: 64 -> 32 (with downsampling)
+            ResidualBlock(64, 128, downsample=True),
+
+            # Layer 3: 32 -> 16 (with downsampling)
+            ResidualBlock(128, 256, downsample=True),
+        )
+
+        # ConvLSTM: process temporal sequence
+        # Input: (B, T, 256, 16, 16)
+        # Output: (B, T, hidden_dim, 16, 16)
+        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
+        )
+
+        # Dropout before latent compression
+        self.dropout = nn.Dropout(0.1)
+
+        # Linear layer to compress spatial latent to fixed size
+        # Input: (B*T, hidden_dim * 16 * 16)
+        # Output: (B*T, latent_size)
+        self.latent_compress = nn.Linear(hidden_dim * 16 * 16, latent_size)
+
+    def forward(self, x):
+        """
+        Args:
+            x: (B, T, 1, H, W) - input video sequence (any size, will be resized to 128x128)
+
+        Returns:
+            z_seq: (B, T, latent_size) - compressed latent sequence
+            z_last: (B, latent_size) - last timestep compressed latent
+        """
+        B, T, C, H, W = x.shape
+
+        # Resize to 128x128 if needed
+        x = x.view(B * T, C, H, W)  # (B*T, 1, H, W)
+        if H != 128 or W != 128:
+            x = torch.nn.functional.interpolate(x, size=(128, 128), mode='bilinear', align_corners=True)
+
+        # Spatial compression: process each frame separately
+        x = self.spatial_cnn(x)      # (B*T, 256, 16, 16)
+        _, C2, H2, W2 = x.shape
+        x = x.view(B, T, C2, H2, W2)  # (B, T, 256, 16, 16)
+
+        # ConvLSTM processes temporal sequence
+        lstm_out, _ = self.convlstm(x)  # list of (B, T, hidden_dim, 16, 16)
+        h_seq = lstm_out[0]             # (B, T, hidden_dim, 16, 16)
+
+        # 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)  # (B*T, hidden_dim * 16 * 16)
+        h_flat = self.dropout(h_flat)  # Apply dropout
+        z_compressed = self.latent_compress(h_flat)  # (B*T, latent_size)
+        z_seq = z_compressed.view(B, T, self.latent_size)  # (B, T, latent_size)
+
+        # Take last timestep
+        z_last = z_seq[:, -1]  # (B, latent_size)
+
+        return z_seq, z_last
+
+
+class Decoder(nn.Module):
+    """
+    Decoder: Linear expansion + ConvLSTM temporal decoding + ConvTranspose spatial reconstruction
+    Input: z_seq (B, T, latent_size)
+    Output: x_rec (B, T, 1, 128, 128)
+    """
+
+    def __init__(self, seq_len, latent_size=4096, latent_dim=256, hidden_dim=128, num_layers=2):
+        super(Decoder, self).__init__()
+        self.seq_len = seq_len
+        self.latent_dim = latent_dim
+        self.latent_size = latent_size
+
+        # Linear layer to expand compressed latent to spatial dimensions
+        # Input: (B*T, latent_size)
+        # Output: (B*T, latent_dim * 16 * 16)
+        self.latent_expand = nn.Linear(latent_size, latent_dim * 16 * 16)
+
+        # ConvLSTM decodes temporal dimension
+        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
+        )
+
+        # Spatial decoding with residual connections: 16x16 -> 32x32 -> 64x64 -> 128x128
+        self.spatial_decoder = nn.Sequential(
+            # 16 -> 32 (with upsampling)
+            ResidualUpBlock(hidden_dim, 128),
+
+            # 32 -> 64 (with upsampling)
+            ResidualUpBlock(128, 64),
+
+            # 64 -> 128 (with upsampling)
+            ResidualUpBlock(64, 32),
+
+            # Final output layer
+            nn.Conv2d(32, 1, kernel_size=3, padding=1),
+            nn.Sigmoid()  # Assume pixels normalized to [0,1]
+        )
+
+    def forward(self, z_seq):
+        """
+        Args:
+            z_seq: (B, T, latent_size) - compressed latent sequence from encoder
+
+        Returns:
+            x_rec: (B, T, 1, 128, 128) - reconstructed video sequence
+        """
+        B, T, L = z_seq.shape
+
+        # Expand compressed latent to spatial dimensions
+        z_flat = z_seq.view(B * T, L)  # (B*T, latent_size)
+        z_expanded = self.latent_expand(z_flat)  # (B*T, latent_dim * 16 * 16)
+        z_spatial = z_expanded.view(B, T, self.latent_dim, 16, 16)  # (B, T, latent_dim, 16, 16)
+
+        # ConvLSTM decodes temporal dimension
+        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)
+
+        # 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, 128, 128)  # (B, T, 1, 128, 128)
+
+        return x_rec
+
+
+class LatentClassifier(nn.Module):
+    """
+    Empty / Non-empty Well Classifier
+    Classifies based on last timestep latent
+    """
+
+    def __init__(self, latent_size=4096, num_classes=2, dropout=0.3):
+        super(LatentClassifier, self).__init__()
+
+        self.head = nn.Sequential(
+            # Classification head - input is already flattened (B, latent_size)
+            nn.Linear(latent_size, 512),
+            nn.BatchNorm1d(512),
+            nn.ReLU(inplace=True),
+            nn.Dropout(dropout),
+
+            nn.Linear(512, 256),
+            nn.BatchNorm1d(256),
+            nn.ReLU(inplace=True),
+            nn.Dropout(dropout),
+
+            nn.Linear(256, num_classes)
+        )
+
+    def forward(self, z_last):
+        """
+        Args:
+            z_last: (B, latent_size) - last timestep compressed latent
+
+        Returns:
+            logits: (B, num_classes) - classification logits
+        """
+        return self.head(z_last)
+
+
+class ConvLSTMAutoencoder(nn.Module, PyTorchModelHubMixin):
+    """
+    Complete ConvLSTM Autoencoder
+    Includes Encoder, Decoder, and optional Classifier
+    Compatible with HuggingFace Hub
+    Works with 128x128 images
+    """
+    def __init__(
+        self,
+        config=None,
+        seq_len=20,
+        input_channels=1,
+        encoder_hidden_dim=256,
+        encoder_layers=2,
+        decoder_hidden_dim=128,
+        decoder_layers=2,
+        latent_size=4096,
+        use_classifier=True,
+        num_classes=2,
+        use_latent_split=False,
+        # Ablation parameters
+        dropout_rate=0.1,
+        use_convlstm=True,
+        use_residual=True,
+        use_batchnorm=True
+        ):
+        super(ConvLSTMAutoencoder, self).__init__()
+        self.seq_len = seq_len
+        self.use_classifier = use_classifier
+        self.encoder_hidden_dim = encoder_hidden_dim
+        self.latent_size = latent_size
+        self.use_latent_split = use_latent_split
+        # Store ablation settings for reproducibility
+        self.dropout_rate = dropout_rate
+        self.use_convlstm = use_convlstm
+        self.use_residual = use_residual
+        self.use_batchnorm = use_batchnorm
+        if(config != None):
+            # Handle config as dict (from HuggingFace) or object
+            if isinstance(config, dict):
+                self.seq_len = config.get('seq_len', seq_len)
+                self.use_classifier = config.get('use_classifier', use_classifier)
+                self.encoder_hidden_dim = config.get('encoder_hidden_dim', encoder_hidden_dim)
+                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_convlstm = config.get('use_convlstm', use_convlstm)
+                self.use_residual = config.get('use_residual', use_residual)
+                self.use_batchnorm = config.get('use_batchnorm', use_batchnorm)
+            else:
+                self.seq_len = config.seq_len
+                self.use_classifier = config.use_classifier
+                self.encoder_hidden_dim = config.encoder_hidden_dim
+                self.latent_size = config.latent_size
+                self.use_latent_split = config.use_latent_split
+                self.dropout_rate = config.dropout_rate
+                self.use_convlstm = config.use_convlstm
+                self.use_residual = config.use_residual
+                self.use_batchnorm = config.use_batchnorm
+
+            # Core components
+        self.encoder = Encoder(
+            input_channels=input_channels,
+            hidden_dim=encoder_hidden_dim,
+            num_layers=encoder_layers,
+            latent_size=latent_size
+        )
+
+        self.decoder = Decoder(
+            seq_len=seq_len,
+            latent_size=latent_size,
+            latent_dim=encoder_hidden_dim,
+            hidden_dim=decoder_hidden_dim,
+            num_layers=decoder_layers
+        )
+
+        # Optional classifier
+        if use_classifier:
+            self.classifier = LatentClassifier(
+                latent_size=latent_size,
+                num_classes=num_classes
+            )
+
+    def forward(self, x, return_all=False):
+        """
+        Args:
+            x: (B, T, 1, H, W) - input video sequence (any size, will be resized internally)
+            return_all: whether to return all intermediate results
+
+        Returns:
+            Tuple of (reconstruction, lat_vec_seq) where:
+                - reconstruction: (B, T, 1, H, W) - reconstructed video (same size as input)
+                - lat_vec_seq: (B, T, latent_size) - compressed latent sequence
+
+            If return_all is True, returns dict with keys:
+                - reconstruction: (B, T, 1, H, W) - reconstructed video
+                - z_seq: (B, T, latent_size) - compressed latent sequence
+                - z_last: (B, latent_size) - last timestep compressed latent
+                - logits: (B, num_classes) - classification logits (if enabled)
+        """
+        B, T, C, orig_H, orig_W = x.shape
+
+        # Encode (will resize to 128x128 internally)
+        z_seq, z_last = self.encoder(x)
+
+        # Decode (outputs 128x128)
+        x_rec = self.decoder(z_seq)
+
+        # Resize back to original input size if needed
+        if orig_H != 128 or orig_W != 128:
+            x_rec_flat = x_rec.view(B * T, C, 128, 128)
+            x_rec_flat = torch.nn.functional.interpolate(x_rec_flat, size=(orig_H, orig_W), mode='bilinear', align_corners=True)
+            x_rec = x_rec_flat.view(B, T, C, orig_H, orig_W)
+
+        if return_all:
+            # Build output dictionary
+            output = {
+                "reconstruction": x_rec,
+                "z_seq": z_seq,
+                "z_last": z_last,
+            }
+
+            # Optional classification
+            if self.use_classifier:
+                logits = self.classifier(z_last)
+                output["logits"] = logits
+
+            return output
+        else:
+            # Return tuple: (reconstruction, latent_vector)
+            return x_rec, z_seq
+
+    def encode(self, x):
+        """Encode only, for extracting latent"""
+        z_seq, z_last = self.encoder(x)
+        return z_seq, z_last
+
+    def decode(self, z_seq):
+        """Decode only, for reconstructing from latent"""
+        return self.decoder(z_seq)