self conditioning step needs to also receive variance normalized noised image

Author: lucidrains

rin_pytorch/rin_pytorch.py

@@ -732,6 +732,8 @@ def forward(self, img, *args, **kwargs):
 
         noised_img = alpha * img + sigma * noise
 
+        noised_img = self.normalize_img_variance(noised_img)
+
         # in the paper, they had to use a really high probability of latent self conditioning, up to 90% of the time
         # slight drawback
 
@@ -745,7 +747,6 @@ def forward(self, img, *args, **kwargs):
 
         # predict and take gradient step
 
-        noised_img = self.normalize_img_variance(noised_img)
         pred = self.model(noised_img, times, self_cond, self_latents)
 
         if self.objective == 'x0':

setup.py

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