address #329

Author: lucidrains

denoising_diffusion_pytorch/classifier_free_guidance.py

@@ -349,8 +349,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_res_block = ResnetBlock(dim * 2, dim, time_emb_dim = time_dim, classes_emb_dim = classes_dim)
-        self.final_conv = nn.Conv2d(dim, self.out_dim, 1)
+        self.final_res_block = ResnetBlock(init_dim * 2, init_dim, time_emb_dim = time_dim, classes_emb_dim = classes_dim)
+        self.final_conv = nn.Conv2d(init_dim, self.out_dim, 1)
 
     def forward_with_cond_scale(
         self,

denoising_diffusion_pytorch/denoising_diffusion_pytorch.py

@@ -380,8 +380,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_res_block = resnet_block(dim * 2, dim)
-        self.final_conv = nn.Conv2d(dim, self.out_dim, 1)
+        self.final_res_block = resnet_block(init_dim * 2, init_dim)
+        self.final_conv = nn.Conv2d(init_dim, self.out_dim, 1)
 
     @property
     def downsample_factor(self):

denoising_diffusion_pytorch/denoising_diffusion_pytorch_1d.py

@@ -343,8 +343,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_res_block = resnet_block(dim * 2, dim)
-        self.final_conv = nn.Conv1d(dim, self.out_dim, 1)
+        self.final_res_block = resnet_block(init_dim * 2, init_dim)
+        self.final_conv = nn.Conv1d(init_dim, self.out_dim, 1)
 
     def forward(self, x, time, x_self_cond = None):
         if self.self_condition:

denoising_diffusion_pytorch/repaint.py

@@ -376,8 +376,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_res_block = ResnetBlock(dim * 2, dim, time_emb_dim = time_dim)
-        self.final_conv = nn.Conv2d(dim, self.out_dim, 1)
+        self.final_res_block = ResnetBlock(init_dim * 2, init_dim, time_emb_dim = time_dim)
+        self.final_conv = nn.Conv2d(init_dim, self.out_dim, 1)
 
     @property
     def downsample_factor(self):

denoising_diffusion_pytorch/simple_diffusion.py

@@ -428,8 +428,8 @@ def __init__(
         default_out_dim = input_channels
         self.out_dim = default(out_dim, default_out_dim)
 
-        self.final_res_block = ResnetBlock(dim * 2, dim, time_emb_dim = time_dim)
-        self.final_conv = nn.Conv2d(dim, self.out_dim, 1)
+        self.final_res_block = ResnetBlock(init_dim * 2, init_dim, time_emb_dim = time_dim)
+        self.final_conv = nn.Conv2d(init_dim, self.out_dim, 1)
 
     def forward(self, x, time):
         x = self.init_img_transform(x)

denoising_diffusion_pytorch/version.py

@@ -1 +1 @@
-__version__ = '2.0.10'
+__version__ = '2.0.12'