misc

GreenWizard2015 · Apr 16, 2024 · b1d2319 · b1d2319
1 parent 523a779
commit b1d2319
Showing 1 changed file with 5 additions and 4 deletions.
diff --git a/NN/RestorationModel/CRestorationModel.py b/NN/RestorationModel/CRestorationModel.py
@@ -84,11 +84,12 @@ def call(self, latents, pos, T, V, residual, R=None, training=False):
   def reverse(self, latents, pos, reverseArgs, training, value, residual, index):
     EPos = self._encodePos(pos, training=training, args=reverseArgs.get('decoder', {}))
     latents = self._addResiduals(latents, residual)
-    fakeR = reverseArgs.get('blurRadius', 0.0)
-    latents = self._addRadius(latents, R=None, fakeR=fakeR, training=training)
 
-    def denoiser(x, t, mask=None):
-      args = dict(condition=latents, coords=EPos, timestep=t, V=x)
+    def denoiser(x, t, mask=None, **kwargs):
+      fakeR = kwargs.get('blurRadius', reverseArgs.get('fakeR', 0.0))
+      latentsPlus = self._addRadius(latents, R=None, fakeR=fakeR, training=training)
+
+      args = dict(condition=latentsPlus, coords=EPos, timestep=t, V=x)
       residuals = residual
       if mask is not None:
         args = {k: masked(v, mask) for k, v in args.items()}