monster/Qwen-Image-Layered-MRP-MLX
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Merge pull request #154 from filipstrand/save-weight-metadata

Browse Source 
Save weight metadata
This commit is contained in:
Filip Strand 2025-03-15 11:50:41 +01:00 committed by GitHub
commit 41cffe8822
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
6 changed files with 57 additions and 22 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

24
README.md
View File

@ -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)
<details>
<summary>Using the community model support, the quantized weights can be also be automatically downloaded</summary>
```sh
mflux-generate \
--model "dhairyashil/FLUX.1-schnell-mflux-v0.6.2-4bit" \
--base-model schnell \
--steps 2 \
--seed 2 \
--prompt "Luxury food photograph"
```
</details>
### 💽 Running a non-quantized model directly from disk

2
src/mflux/dreambooth/lora_layers/lora_layers.py
View File

@ -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

7
src/mflux/weights/model_saver.py
View File

@ -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

36
src/mflux/weights/weight_handler.py
View File

@ -14,7 +14,7 @@ class MetaData:
quantization_level: int | None = None
scale: float | None = None
is_lora: bool = False
is_mflux: bool = False
mflux_version: str | None = None
class WeightHandler:
@ -40,10 +40,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 +54,7 @@ class WeightHandler:
quantization_level=quantization_level,
scale=None,
is_lora=False,
is_mflux=False,
mflux_version=mflux_version,
),
)
@ -64,17 +65,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 +94,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 +125,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 +139,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 +157,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:

2
src/mflux/weights/weight_handler_lora.py
View File

@ -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)

8
tests/model_saving/test_model_saving.py
View File

@ -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(),