From 03b929e7324d7089f25fede43c77d8a61728633a Mon Sep 17 00:00:00 2001 From: Sujip Maharjan Date: Tue, 3 Sep 2024 18:39:01 +0545 Subject: [PATCH] feat: Add --lora-scale parameter to adjust the scaling of LoRA weights --- main.py | 4 +++- src/flux_1/flux.py | 7 +++---- src/flux_1/weights/weight_handler.py | 29 ++++++++++++++-------------- 3 files changed, 21 insertions(+), 19 deletions(-) diff --git a/main.py b/main.py index d6a1790..a27af81 100644 --- a/main.py +++ b/main.py @@ -24,6 +24,7 @@ def main(): parser.add_argument('--quantize', "-q", type=int, choices=[4, 8], default=None, help='Quantize the model (4 or 8, Default is None)') parser.add_argument('--path', type=str, default=None, help='Local path for loading a model from disk') parser.add_argument('--apply-lora', type=str, default=None, help='Local safetensors for applying LORA from disk') + parser.add_argument('--lora-scale', type=float, default=1.0, help='Scaling factor to adjust the impact of LoRA weights on the model. A value of 1.0 applies the LoRA weights as they are.') args = parser.parse_args() @@ -37,7 +38,8 @@ def main(): model_config=ModelConfig.from_alias(args.model), quantize_full_weights=args.quantize, local_path=args.path, - lora_path=args.apply_lora + lora_path=args.apply_lora, + lora_scale=args.lora_scale ) image = flux.generate_image( diff --git a/src/flux_1/flux.py b/src/flux_1/flux.py index 08e1036..298be09 100644 --- a/src/flux_1/flux.py +++ b/src/flux_1/flux.py @@ -19,8 +19,6 @@ from flux_1.tokenizer.clip_tokenizer import TokenizerCLIP from flux_1.tokenizer.t5_tokenizer import TokenizerT5 from flux_1.tokenizer.tokenizer_handler import TokenizerHandler from flux_1.weights.weight_handler import WeightHandler -import safetensors -from safetensors import safe_open @@ -31,7 +29,8 @@ class Flux1: model_config: ModelConfig, quantize_full_weights: int | None = None, local_path: str | None = None, - lora_path: str | None = None + lora_path: str | None = None, + lora_scale: float =1.0, ): self.model_config = model_config self.quantize_full_weights = quantize_full_weights @@ -48,7 +47,7 @@ class Flux1: self.clip_text_encoder = CLIPEncoder() # Load the weights from disk, huggingface cache, or download from huggingface - weights = WeightHandler(repo_id=model_config.model_name, local_path=local_path,lora_path=lora_path) + weights = WeightHandler(repo_id=model_config.model_name, local_path=local_path,lora_path=lora_path,lora_scale=lora_scale) # Set the loaded weights if they are not quantized if weights.quantization_level is None: diff --git a/src/flux_1/weights/weight_handler.py b/src/flux_1/weights/weight_handler.py index f4247f2..6d312cd 100644 --- a/src/flux_1/weights/weight_handler.py +++ b/src/flux_1/weights/weight_handler.py @@ -6,22 +6,19 @@ from mlx.utils import tree_unflatten from safetensors import safe_open from flux_1.config.config import Config -import json from mlx.utils import tree_flatten -from functools import reduce +import logging + +log = logging.getLogger(__name__) -class NumpyEncoder(json.JSONEncoder): - def default(self, obj): - if isinstance(obj, np.ndarray): - return obj.tolist() - return super(NumpyEncoder, self).default(obj) class WeightHandler: def __init__( self, repo_id: str | None = None, local_path: str | None = None, - lora_path: str | None = None + lora_path: str | None = None, + lora_scale: float = 1.0 ): root_path = Path(local_path) if local_path else WeightHandler._download_or_get_cached_weights(repo_id) @@ -30,12 +27,16 @@ class WeightHandler: self.vae, _ = WeightHandler._vae(root_path=root_path) self.transformer, self.quantization_level = WeightHandler._transformer(root_path=root_path) if(lora_path is not None): - self.lora_transformer,self.lora_quantization_level= WeightHandler._lora_transformer(lora_path=lora_path) - if 'transformer' not in self.lora_transformer: - pass - self._apply_transformer(self.transformer,self.lora_transformer['transformer']) + try: + self.lora_transformer,self.lora_quantization_level= WeightHandler._lora_transformer(lora_path=lora_path) + if 'transformer' not in self.lora_transformer: + 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.") + self._apply_transformer(self.transformer,self.lora_transformer['transformer'],lora_scale) + except Exception as e: + log.error(f"Error loading the LoRA safetensors file: {e}") + - def _apply_transformer(self,transformer,lora_transformer): + def _apply_transformer(self,transformer,lora_transformer,lora_scale): lora_weights = tree_flatten(lora_transformer) visited={} @@ -82,7 +83,7 @@ class WeightHandler: lora_a=visited[parentKey]['lora_A'] lora_b=visited[parentKey]['lora_B'] transWeight=target['weight'] - weight=transWeight + lora_b @lora_a + weight=transWeight + lora_scale* (lora_b @lora_a) target['weight']=weight