Add ability to save model with applied LoRA

README.md

@@ -264,6 +264,20 @@ mflux-save \
 
 *Note that when saving a quantized version, you will need the original huggingface weights.*
 
+It is also possible to specify [LoRA](#lora) adapters when saving the model, e.g 
+
+```sh
+mflux-save \
+    --path "/Users/filipstrand/Desktop/schnell_8bit" \
+    --model schnell \
+    --quantize 8 \
+    --lora-paths "/path/to/lora.safetensors" \
+    --lora-scales 0.7
+```
+
+When generating images with a model like this, no LoRA adapter is needed to be specified since
+it is already baked into the saved quantized weights.
+
 #### Loading and running a quantized version from disk
 
 To generate a new image from the quantized model, simply provide a `--path` to where it was saved:

src/mflux/save.py

@@ -9,6 +9,8 @@ def main():
     parser.add_argument("--path", type=str, required=True, help="Local path for loading a model from disk")
     parser.add_argument("--model", "-m", type=str, required=True, choices=["dev", "schnell"], help="The model to use (\"schnell\" or \"dev\").")
     parser.add_argument("--quantize", "-q", type=int, choices=[4, 8], default=8, help="Quantize the model (4 or 8, Default is 8)")
+    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.")
     # fmt: on
 
     args = parser.parse_args()
@@ -18,6 +20,8 @@ def main():
     flux = Flux1(
         model_config=ModelConfig.from_alias(args.model),
         quantize=args.quantize,
+        lora_paths=args.lora_paths,
+        lora_scales=args.lora_scales,
     )
 
     flux.save_model(args.path)