From 6d35db5e9a28282ee1bf46d4f5b1a71356bcf59b Mon Sep 17 00:00:00 2001 From: filipstrand Date: Wed, 4 Sep 2024 21:05:19 +0200 Subject: [PATCH] Use `mx.load` instead of `safe_open` --- src/flux_1/weights/weight_handler.py | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/src/flux_1/weights/weight_handler.py b/src/flux_1/weights/weight_handler.py index c48ba29..0132746 100644 --- a/src/flux_1/weights/weight_handler.py +++ b/src/flux_1/weights/weight_handler.py @@ -3,7 +3,6 @@ from pathlib import Path import mlx.core as mx from huggingface_hub import snapshot_download from mlx.utils import tree_unflatten -from safetensors import safe_open from flux_1.weights.lora_util import LoraUtil from flux_1.weights.weight_util import WeightUtil @@ -108,7 +107,7 @@ class WeightHandler: weights = [] quantization_level = None for file in sorted(root_path.glob(model_name + "/*.safetensors")): - quantization_level = safe_open(file, framework="pt").metadata().get("quantization_level") + quantization_level = mx.load(str(file), return_metadata=True)[1].get("quantization_level") weight = list(mx.load(str(file)).items()) weights.extend(weight)