add one more residual

denoising_diffusion_pytorch/denoising_diffusion_pytorch.py

@@ -243,7 +243,7 @@ class Unet(nn.Module):
 
         self.channels = channels
 
-        init_dim = default(init_dim, dim // 3 * 2)
+        init_dim = default(init_dim, dim)
         self.init_conv = nn.Conv2d(channels, init_dim, 7, padding = 3)
 
         dims = [init_dim, *map(lambda m: dim * m, dim_mults)]
@@ -288,8 +288,8 @@ class Unet(nn.Module):
         self.mid_attn = Residual(PreNorm(mid_dim, Attention(mid_dim)))
         self.mid_block2 = block_klass(mid_dim, mid_dim, time_emb_dim = time_dim)
 
-        for ind, (dim_in, dim_out) in enumerate(reversed(in_out[1:])):
-            is_last = ind >= (num_resolutions - 1)
+        for ind, (dim_in, dim_out) in enumerate(reversed(in_out)):
+            is_last = ind == (len(in_out) - 1)
 
             self.ups.append(nn.ModuleList([
                 block_klass(dim_out * 2, dim_in, time_emb_dim = time_dim),
@@ -308,6 +308,8 @@ class Unet(nn.Module):
 
     def forward(self, x, time):
         x = self.init_conv(x)
+        r = x.clone()
+
         t = self.time_mlp(time)
 
         h = []
@@ -330,7 +332,7 @@ class Unet(nn.Module):
             x = attn(x)
             x = upsample(x)
 
-        x = torch.cat((x, h.pop()), dim = 1)
+        x = torch.cat((x, r), dim = 1)
         return self.final_conv(x)
 
 # gaussian diffusion trainer class

setup.py

@@ -3,7 +3,7 @@ from setuptools import setup, find_packages
 setup(
   name = 'denoising-diffusion-pytorch',
   packages = find_packages(),
-  version = '0.19.0',
+  version = '0.19.2',
   license='MIT',
   description = 'Denoising Diffusion Probabilistic Models - Pytorch',
   author = 'Phil Wang',