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

Use mx.load instead of safe_open

Browse Source
This commit is contained in:
filipstrand 2024-09-04 21:05:19 +02:00
parent 5ad5352481
commit 6d35db5e9a
1 changed files with 1 additions and 2 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

3
src/flux_1/weights/weight_handler.py
View File

@ -3,7 +3,6 @@ from pathlib import Path
import mlx.core as mx import mlx.core as mx
from huggingface_hub import snapshot_download from huggingface_hub import snapshot_download
from mlx.utils import tree_unflatten from mlx.utils import tree_unflatten
from safetensors import safe_open
from flux_1.weights.lora_util import LoraUtil from flux_1.weights.lora_util import LoraUtil
from flux_1.weights.weight_util import WeightUtil from flux_1.weights.weight_util import WeightUtil
@ -108,7 +107,7 @@ class WeightHandler:
weights = [] weights = []
quantization_level = None quantization_level = None
for file in sorted(root_path.glob(model_name + "/*.safetensors")): 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()) weight = list(mx.load(str(file)).items())
weights.extend(weight) weights.extend(weight)