From 482a2d06c930549e0d04328355586e179d7681ea Mon Sep 17 00:00:00 2001 From: Filip Strand Date: Sun, 3 Aug 2025 12:46:31 +0200 Subject: [PATCH] Fix bug when loading in weights after a resumed training run (#245) --- CHANGELOG.md | 6 +++- .../dreambooth/lora_layers/lora_layers.py | 35 +++++++++++++++---- tests/dreambooth/test_resume_training.py | 30 ++++++++++++++++ 3 files changed, 63 insertions(+), 8 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 72c4894..189f808 100644 --- a/CHANGELOG.md +++ b/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 diff --git a/src/mflux/dreambooth/lora_layers/lora_layers.py b/src/mflux/dreambooth/lora_layers/lora_layers.py index f794d4a..493110f 100644 --- a/src/mflux/dreambooth/lora_layers/lora_layers.py +++ b/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): diff --git a/tests/dreambooth/test_resume_training.py b/tests/dreambooth/test_resume_training.py index 7423372..c25ca28 100644 --- a/tests/dreambooth/test_resume_training.py +++ b/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")