more fixes

Author: prishajain1

src/maxdiffusion/generate_wan.py

@@ -18,10 +18,8 @@
 import os
 from maxdiffusion.pipelines.wan.wan_pipeline import WanPipeline
 from maxdiffusion.checkpointing.wan_checkpointer import WanCheckpointer
-from functools import partial
 from maxdiffusion import pyconfig, max_logging, max_utils
 from absl import app
-from absl import flags
 from maxdiffusion.utils import export_to_video
 from google.cloud import storage
 import flax
@@ -127,7 +125,7 @@ def run(config, pipeline=None, filename_prefix=""):
   # Using global_batch_size_to_train_on so not to create more config variables
   prompt = [config.prompt] * config.global_batch_size_to_train_on
   negative_prompt = [config.negative_prompt] * config.global_batch_size_to_train_on
-  
+
   max_logging.log(
       f"Num steps: {config.num_inference_steps}, height: {config.height}, width: {config.width}, frames: {config.num_frames}"
   )
@@ -162,4 +160,4 @@ def main(argv: Sequence[str]) -> None:
 
 
 if __name__ == "__main__":
-  app.run(main)
+  app.run(main)

src/maxdiffusion/pipelines/wan/wan_pipeline.py

@@ -12,7 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from abc import ABC, abstractmethod
+from abc import abstractmethod
 from typing import List, Union, Optional, Type
 from functools import partial
 import numpy as np
@@ -466,7 +466,7 @@ def _denormalize_latents(self, latents: jax.Array) -> jax.Array:
       latents = latents / latents_std + latents_mean
       latents = latents.astype(jnp.float32)
       return latents
-    
+
   def _decode_latents_to_video(self, latents: jax.Array) -> np.ndarray:
       """Decodes latents to video frames and postprocesses."""
       with self.mesh, nn_partitioning.axis_rules(self.config.logical_axis_rules):
@@ -508,7 +508,7 @@ def _get_subclass(cls, model_key: str) -> Type['WanPipeline']:
             f"Supported keys are: {list(cls._SUBCLASS_MAP.keys())}"
         )
     return subclass
-  
+
   @classmethod
   def from_checkpoint(cls, model_key: str, config: HyperParameters, restored_checkpoint=None, vae_only=False, load_transformer=True):
     subclass = cls._get_subclass(model_key)
@@ -708,7 +708,6 @@ def _load_and_init(cls, config, restored_checkpoint=None, vae_only=False, load_t
     common_components = cls._create_common_components(config, vae_only)
     low_noise_transformer, high_noise_transformer = None, None
     if not vae_only and load_transformer:
-        rngs = nnx.Rngs(jax.random.key(config.seed))
         low_noise_transformer = super().load_transformer(
             devices_array=common_components["devices_array"],
             mesh=common_components["mesh"],

src/maxdiffusion/tests/wan_checkpointer_test.py

@@ -14,12 +14,7 @@
 import unittest
 from unittest.mock import patch, MagicMock
 
-from maxdiffusion.checkpointing.wan_checkpointer import (
-    WanCheckpointer2_1, 
-    WanCheckpointer2_2, 
-    WAN_CHECKPOINT
-)
-
+from maxdiffusion.checkpointing.wan_checkpointer import WanCheckpointer2_1, WanCheckpointer2_2
 
 class WanCheckpointer2_1Test(unittest.TestCase):
   """Tests for WAN 2.1 checkpointer."""
@@ -240,15 +235,14 @@ def test_load_checkpoint_with_optimizer_in_high_noise(self, mock_wan_pipeline, m
     self.assertEqual(step, 1)
 
 
-class WanCheckpointerFactoryTest(unittest.TestCase):
-  """Tests for checkpointer factory/selection logic."""
+class WanCheckpointerEdgeCasesTest(unittest.TestCase):
+  """Tests for edge cases and error handling."""
 
   def setUp(self):
     self.config = MagicMock()
-    self.config.checkpoint_dir = "/tmp/wan_checkpoint_factory_test"
+    self.config.checkpoint_dir = "/tmp/wan_checkpoint_edge_test"
     self.config.dataset_type = "test_dataset"
 
-
   @patch("maxdiffusion.checkpointing.wan_checkpointer.create_orbax_checkpoint_manager")
   @patch("maxdiffusion.checkpointing.wan_checkpointer.WanPipeline2_1")
   def test_load_checkpoint_with_explicit_none_step(self, mock_wan_pipeline, mock_create_manager):