working ldm

2 files changed, 58 insertions, 46 deletions

diff --git a/autoencoder.py b/autoencoder.py
index 645f628..58326fe 100644
--- a/autoencoder.py
+++ b/autoencoder.py
@@ -21,9 +21,8 @@ class Autoencoder(torch.nn.Module):
         super(Autoencoder, self).__init__()
         self.input_dim = input_dim
         self.latent_dim = latent_dim
-        self.latent_size = torch.prod(torch.tensor(latent_dim))
         self.encoder = torch.nn.Sequential(
-            torch.nn.Conv2d(input_dim[0], 16, kernel_size=3, padding=1),
+            torch.nn.Conv2d(input_dim, 16, kernel_size=3, padding=1),
             torch.nn.ReLU(),
             torch.nn.Conv2d(16, 16, kernel_size=3, padding=1),
             torch.nn.ReLU(),
@@ -37,16 +36,19 @@ class Autoencoder(torch.nn.Module):
             torch.nn.ReLU(),
             torch.nn.Conv2d(64, 64, kernel_size=3, padding=1),
             torch.nn.ReLU(),
-            torch.nn.Flatten(),
-            torch.nn.Linear(64 * 7 * 7, self.latent_size),
+            torch.nn.Conv2d(64, latent_dim, kernel_size=3, padding=1),
+            torch.nn.ReLU(),
+            # 1x7x7 to 1x8x8
+            torch.nn.Conv2d(latent_dim, latent_dim, kernel_size=2, padding=1),
             torch.nn.Sigmoid(),
-            torch.nn.Unflatten(1, latent_dim),
         )
         self.decoder = torch.nn.Sequential(
-            torch.nn.Flatten(),
-            torch.nn.Linear(self.latent_size, 64 * 7 * 7),
+            # 1x8x8 to 1x7x7
+            torch.nn.Conv2d(latent_dim, latent_dim, kernel_size=2, padding=0),
+            torch.nn.ReLU(),
+            # main decoder
+            torch.nn.Conv2d(latent_dim, 64, kernel_size=3, padding=1),
             torch.nn.ReLU(),
-            torch.nn.Unflatten(1, (64, 7, 7)),
             torch.nn.Conv2d(64, 64, kernel_size=3, padding=1),
             torch.nn.ReLU(),
             torch.nn.ConvTranspose2d(64, 32, kernel_size=2, stride=2),
@@ -57,7 +59,7 @@ class Autoencoder(torch.nn.Module):
             torch.nn.ReLU(),
             torch.nn.Conv2d(16, 16, kernel_size=3, padding=1),
             torch.nn.ReLU(),
-            torch.nn.Conv2d(16, input_dim[0], kernel_size=3, padding=1),
+            torch.nn.Conv2d(16, input_dim, kernel_size=3, padding=1),
             torch.nn.Sigmoid(),
         )
 
@@ -102,14 +104,14 @@ def main():
                             num_workers=4, persistent_workers=True)
 
     # Initialize model
-    model = Autoencoder(input_dim=(1, 28, 28), latent_dim=(1, 8, 8))
-    model.load_state_dict(torch.load('autoencoder.pth'))
+    model = Autoencoder(input_dim=1, latent_dim=1)
+    model.load_state_dict(torch.load('mnist_autoencoder.pth'))
     model.to(device)
 
     # Train model
     trainer = Trainer()
     lr = 1e-3
-    epochs = 10
+    epochs = 1
     optimizer = torch.optim.Adam(model.parameters(), lr=lr)
     criterion = torch.nn.functional.binary_cross_entropy
     trainer.train(model, dataloader, epochs, optimizer, criterion)
@@ -120,17 +122,31 @@ def main():
     # Visualize results
     n = 10
     with torch.no_grad():
-        plt.figure(figsize=(2*n, 4))
+        plt.figure(figsize=(2*n, 6))
         for i, j in enumerate(torch.randint(0, len(dataset), (n,))):
             x, _ = dataset[j]
             x = x.unsqueeze(0)
-            x_hat = model(x)
-            plt.subplot(2, n, i + 1)
-            plt.imshow(x.cpu().squeeze().numpy(), cmap='gray')
+            x_latent = model.encode(x)
+            x_hat = model.decode(x_latent)
+            plt.subplot(3, n+1, i + 2)
+            plt.imshow(x.cpu().squeeze().numpy())
+            plt.axis('off')
+            plt.subplot(3, n+1, i + n + 3)
+            plt.imshow(x_latent.cpu().squeeze().numpy())
             plt.axis('off')
-            plt.subplot(2, n, i + n + 1)
-            plt.imshow(x_hat.cpu().squeeze().numpy(), cmap='gray')
+            plt.subplot(3, n+1, i + 2*n + 4)
+            plt.imshow(x_hat.cpu().squeeze().numpy())
             plt.axis('off')
+        plt.subplot(3, n+1, 1)
+        plt.text(0.5, 0.5, 'Original', horizontalalignment='center', fontsize=12)
+        plt.axis('off')
+        plt.subplot(3, n+1, n + 2)
+        plt.text(0.5, 0.5, 'Latent', horizontalalignment='center', fontsize=12)
+        plt.axis('off')
+        plt.subplot(3, n+1, 2*n + 3)
+        plt.text(0.5, 0.5, 'Reconstructed', horizontalalignment='center', fontsize=12)
+        plt.axis('off')
+        plt.suptitle('Autoencoder')
         plt.tight_layout()
         plt.savefig('autoencoder.tmp.png')
 
diff --git a/main.py b/main.py
index 444292f..72ea15e 100644
--- a/main.py
+++ b/main.py
@@ -142,15 +142,20 @@ def p_xt_1_xt(model, xt, t, vec):
 
 
 class DiffusionDataset(Dataset):
-    def __init__(self, dataset):
+    def __init__(self, dataset, autoencoder=None):
         super().__init__()
         self.dataset = dataset
+        self.autoencoder = autoencoder
 
     def __getitem__(self, index):
         # Get the image and the label.
         img, label = self.dataset[index]
         img = img.to(DEVICE)
 
+        # Encode the image.
+        if self.autoencoder is not None:
+            img = self.autoencoder.encode(img.unsqueeze(0)).squeeze(0)
+
         # Normalize the image.
         img = img * 2 - 1
 
@@ -177,22 +182,6 @@ class DiffusionDataset(Dataset):
         return len(self.dataset)
 
 
-class LatentDataset(Dataset):
-    def __init__(self, dataset, autoencoder):
-        super().__init__()
-        self.dataset = dataset
-        self.autoencoder = autoencoder
-
-    def __getitem__(self, index):
-        img, label = self.dataset[index]
-        img = img.to(DEVICE)
-        latent = self.autoencoder.encode(img.unsqueeze(0)).squeeze(0)
-        return latent, label
-
-    def __len__(self):
-        return len(self.dataset)
-
-
 def loss(y_pred, y_true):
     return nn.MSELoss()(y_pred, y_true)
 
@@ -238,16 +227,18 @@ dataset = datasets.MNIST(
 # dataset = FolderDataset('data/lfwcrop_color/faces')
 # dataset = FolderDataset('data/edface')
 
-autoencoder = Autoencoder((1, 28, 28), (1, 8, 8)).to(DEVICE)
+autoencoder = None
+autoencoder = Autoencoder(1, 1).to(DEVICE)
 autoencoder.load_state_dict(torch.load('autoencoder.pth'))
 autoencoder.eval()
-dataset = LatentDataset(dataset, autoencoder)
 
-img = dataset[0][0]
+img = dataset[0][0].to(DEVICE)
+if autoencoder is not None:
+    img = autoencoder.encode(img.unsqueeze(0)).squeeze(0)
 NB_CHANNEL, IMG_SIZE, _ = img.shape
 NB_LABEL = 10
 
-EPOCHS = 0
+EPOCHS = 1
 LEARNING_RATE = 2e-4
 
 
@@ -262,7 +253,7 @@ if __name__ == '__main__':
     # Load the model.
     model = UNet().to(DEVICE)
     try:
-        model.load_state_dict(torch.load('model.pth'))
+        model.load_state_dict(torch.load('mnist_latent_model.pth'))
     except FileNotFoundError:
         print("No model found, training a new one.")
         pass
@@ -271,7 +262,7 @@ if __name__ == '__main__':
     optimizer = torch.optim.Adam(model.parameters(), lr=LEARNING_RATE)
 
     # Define the training dataset.
-    train_dataset = DiffusionDataset(dataset)
+    train_dataset = DiffusionDataset(dataset, autoencoder)
     train_loader = DataLoader(train_dataset, batch_size=64, shuffle=True,
                               num_workers=4, persistent_workers=True)
     trainer = Trainer()
@@ -291,7 +282,10 @@ if __name__ == '__main__':
     with torch.no_grad():
         # Forward diffusion
         img, label = dataset[np.random.randint(0, len(dataset))]
-        img = img.to(DEVICE) * 2 - 1
+        img = img.to(DEVICE)
+        if autoencoder is not None:
+            img = autoencoder.encode(img.unsqueeze(0)).squeeze(0)
+        img = img * 2 - 1
 
         nb_plots = 10
         plots_id = [i for i in np.linspace(1, DIFFU_STEPS, nb_plots, dtype=int)]
@@ -322,7 +316,6 @@ if __name__ == '__main__':
         plt.suptitle("Forward diffusion")
         plt.tight_layout()
         plt.savefig("forward_diffusion.tmp.png")
-        exit()
 
 
         # Backward diffusion
@@ -345,7 +338,7 @@ if __name__ == '__main__':
                     plt.subplot(n_classes, nb_plots, t_plot_i + nb_plots * class_i + 1)
                     if class_i == 0:
                         plt.title(f"t={t}")
-                    plt.imshow(tensor_to_image(dataset.autoencoder.decode(x[class_i].unsqueeze(0)).squeeze(0)))
+                    plt.imshow(tensor_to_image(x[class_i]))
                     plt.axis("off")
         plt.tight_layout()
         plt.savefig("backward_diffusion.tmp.png")
@@ -368,10 +361,13 @@ if __name__ == '__main__':
         for i in range(nb_plots):
             for j in range(n_classes):
                 id = i * n_classes + j
+                img = x[id]
+                if autoencoder is not None:
+                    img = train_dataset.autoencoder.decode(img.unsqueeze(0)).squeeze(0)
                 plt.subplot(n_classes, nb_plots, id + 1)
-                plt.imshow(tensor_to_image(dataset.autoencoder.decode(x[id].unsqueeze(0)).squeeze(0)))
+                plt.imshow(tensor_to_image(img))
                 plt.axis("off")
         plt.tight_layout()
         plt.savefig("benchmark.tmp.png")
 
-        # plt.show()
\ No newline at end of file
+        # plt.show()