scheduler reverted to original

Author: prishajain1

src/maxdiffusion/pipelines/wan/wan_pipeline_i2v_2p1.py

@@ -25,7 +25,7 @@
 import jax
 import jax.numpy as jnp
 from jax.sharding import NamedSharding, PartitionSpec as P
-from ...schedulers.scheduling_flow_match_flax import FlaxFlowMatchScheduler
+from ...schedulers.scheduling_unipc_multistep_flax import FlaxUniPCMultistepScheduler
 from ...max_utils import randn_tensor
 
 class WanPipelineI2V_2_1(WanPipeline):
@@ -65,19 +65,6 @@ def _load_and_init(cls, config, restored_checkpoint=None, vae_only=False, load_t
       config=config,
     )
     return pipeline, transformer
-  
-  @classmethod
-  def load_scheduler(cls, config):
-      """Overrides the base scheduler loader to use Flow Matching for I2V."""
-      # Wan 2.1 I2V requires Flow Matching with these specific settings:
-      # shift=1.0, num_train_timesteps=1000, and usually reverse_sigmas=True (1.0 -> 0.0)
-      scheduler, scheduler_state = FlaxFlowMatchScheduler.from_pretrained(
-          config.pretrained_model_name_or_path,
-          subfolder="scheduler",
-          shift=1.0,
-          reverse_sigmas=True, 
-      )
-      return scheduler, scheduler_state
 
   @classmethod
   def from_pretrained(cls, config: HyperParameters, vae_only=False, load_transformer=True):
@@ -230,12 +217,12 @@ def __call__(
         self.scheduler_state, num_inference_steps=num_inference_steps, shape=latents.shape
     )
 
-    # if self.scheduler_state.last_sample is None or self.scheduler_state.step_index is None:
-    #   t0 = jnp.array(scheduler_state.timesteps, dtype=jnp.int32)[0]
-    #   dummy_noise = jnp.zeros_like(latents)
-    #   # This call initializes the internal state arrays
-    #   step_output = self.scheduler.step(scheduler_state, dummy_noise, t0, latents)
-    #   scheduler_state = step_output.state
+    if self.scheduler_state.last_sample is None or self.scheduler_state.step_index is None:
+      t0 = jnp.array(scheduler_state.timesteps, dtype=jnp.int32)[0]
+      dummy_noise = jnp.zeros_like(latents)
+      # This call initializes the internal state arrays
+      step_output = self.scheduler.step(scheduler_state, dummy_noise, t0, latents)
+      scheduler_state = step_output.state
     graphdef, state, rest_of_state = nnx.split(self.transformer, nnx.Param, ...)
     data_sharding = NamedSharding(self.mesh, P(*self.config.data_sharding))
     latents = jax.device_put(latents, data_sharding)
@@ -288,7 +275,7 @@ def run_inference_2_1_i2v(
     image_embeds: jnp.array,
     guidance_scale: float,
     num_inference_steps: int,
-    scheduler: FlaxFlowMatchScheduler,
+    scheduler: FlaxUniPCMultistepScheduler,
     scheduler_state,
     rng: jax.Array,
     expand_timesteps: bool,
@@ -336,16 +323,7 @@ def loop_body(step, vals):
                     s=step,
                     std=jnp.std(latents),
                     mean=jnp.mean(latents))
-    
-    step_output = scheduler.step(
-        state=scheduler_state,
-        model_output=noise_pred,
-        timestep=t,
-        sample=latents,
-        return_dict=True
-    )
-    
-    latents = step_output.prev_sample
+    latents, scheduler_state = scheduler.step(scheduler_state, noise_pred, t, latents).to_tuple()
     jax.debug.print("Step {s}: latents_next std={std}, mean={mean}",
                     s=step,
                     std=jnp.std(latents),