From 96848e320ffafe5a5b4b1d86c3edcf3a4582afac Mon Sep 17 00:00:00 2001 From: filipstrand Date: Thu, 5 Sep 2024 19:37:50 +0200 Subject: [PATCH] Add support for multiple LoRAs again --- main.py | 8 ++++---- src/flux_1/flux.py | 6 +++--- src/flux_1/weights/lora_util.py | 21 ++++++++++++++++++--- src/flux_1/weights/weight_handler.py | 8 ++++---- 4 files changed, 29 insertions(+), 14 deletions(-) diff --git a/main.py b/main.py index 92dc102..e1b6e29 100644 --- a/main.py +++ b/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( diff --git a/src/flux_1/flux.py b/src/flux_1/flux.py index bbae938..9a00d46 100644 --- a/src/flux_1/flux.py +++ b/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: diff --git a/src/flux_1/weights/lora_util.py b/src/flux_1/weights/lora_util.py index 12f066d..6706d68 100644 --- a/src/flux_1/weights/lora_util.py +++ b/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: diff --git a/src/flux_1/weights/weight_handler.py b/src/flux_1/weights/weight_handler.py index fbc2a88..0314437 100644 --- a/src/flux_1/weights/weight_handler.py +++ b/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):