some changes reverted

Author: prishajain1

src/maxdiffusion/configs/base_wan_i2v_14b.yml

@@ -113,11 +113,10 @@ timestep_bias: {
 
 # Override parameters from checkpoints's scheduler.
 diffusion_scheduler_config: {
-  _class_name: 'FlaxUniPCMultistepScheduler',
+  _class_name: 'FlaxEulerDiscreteScheduler',
   prediction_type: 'epsilon',
   rescale_zero_terminal_snr: False,
-  timestep_spacing: 'trailing',
-  final_sigmas_type: 'sigma_min'
+  timestep_spacing: 'trailing'
 }
 
 # Output directory

src/maxdiffusion/pipelines/wan/wan_pipeline_i2v_2p1.py

@@ -189,7 +189,7 @@ def __call__(
         height=height,
         width=width,
         num_frames=num_frames,
-        dtype=jnp.float32,
+        dtype=image_embeds.dtype,
         rng=latents_rng,
         latents=latents,
         last_image=last_image_tensor,

src/maxdiffusion/schedulers/scheduling_unipc_multistep_flax.py

@@ -382,9 +382,9 @@ def multistep_uni_p_bh_update(
     check_nan_jit(alpha_s0, "P alpha_s0", step)
     check_nan_jit(sigma_s0, "P sigma_s0", step)
 
-    lambda_t = jnp.log(alpha_t + 1e-5) - jnp.log(sigma_t + 1e-5)
+    lambda_t = jnp.log(alpha_t + 1e-10) - jnp.log(sigma_t + 1e-10)
     check_nan_jit(lambda_t, "P lambda_t", step)
-    lambda_s0 = jnp.log(alpha_s0 + 1e-5) - jnp.log(sigma_s0 + 1e-5)
+    lambda_s0 = jnp.log(alpha_s0 + 1e-10) - jnp.log(sigma_s0 + 1e-10)
     check_nan_jit(lambda_s0, "P lambda_s0", step)
 
     h = lambda_t - lambda_s0