add centered back

patil-suraj · patil-suraj · commit 0dbc4779c8bf · 2022-07-01T12:50:34.000+02:00
diff --git a/src/diffusers/models/unet_sde_score_estimation.py b/src/diffusers/models/unet_sde_score_estimation.py
@@ -229,6 +229,7 @@ def __init__(
         self,
         image_size=1024,
         num_channels=3,
+        centered=False,
         attn_resolutions=(16,),
         ch_mult=(1, 2, 4, 8, 16, 32, 32, 32),
         conditional=True,
@@ -253,6 +254,7 @@ def __init__(
         self.register_to_config(
             image_size=image_size,
             num_channels=num_channels,
+            centered=centered,
             attn_resolutions=attn_resolutions,
             ch_mult=ch_mult,
             conditional=conditional,
@@ -457,7 +459,8 @@ def forward(self, x, timesteps, sigmas=None):
             temb = None
 
         # If input data is in [0, 1]
-        x = 2 * x - 1.0
+        if not self.config.centered:
+            x = 2 * x - 1.0
 
         # Downsampling block
         input_pyramid = None
