k_proj bias is needed for qk norm

Author: arda-argmax

python/src/diffusionkit/mlx/mmdit.py

@@ -693,7 +693,7 @@ def __init__(
         self.kv_proj_embed_dim = self.per_head_dim * n_heads
 
         # Note: key bias is redundant due to softmax invariance
-        self.k_proj = nn.Linear(embed_dim, self.kv_proj_embed_dim, bias=False)
+        self.k_proj = nn.Linear(embed_dim, self.kv_proj_embed_dim)
         self.q_proj = nn.Linear(embed_dim, embed_dim)
         self.v_proj = nn.Linear(embed_dim, self.kv_proj_embed_dim)
         self.o_proj = nn.Linear(embed_dim, embed_dim)

python/src/diffusionkit/mlx/model_io.py

@@ -248,9 +248,6 @@ def flux_state_dict_adjustments(state_dict, prefix="", hidden_size=3072, mlp_rat
         for k, v in state_dict.items()
     }
 
-    # Remove k_proj bias
-    state_dict = {k: v for k, v in state_dict.items() if "k_proj.bias" not in k}
-
     state_dict["x_embedder.proj.weight"] = mx.expand_dims(
         mx.expand_dims(state_dict["x_embedder.proj.weight"], axis=1), axis=1
     )