Fix bug when loading in weights after a resumed training run (#245)

CHANGELOG.md

@@ -5,7 +5,7 @@ All notable changes to this project will be documented in this file.
 The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/),
 and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0.html).
 
-## [0.10.0] - 2025-08-02
+## [0.10.0] - 2025-08-03
 
 # MFLUX v.0.10.0 Release Notes
 
@@ -24,6 +24,10 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 - **Enhanced Default Inference Steps**: Increased default inference steps for dev models from 14 to 25 for improved image quality
 - **Multiple Model Aliases Support**: Improved model configuration system to properly support multiple aliases per model, making model selection more flexible and robust
 
+### 🐛 Bug Fixes
+
+- **LoRA Resume Training**: Fixed critical bug where adapters created after training interruption would fail to load for generation with `AttributeError: 'list' object has no attribute 'weight'`. The issue occurred because the resume loading logic wasn't properly handling layers that are legitimately lists in the transformer architecture (like `attn.to_out`). (see [#224](https://github.com/filipstrand/mflux/issues/224))
+
 ### 🔧 Technical Requirements
 
 - **MLX Compatibility**: This release assumes MLX 0.27.0 and upwards for optimal performance and compatibility

src/mflux/dreambooth/lora_layers/lora_layers.py

@@ -140,17 +140,29 @@ class LoRALayers:
         if module_name == "attn" and attr_name == "to_out" and len(parts) >= 6 and parts[5] == "0":
             return module.to_out[0], f"transformer.transformer_blocks.{block_idx}.attn.to_out"
 
-        # Default case
+        # Default case - get the attribute and check if it's actually a list
         original_layer = getattr(module, attr_name)
-        if len(parts) == 6:
-            original_layer = original_layer[0]
+
+        # Some layers (like attn.to_out) are legitimately lists in the transformer architecture
+        # In such cases, we need to extract the first element for LoRA creation
+        if isinstance(original_layer, list):
+            # For list layers, we use the first element for LoRA creation
+            # but we need to return the list itself for proper storage
+            return original_layer, ".".join(parts)
+
         return original_layer, ".".join(parts)
 
     @staticmethod
     def _create_lora_layer(weights: dict, base_path: str, original_layer, scale: float):
         lora_A = weights[f"{base_path}.lora_A"]
         rank = lora_A.shape[1]
-        lora_layer = LoRALinear.from_linear(linear=original_layer, r=rank, scale=scale)
+
+        # Handle the case where original_layer is a list (like attn.to_out)
+        # In such cases, use the first element for LoRA creation (same as construction code)
+        is_list = isinstance(original_layer, list)
+        layer_for_lora = original_layer[0] if is_list else original_layer
+
+        lora_layer = LoRALinear.from_linear(linear=layer_for_lora, r=rank, scale=scale)
         lora_layer.lora_A = lora_A
         lora_layer.lora_B = weights[f"{base_path}.lora_B"]
         return lora_layer
@@ -177,7 +189,10 @@ class LoRALayers:
 
         # Create and store LoRA layer
         lora_layer = LoRALayers._create_lora_layer(weights, base_path, original_layer, scale)
-        lora_layers[storage_path] = lora_layer
+
+        # Store as list if original layer was a list (matching construction logic)
+        is_list = isinstance(original_layer, list)
+        lora_layers[storage_path] = [lora_layer] if is_list else lora_layer
 
     @staticmethod
     def _handle_single_transformer_blocks(
@@ -197,7 +212,10 @@ class LoRALayers:
 
         # Create and store LoRA layer
         lora_layer = LoRALayers._create_lora_layer(weights, base_path, original_layer, scale)
-        lora_layers[base_path] = lora_layer
+
+        # Store as list if original layer was a list (matching construction logic)
+        is_list = isinstance(original_layer, list)
+        lora_layers[base_path] = [lora_layer] if is_list else lora_layer
 
     @staticmethod
     def _handle_top_level_component(
@@ -205,7 +223,10 @@ class LoRALayers:
     ):
         original_layer = getattr(transformer, component_name)
         lora_layer = LoRALayers._create_lora_layer(weights, base_path, original_layer, scale)
-        lora_layers[base_path] = lora_layer
+
+        # Store as list if original layer was a list (matching construction logic)
+        is_list = isinstance(original_layer, list)
+        lora_layers[base_path] = [lora_layer] if is_list else lora_layer
 
     @staticmethod
     def set_transformer_block(transformer_block, dictionary: dict):

tests/dreambooth/test_resume_training.py

@@ -2,14 +2,20 @@ import os
 import shutil
 
 import mlx.core as mx
+import numpy as np
 
+from mflux.config.config import Config
+from mflux.config.model_config import ModelConfig
 from mflux.dreambooth.dreambooth import DreamBooth
 from mflux.dreambooth.dreambooth_initializer import DreamBoothInitializer
 from mflux.dreambooth.state.zip_util import ZipUtil
+from mflux.flux.flux import Flux1
 
 CHECKPOINT_3 = "tests/dreambooth/tmp/_checkpoints/0000003_checkpoint.zip"
 CHECKPOINT_4 = "tests/dreambooth/tmp/_checkpoints/0000004_checkpoint.zip"
 CHECKPOINT_5 = "tests/dreambooth/tmp/_checkpoints/0000005_checkpoint.zip"
+OUTPUT_DIR = "tests/dreambooth/tmp/_checkpoints/0000005_checkpoint"
+LORA_FILE = "tests/dreambooth/tmp/_checkpoints/0000005_checkpoint/0000005_adapter.safetensors"
 
 
 class TestResumeTraining:
@@ -66,6 +72,30 @@ class TestResumeTraining:
                     array1 = adapter_after_5_steps[key]
                     array2 = adapter_after_5_steps_resumed[key]
                     assert mx.array_equal(array1, array2)
+
+            # And: We want to confirm that the resumed weights can actually be loaded into Flux for generation...
+            ZipUtil.extract_all(zip_path=CHECKPOINT_5, output_dir=OUTPUT_DIR)
+            flux_with_resumed_lora = Flux1(
+                model_config=ModelConfig.dev(),
+                quantize=4,
+                lora_paths=[LORA_FILE],
+                lora_scales=[1.0],
+            )
+
+            # ...and we should be able to generate with the resumed LoRA weights (in this case, the actual image is nonsense and doesn't matter)
+            image = flux_with_resumed_lora.generate_image(
+                seed=42,
+                prompt="test",
+                config=Config(
+                    num_inference_steps=1,
+                    height=128,
+                    width=128,
+                ),
+            )
+
+            # Basic sanity check that we got a valid image
+            assert image.image is not None
+            assert np.array(image.image).shape == (128, 128, 3)
         finally:
             # cleanup
             TestResumeTraining.delete_folder("tests/dreambooth/tmp")