residual conditional

NN/RestorationModel/CRepeatedRestorator.py

@@ -48,6 +48,7 @@ def train_step(self, x0, latents, positions, params, xT=None):
     xT = tf.stop_gradient(trainStep['value'])
 
     for i in range(1, self._N):
+      params['residual'] = xT
       trainStep = self._restorator.train_step(
         x0=x0,
         xT=xT,

NN/RestorationModel/CRestorationModel.py

@@ -9,6 +9,7 @@ class CRestorationModel(tf.keras.Model):
   def __init__(self,
     decoder, restorator,
     posEncoder, timeEncoder,
+    residualCondition=False,
     **kwargs
   ):
     assert posEncoder is not None, "posEncoder is not provided"
@@ -18,6 +19,7 @@ def __init__(self,
     self._restorator = restorator
     self._posEncoder = posEncoder
     self._timeEncoder = timeEncoder
+    self._residualCondition = residualCondition
     return
 
   def _encodeTime(self, t, training):
@@ -60,6 +62,13 @@ def denoiser(x, t, mask=None):
         args = {k: masked(v, mask) for k, v in args.items()}
         residuals = masked(residual, mask)
 
+      if self._residualCondition: # add residuals to the condition
+        B = tf.shape(residuals)[0]
+        args['condition'] = tf.reshape( # ensure that we know the shape of the condition
+          tf.concat([args['condition'], residuals], axis=-1),
+          [B, -1]
+        )
+
       res = self._decoder(**args, training=training)
       return self._withResidual(res, residuals)
 

NN/RestorationModel/CSequentialRestorator.py

@@ -35,6 +35,7 @@ def train_step(self, x0, latents, positions, params, xT=None):
     xT = tf.stop_gradient(trainStep['value'])
     # the rest of the restorators are trained sequentially
     for restorator in self._restorators:
+      params['residual'] = xT
       trainStep = restorator.train_step(
         x0=x0,
         xT=xT,

NN/RestorationModel/__init__.py

@@ -31,6 +31,7 @@ def restorationModel_from_config(config):
       restorator=restorator,
       posEncoder=encoding_from_config(config['position encoding']),
       timeEncoder=encoding_from_config(config['time encoding']),
+      residualCondition=config.get('residual condition', False),
     )
 
   if 'repeated' == name: