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

Point out the .safetensors file again

Browse Source
This commit is contained in:
filipstrand 2024-09-04 22:42:29 +02:00
parent a3567bddd7
commit f28fffe0cf
2 changed files with 14 additions and 8 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

2
src/flux_1/weights/lora_util.py
View File

@ -15,7 +15,7 @@ class LoraUtil:
raise Exception(f"Invalid scale {lora_scale} provided for {lora_file}. Valid Range [0.0 - 1.0] ")
try:
from flux_1.weights.weight_handler import WeightHandler
lora_transformer, _ = WeightHandler.load_transformer(Path(lora_file), is_lora=True)
lora_transformer, _ = WeightHandler.load_transformer(lora_path=lora_file)
LoraUtil._apply_transformer(transformer, lora_transformer, lora_scale)
except Exception as e:
log.error(f"Error loading the LoRA safetensors file: {e}")

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

@ -60,10 +60,10 @@ class WeightHandler:
return weights, quantization_level
@staticmethod
def load_transformer(root_path: Path, is_lora: bool = False) -> (dict, int):
weights, quantization_level = WeightHandler._get_weights("transformer", root_path)
def load_transformer(root_path: Path | None = None, lora_path: str | None = None) -> (dict, int):
weights, quantization_level = WeightHandler._get_weights("transformer", root_path, lora_path)
if is_lora:
if lora_path:
if 'transformer' not in weights:
raise Exception("The key `transformer` is missing in the LoRA safetensors file. Please ensure that the file is correctly formatted and contains the expected keys.")
weights = weights["transformer"]
@ -104,12 +104,18 @@ class WeightHandler:
return weights, quantization_level
@staticmethod
def _get_weights(model_name: str, root_path: Path) -> (dict, int):
def _get_weights(model_name: str, root_path: Path | None = None, lora_path: str | None = None) -> (dict, int):
weights = []
quantization_level = None
for file in sorted(root_path.glob(model_name + "/*.safetensors")):
quantization_level = mx.load(str(file), return_metadata=True)[1].get("quantization_level")
weight = list(mx.load(str(file)).items())
if root_path is not None:
for file in sorted(root_path.glob(model_name + "/*.safetensors")):
quantization_level = mx.load(str(file), return_metadata=True)[1].get("quantization_level")
weight = list(mx.load(str(file)).items())
weights.extend(weight)
if lora_path and root_path is None:
weight = list(mx.load(lora_path).items())
weights.extend(weight)
# Non huggingface weights (i.e. ones exported from this project) don't need any reshaping.