diff --git a/CHANGELOG.md b/CHANGELOG.md index e02b8ab..9b3c9f0 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -5,6 +5,15 @@ 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.13.3] - 2025-12-06 + +### 🐛 Bug Fixes + +- **LoRA save bloat prevention**: Bake and strip LoRA wrappers before sharding to avoid exploding shard counts/sizes when saving quantized models with multiple/mismatched LoRAs (see [issue #217 comment](https://github.com/filipstrand/mflux/issues/217#issuecomment-3615321206)). +- **Regression test hardening**: LoRA model-saving tests now include size guardrails (5% tolerance) while using the bundled local LoRA fixtures to catch shard bloat regressions early. + +--- + ## [0.13.2] - 2025-12-05 ### ✨ Improvements diff --git a/pyproject.toml b/pyproject.toml index 9de51c6..90a2296 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -14,7 +14,7 @@ source-exclude = [ [project] name = "mflux" -version = "0.13.2" +version = "0.13.3" description = "A MLX port of FLUX based on the Huggingface Diffusers implementation." readme = "README.md" keywords = ["diffusers", "flux", "mlx"] diff --git a/src/mflux/models/common/lora/mapping/lora_saver.py b/src/mflux/models/common/lora/mapping/lora_saver.py new file mode 100644 index 0000000..0a6399c --- /dev/null +++ b/src/mflux/models/common/lora/mapping/lora_saver.py @@ -0,0 +1,82 @@ +import mlx.core as mx +import mlx.nn as nn + +from mflux.models.common.lora.layer.fused_linear_lora_layer import FusedLoRALinear +from mflux.models.common.lora.layer.linear_lora_layer import LoRALinear + + +class LoRASaver: + @staticmethod + def bake_and_strip_lora(module: nn.Module) -> nn.Module: + def _assign(parent, attr_name, idx, new_child): + if parent is None: + return + if isinstance(parent, list) and idx is not None: + parent[idx] = new_child + elif isinstance(parent, dict) and attr_name is not None: + parent[attr_name] = new_child + elif attr_name is not None: + setattr(parent, attr_name, new_child) + + def _bake_single(lora_layer: LoRALinear) -> nn.Module: + base_linear = lora_layer.linear + LoRASaver._apply_lora_delta(base_linear, lora_layer) + return base_linear + + def _bake_fused(fused_layer: FusedLoRALinear) -> nn.Module: + base_linear = fused_layer.base_linear + for lora in fused_layer.loras: + if isinstance(lora, LoRALinear): + LoRASaver._apply_lora_delta(base_linear, lora) + return base_linear + + def _walk(obj, parent=None, attr_name=None, idx=None): + # Replace wrappers first + if isinstance(obj, FusedLoRALinear): + new_child = _bake_fused(obj) + _assign(parent, attr_name, idx, new_child) + obj = new_child + elif isinstance(obj, LoRALinear): + new_child = _bake_single(obj) + _assign(parent, attr_name, idx, new_child) + obj = new_child + + # Recurse into containers/modules + if isinstance(obj, list): + for i, child in enumerate(list(obj)): + _walk(child, obj, None, i) + elif isinstance(obj, tuple): + temp_list = list(obj) + for i, child in enumerate(temp_list): + _walk(child, temp_list, None, i) + if parent is not None: + _assign(parent, attr_name, idx, type(obj)(temp_list)) + elif isinstance(obj, dict): + for key, child in list(obj.items()): + _walk(child, obj, key, None) + elif isinstance(obj, nn.Module): + for name, child in vars(obj).items(): + if isinstance(child, (nn.Module, list, tuple, dict)): + _walk(child, obj, name, None) + + _walk(module, None, None, None) + return module + + @staticmethod + def _apply_lora_delta(base_linear: nn.Module, lora_layer: LoRALinear) -> None: + if not hasattr(base_linear, "weight"): + return + + weight = base_linear.weight + delta = mx.matmul(lora_layer.lora_A, lora_layer.lora_B) # shape: [in, out] + delta = mx.transpose(delta) # shape: [out, in] + delta = lora_layer.scale * delta + + if weight.shape != delta.shape: + print(f"⚠️ Skipping LoRA bake due to shape mismatch: weight {weight.shape} vs delta {delta.shape}") + return + + try: + base_linear.weight = weight + delta.astype(weight.dtype) + except Exception as e: # noqa: BLE001 + print(f"⚠️ Failed to bake LoRA into base layer: {e}") diff --git a/src/mflux/models/common/weights/saving/model_saver.py b/src/mflux/models/common/weights/saving/model_saver.py index 6ebeba8..aad6d9a 100644 --- a/src/mflux/models/common/weights/saving/model_saver.py +++ b/src/mflux/models/common/weights/saving/model_saver.py @@ -8,6 +8,7 @@ from mlx.utils import tree_flatten from tqdm import tqdm from transformers import PreTrainedTokenizer +from mflux.models.common.lora.mapping.lora_saver import LoRASaver from mflux.utils.version_util import VersionUtil if TYPE_CHECKING: @@ -35,6 +36,8 @@ class ModelSaver: for attr_name, subdir in tqdm(components, desc="Saving components", unit="component"): component = getattr(model, attr_name, None) if component is not None: + # Bake and strip any LoRA wrappers to avoid duplicating shared weights + LoRASaver.bake_and_strip_lora(component) ModelSaver._save_weights(base_path, bits, component, subdir) @staticmethod diff --git a/tests/model_saving/test_model_saving_lora.py b/tests/model_saving/test_model_saving_lora.py index 7fb932e..74cfe72 100644 --- a/tests/model_saving/test_model_saving_lora.py +++ b/tests/model_saving/test_model_saving_lora.py @@ -9,15 +9,25 @@ from mflux.models.common.config import ModelConfig from mflux.models.flux.variants.txt2img.flux import Flux1 PATH = "tests/4bit/" +SIZE_TOLERANCE_RATIO = 0.05 # allow small metadata/header differences class TestModelSavingLora: + LORA_FILES = [ + "FLUX-dev-lora-MiaoKa-Yarn-World.safetensors", + "Flux_-_Renaissance_art_style.safetensors", + ] + LORA_SCALES = [0.6, 0.4] + @pytest.mark.slow def test_save_and_load_4bit_model_with_lora(self): # Clean up any existing temporary directories from previous test runs - TestModelSavingLora.delete_folder_if_exists(PATH) + TestModelSavingLora._delete_folder_if_exists(PATH) try: + lora_paths = TestModelSavingLora._get_lora_paths() + assert len(lora_paths) == len(TestModelSavingLora.LORA_SCALES) + # given a saved quantized model on disk (without LoRA)... fluxA = Flux1( model_config=ModelConfig.schnell(), @@ -30,8 +40,8 @@ class TestModelSavingLora: fluxB = Flux1( model_config=ModelConfig.schnell(), quantize=4, - lora_paths=TestModelSavingLora.get_lora_path(), - lora_scales=[1.0], + lora_paths=lora_paths, + lora_scales=TestModelSavingLora.LORA_SCALES, ) image1 = fluxB.generate_image( seed=44, @@ -46,8 +56,8 @@ class TestModelSavingLora: fluxC = Flux1( model_config=ModelConfig.schnell(), model_path=PATH, - lora_paths=TestModelSavingLora.get_lora_path(), - lora_scales=[1.0], + lora_paths=lora_paths, + lora_scales=TestModelSavingLora.LORA_SCALES, ) # ...and generating the identical image @@ -68,14 +78,51 @@ class TestModelSavingLora: finally: # cleanup - TestModelSavingLora.delete_folder(PATH) + TestModelSavingLora._delete_folder_if_exists(PATH) + + @pytest.mark.slow + def test_save_with_lora_has_same_shard_count_as_base(self): + base_path = "tests/4bit_base/" + lora_path = "tests/4bit_with_lora/" + + TestModelSavingLora._delete_folder_if_exists(base_path) + TestModelSavingLora._delete_folder_if_exists(lora_path) + + try: + flux_base = Flux1( + model_config=ModelConfig.schnell(), + quantize=4, + ) + flux_base.save_model(base_path) + del flux_base + + flux_lora = Flux1( + model_config=ModelConfig.schnell(), + quantize=4, + lora_paths=TestModelSavingLora._get_lora_paths(), + lora_scales=TestModelSavingLora.LORA_SCALES, + ) + flux_lora.save_model(lora_path) + del flux_lora + + base_shards = list((Path(base_path) / "transformer").glob("*.safetensors")) + lora_shards = list((Path(lora_path) / "transformer").glob("*.safetensors")) + + assert len(base_shards) == len(lora_shards), "LoRA save should not inflate transformer shard count" + + # Also assert the total saved size stays within a tight bound to catch shard bloat + base_size = TestModelSavingLora._dir_size_bytes(base_path) + lora_size = TestModelSavingLora._dir_size_bytes(lora_path) + + # LoRA baking should leave sizes effectively unchanged; allow a small tolerance for metadata + max_allowed = base_size * (1 + SIZE_TOLERANCE_RATIO) + assert lora_size <= max_allowed, f"LoRA save size grew unexpectedly: base={base_size}B vs lora={lora_size}B" + finally: + TestModelSavingLora._delete_folder_if_exists(base_path) + TestModelSavingLora._delete_folder_if_exists(lora_path) @staticmethod - def delete_folder(path: str) -> None: - return shutil.rmtree(path) - - @staticmethod - def delete_folder_if_exists(path: str) -> None: + def _delete_folder_if_exists(path: str) -> None: if os.path.exists(path): shutil.rmtree(path) print(f"Deleted folder: {path}") @@ -83,12 +130,23 @@ class TestModelSavingLora: print(f"Folder does not exist: {path}") @staticmethod - def resolve_path(path) -> Path | None: + def _resolve_path(path) -> Path | None: if path is None: return None return Path(__file__).parent.parent / "resources" / path @staticmethod - def get_lora_path() -> list[str]: - path = TestModelSavingLora.resolve_path("FLUX-dev-lora-MiaoKa-Yarn-World.safetensors") - return [str(path)] + def _get_lora_paths() -> list[str]: + resolved_paths = [TestModelSavingLora._resolve_path(path) for path in TestModelSavingLora.LORA_FILES] + missing = [p for p in resolved_paths if p is None or not p.exists()] + if missing: + missing_names = ", ".join(sorted(p.name if p else "unknown" for p in missing)) + pytest.skip(f"Missing local LoRA test asset(s): {missing_names}") + return [str(path) for path in resolved_paths if path is not None] + + @staticmethod + def _dir_size_bytes(path: str | Path) -> int: + p = Path(path) + if not p.exists(): + return 0 + return sum(f.stat().st_size for f in p.rglob("*") if f.is_file())