Linting and styling changes based on All checks passed!

Author: Perseus14

src/maxdiffusion/loaders/lora_conversion_utils.py

@@ -391,7 +391,7 @@ def _convert_to_ai_toolkit_cat(sds_sd, ait_sd, sds_key, ait_keys, dims=None):
     ait_up_keys = [k + ".lora_B.weight" for k in ait_keys]
     if not is_sparse:
       # down_weight is copied to each split
-      ait_sd.update({k: down_weight for k in ait_down_keys})
+      ait_sd.update(dict.fromkeys(ait_down_keys, down_weight))
 
       # up_weight is split to each split
       ait_sd.update({k: v for k, v in zip(ait_up_keys, torch.split(up_weight, dims, dim=0))})  # noqa: C416
@@ -534,7 +534,7 @@ def handle_qkv(sds_sd, ait_sd, sds_key, ait_keys, dims=None):
     ait_up_keys = [k + ".lora_B.weight" for k in ait_keys]
 
     # down_weight is copied to each split
-    ait_sd.update({k: down_weight for k in ait_down_keys})
+    ait_sd.update(dict.fromkeys(ait_down_keys, down_weight))
 
     # up_weight is split to each split
     ait_sd.update({k: v for k, v in zip(ait_up_keys, torch.split(up_weight, dims, dim=0))})  # noqa: C416
@@ -615,7 +615,7 @@ def translate_wan_nnx_path_to_diffusers_lora(nnx_path_str, scan_layers=False):
     Translates WAN NNX path to Diffusers/LoRA keys.
     Verified against wan_utils.py mappings.
     """
-    
+
     # --- 1. Embeddings (Exact Matches) ---
     if nnx_path_str == 'condition_embedder.text_embedder.linear_1':
         return 'diffusion_model.text_embedding.0'
@@ -640,11 +640,11 @@ def translate_wan_nnx_path_to_diffusers_lora(nnx_path_str, scan_layers=False):
         "attn1.key":       "self_attn.k",
         "attn1.value":     "self_attn.v",
         "attn1.proj_attn": "self_attn.o",
-        
+
         # Self Attention Norms (QK Norm) - Added per your request
         "attn1.norm_q":    "self_attn.norm_q",
         "attn1.norm_k":    "self_attn.norm_k",
-        
+
         # Cross Attention (attn2)
         "attn2.query":     "cross_attn.q",
         "attn2.key":       "cross_attn.k",
@@ -654,22 +654,22 @@ def translate_wan_nnx_path_to_diffusers_lora(nnx_path_str, scan_layers=False):
         # Cross Attention Norms (QK Norm) - Added per your request
         "attn2.norm_q":    "cross_attn.norm_q",
         "attn2.norm_k":    "cross_attn.norm_k",
-        
+
         # Feed Forward (ffn)
         "ffn.act_fn.proj": "ffn.0",  # Up proj
         "ffn.proj_out":    "ffn.2",  # Down proj
-        
+
         # Global Norms & Modulation
-        "norm2.layer_norm": "norm3", 
+        "norm2.layer_norm": "norm3",
         "scale_shift_table": "modulation",
-        "proj_out": "head.head" 
+        "proj_out": "head.head"
     }
 
     # --- 3. Translation Logic ---
     if scan_layers:
         # Scanned Pattern: "blocks.attn1.query" -> "diffusion_model.blocks.{}.self_attn.q"
         if nnx_path_str.startswith("blocks."):
-            inner_suffix = nnx_path_str[len("blocks."):] 
+            inner_suffix = nnx_path_str[len("blocks."):]
             if inner_suffix in suffix_map:
                 return f"diffusion_model.blocks.{{}}.{suffix_map[inner_suffix]}"
     else:
@@ -679,6 +679,6 @@ def translate_wan_nnx_path_to_diffusers_lora(nnx_path_str, scan_layers=False):
             idx, inner_suffix = m.group(1), m.group(2)
             if inner_suffix in suffix_map:
                 return f"diffusion_model.blocks.{idx}.{suffix_map[inner_suffix]}"
-    
+
     return None
 

src/maxdiffusion/loaders/wan_lora_nnx_loader.py

@@ -15,8 +15,6 @@
 """NNX-based LoRA loader for WAN models."""
 
 from flax import nnx
-import jax
-import re
 from .lora_base import LoRABaseMixin
 from .lora_pipeline import StableDiffusionLoraLoaderMixin
 from ..models import lora_nnx
@@ -46,7 +44,8 @@ def load_lora_weights(
     lora_loader = StableDiffusionLoraLoaderMixin()
 
     merge_fn = lora_nnx.merge_lora_for_scanned if scan_layers else lora_nnx.merge_lora
-    translate_fn = lambda nnx_path_str: lora_conversion_utils.translate_wan_nnx_path_to_diffusers_lora(nnx_path_str, scan_layers=scan_layers)
+    def translate_fn(nnx_path_str):
+      return lora_conversion_utils.translate_wan_nnx_path_to_diffusers_lora(nnx_path_str, scan_layers=scan_layers)
 
     # Handle high noise model
     if hasattr(pipeline, "transformer") and transformer_weight_name:
@@ -84,7 +83,8 @@ def load_lora_weights(
     lora_loader = StableDiffusionLoraLoaderMixin()
 
     merge_fn = lora_nnx.merge_lora_for_scanned if scan_layers else lora_nnx.merge_lora
-    translate_fn = lambda nnx_path_str: lora_conversion_utils.translate_wan_nnx_path_to_diffusers_lora(nnx_path_str, scan_layers=scan_layers)
+    def translate_fn(nnx_path_str: str):
+      return lora_conversion_utils.translate_wan_nnx_path_to_diffusers_lora(nnx_path_str, scan_layers=scan_layers)
 
     # Handle high noise model
     if hasattr(pipeline, "high_noise_transformer") and high_noise_weight_name: