new line added

Author: prishajain1

src/maxdiffusion/checkpointing/wan_checkpointer2_2.py

@@ -24,7 +24,6 @@
 from ..pipelines.wan.wan_pipeline2_2 import WanPipeline
 from .. import max_logging, max_utils
 import orbax.checkpoint as ocp
-from etils import epath
 
 WAN_CHECKPOINT = "WAN_CHECKPOINT"
 
@@ -206,4 +205,4 @@ def config_to_json(model_or_config):
 
   # Save the checkpoint
   self.checkpoint_manager.save(train_step, args=save_args)
-  max_logging.log(f"Checkpoint for step {train_step} saved.")
+  max_logging.log(f"Checkpoint for step {train_step} saved.")

src/maxdiffusion/generate_wan.py

@@ -18,7 +18,6 @@
 import os
 from maxdiffusion.pipelines.wan.wan_pipeline import WanPipeline as WanPipeline2_1
 from maxdiffusion.pipelines.wan.wan_pipeline2_2 import WanPipeline as WanPipeline2_2
-from functools import partial
 from maxdiffusion import pyconfig, max_logging, max_utils
 from absl import app
 from absl import flags
@@ -143,7 +142,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}"
   )

src/maxdiffusion/pipelines/wan/wan_pipeline.py

@@ -387,7 +387,7 @@ def from_checkpoint(cls, config: HyperParameters, restored_checkpoint=None, vae_
           transformer = cls.load_transformer(
               devices_array=devices_array, mesh=mesh, rngs=rngs, config=config, restored_checkpoint=restored_checkpoint, subfolder="transformer"
           )
-          
+
       text_encoder = cls.load_text_encoder(config=config)
       tokenizer = cls.load_tokenizer(config=config)
 

src/maxdiffusion/pipelines/wan/wan_pipeline2_2.py

@@ -722,4 +722,4 @@ def high_noise_branch(operands):
     )
 
     latents, scheduler_state = scheduler.step(scheduler_state, noise_pred, t, latents).to_tuple()
-  return latents
+  return latents