diff --git a/README.md b/README.md index 0bb24a4..e7cf7cd 100644 --- a/README.md +++ b/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. diff --git a/main.py b/main.py index e1b6e29..4660c06 100644 --- a/main.py +++ b/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 ) diff --git a/src/flux_1/weights/lora_util.py b/src/flux_1/weights/lora_util.py index 6706d68..afb539d 100644 --- a/src/flux_1/weights/lora_util.py +++ b/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: