dfhsfhs

1 files changed, 28 insertions, 12 deletions

diff --git a/main.py b/main.py
index 1757574..fca0b40 100644
--- a/main.py
+++ b/main.py
@@ -87,6 +87,7 @@ SIGMA2 = BETA
 
 def q_xt_xt_1(xt_1, t):
     t_ind = t.to(dtype=torch.long) if isinstance(t, torch.Tensor) else t
+    t_ind -= 1
     beta = BETA[t_ind]
     mean = torch.sqrt(1 - beta) * xt_1
     std = beta
@@ -95,6 +96,7 @@ def q_xt_xt_1(xt_1, t):
 
 def q_xt_x0(x0, t):
     t_ind = t.to(dtype=torch.long) if isinstance(t, torch.Tensor) else t
+    t_ind -= 1
     alpha_bar = ALPHA_BAR[t_ind]
     mean = torch.sqrt(alpha_bar) * x0
     std = 1 - alpha_bar
@@ -103,6 +105,7 @@ def q_xt_x0(x0, t):
 
 def p_xt_1_xt(model, xt, t, vec):
     t_ind = t.to(dtype=torch.long) if isinstance(t, torch.Tensor) else t
+    t_ind -= 1
     alpha_bar_t = ALPHA_BAR[t_ind]
     alpha_bar_t_1 = ALPHA_BAR[t_ind-1]
     alpha_t = ALPHA[t_ind]
@@ -193,36 +196,49 @@ model = UNet().to(DEVICE)
 model.load_state_dict(torch.load('model.pth'))
 
 # Define the optimizer.
-optimizer = torch.optim.Adam(model.parameters(), lr=5e-5)
+optimizer = torch.optim.Adam(model.parameters(), lr=3e-5)
 
 # Define the training dataset.
 train_dataset = MNISTDiffusionDataset(train=True)
-train_loader = DataLoader(train_dataset, batch_size=64, shuffle=True, num_workers=4)
+train_loader = DataLoader(train_dataset, batch_size=64, shuffle=True, num_workers=0)
 trainer = Trainer()
 criterion = loss
 epochs = 5
 
-# # Train the model.
-# trainer.train(model, train_loader, epochs, optimizer, criterion)
-
-# # Save the model.
-# torch.save(model.state_dict(), "model.pth")
+# Train the model.
+trainer.train(model, train_loader, epochs, optimizer, criterion)
 
+# Save the model.
+torch.save(model.state_dict(), "model.pth")
 
 ##############
 # Evaluation #
 ##############
 
 x = torch.randn(1, 1, 28, 28).to(DEVICE)
-vec = torch.nn.functional.one_hot(torch.randint(0, 10, (1, 1)), num_classes=10).to(device=DEVICE, dtype=torch.float32)
-print(vec)
+n = torch.randint(0, 10, (1, 1))
+vec = torch.nn.functional.one_hot(n, num_classes=10).to(device=DEVICE, dtype=torch.float32)
+img_vec = model.encodevec(vec)
+img_vec = F.relu(img_vec)
+img_vec = img_vec.view(-1, 1, 28, 28)
+img_vec = img_vec.cpu().detach().numpy()
+fig = plt.figure()
+plt.imshow(img_vec[0, 0], cmap="gray")
+plt.axis("off")
+plt.savefig("vec.tmp.png")
 
 fig = plt.figure(figsize=(DIFFU_STEPS, 2))
-for ti in range(1, DIFFU_STEPS):
-    t = torch.tensor([[ti]], device=DEVICE, dtype=torch.float32)
+for ti in range(DIFFU_STEPS):
+    t = torch.tensor([[ti+1]], device=DEVICE, dtype=torch.float32)
+    eps_theta = model(x, t, vec)
+    x_theta = x - eps_theta
     x = p_xt_1_xt(model, x, t, vec).sample()
-    ax = fig.add_subplot(1, DIFFU_STEPS, ti)
+    ax = fig.add_subplot(2, DIFFU_STEPS, ti+1)
     ax.imshow(x[0, 0].detach().cpu(), cmap="gray")
     ax.axis("off")
+    ax = fig.add_subplot(2, DIFFU_STEPS, ti+1+DIFFU_STEPS)
+    ax.imshow(x_theta[0, 0].detach().cpu(), cmap="gray")
+    ax.axis("off")
 fig.tight_layout()
+plt.title(f'{n.item()}')
 fig.savefig("diffused.tmp.png")
\ No newline at end of file