trying fix for NaN

Author: prishajain1

src/maxdiffusion/configs/base_wan_i2v_14b.yml

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

src/maxdiffusion/models/embeddings_flax.py

@@ -287,7 +287,6 @@ def __call__(self, encoder_hidden_states_image: jax.Array) -> jax.Array:
         padding_size = target_seq_len - current_seq_len
         padding = jnp.zeros((B, padding_size, D_out), dtype=hidden_states.dtype)
         hidden_states = jnp.concatenate([hidden_states, padding], axis=1)
-        print(f"[DEBUG EMB] Padded image embeds from {current_seq_len} to {target_seq_len}. New shape: {hidden_states.shape}")
 
     return hidden_states
 

src/maxdiffusion/pipelines/wan/wan_pipeline_2_1.py

@@ -63,7 +63,7 @@ def _load_and_init(cls, config, restored_checkpoint=None, vae_only=False, load_t
   @classmethod
   def from_pretrained(cls, config: HyperParameters, vae_only=False, load_transformer=True):
     pipeline , transformer = cls._load_and_init(config, None, vae_only, load_transformer)
-    transformer = cls.quantize_transformer(config, transformer, pipeline, pipeline.mesh)
+    pipeline.transformer = cls.quantize_transformer(config, transformer, pipeline, pipeline.mesh)
     return pipeline
 
   @classmethod

src/maxdiffusion/pipelines/wan/wan_pipeline_2_2.py

@@ -72,8 +72,8 @@ def _load_and_init(cls, config, restored_checkpoint=None, vae_only=False, load_t
   @classmethod
   def from_pretrained(cls, config: HyperParameters, vae_only=False, load_transformer=True):
     pipeline, low_noise_transformer, high_noise_transformer = cls._load_and_init(config, None, vae_only, load_transformer)
-    low_noise_transformer = cls.quantize_transformer(config, low_noise_transformer, pipeline, pipeline.mesh)
-    high_noise_transformer = cls.quantize_transformer(config, high_noise_transformer, pipeline, pipeline.mesh)
+    pipeline.low_noise_transformer = cls.quantize_transformer(config, low_noise_transformer, pipeline, pipeline.mesh)
+    pipeline.high_noise_transformer = cls.quantize_transformer(config, high_noise_transformer, pipeline, pipeline.mesh)
     return pipeline
 
   @classmethod

src/maxdiffusion/pipelines/wan/wan_pipeline_i2v_2p1.py

@@ -68,7 +68,8 @@ def _load_and_init(cls, config, restored_checkpoint=None, vae_only=False, load_t
 
   @classmethod
   def from_pretrained(cls, config: HyperParameters, vae_only=False, load_transformer=True):
-    pipeline , _ = cls._load_and_init(config, None, vae_only, load_transformer)
+    pipeline , transformer = cls._load_and_init(config, None, vae_only, load_transformer)
+    pipeline.transformer = cls.quantize_transformer(config, transformer, pipeline, pipeline.mesh)
     return pipeline
 
   @classmethod

src/maxdiffusion/pipelines/wan/wan_pipeline_i2v_2p2.py

@@ -65,7 +65,9 @@ def _load_and_init(cls, config, restored_checkpoint=None, vae_only=False, load_t
 
   @classmethod
   def from_pretrained(cls, config: HyperParameters, vae_only=False, load_transformer=True):
-    pipeline, _, _ = cls._load_and_init(config, None, vae_only, load_transformer)
+    pipeline, low_noise_transformer, high_noise_transformer = cls._load_and_init(config, None, vae_only, load_transformer)
+    pipeline.low_noise_transformer = cls.quantize_transformer(config, low_noise_transformer, pipeline, pipeline.mesh)
+    pipeline.high_noise_transformer = cls.quantize_transformer(config, high_noise_transformer, pipeline, pipeline.mesh)
     return pipeline
 
   @classmethod

src/maxdiffusion/schedulers/scheduling_unipc_multistep_flax.py

@@ -518,13 +518,13 @@ def solve_for_rhos_p(R_mat, b_vec, current_order):
     check_nan_jit(pred_res, "P pred_res", step)
 
     if self.config.predict_x0:
-      x_t_ = sigma_t / (sigma_s0 + 1e-8) * x - alpha_t * h_phi_1 * m0
+      x_t_ = sigma_t / (sigma_s0) * x - alpha_t * h_phi_1 * m0
       check_nan_jit(x_t_, "P x_t_ term", step)
       term2 = alpha_t * B_h * pred_res
       check_nan_jit(term2, "P term2", step)
       x_t = x_t_ - term2
     else:  # Predict epsilon
-      x_t_ = alpha_t / (alpha_s0 + 1e-8) * x - sigma_t * h_phi_1 * m0
+      x_t_ = alpha_t / (alpha_s0) * x - sigma_t * h_phi_1 * m0
       check_nan_jit(x_t_, "P x_t_ term eps", step)
       term2 = sigma_t * B_h * pred_res
       check_nan_jit(term2, "P term2 eps", step)