Add support for multiple LoRAs again

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=None, help='Local safetensors for applying LORA from disk')
+    parser.add_argument('--lora-scales', type=float, nargs='*', default=None, 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()
 
@@ -35,8 +35,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_scale=args.lora_scale
+        lora_paths=args.apply_lora,
+        lora_scales=args.lora_scales
     )
 
     image = flux.generate_image(

src/flux_1/flux.py

@@ -28,8 +28,8 @@ 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_paths: list[str] | None = None,
+            lora_scales: list[float] | None = None,
     ):
         self.model_config = model_config
         self.quantize_full_weights = quantize_full_weights
@@ -46,7 +46,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_paths=lora_paths, lora_scales=lora_scales)
 
         # Set the loaded weights if they are not quantized
         if weights.quantization_level is None:

src/flux_1/weights/lora_util.py

@@ -1,16 +1,31 @@
 import logging
-from pathlib import Path
 
 from mlx.utils import tree_flatten
 
-
 log = logging.getLogger(__name__)
 
 
 class LoraUtil:
 
     @staticmethod
-    def apply_lora(transformer: dict, lora_file: str, lora_scale: float) -> None:
+    def apply_loras(transformer: dict, lora_files: list[str], lora_scales: list[float] | None = None) -> None:
+        lora_scales = LoraUtil._validate_lora_scales(lora_files, lora_scales)
+
+        for lora_file, lora_scale in zip(lora_files, lora_scales):
+            LoraUtil._apply_lora(transformer, lora_file, lora_scale)
+
+    @staticmethod
+    def _validate_lora_scales(lora_files: list[str], lora_scales: list[float]):
+        if len(lora_files) == 1:
+            if not lora_scales:
+                lora_scales = [1.0]
+        elif len(lora_files) > 1:
+            if len(lora_files) != len(lora_scales):
+                raise ValueError("When providing multiple LoRAs, be sure to specify a scale for each one respectively")
+        return lora_scales
+
+    @staticmethod
+    def _apply_lora(transformer: dict, lora_file: str, lora_scale: float) -> None:
         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:

src/flux_1/weights/weight_handler.py

@@ -14,8 +14,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_paths: list[str] | None = None,
+            lora_scales: list[float] | None = None,
     ):
         root_path = Path(local_path) if local_path else WeightHandler._download_or_get_cached_weights(repo_id)
 
@@ -24,8 +24,8 @@ class WeightHandler:
         self.vae, _ = WeightHandler.load_vae(root_path=root_path)
         self.transformer, self.quantization_level = WeightHandler.load_transformer(root_path=root_path)
 
-        if lora_path:
-            LoraUtil.apply_lora(self.transformer, lora_path, lora_scale)
+        if lora_paths:
+            LoraUtil.apply_loras(self.transformer, lora_paths, lora_scales)
 
     @staticmethod
     def load_clip_encoder(root_path: Path) -> (dict, int):