From 49af3504e6a054841fba33c5880d9b111ae1ec34 Mon Sep 17 00:00:00 2001 From: Anthony Wu <462072+anthonywu@users.noreply.github.com> Date: Mon, 8 Sep 2025 05:30:33 -0700 Subject: [PATCH] Fix: loading non-quantized saved models (#261) Co-authored-by: Anthony Wu --- src/mflux/weights/model_saver.py | 4 ++++ src/mflux/weights/quantization_util.py | 7 +++++++ 2 files changed, 11 insertions(+) diff --git a/src/mflux/weights/model_saver.py b/src/mflux/weights/model_saver.py index 2f94e8a..deb8e6d 100644 --- a/src/mflux/weights/model_saver.py +++ b/src/mflux/weights/model_saver.py @@ -34,8 +34,12 @@ class ModelSaver: weights = ModelSaver._split_weights(base_path, dict(tree_flatten(model.parameters()))) for i, weight in enumerate(weights): mx.save_safetensors( + # usage: save_safetensors(file: str, arrays: dict[str, array], metadata: Optional[dict[str, str]] = None) str(path / f"{i}.safetensors"), + # arrays (dict(str, array)): The dictionary of names to arrays to be saved. weight, + # [save_safetensors] Metadata must be a dictionary with string keys and values + # i.e. 'None' and other special values are string-ified and need to be parsed by readers { "quantization_level": str(bits), "mflux_version": VersionUtil.get_mflux_version(), diff --git a/src/mflux/weights/quantization_util.py b/src/mflux/weights/quantization_util.py index a530160..e4b763a 100644 --- a/src/mflux/weights/quantization_util.py +++ b/src/mflux/weights/quantization_util.py @@ -19,6 +19,13 @@ class QuantizationUtil: ) -> None: q_level = weights.meta_data.quantization_level + # mx.save_tensors saves metadata dict kv 'quantization_level': 'None' as a str: str mapping + # we coerce both configs to NoneType to help users use non-quantized saved model files + if q_level == "None": + q_level = None + if quantize == "None": + quantize = None + if quantize is not None or q_level is not None: bits = int(q_level) if q_level is not None else quantize nn.quantize(vae, bits=bits)