conditioning on final resnet block

andrewmoise · Jun 17, 2022 · b4fb880 · b4fb880
1 parent 9fd05f1
commit b4fb880
Show file tree

Hide file tree

Showing 2 changed files with 5 additions and 5 deletions.
diff --git a/denoising_diffusion_pytorch/denoising_diffusion_pytorch.py b/denoising_diffusion_pytorch/denoising_diffusion_pytorch.py
@@ -314,10 +314,8 @@ def __init__(
         default_out_dim = channels * (1 if not learned_variance else 2)
         self.out_dim = default(out_dim, default_out_dim)
 
-        self.final_conv = nn.Sequential(
-            block_klass(dim * 2, dim),
-            nn.Conv2d(dim, self.out_dim, 1)
-        )
+        self.final_res_block = block_klass(dim * 2, dim, time_emb_dim = time_dim)
+        self.final_conv = nn.Conv2d(dim, self.out_dim, 1)
 
     def forward(self, x, time):
         x = self.init_conv(x)
@@ -346,6 +344,8 @@ def forward(self, x, time):
             x = upsample(x)
 
         x = torch.cat((x, r), dim = 1)
+
+        x = self.final_res_block(x, t)
         return self.final_conv(x)
 
 # gaussian diffusion trainer class

diff --git a/setup.py b/setup.py
@@ -3,7 +3,7 @@
 setup(
   name = 'denoising-diffusion-pytorch',
   packages = find_packages(),
-  version = '0.20.0',
+  version = '0.20.1',
   license='MIT',
   description = 'Denoising Diffusion Probabilistic Models - Pytorch',
   author = 'Phil Wang',