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

Save weight metadata

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)
+- For mflux < v.0.6.0:
-- [madroid/flux.1-dev-mflux-4bit](https://huggingface.co/madroid/flux.1-dev-mflux-4bit)
+  - [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
 

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
 

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

src/mflux/weights/weight_handler.py

@@ -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)
+        clip_encoder, _, _ = WeightHandler._load_clip_encoder(root_path=root_path)
-        t5_encoder, _ = WeightHandler._load_t5_encoder(root_path=root_path)
+        t5_encoder, _, _ = WeightHandler._load_t5_encoder(root_path=root_path)
-        vae, _ = WeightHandler._load_vae(root_path=root_path)
+        vae, _, _ = WeightHandler._load_vae(root_path=root_path)
-        transformer, quantization_level, _ = WeightHandler.load_transformer(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):
+    def _load_clip_encoder(root_path: Path) -> (dict, int, str | None):
-        weights, quantization_level, _ = WeightHandler._get_weights("text_encoder", root_path)
+        weights, quantization_level, mflux_version = WeightHandler._get_weights("text_encoder", root_path)
-        return weights, quantization_level
+        return weights, quantization_level, mflux_version
 
     @staticmethod
-    def _load_t5_encoder(root_path: Path) -> (dict, int):
+    def _load_t5_encoder(root_path: Path) -> (dict, int, str | None):
-        weights, quantization_level, _ = WeightHandler._get_weights("text_encoder_2", root_path)
+        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):
+    def _load_vae(root_path: Path) -> (dict, int, str | None):
-        weights, quantization_level, _ = WeightHandler._get_weights("vae", root_path)
+        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:

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)

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