Added randn_tensors

Author: prishajain1

src/maxdiffusion/pipelines/wan/wan_pipeline_i2v_2p1.py

@@ -218,7 +218,7 @@ def run_inference_2_1_i2v(
     num_inference_steps: int,
     scheduler: FlaxUniPCMultistepScheduler,
     scheduler_state,
-    rng: jax.random.KeyArray,
+    rng: jax.Array,
     expand_timesteps: bool,
     first_frame_mask: Optional[jnp.array],
 ):

src/maxdiffusion/pipelines/wan/wan_pipeline_i2v_2p2.py

@@ -143,7 +143,7 @@ def __call__(
     image_embeds: Optional[jax.Array] = None,
     last_image: Optional[PipelineImageInput] = None,
     output_type: Optional[str] = "np",
-    rng: Optional[jax.random.KeyArray] = None,
+    rng: Optional[jax.Array] = None,
   ):
     height = height or self.config.height
     width = width or self.config.width
@@ -236,7 +236,7 @@ def run_inference_2_2_i2v(
     num_inference_steps: int,
     scheduler: FlaxUniPCMultistepScheduler,
     scheduler_state,
-    rng: jax.random.KeyArray,
+    rng: jax.Array,
     expand_timesteps: bool,
 ):
   do_classifier_free_guidance = guidance_scale > 1.0