feat: Add --lora-scale parameter to adjust the scaling of LoRA weights

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(

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:

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