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] ") raise Exception(f"Invalid scale {lora_scale} provided for {lora_file}. Valid Range [0.0 - 1.0] ")
try: try:
from flux_1.weights.weight_handler import WeightHandler 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) LoraUtil._apply_transformer(transformer, lora_transformer, lora_scale)
except Exception as e: except Exception as e:
log.error(f"Error loading the LoRA safetensors file: {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 return weights, quantization_level
@staticmethod @staticmethod
def load_transformer(root_path: Path, is_lora: bool = False) -> (dict, int): def load_transformer(root_path: Path | None = None, lora_path: str | None = None) -> (dict, int):
weights, quantization_level = WeightHandler._get_weights("transformer", root_path) weights, quantization_level = WeightHandler._get_weights("transformer", root_path, lora_path)
if is_lora: if lora_path:
if 'transformer' not in weights: 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.") 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"] weights = weights["transformer"]
@ -104,12 +104,18 @@ class WeightHandler:
return weights, quantization_level return weights, quantization_level
@staticmethod @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 = [] weights = []
quantization_level = None 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") if root_path is not None:
weight = list(mx.load(str(file)).items()) 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) weights.extend(weight)
# Non huggingface weights (i.e. ones exported from this project) don't need any reshaping. # Non huggingface weights (i.e. ones exported from this project) don't need any reshaping.