From 989834728791abe50c14aad4294e1ef10d9bbf35 Mon Sep 17 00:00:00 2001 From: Francois Fleuret Date: Mon, 15 Aug 2022 09:27:47 +0200 Subject: [PATCH] Tried to make the source clearer, added the TimeAppender Module. --- minidiffusion.py | 37 ++++++++++++++++++++++++++++--------- 1 file changed, 28 insertions(+), 9 deletions(-) diff --git a/minidiffusion.py b/minidiffusion.py index 65ca947..27842d9 100755 --- 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) ###################################################################### -- 2.39.5