Tried to make the source clearer, added the TimeAppender Module.

diff --git a/minidiffusion.py b/minidiffusion.py
index 65ca947..27842d9 100755 (executable)
--- a/minidiffusion.py
+++ b/minidiffusion.py
@@ -146,6 +146,20 @@ class EMA:
 
 ######################################################################
 
+# Gets a pair (x, t) and appends t (scalar or 1d tensor) to x as an
+# additional dimension / channel
+
+class TimeAppender(nn.Module):
+    def __init__(self):
+        super().__init__()
+
+    def forward(self, u):
+        x, t = u
+        if not torch.is_tensor(t):
+            t = x.new_full((x.size(0),), t)
+        t = t.view((-1,) + (1,) * (x.dim() - 1)).expand_as(x[:,:1])
+        return torch.cat((x, t), 1)
+
 class ConvNet(nn.Module):
     def __init__(self, in_channels, out_channels):
         super().__init__()
@@ -153,7 +167,8 @@ class ConvNet(nn.Module):
         ks, nc = 5, 64
 
         self.core = nn.Sequential(
-            nn.Conv2d(in_channels, nc, ks, padding = ks//2),
+            TimeAppender(),
+            nn.Conv2d(in_channels + 1, nc, ks, padding = ks//2),
             nn.ReLU(),
             nn.Conv2d(nc, nc, ks, padding = ks//2),
             nn.ReLU(),
@@ -166,8 +181,8 @@ class ConvNet(nn.Module):
             nn.Conv2d(nc, out_channels, ks, padding = ks//2),
         )
 
-    def forward(self, x):
-        return self.core(x)
+    def forward(self, u):
+        return self.core(u)
 
 ######################################################################
 # Data
@@ -187,6 +202,7 @@ if train_input.dim() == 2:
     nh = 256
 
     model = nn.Sequential(
+        TimeAppender(),
         nn.Linear(train_input.size(1) + 1, nh),
         nn.ReLU(),
         nn.Linear(nh, nh),
@@ -198,7 +214,7 @@ if train_input.dim() == 2:
 
 elif train_input.dim() == 4:
 
-    model = ConvNet(train_input.size(1) + 1, train_input.size(1))
+    model = ConvNet(train_input.size(1), train_input.size(1))
 
 model.to(device)
 
@@ -214,10 +230,10 @@ def generate(size, alpha, alpha_bar, sigma, model, train_mean, train_std):
         x = torch.randn(size, device = device)
 
         for t in range(T-1, -1, -1):
+            output = model((x, t / (T - 1) - 0.5))
             z = torch.zeros_like(x) if t == 0 else torch.randn_like(x)
-            input = torch.cat((x, torch.full_like(x[:,:1], t / (T - 1) - 0.5)), 1)
             x = 1/torch.sqrt(alpha[t]) \
-                * (x - (1-alpha[t]) / torch.sqrt(1-alpha_bar[t]) * model(input)) \
+                * (x - (1-alpha[t]) / torch.sqrt(1-alpha_bar[t]) * output) \
                 + sigma[t] * z
 
         x = x * train_std + train_mean
@@ -245,8 +261,8 @@ for k in range(args.nb_epochs):
         t = torch.randint(T, (x0.size(0),) + (1,) * (x0.dim() - 1), device = x0.device)
         eps = torch.randn_like(x0)
         xt = torch.sqrt(alpha_bar[t]) * x0 + torch.sqrt(1 - alpha_bar[t]) * eps
-        input = torch.cat((xt, t.expand_as(x0[:,:1]) / (T - 1) - 0.5), 1)
-        loss = (eps - model(input)).pow(2).mean()
+        output = model((xt, t / (T - 1) - 0.5))
+        loss = (eps - output).pow(2).mean()
         acc_loss += loss.item() * x0.size(0)
 
         optimizer.zero_grad()
@@ -327,6 +343,9 @@ elif train_input.dim() == 4:
     x = generate((128,) + train_input.size()[1:], alpha, alpha_bar, sigma,
                  model, train_mean, train_std)
     x = 1 - x.clamp(min = 0, max = 255) / 255
-    torchvision.utils.save_image(x, f'diffusion_{args.data}.png', nrow = 16, pad_value = 0.8)
+
+    filename = f'diffusion_{args.data}.png'
+    print(f'saving {filename}')
+    torchvision.utils.save_image(x, filename, nrow = 16, pad_value = 0.8)
 
 ######################################################################