Fix

Author: Perseus14

src/maxdiffusion/models/lora_nnx.py

@@ -253,7 +253,16 @@ def merge_lora(model: nnx.Module, state_dict: dict, scale: float, translate_fn=N
             w_diff = weights.get("diff", None)
             b_diff = weights.get("diff_b", None)
 
-            if w_diff is not None: w_diff = np.array(w_diff)
+            if w_diff is not None:
+                w_diff = np.array(w_diff)
+                # Transpose weights from PyTorch OIHW/OIDHW to Flax HWIO/DHWIO if needed.
+                if isinstance(module, nnx.Conv):
+                    if w_diff.ndim == 5:
+                        w_diff = w_diff.transpose((2,3,4,1,0))
+                    elif w_diff.ndim == 4:
+                        w_diff = w_diff.transpose((2,3,1,0))
+                elif isinstance(module, nnx.Linear) and w_diff.ndim == 2:
+                     w_diff = w_diff.transpose((1,0))
             if b_diff is not None: b_diff = np.array(b_diff)
 
             # Check for Bias existence
@@ -351,9 +360,14 @@ def merge_lora_for_scanned(model: nnx.Module, state_dict: dict, scale: float, tr
                     if stack_w_diff is None: 
                         stack_w_diff = np.zeros(module.kernel.shape, dtype=np.float32)
                     wd = np.array(w["diff"])
-                    # Reshape if 1x1 conv diff (squeeze spatial dims if needed, or broadcast)
-                    if is_conv and wd.ndim != 5: wd = wd.reshape(1, 1, 1, in_feat, out_feat)
-                    elif is_linear and wd.ndim != 2: wd = wd.reshape(in_feat, out_feat)
+                    # Transpose weights from PyTorch OIHW/OIDHW to Flax HWIO/DHWIO if needed.
+                    if is_conv:
+                        if wd.ndim == 5:
+                            wd = wd.transpose((2,3,4,1,0))
+                        elif wd.ndim == 4:
+                            wd = wd.transpose((2,3,1,0))
+                    elif is_linear and wd.ndim == 2:
+                        wd = wd.transpose((1,0))
 
                     stack_w_diff[i] = wd
                     has_diff = True