class Decoder2(nn.Module):
    def __init__(self, embedding_dim=256):
        super(Decoder2, self).__init__()
        # Project embedding and expand to feature map
        # Initial convolution to process the first frame
        self.first_frame_conv = nn.Sequential(
            nn.Conv2d(3, 32, kernel_size=7, stride=1, padding=3),  # Output: 32 x 512 x 512
            nn.BatchNorm2d(32),
            nn.ReLU(True),
            nn.MaxPool2d(kernel_size=2, stride=2, padding=0),  # Output: 32 x 256 x 256

            nn.Conv2d(32, 64, kernel_size=5, stride=1, padding=2),  # Output: 64 x 256 x 256
            nn.BatchNorm2d(64),
            nn.ReLU(True),
            nn.MaxPool2d(kernel_size=2, stride=2, padding=0),  # Output: 64 x 128 x 128

            nn.Conv2d(64, 128, kernel_size=3, stride=1, padding=1),  # Output: 128 x 128 x 128
            nn.BatchNorm2d(128),
            nn.ReLU(True),
            nn.MaxPool2d(kernel_size=2, stride=2, padding=0),  # Output: 128 x 64 x 64
        )

        self.attention1 = SelfAttention(128)
        
        self.second_frame_conv = nn.Sequential(
            nn.Conv2d(128, 256, kernel_size=3, stride=1, padding=1),  # Output: 256 x 64 x 64
            nn.BatchNorm2d(256),
            nn.ReLU(True),
            nn.MaxPool2d(kernel_size=2, stride=2, padding=0),  # Output: 256 x 32 x 32

            nn.Conv2d(256, 512, kernel_size=3, stride=1, padding=1),  # Output: 512 x 32 x 32
            nn.BatchNorm2d(512),
            nn.ReLU(True),
            nn.MaxPool2d(kernel_size=2, stride=2, padding=0),  # Output: 512 x 16 x 16

            nn.Dropout2d(0.5),  # Adding dropout for regularization

            nn.AdaptiveAvgPool2d((4, 4))  # Output: 512 x 4 x 4
        )

        self.attention2 = SelfAttention(512)
        # Decoder network
        self.decoder = nn.Sequential(
            nn.ConvTranspose2d(512, 512, kernel_size=4, stride=2, padding=1),  # Output: 512 x 8 x 8
            nn.BatchNorm2d(512),
            nn.ReLU(True),
            nn.ConvTranspose2d(512, 256, kernel_size=4, stride=2, padding=1),  # Output: 256 x 16 x 16
            nn.BatchNorm2d(256),
            nn.ReLU(True),
            nn.ConvTranspose2d(256, 128, kernel_size=4, stride=2, padding=1),  # Output: 128 x 32 x 32
            nn.BatchNorm2d(128),
            nn.ReLU(True),
            nn.ConvTranspose2d(128, 64, kernel_size=4, stride=2, padding=1),  # Output: 64 x 64 x 64
            nn.BatchNorm2d(64),
            nn.ReLU(True),
            nn.ConvTranspose2d(64, 32, kernel_size=4, stride=2, padding=1),  # Output: 32 x 128 x 128
            nn.BatchNorm2d(32),
            nn.ReLU(True),
            nn.ConvTranspose2d(32, 16, kernel_size=4, stride=2, padding=1),  # Output: 16 x 256 x 256
            nn.BatchNorm2d(16),
            nn.ReLU(True),
            nn.ConvTranspose2d(16, 3, kernel_size=4, stride=2, padding=1),  # Output: 3 x 512 x 512
            nn.Tanh()  # Output values between -1 and 1
        ) # torch.Size([32, 3, 128, 128])

    def forward(self, first_frame):
        # Process first frame
        first_frame_features = self.first_frame_conv(first_frame)  # Output: [batch_size, 512, 4, 4]
        first_frame_features = self.attention1(first_frame_features)
        first_frame_features = self.second_frame_conv(first_frame_features)
        first_frame_features = self.attention2(first_frame_features)
        # Process embedding
        # Concatenate along channel dimensiontput: [batch_size, 1024, 4, 4]
        x = self.decoder(first_frame_features)
        #print(x.shape)
        return x


make it have lower loss