Revert "replace dropout_prob by dropout in vae"

This reverts commit cd7fb4f.

Author: younesbelkada

src/diffusers/models/vae_flax.py

@@ -127,7 +127,7 @@ class FlaxResnetBlock2D(nn.Module):
 
     in_channels: int
     out_channels: int = None
-    dropout: float = 0.0
+    dropout_prob: float = 0.0
     use_nin_shortcut: bool = None
     dtype: jnp.dtype = jnp.float32
 
@@ -144,7 +144,7 @@ def setup(self):
         )
 
         self.norm2 = nn.GroupNorm(num_groups=32, epsilon=1e-6)
-        self.dropout_layer = nn.Dropout(self.dropout)
+        self.dropout = nn.Dropout(self.dropout_prob)
         self.conv2 = nn.Conv(
             out_channels,
             kernel_size=(3, 3),
@@ -173,7 +173,7 @@ def __call__(self, hidden_states, deterministic=True):
 
         hidden_states = self.norm2(hidden_states)
         hidden_states = nn.swish(hidden_states)
-        hidden_states = self.dropout_layer(hidden_states, deterministic)
+        hidden_states = self.dropout(hidden_states, deterministic)
         hidden_states = self.conv2(hidden_states)
 
         if self.conv_shortcut is not None:
@@ -284,7 +284,7 @@ def setup(self):
             res_block = FlaxResnetBlock2D(
                 in_channels=in_channels,
                 out_channels=self.out_channels,
-                dropout=self.dropout,
+                dropout_prob=self.dropout,
                 dtype=self.dtype,
             )
             resnets.append(res_block)
@@ -335,7 +335,7 @@ def setup(self):
             res_block = FlaxResnetBlock2D(
                 in_channels=in_channels,
                 out_channels=self.out_channels,
-                dropout=self.dropout,
+                dropout_prob=self.dropout,
                 dtype=self.dtype,
             )
             resnets.append(res_block)
@@ -383,7 +383,7 @@ def setup(self):
             FlaxResnetBlock2D(
                 in_channels=self.in_channels,
                 out_channels=self.in_channels,
-                dropout=self.dropout,
+                dropout_prob=self.dropout,
                 dtype=self.dtype,
             )
         ]
@@ -399,7 +399,7 @@ def setup(self):
             res_block = FlaxResnetBlock2D(
                 in_channels=self.in_channels,
                 out_channels=self.in_channels,
-                dropout=self.dropout,
+                dropout_prob=self.dropout,
                 dtype=self.dtype,
             )
             resnets.append(res_block)