feat: --apply-lora a.safetensors b.safetensors c.safetensor --lora-scales 0.5 1.0 0.5

README.md

@@ -83,6 +83,12 @@ python main.py --model dev --prompt "Luxury food photograph" --steps 25 --seed 2
 
 - **`--quantize`** or **`-q`** (optional, `int`, default: `None`): [Quantization](#quantization) (choose between `4` or `8`).
 
+- **`--apply-lora`** (optional, `[str]`, default: `[]`): [Lora Safetensors file]
+
+- **`--lora-scales`** (optional, `[float]`, default: `[1.0]`): [Scaling factor for each LoRA files]
+
+### Note:
+Ensure that the safetensors file provided is compatible with the model's architecture and that the LoRA keys correctly map to the model's layers. The missing lora-scales shall be treated as 1.0 by default.
 
 Or make a new separate script like the following
 

main.py

@@ -23,8 +23,8 @@ def main():
     parser.add_argument('--guidance', type=float, default=3.5, help='Guidance Scale (Default is 3.5)')
     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.')
+    parser.add_argument('--apply-lora', type=str, nargs='*', default=[], help='Local safetensors for applying LORA from disk')
+    parser.add_argument('--lora-scales', type=float,nargs='*', 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()
@@ -33,13 +33,12 @@ def main():
         parser.error("--model must be specified when using --path")
 
     seed = int(time.time()) if args.seed is None else args.seed
-
     flux = Flux1(
         model_config=ModelConfig.from_alias(args.model),
         quantize_full_weights=args.quantize,
         local_path=args.path,
-        lora_path=args.apply_lora,
-        lora_scale=args.lora_scale
+        lora_files=args.apply_lora,
+        lora_scales=args.lora_scales
     )
 
     image = flux.generate_image(

src/flux_1/flux.py

@@ -29,12 +29,12 @@ class Flux1:
             model_config: ModelConfig,
             quantize_full_weights: int | None = None,
             local_path: str | None = None,
-            lora_path: str | None = None,
-            lora_scale: float =1.0,
+            lora_files: [str] =[], 
+            lora_scales: [float] = [1.0]
     ):
         self.model_config = model_config
         self.quantize_full_weights = quantize_full_weights
-        self.lora_path = lora_path
+        self.lora_files = lora_files
 
         # Load and initialize the tokenizers from disk, huggingface cache, or download from huggingface
         tokenizers = TokenizerHandler(model_config.model_name, self.model_config.max_sequence_length, local_path)
@@ -47,8 +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,lora_scale=lora_scale)
-        
+        weights = WeightHandler(repo_id=model_config.model_name, local_path=local_path,lora_files=lora_files,lora_scales=lora_scales)
         # Set the loaded weights if they are not quantized
         if weights.quantization_level is None:
             self._set_model_weights(weights)

src/flux_1/weights/weight_handler.py

@@ -17,8 +17,8 @@ class WeightHandler:
             self,
             repo_id: str | None = None,
             local_path: str | None = None,
-            lora_path: str | None = None,
-            lora_scale: float = 1.0
+            lora_files: [str] =[], 
+            lora_scales: [float] = [1.0]
     ):
         root_path = Path(local_path) if local_path else WeightHandler._download_or_get_cached_weights(repo_id)
 
@@ -26,14 +26,22 @@ class WeightHandler:
         self.t5_encoder, _ = WeightHandler._t5_encoder(root_path=root_path)
         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):
-            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}")
+        if(lora_files):
+            if(len(lora_files)< len(lora_scales)):
+                lora_scales=lora_scales[0:len(lora_files)]
+            if(len(lora_scales)<len(lora_files)):
+                lora_scales= lora_scales + (len(lora_files)-len(lora_scales)) * [1.0]
+            for lora_file, lora_scale in zip(lora_files, lora_scales):
+                if( lora_scale<0.0 or lora_scale>1.0):
+                    raise Exception(f"Invalid scale {lora_scale} provided for {lora_file}. Valid Range [0.0-1.0] ")
+
+                try:
+                    lora_transformer,_ = WeightHandler._lora_transformer(lora_file=lora_file)
+                    if 'transformer' not in 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,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,lora_scale):
@@ -91,9 +99,9 @@ class WeightHandler:
     
 
     @staticmethod
-    def _lora_transformer(lora_path: Path) -> (dict, int):
-        quantization_level = safe_open(lora_path, framework="pt").metadata().get("quantization_level")
-        weights = list(mx.load(str(lora_path)).items())
+    def _lora_transformer(lora_file: Path) -> (dict, int):
+        quantization_level = safe_open(lora_file, framework="pt").metadata().get("quantization_level")
+        weights = list(mx.load(str(lora_file)).items())
         weights = [WeightHandler._reshape_weights(k, v) for k, v in weights]
         weights = WeightHandler._flatten(weights)
         unflatten = tree_unflatten(weights)