Add checkpoints its still not clear

checkpoints/model_Epoch10.pth

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:bf1f268fa95855fd879e4690b920b7820cf6e80848850dbacc0b479d69ff1953
+size 6012986

checkpoints/model_Epoch20.pth

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:55ff29fd4bbe990d286d489a88c712aef197713e5e61c536664b3c3b67232e36
+size 6012986

checkpoints/model_Epoch30.pth

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:a323611d36eb7c639bb5f1de6f4d913e7889be5fe4303a8ea4212f5fcc213475
+size 6012986

checkpoints/model_Epoch40.pth

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:f8d804cb3d2292dc9c4242c7a5914302026a65a1bf1b70fc1ac2bd1f73e174d8
+size 6012986

checkpoints/model_Epoch50.pth

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:b351a9568c0eab7d0b76a090cf7827aaf9e9b54f7ffed21442619695cc4d9be1
+size 6012986

models.py

@@ -43,21 +43,27 @@ class ResidualBlock(nn.Module):
 
 
 class UnetUp(nn.Module):
-    def __init__(self, in_channels, out_channels) -> None:
-        super(UnetUp,self).__init__()
-
-        self.model = nn.Sequential(
-            nn.ConvTranspose2d(in_channels,out_channels,2,2),
-            ResidualBlock(out_channels,out_channels),
-            ResidualBlock(out_channels,out_channels),
-        )
+    def __init__(self, in_channels, out_channels):
+        super(UnetUp, self).__init__()
+        
+        # Create a list of layers for the upsampling block
+        # The block consists of a ConvTranspose2d layer for upsampling, followed by two ResidualConvBlock layers
+        layers = [
+            nn.ConvTranspose2d(in_channels, out_channels, 2, 2),
+            ResidualBlock(out_channels, out_channels),
+            ResidualBlock(out_channels, out_channels),
+        ]
+        
+        # Use the layers to create a sequential model
+        self.model = nn.Sequential(*layers)
 
     def forward(self, x, skip):
-        x = torch.cat([x,skip],1)
-
+        # Concatenate the input tensor x with the skip connection tensor along the channel dimension
+        x = torch.cat((x, skip), 1)
+        
+        # Pass the concatenated tensor through the sequential model and return the output
         x = self.model(x)
-        return x
-    
+        return x 
 class UnetDown(nn.Module):
     def __init__(self, input_channels, out_channels) -> None:
         super(UnetDown,self).__init__()
@@ -106,9 +112,9 @@ class ContextUnet(nn.Module):
         self.to_vec = nn.Sequential(nn.AvgPool2d((4)),nn.GELU())
 
         self.timeembed1 = EmbedFC(1, 2 *n_feat)
-        self.timeembed2 = EmbedFC(1,n_feat)
+        self.timeembed2 = EmbedFC(1,embed_dm=1*n_feat)
         self.contextembed1 = EmbedFC(n_cfeat,2 * n_feat)
-        self.contextembed2 = EmbedFC(n_cfeat,n_feat)
+        self.contextembed2 = EmbedFC(n_cfeat,1*n_feat)
 
         self.up0 = nn.Sequential(
             nn.ConvTranspose2d(2 * n_feat,2*n_feat,self.h // 4,self.h // 4),
@@ -137,15 +143,13 @@ class ContextUnet(nn.Module):
         if c is None:
             c = torch.zeros(x.shape[0],self.n_cfeat).to(x)
         
-        cemb1 = self.contextembed1(c).view(-1,self.n_cfeat*2,1,1)
-        temb1 = self.timeembed1(t).view(-1,self.n_cfeat * 2,1,1)
-        cemb2 = self.contextembed2(c).view(-1,self.n_cfeat,1,1)
-        temb2 = self.timeembed2(t).view(-1,self.n_cfeat,1,1)
-
-        up0 = self.up0(hidden_vec)
-        up1 =self.up1(up0*cemb1 + temb1,down2)
-        up2 = self.up2(up1*cemb2+temb2,down1)
-
-        out = self.out(torch.cat((up2,x),1))
-
+        cemb1 = self.contextembed1(c).view(-1,self.n_feat*2,1,1)
+        temb1 = self.timeembed1(t).view(-1,self.n_feat * 2,1,1)
+        cemb2 = self.contextembed2(c).view(-1,self.n_feat,1,1)
+        temb2 = self.timeembed2(t).view(-1,self.n_feat,1,1)
+
+        up1 = self.up0(hidden_vec)
+        up2 = self.up1(cemb1*up1 + temb1, down2)  # add and multiply embeddings
+        up3 = self.up2(cemb2*up2 + temb2, down1)
+        out = self.out(torch.cat((up3, x), 1))
         return out

train.py

@@ -3,6 +3,8 @@ from utils import *
 from torch.utils.data import DataLoader
 from models import *
 from tqdm.auto import tqdm
+import os
+import torch.nn.functional as F
 
 timesteps = 500
 beta1 = 1e-4
@@ -15,7 +17,7 @@ height = 16
 save_dir="./checkpoints"
 
 batch_size = 100
-n_epoch = 40
+n_epoch = 60
 lrate = 1e-3
 
 
@@ -25,11 +27,17 @@ a_bt = torch.cumsum(a_t.log(),0).exp()
 a_bt[0] = 1
 
 
+
+transform = transforms.Compose([
+    transforms.ToTensor(),                # from [0,255] to range [0.0,1.0]
+    transforms.Normalize((0.5,), (0.5,))  # range [-1,1]
+
+])
 dataset = CustomDataset("./sprites_1788_16x16.npy", "./sprite_labels_nc_1788_16x16.npy", transform, null_context=False)
 dataloader = DataLoader(dataset, batch_size=batch_size, shuffle=True, num_workers=1)
 
 
-nn_model = ContextUnet(3,n_feat,n_cfeat,height)
+nn_model = ContextUnet(3,n_feat,n_cfeat,height).to(device)
 optim = torch.optim.Adam(nn_model.parameters(),lrate)
 
 def perturb_input(x, t, noise):
@@ -46,7 +54,7 @@ for epoch in range(n_epoch):
 
         x = x.to(device)
 
-        t = torch.randint(1,timesteps+1,x.shape[0]).to(device)
+        t = torch.randint(1,timesteps+1,(x.shape[0],)).to(device)
         noise = torch.randn_like(x)
         x_pert = perturb_input(x,t,noise)
 
@@ -56,8 +64,9 @@ for epoch in range(n_epoch):
         loss.backward()
         optim.step()
 
-    if epoch % 1 == 0 and epoch >0:
+    if epoch % 10 == 0 and epoch > 0:
+        print(f"Epoch: {epoch} | Loss: {loss}")
         if not os.path.exists(save_dir):
             os.mkdir(save_dir)
-        torch.save(nn_model,save_dir + f"model_Epoch{epoch}.pth")
+        torch.save(nn_model,save_dir + f"/model_Epoch{epoch}.pth")
         print("Saved model")

utils.py

@@ -66,16 +66,10 @@ def plot_sample(x_gen_store,n_sample,nrows,save_dir, fn,  w, save=False):
     return ani
 
 
-transform = transforms.Compose([
-    transforms.ToTensor(),                # from [0,255] to range [0.0,1.0]
-    transforms.Normalize((0.5,), (0.5,))  # range [-1,1]
-
-])
-
 class CustomDataset(Dataset):
     def __init__(self, sfilename, lfilename, transform, null_context=False):
-        self.sprites = np.load(sfilename)
-        self.slabels = np.load(lfilename)
+        self.sprites = np.load(sfilename,allow_pickle=True,fix_imports=True,encoding='latin1')
+        self.slabels = np.load(lfilename,allow_pickle=True,fix_imports=True,encoding='latin1')
         print(f"sprite shape: {self.sprites.shape}")
         print(f"labels shape: {self.slabels.shape}")
         self.transform = transform
@@ -83,10 +77,13 @@ class CustomDataset(Dataset):
         self.sprites_shape = self.sprites.shape
         self.slabel_shape = self.slabels.shape
                 
+    # Return the number of images in the dataset
     def __len__(self):
         return len(self.sprites)
     
+    # Get the image and label at a given index
     def __getitem__(self, idx):
+        # Return the image and label as a tuple
         if self.transform:
             image = self.transform(self.sprites[idx])
             if self.null_context:
@@ -94,3 +91,7 @@ class CustomDataset(Dataset):
             else:
                 label = torch.tensor(self.slabels[idx]).to(torch.int64)
         return (image, label)
+
+    def getshapes(self):
+        # return shapes of data and labels
+        return self.sprites_shape, self.slabel_shape