From 4dfb3aa892c4a4c0d287ea703720f46947104b61 Mon Sep 17 00:00:00 2001 From: elitexp Date: Tue, 3 Sep 2024 18:19:21 +0545 Subject: [PATCH] Clean up weight handler @xuzzo Cleanup code Co-authored-by: Fabio Peruzzo <39661245+Xuzzo@users.noreply.github.com> --- src/flux_1/weights/weight_handler.py | 4 +--- 1 file changed, 1 insertion(+), 3 deletions(-) diff --git a/src/flux_1/weights/weight_handler.py b/src/flux_1/weights/weight_handler.py index 719f753..f4247f2 100644 --- a/src/flux_1/weights/weight_handler.py +++ b/src/flux_1/weights/weight_handler.py @@ -91,10 +91,8 @@ class WeightHandler: @staticmethod def _lora_transformer(lora_path: Path) -> (dict, int): - weights = [] quantization_level = safe_open(lora_path, framework="pt").metadata().get("quantization_level") - weight = list(mx.load(str(lora_path)).items()) - weights.extend(weight) + weights = list(mx.load(str(lora_path)).items()) weights = [WeightHandler._reshape_weights(k, v) for k, v in weights] weights = WeightHandler._flatten(weights) unflatten = tree_unflatten(weights)