From 3d8d5563bb3c902e3dd6fb0a480dcb9221a29bf2 Mon Sep 17 00:00:00 2001 From: gdamms Date: Tue, 28 May 2024 17:03:04 +0200 Subject: testing with latent diffusion --- main.py | 38 +++++++++++++++++++++++++++++--------- 1 file changed, 29 insertions(+), 9 deletions(-) (limited to 'main.py') diff --git a/main.py b/main.py index 6ba7eb6..7bbbc33 100644 --- a/main.py +++ b/main.py @@ -12,6 +12,7 @@ import os import cv2 from trainer import Trainer +from autoencoder import Autoencoder class UNet(nn.Module): @@ -176,6 +177,22 @@ 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(self.autoencoder.device) + latent = self.autoencoder.encode(img.unsqueeze(0)) + return latent, label + + def __len__(self): + return len(self.dataset) + + def loss(y_pred, y_true): return nn.MSELoss()(y_pred, y_true) @@ -204,12 +221,12 @@ BETA = torch.cat((torch.tensor([0.], device=DEVICE), BETA)) ALPHA = 1 - BETA ALPHA_BAR = torch.cumprod(ALPHA, dim=0) -# dataset = datasets.MNIST( -# root="./data", -# train=True, -# download=True, -# transform=transforms.ToTensor(), -# ) +dataset = datasets.MNIST( + root="./data", + train=True, + download=True, + transform=transforms.ToTensor(), +) # dataset = datasets.LFWPeople( # root="./data", # download=True, @@ -219,13 +236,16 @@ ALPHA_BAR = torch.cumprod(ALPHA, dim=0) # ]), # ) # dataset = FolderDataset('data/lfwcrop_color/faces') -dataset = FolderDataset('data/edface') +# dataset = FolderDataset('data/edface') + +autoencoder = Autoencoder(1, 64).to(DEVICE) +dataset = LatentDataset(dataset, autoencoder) img = dataset[0][0] NB_CHANNEL, IMG_SIZE, _ = img.shape -NB_LABEL = 1 +NB_LABEL = 10 -EPOCHS = 100 +EPOCHS = 10 LEARNING_RATE = 2e-4 -- cgit v1.3.1