From cb09eacb6913e026f06bf6a344fbf3cf9386e1ae Mon Sep 17 00:00:00 2001 From: filipstrand Date: Sat, 15 Mar 2025 10:00:47 +0100 Subject: [PATCH 1/3] Update readme with new 4-bit download location --- README.md | 24 ++++++++++++++++++++++-- 1 file changed, 22 insertions(+), 2 deletions(-) diff --git a/README.md b/README.md index 38b09c7..ec2cf30 100644 --- a/README.md +++ b/README.md @@ -503,9 +503,29 @@ mflux-generate \ *Also Note: Once we have a local model (quantized [or not](#-running-a-non-quantized-model-directly-from-disk)) specified via the `--path` argument, the huggingface cache models are not required to launch the model. In other words, you can reclaim the 34GB diskspace (per model) by deleting the full 16-bit model from the [Huggingface cache](#%EF%B8%8F-generating-an-image) if you choose.* +⚠️ * Quantized models saved with mflux < v.0.6.0 will not work with v.0.6.0 and later due to updated implementation. The solution is to [save a new quantized local copy](https://github.com/filipstrand/mflux/issues/149) + *If you don't want to download the full models and quantize them yourself, the 4-bit weights are available here for a direct download:* -- [madroid/flux.1-schnell-mflux-4bit](https://huggingface.co/madroid/flux.1-schnell-mflux-4bit) -- [madroid/flux.1-dev-mflux-4bit](https://huggingface.co/madroid/flux.1-dev-mflux-4bit) +- For mflux < v.0.6.0: + - [madroid/flux.1-schnell-mflux-4bit](https://huggingface.co/madroid/flux.1-schnell-mflux-4bit) + - [madroid/flux.1-dev-mflux-4bit](https://huggingface.co/madroid/flux.1-dev-mflux-4bit) +- For mflux >= v.0.6.0: + - [dhairyashil/FLUX.1-schnell-mflux-v0.6.2-4bit](https://huggingface.co/dhairyashil/FLUX.1-schnell-mflux-v0.6.2-4bit) + - [dhairyashil/FLUX.1-dev-mflux-4bit](https://huggingface.co/dhairyashil/FLUX.1-dev-mflux-4bit) + +
+Using the community model support, the quantized weights can be also be automatically downloaded + +```sh +mflux-generate \ + --model "dhairyashil/FLUX.1-schnell-mflux-v0.6.2-4bit" \ + --base-model schnell \ + --steps 2 \ + --seed 2 \ + --prompt "Luxury food photograph" +``` + +
### 💽 Running a non-quantized model directly from disk From 696762af2e44b505b274770aa3e0b9cd763d0ab6 Mon Sep 17 00:00:00 2001 From: filipstrand Date: Sat, 15 Mar 2025 10:07:38 +0100 Subject: [PATCH 2/3] Save mflux versions in weight metadata --- src/mflux/weights/model_saver.py | 7 ++++- src/mflux/weights/weight_handler.py | 35 ++++++++++++++----------- tests/model_saving/test_model_saving.py | 8 ++++++ 3 files changed, 33 insertions(+), 17 deletions(-) diff --git a/src/mflux/weights/model_saver.py b/src/mflux/weights/model_saver.py index da133e9..f831ae0 100644 --- a/src/mflux/weights/model_saver.py +++ b/src/mflux/weights/model_saver.py @@ -5,6 +5,8 @@ from mlx import nn from mlx.utils import tree_flatten from transformers import CLIPTokenizer, T5Tokenizer +from mflux.post_processing.generated_image import GeneratedImage + class ModelSaver: @staticmethod @@ -34,7 +36,10 @@ class ModelSaver: mx.save_safetensors( str(path / f"{i}.safetensors"), weight, - {"quantization_level": str(bits)}, + { + "quantization_level": str(bits), + "mflux_version": GeneratedImage.get_version(), + }, ) @staticmethod diff --git a/src/mflux/weights/weight_handler.py b/src/mflux/weights/weight_handler.py index 9679bb4..4b162f0 100644 --- a/src/mflux/weights/weight_handler.py +++ b/src/mflux/weights/weight_handler.py @@ -15,6 +15,7 @@ class MetaData: scale: float | None = None is_lora: bool = False is_mflux: bool = False + mflux_version: str | None = None class WeightHandler: @@ -40,10 +41,11 @@ class WeightHandler: # Load the weights from disk, huggingface cache, or download from huggingface root_path = Path(local_path) if local_path else WeightHandler._download_or_get_cached_weights(repo_id) - clip_encoder, _ = WeightHandler._load_clip_encoder(root_path=root_path) - t5_encoder, _ = WeightHandler._load_t5_encoder(root_path=root_path) - vae, _ = WeightHandler._load_vae(root_path=root_path) - transformer, quantization_level, _ = WeightHandler.load_transformer(root_path=root_path) + clip_encoder, _, _ = WeightHandler._load_clip_encoder(root_path=root_path) + t5_encoder, _, _ = WeightHandler._load_t5_encoder(root_path=root_path) + vae, _, _ = WeightHandler._load_vae(root_path=root_path) + transformer, quantization_level, mflux_version = WeightHandler.load_transformer(root_path=root_path) + return WeightHandler( clip_encoder=clip_encoder, t5_encoder=t5_encoder, @@ -53,7 +55,7 @@ class WeightHandler: quantization_level=quantization_level, scale=None, is_lora=False, - is_mflux=False, + mflux_version=mflux_version, ), ) @@ -64,17 +66,17 @@ class WeightHandler: return len(self.transformer["single_transformer_blocks"]) @staticmethod - def _load_clip_encoder(root_path: Path) -> (dict, int): - weights, quantization_level, _ = WeightHandler._get_weights("text_encoder", root_path) - return weights, quantization_level + def _load_clip_encoder(root_path: Path) -> (dict, int, str | None): + weights, quantization_level, mflux_version = WeightHandler._get_weights("text_encoder", root_path) + return weights, quantization_level, mflux_version @staticmethod - def _load_t5_encoder(root_path: Path) -> (dict, int): - weights, quantization_level, _ = WeightHandler._get_weights("text_encoder_2", root_path) + def _load_t5_encoder(root_path: Path) -> (dict, int, str | None): + weights, quantization_level, mflux_version = WeightHandler._get_weights("text_encoder_2", root_path) # Quantized weights (i.e. ones exported from this project) don't need any post-processing. if quantization_level is not None: - return weights, quantization_level + return weights, quantization_level, mflux_version # Reshape and process the huggingface weights weights["final_layer_norm"] = weights["encoder"]["final_layer_norm"] @@ -93,7 +95,7 @@ class WeightHandler: block["attention"]["SelfAttention"]["relative_attention_bias"] = relative_attention_bias weights.pop("encoder") - return weights, quantization_level + return weights, quantization_level, mflux_version @staticmethod def load_transformer(root_path: Path | None = None, lora_path: str | None = None) -> (dict, int, str | None): @@ -124,12 +126,12 @@ class WeightHandler: return weights, quantization_level, mflux_version @staticmethod - def _load_vae(root_path: Path) -> (dict, int): - weights, quantization_level, _ = WeightHandler._get_weights("vae", root_path) + def _load_vae(root_path: Path) -> (dict, int, str | None): + weights, quantization_level, mflux_version = WeightHandler._get_weights("vae", root_path) # Quantized weights (i.e. ones exported from this project) don't need any post-processing. if quantization_level is not None: - return weights, quantization_level + return weights, quantization_level, mflux_version # Reshape and process the huggingface weights weights["decoder"]["conv_in"] = {"conv2d": weights["decoder"]["conv_in"]} @@ -138,7 +140,7 @@ class WeightHandler: weights["encoder"]["conv_in"] = {"conv2d": weights["encoder"]["conv_in"]} weights["encoder"]["conv_out"] = {"conv2d": weights["encoder"]["conv_out"]} weights["encoder"]["conv_norm_out"] = {"norm": weights["encoder"]["conv_norm_out"]} - return weights, quantization_level + return weights, quantization_level, mflux_version @staticmethod def _get_weights( @@ -156,6 +158,7 @@ class WeightHandler: weight = list(data[0].items()) if len(data) > 1: quantization_level = data[1].get("quantization_level") + mflux_version = data[1].get("mflux_version") weights.extend(weight) if lora_path and root_path is None: diff --git a/tests/model_saving/test_model_saving.py b/tests/model_saving/test_model_saving.py index 96cc696..ef512c1 100644 --- a/tests/model_saving/test_model_saving.py +++ b/tests/model_saving/test_model_saving.py @@ -1,9 +1,12 @@ import os import shutil +from pathlib import Path import numpy as np from mflux import Config, Flux1, ModelConfig +from mflux.post_processing.generated_image import GeneratedImage +from mflux.weights.weight_handler import WeightHandler PATH = "tests/4bit/" @@ -31,6 +34,11 @@ class TestModelSaving: fluxA.save_model(PATH) del fluxA + # Verify that the mflux version is correctly saved in the model's metadata + _, quantization_level, mflux_version = WeightHandler._load_vae(root_path=Path(PATH)) + assert mflux_version == GeneratedImage.get_version(), "mflux version not correctly saved in metadata" # fmt: off + assert quantization_level == "4", "quantization level not correctly saved in metadata" # fmt: off + # when loading the quantized model (also without specifying bits) fluxB = Flux1( model_config=ModelConfig.dev(), From 273431f0106c28850fe0fd586ff341e7bf0e544d Mon Sep 17 00:00:00 2001 From: filipstrand Date: Sat, 15 Mar 2025 10:22:36 +0100 Subject: [PATCH 3/3] Remove redundant is_mflux weight attribute --- src/mflux/dreambooth/lora_layers/lora_layers.py | 2 +- src/mflux/weights/weight_handler.py | 1 - src/mflux/weights/weight_handler_lora.py | 2 +- 3 files changed, 2 insertions(+), 3 deletions(-) diff --git a/src/mflux/dreambooth/lora_layers/lora_layers.py b/src/mflux/dreambooth/lora_layers/lora_layers.py index 2ad62d4..edaa5cb 100644 --- a/src/mflux/dreambooth/lora_layers/lora_layers.py +++ b/src/mflux/dreambooth/lora_layers/lora_layers.py @@ -58,7 +58,7 @@ class LoRALayers: lora_layers = {**transformer_lora_layers, **single_transformer_lora_layers} weights = WeightHandler( - meta_data=MetaData(is_mflux=True), + meta_data=MetaData(mflux_version=GeneratedImage.get_version()), transformer=mlx.utils.tree_unflatten(list(lora_layers.items()))['transformer'], ) # fmt:off diff --git a/src/mflux/weights/weight_handler.py b/src/mflux/weights/weight_handler.py index 4b162f0..830b5e5 100644 --- a/src/mflux/weights/weight_handler.py +++ b/src/mflux/weights/weight_handler.py @@ -14,7 +14,6 @@ class MetaData: quantization_level: int | None = None scale: float | None = None is_lora: bool = False - is_mflux: bool = False mflux_version: str | None = None diff --git a/src/mflux/weights/weight_handler_lora.py b/src/mflux/weights/weight_handler_lora.py index 22edb7d..da514f6 100644 --- a/src/mflux/weights/weight_handler_lora.py +++ b/src/mflux/weights/weight_handler_lora.py @@ -38,7 +38,7 @@ class WeightHandlerLoRA: quantization_level=None, scale=lora_scale, is_lora=True, - is_mflux=True if mflux_version is not None else False, + mflux_version=mflux_version, ), ) lora_weights.append(weights)