lint.

Author: jfacevedo-google

src/maxdiffusion/checkpointing/wan_checkpointer.py

@@ -15,35 +15,35 @@
 """
 
 from abc import ABC
-import jax
 from flax import nnx
 from maxdiffusion.checkpointing.checkpointing_utils import (create_orbax_checkpoint_manager)
 from ..pipelines.wan.wan_pipeline import WanPipeline
 from .. import max_logging, max_utils
 
 WAN_CHECKPOINT = "WAN_CHECKPOINT"
 
+
 class WanCheckpointer(ABC):
+
   def __init__(self, config, checkpoint_type):
     self.config = config
     self.checkpoint_type = checkpoint_type
 
     self.checkpoint_manager = create_orbax_checkpoint_manager(
-      self.config.checkpoint_dir,
-      enable_checkpointing=True,
-      save_interval_steps=1,
-      checkpoint_type=checkpoint_type,
-      dataset_type=config.dataset_type
+        self.config.checkpoint_dir,
+        enable_checkpointing=True,
+        save_interval_steps=1,
+        checkpoint_type=checkpoint_type,
+        dataset_type=config.dataset_type,
     )
-  
+
   def _create_optimizer(self, model, config, learning_rate):
     learning_rate_scheduler = max_utils.create_learning_rate_schedule(
-      learning_rate, config.learning_rate_schedule_steps, config.warmup_steps_fraction, config.max_train_steps
+        learning_rate, config.learning_rate_schedule_steps, config.warmup_steps_fraction, config.max_train_steps
     )
     tx = max_utils.create_optimizer(config, learning_rate_scheduler)
     return nnx.Optimizer(model, tx), learning_rate_scheduler
 
-
   def load_wan_configs_from_orbax(self, step):
     max_logging.log("Restoring stable diffusion configs")
     if step is None:
@@ -59,8 +59,8 @@ def load_checkpoint(self, step=None):
     model_configs = self.load_wan_configs_from_orbax(step)
 
     if model_configs:
-      raise NotImplemented("model configs should not exist in orbax")
+      raise NotImplementedError("model configs should not exist in orbax")
     else:
       pipeline = self.load_diffusers_checkpoint()
-    
-    return pipeline
+
+    return pipeline

src/maxdiffusion/generate_wan.py

@@ -20,38 +20,21 @@
 from absl import app
 from maxdiffusion.utils import export_to_video
 
+
 def run(config):
   print("seed: ", config.seed)
   pipeline = WanPipeline.from_pretrained(config)
   s0 = time.perf_counter()
-  
+
   # Skip layer guidance
   slg_layers = config.slg_layers
   slg_start = config.slg_start
   slg_end = config.slg_end
 
   prompt = [config.prompt] * jax.device_count()
-  negative_prompt= [config.negative_prompt] * jax.device_count()
-   
-  videos = pipeline(
-    prompt=prompt,
-    negative_prompt=negative_prompt,
-    height=config.height,
-    width=config.width,
-    num_frames=config.num_frames,
-    num_inference_steps=config.num_inference_steps,
-    guidance_scale=config.guidance_scale,
-    slg_layers=slg_layers,
-    slg_start=slg_start,
-    slg_end=slg_end
-  )
+  negative_prompt = [config.negative_prompt] * jax.device_count()
 
-  print("compile time: ", (time.perf_counter() - s0))
-  for i in range(len(videos)):
-    export_to_video(videos[i], f"wan_output_{config.seed}_{i}.mp4", fps=config.fps)
-  s0 = time.perf_counter()
-  with jax.profiler.trace("/tmp/trace/"):
-    videos = pipeline(
+  videos = pipeline(
       prompt=prompt,
       negative_prompt=negative_prompt,
       height=config.height,
@@ -61,7 +44,25 @@ def run(config):
       guidance_scale=config.guidance_scale,
       slg_layers=slg_layers,
       slg_start=slg_start,
-      slg_end=slg_end
+      slg_end=slg_end,
+  )
+
+  print("compile time: ", (time.perf_counter() - s0))
+  for i in range(len(videos)):
+    export_to_video(videos[i], f"wan_output_{config.seed}_{i}.mp4", fps=config.fps)
+  s0 = time.perf_counter()
+  with jax.profiler.trace("/tmp/trace/"):
+    videos = pipeline(
+        prompt=prompt,
+        negative_prompt=negative_prompt,
+        height=config.height,
+        width=config.width,
+        num_frames=config.num_frames,
+        num_inference_steps=config.num_inference_steps,
+        guidance_scale=config.guidance_scale,
+        slg_layers=slg_layers,
+        slg_start=slg_start,
+        slg_end=slg_end,
     )
   print("generation time: ", (time.perf_counter() - s0))
   for i in range(len(videos)):
@@ -74,4 +75,4 @@ def main(argv: Sequence[str]) -> None:
 
 
 if __name__ == "__main__":
-  app.run(main)
+  app.run(main)

src/maxdiffusion/maxdiffusion_utils.py

@@ -286,21 +286,23 @@ def get_dummy_flux_inputs(config, pipeline, batch_size):
 
   return (latents, timesteps, latents_ids, guidance_vec, t5_hidden_states, t5_ids, clip_hidden_states)
 
+
 def get_dummy_wan_inputs(config, pipeline, batch_size):
   latents = pipeline.prepare_latents(
-    batch_size,
-    vae_scale_factor_temporal=pipeline.vae_scale_factor_temporal,
-    vae_scale_factor_spatial=pipeline.vae_scale_factor_spatial,
-    height=config.height,
-    width=config.width,
-    num_frames=config.num_frames,
-    num_channels_latents=pipeline.transformer.config.in_channels
+      batch_size,
+      vae_scale_factor_temporal=pipeline.vae_scale_factor_temporal,
+      vae_scale_factor_spatial=pipeline.vae_scale_factor_spatial,
+      height=config.height,
+      width=config.width,
+      num_frames=config.num_frames,
+      num_channels_latents=pipeline.transformer.config.in_channels,
   )
   bsz = latents.shape[0]
   prompt_embeds = jax.random.normal(jax.random.key(config.seed), (batch_size, 512, 4096))
   timesteps = jnp.array([0] * bsz, dtype=jnp.int32)
   return (latents, prompt_embeds, timesteps)
 
+
 def calculate_wan_tflops(config, pipeline, batch_size, rngs, train):
   """
   Calculates jflux tflops.
@@ -309,10 +311,10 @@ def calculate_wan_tflops(config, pipeline, batch_size, rngs, train):
   """
   (latents, prompt_embeds, timesteps) = get_dummy_wan_inputs(config, pipeline, batch_size)
   return max_utils.calculate_model_tflops(
-    pipeline.transformer,
-    
+      pipeline.transformer,
   )
 
+
 def calculate_flux_tflops(config, pipeline, batch_size, rngs, train):
   """
   Calculates jflux tflops.