[Flax DDPM] Make key optional so default pipelines don't fail (huggingface#2176)

Make `key` optional so default pipelines don't fail.

Author: pcuenca

src/diffusers/schedulers/scheduling_ddpm_flax.py

@@ -198,7 +198,7 @@ def step(
         model_output: jnp.ndarray,
         timestep: int,
         sample: jnp.ndarray,
-        key: jax.random.KeyArray,
+        key: Optional[jax.random.KeyArray] = None,
         return_dict: bool = True,
     ) -> Union[FlaxDDPMSchedulerOutput, Tuple]:
         """
@@ -221,6 +221,9 @@ def step(
         """
         t = timestep
 
+        if key is None:
+            key = jax.random.PRNGKey(0)
+
         if model_output.shape[1] == sample.shape[1] * 2 and self.config.variance_type in ["learned", "learned_range"]:
             model_output, predicted_variance = jnp.split(model_output, sample.shape[1], axis=1)
         else: