Fix

Author: Perseus14

src/maxdiffusion/models/lora_nnx.py

@@ -341,14 +341,46 @@ def merge_lora_for_scanned(model: nnx.Module, state_dict: dict, scale: float, tr
                 if lora_found_in_module:
                     module.kernel.value = kernel_value_updated
                     assigned_count += 1
-                    max_logging.log(f"Merged LoRA into scanned layer {nnx_path_str}")
                 else:
                     max_logging.log(f"Scanned layer {nnx_path_str} matched template but no LoRA weights found for any block.")
             else:
                  max_logging.log(f"Scanned NNX layer '{nnx_path_str}' could not be translated to a LoRA key template.")
 
         # Handle scanned Conv layers (ndim=5)
         elif isinstance(module, nnx.Conv) and module.kernel.ndim == 5:
-             max_logging.log(f"Merging LoRA into scanned Conv layers not implemented: {nnx_path_str}")
+            if module.kernel_size != (1, 1):
+                max_logging.warn(f"Skipping merge for scanned Conv layer {nnx_path_str} with kernel size {module.kernel_size}, only 1x1 is supported for merging.")
+                continue
+
+            lora_key_template = translate_fn(nnx_path_str) if translate_fn else None
+            if lora_key_template:
+                num_layers, _, _, in_features, out_features = module.kernel.shape
+                kernel_value_updated = module.kernel.value
+                lora_found_in_module = False
+                for i in range(num_layers):
+                    lora_key = lora_key_template.format(i)
+                    if lora_key in lora_params and "down" in lora_params[lora_key] and "up" in lora_params[lora_key]:
+                        weights = lora_params[lora_key]
+                        down_w, up_w = weights["down"], weights["up"]
+
+                        if down_w.ndim == 4:
+                            down_w = jnp.squeeze(down_w)
+                        if up_w.ndim == 4:
+                            up_w = jnp.squeeze(up_w)
+                        
+                        rank = down_w.shape[0]
+                        alpha = weights.get("alpha", rank)
+                        current_scale = scale * alpha / rank
+                        delta_i = (down_w.T @ up_w.T).reshape(1, 1, in_features, out_features) * current_scale
+                        kernel_value_updated = kernel_value_updated.at[i].add(delta_i)
+                        lora_found_in_module = True
+
+                if lora_found_in_module:
+                    module.kernel.value = kernel_value_updated
+                    assigned_count += 1
+                else:
+                    max_logging.log(f"Scanned 1x1 Conv layer {nnx_path_str} matched template but no LoRA weights found for any block.")
+            else:
+                 max_logging.log(f"Scanned 1x1 Conv layer '{nnx_path_str}' could not be translated to a LoRA key template.")
 
     max_logging.log(f"Merged weights into {assigned_count} scanned layers in {type(model).__name__}.")