fix(flux): support .weight variant for attention norm parameter mappings (invoke-ai#9190)

Author: Fortyseven

invokeai/backend/model_manager/load/model_loaders/flux.py

@@ -984,13 +984,13 @@ def _convert_double_block_key(self, key: str, tensor: torch.Tensor, converted: d
             return f"{prefix}.attn.to_add_out.weight"
 
         # Attention norms
-        if "img_attn.norm.query_norm.scale" in rest:
+        if "img_attn.norm.query_norm.scale" in rest or "img_attn.norm.query_norm.weight" in rest:
             return f"{prefix}.attn.norm_q.weight"
-        elif "img_attn.norm.key_norm.scale" in rest:
+        elif "img_attn.norm.key_norm.scale" in rest or "img_attn.norm.key_norm.weight" in rest:
             return f"{prefix}.attn.norm_k.weight"
-        elif "txt_attn.norm.query_norm.scale" in rest:
+        elif "txt_attn.norm.query_norm.scale" in rest or "txt_attn.norm.query_norm.weight" in rest:
             return f"{prefix}.attn.norm_added_q.weight"
-        elif "txt_attn.norm.key_norm.scale" in rest:
+        elif "txt_attn.norm.key_norm.scale" in rest or "txt_attn.norm.key_norm.weight" in rest:
             return f"{prefix}.attn.norm_added_k.weight"
 
         # MLP layers
@@ -1020,9 +1020,9 @@ def _convert_single_block_key(self, key: str, tensor: torch.Tensor, converted: d
             return f"{prefix}.attn.to_out.weight"
 
         # Norms
-        if "norm.query_norm.scale" in rest:
+        if "norm.query_norm.scale" in rest or "norm.query_norm.weight" in rest:
             return f"{prefix}.attn.norm_q.weight"
-        elif "norm.key_norm.scale" in rest:
+        elif "norm.key_norm.scale" in rest or "norm.key_norm.weight" in rest:
             return f"{prefix}.attn.norm_k.weight"
 
         return key
@@ -1328,13 +1328,13 @@ def _convert_double_block_key(self, key: str, tensor, converted: dict) -> str |
         elif "txt_attn.proj.weight" in rest:
             return f"{prefix}.attn.to_add_out.weight"
 
-        if "img_attn.norm.query_norm.scale" in rest:
+        if "img_attn.norm.query_norm.scale" in rest or "img_attn.norm.query_norm.weight" in rest:
             return f"{prefix}.attn.norm_q.weight"
-        elif "img_attn.norm.key_norm.scale" in rest:
+        elif "img_attn.norm.key_norm.scale" in rest or "img_attn.norm.key_norm.weight" in rest:
             return f"{prefix}.attn.norm_k.weight"
-        elif "txt_attn.norm.query_norm.scale" in rest:
+        elif "txt_attn.norm.query_norm.scale" in rest or "txt_attn.norm.query_norm.weight" in rest:
             return f"{prefix}.attn.norm_added_q.weight"
-        elif "txt_attn.norm.key_norm.scale" in rest:
+        elif "txt_attn.norm.key_norm.scale" in rest or "txt_attn.norm.key_norm.weight" in rest:
             return f"{prefix}.attn.norm_added_k.weight"
 
         if "img_mlp.0.weight" in rest:
@@ -1359,9 +1359,9 @@ def _convert_single_block_key(self, key: str, tensor, converted: dict) -> str |
         elif "linear2.weight" in rest:
             return f"{prefix}.attn.to_out.weight"
 
-        if "norm.query_norm.scale" in rest:
+        if "norm.query_norm.scale" in rest or "norm.query_norm.weight" in rest:
             return f"{prefix}.attn.norm_q.weight"
-        elif "norm.key_norm.scale" in rest:
+        elif "norm.key_norm.scale" in rest or "norm.key_norm.weight" in rest:
             return f"{prefix}.attn.norm_k.weight"
 
         return key