Styling changes

Author: Perseus14

src/maxdiffusion/generate_wan.py

@@ -193,7 +193,12 @@ def run(config, pipeline=None, filename_prefix=""):
     pipeline, _, _ = checkpoint_loader.load_checkpoint()
 
   # If LoRA is specified, inject layers and load weights.
-  if config.enable_lora and hasattr(config, "lora_config") and config.lora_config and config.lora_config["lora_model_name_or_path"]:
+  if (
+      config.enable_lora
+      and hasattr(config, "lora_config")
+      and config.lora_config
+      and config.lora_config["lora_model_name_or_path"]
+  ):
     if model_key == WAN2_1:
       lora_loader = Wan2_1NNXLoraLoader()
       lora_config = config.lora_config

src/maxdiffusion/models/lora_nnx.py

@@ -267,7 +267,9 @@ def merge_lora(model: nnx.Module, state_dict: dict, rank: int, scale: float, tra
     )
 
 
-def merge_lora_for_scanned(model: nnx.Module, state_dict: dict, rank: int, scale: float, translate_fn=None, dtype: str = "float32"):
+def merge_lora_for_scanned(
+    model: nnx.Module, state_dict: dict, rank: int, scale: float, translate_fn=None, dtype: str = "float32"
+):
   """
   Device-Side Optimized Merge for Scanned Layers.
   Now supports diff and diff_b.