add a missing ff block after latents attending to patches, addressing #3

Author: lucidrains

rin_pytorch/rin_pytorch.py

@@ -285,6 +285,7 @@ def __init__(
         dim_latent = default(dim_latent, dim)
 
         self.latents_attend_to_patches = Attention(dim_latent, dim_context = dim, norm = True, norm_context = True, **attn_kwargs)
+        self.latents_cross_attn_ff = FeedForward(dim_latent)
 
         self.latent_self_attns = nn.ModuleList([])
         for _ in range(latent_self_attn_depth):
@@ -309,6 +310,8 @@ def forward(self, patches, latents, t):
 
         latents = self.latents_attend_to_patches(latents, patches, time = t) + latents
 
+        latents = self.latents_cross_attn_ff(latents, time = t) + latents
+
         # latent self attention
 
         for attn, ff in self.latent_self_attns:
@@ -583,7 +586,7 @@ def ddpm_sample(self, shape, time_difference = None):
                 x_start = model_output
 
             elif self.objective == 'eps':
-                x_start = (img - sigma * model_output) / alpha
+                x_start = (img - sigma * model_output) / alpha.clamp(min = 1e-8)
 
             # clip x0
 
@@ -648,7 +651,7 @@ def ddim_sample(self, shape, time_difference = None):
                 x_start = model_output
 
             elif self.objective == 'eps':
-                x_start = (img - sigma * model_output) / alpha
+                x_start = (img - sigma * model_output) / alpha.clamp(min = 1e-8)
 
             # clip x0
 

setup.py

@@ -3,7 +3,7 @@
 setup(
   name = 'RIN-pytorch',
   packages = find_packages(exclude=[]),
-  version = '0.2.0',
+  version = '0.3.0',
   license='MIT',
   description = 'RIN - Recurrent Interface Network - Pytorch',
   author = 'Phil Wang',