default functions for normalizing image data from 0 to 1 to -1 to 1 and back when sampling

Author: lucidrains

rectified_flow_pytorch/rectified_flow.py

@@ -30,6 +30,14 @@ def append_dims(t, ndims):
     shape = t.shape
     return t.reshape(*shape, *((1,) * ndims))
 
+# normalizing helpers
+
+def normalize_to_neg_one_to_one(img):
+    return img * 2 - 1
+
+def unnormalize_to_zero_to_one(t):
+    return (t + 1) * 0.5
+
 # losses
 
 class LPIPSLoss(Module):
@@ -100,7 +108,9 @@ def __init__(
         loss_fn: Literal['mse', 'pseudo_huber'] | Module = 'mse',
         loss_fn_kwargs: dict = dict(),
         data_shape: Tuple[int, ...] | None = None,
-        immiscible = False
+        immiscible = False,
+        data_normalize_fn = normalize_to_neg_one_to_one,
+        data_unnormalize_fn = unnormalize_to_zero_to_one
     ):
         super().__init__()
         self.model = model
@@ -135,6 +145,11 @@ def __init__(
 
         self.immiscible = immiscible
 
+        # normalizing fn
+
+        self.data_normalize_fn = data_normalize_fn
+        self.data_unnormalize_fn = data_unnormalize_fn
+
     @property
     def device(self):
         return next(self.model.parameters()).device
@@ -177,7 +192,8 @@ def ode_fn(t, x):
         sampled_data = trajectory[-1]
 
         self.train(was_training)
-        return sampled_data
+
+        return self.data_unnormalize_fn(sampled_data)
 
     def forward(
         self,
@@ -187,6 +203,8 @@ def forward(
     ):
         batch, *data_shape = data.shape
 
+        data = self.data_normalize_fn(data)
+
         self.data_shape = default(self.data_shape, data_shape)
 
         # x0 - gaussian noise, x1 - data