lint/format files.

Author: jfacevedo-google

src/maxdiffusion/models/wan/wan_utils.py

@@ -28,6 +28,7 @@ def rename_for_nnx(key):
     new_key = key[:-1] + ("scale",)
   return new_key
 
+
 def rename_for_custom_trasformer(key):
   renamed_pt_key = key.replace("model.diffusion_model.", "")
 
@@ -53,6 +54,7 @@ def rename_for_custom_trasformer(key):
 
   return renamed_pt_key
 
+
 def load_fusionx_transformer(pretrained_model_name_or_path: str, eval_shapes: dict, device: str, hf_download: bool = True):
   device = jax.devices(device)[0]
   with jax.default_device(device):
@@ -74,7 +76,7 @@ def load_fusionx_transformer(pretrained_model_name_or_path: str, eval_shapes: di
         random_flax_state_dict[string_tuple] = flattened_dict[key]
       for pt_key, tensor in tensors.items():
         renamed_pt_key = rename_key(pt_key)
-        
+
         renamed_pt_key = rename_for_custom_trasformer(renamed_pt_key)
 
         pt_tuple_key = tuple(renamed_pt_key.split("."))
@@ -89,6 +91,7 @@ def load_fusionx_transformer(pretrained_model_name_or_path: str, eval_shapes: di
       jax.clear_caches()
       return flax_state_dict
 
+
 def load_causvid_transformer(pretrained_model_name_or_path: str, eval_shapes: dict, device: str, hf_download: bool = True):
   device = jax.devices(device)[0]
   with jax.default_device(device):

src/maxdiffusion/pyconfig.py

@@ -118,7 +118,10 @@ def wan_init(raw_keys):
       transformer_pretrained_model_name_or_path = raw_keys["wan_transformer_pretrained_model_name_or_path"]
       if transformer_pretrained_model_name_or_path == "":
         raw_keys["wan_transformer_pretrained_model_name_or_path"] = raw_keys["pretrained_model_name_or_path"]
-      elif transformer_pretrained_model_name_or_path == CAUSVID_TRANSFORMER_MODEL_NAME_OR_PATH or transformer_pretrained_model_name_or_path == WAN_21_FUSION_X_MODEL_NAME_OR_PATH:
+      elif (
+          transformer_pretrained_model_name_or_path == CAUSVID_TRANSFORMER_MODEL_NAME_OR_PATH
+          or transformer_pretrained_model_name_or_path == WAN_21_FUSION_X_MODEL_NAME_OR_PATH
+      ):
         # Set correct parameters for CausVid in case of user error.
         raw_keys["guidance_scale"] = 1.0
         num_inference_steps = raw_keys["num_inference_steps"]