Rename apply-lora arg

README.md

@@ -83,12 +83,9 @@ 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-paths`** (optional, `[str]`, default: `None`): The paths to the [LoRA](#LoRA) weights.
 
-- **`--lora-scale`** (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.
+- **`--lora-scales`** (optional, `[float]`, default: `None`): The scale for each respective [LoRA](#LoRA) (will default to `1.0` if not specified and only one LoRA weight is loaded.)
 
 Or make a new separate script like the following
 
@@ -301,6 +298,11 @@ when loading a model directly from disk, we require the downloaded models to loo
 ```
 This mirrors how the resources are placed in the [HuggingFace Repo](https://huggingface.co/black-forest-labs/FLUX.1-schnell/tree/main) for FLUX.1.
 
+### LoRA
+
+MFLUX support loading precomputed [LoRA](https://huggingface.co/docs/diffusers/en/training/lora) weights (training support is coming soon). 
+By pointing out 
+
 ### Current limitations
 
 - Images are generated one by one.

main.py

@@ -23,7 +23,7 @@ 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, nargs='*', default=None, help='Local safetensors for applying LORA from disk')
+    parser.add_argument('--lora-paths', 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,7 +35,7 @@ def main():
         model_config=ModelConfig.from_alias(args.model),
         quantize_full_weights=args.quantize,
         local_path=args.path,
-        lora_paths=args.apply_lora,
+        lora_paths=args.lora_paths,
         lora_scales=args.lora_scales
     )
 

src/flux_1/weights/lora_util.py

@@ -15,10 +15,12 @@ class LoraUtil:
             LoraUtil._apply_lora(transformer, lora_file, lora_scale)
 
     @staticmethod
-    def _validate_lora_scales(lora_files: list[str], lora_scales: list[float]):
+    def _validate_lora_scales(lora_files: list[str], lora_scales: list[float]) -> list[float]:
         if len(lora_files) == 1:
             if not lora_scales:
                 lora_scales = [1.0]
+            if len(lora_scales) > 1:
+                raise ValueError("Please provide a single scale for the LoRA, or skip it to default to 1")
         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")
@@ -28,12 +30,10 @@ class LoraUtil:
     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:
-            from flux_1.weights.weight_handler import WeightHandler
-            lora_transformer, _ = WeightHandler.load_transformer(lora_path=lora_file)
-            LoraUtil._apply_transformer(transformer, lora_transformer, lora_scale)
-        except Exception as e:
-            log.error(f"Error loading the LoRA safetensors file: {e}")
+
+        from flux_1.weights.weight_handler import WeightHandler
+        lora_transformer, _ = WeightHandler.load_transformer(lora_path=lora_file)
+        LoraUtil._apply_transformer(transformer, lora_transformer, lora_scale)
 
     @staticmethod
     def _apply_transformer(transformer: dict, lora_transformer: dict, lora_scale: float) -> None: