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

Fix: loading non-quantized saved models (#261)

Browse Source 
Co-authored-by: Anthony Wu <pls-file-gh-issue@users.noreply.github.com>
This commit is contained in:
Anthony Wu 2025-09-08 05:30:33 -07:00 committed by GitHub
parent 0d9af4f380
commit 49af3504e6
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
2 changed files with 11 additions and 0 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

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

@ -34,8 +34,12 @@ class ModelSaver:
weights = ModelSaver._split_weights(base_path, dict(tree_flatten(model.parameters()))) weights = ModelSaver._split_weights(base_path, dict(tree_flatten(model.parameters())))
for i, weight in enumerate(weights): for i, weight in enumerate(weights):
mx.save_safetensors( mx.save_safetensors(
# usage: save_safetensors(file: str, arrays: dict[str, array], metadata: Optional[dict[str, str]] = None)
str(path / f"{i}.safetensors"), str(path / f"{i}.safetensors"),
# arrays (dict(str, array)): The dictionary of names to arrays to be saved.
weight, 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), "quantization_level": str(bits),
"mflux_version": VersionUtil.get_mflux_version(), "mflux_version": VersionUtil.get_mflux_version(),

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

@ -19,6 +19,13 @@ class QuantizationUtil:
) -> None: ) -> None:
q_level = weights.meta_data.quantization_level 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: if quantize is not None or q_level is not None:
bits = int(q_level) if q_level is not None else quantize bits = int(q_level) if q_level is not None else quantize
nn.quantize(vae, bits=bits) nn.quantize(vae, bits=bits)