From a3567bddd718aa107acff411582b8f0f04e87499 Mon Sep 17 00:00:00 2001 From: filipstrand Date: Wed, 4 Sep 2024 21:32:10 +0200 Subject: [PATCH] Start supporting only single LoRA (and some other simplifications) --- main.py | 10 +++--- src/flux_1/flux.py | 6 ++-- src/flux_1/weights/lora_util.py | 50 ++++++-------------------- src/flux_1/weights/weight_handler.py | 53 ++++++++++++++-------------- 4 files changed, 45 insertions(+), 74 deletions(-) diff --git a/main.py b/main.py index eaf693d..92dc102 100644 --- a/main.py +++ b/main.py @@ -1,6 +1,6 @@ +import argparse import os import sys -import argparse import time sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), 'src'))) @@ -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, nargs='*', default=[], help='Local safetensors for applying LORA from disk') - parser.add_argument('--lora-scales', type=float,nargs='*', 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, 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.') 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_files=args.apply_lora, - lora_scales=args.lora_scales + lora_path=args.apply_lora, + lora_scale=args.lora_scale ) image = flux.generate_image( diff --git a/src/flux_1/flux.py b/src/flux_1/flux.py index 7f53a80..92aa750 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_files: [str] = [], - lora_scales: [float] = [1.0] + lora_path: str | None = None, + lora_scale: str | 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_files=lora_files, lora_scales=lora_scales) + weights = WeightHandler(repo_id=model_config.model_name, local_path=local_path, lora_path=lora_path, lora_scale=lora_scale) # 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 d82df22..25df752 100644 --- a/src/flux_1/weights/lora_util.py +++ b/src/flux_1/weights/lora_util.py @@ -1,12 +1,8 @@ import logging from pathlib import Path -import mlx.core as mx from mlx.utils import tree_flatten -from mlx.utils import tree_unflatten -from safetensors import safe_open -from flux_1.weights.weight_util import WeightUtil log = logging.getLogger(__name__) @@ -14,25 +10,18 @@ log = logging.getLogger(__name__) class LoraUtil: @staticmethod - def apply_lora(transformer, lora_files, lora_scales): - if lora_files: - if len(lora_files) < len(lora_scales): - lora_scales = lora_scales[0:len(lora_files)] - if len(lora_scales) < len(lora_files): - lora_scales = lora_scales + (len(lora_files) - len(lora_scales)) * [1.0] - for lora_file, lora_scale in zip(lora_files, lora_scales): - 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: - lora_transformer, _ = LoraUtil._lora_transformer(lora_file=lora_file) - if 'transformer' not in lora_transformer: - raise Exception("The key `transformer` is missing in the LoRA safetensors file. Please ensure that the file is correctly formatted and contains the expected keys.") - LoraUtil._apply_transformer(transformer, lora_transformer['transformer'], lora_scale) - except Exception as e: - log.error(f"Error loading the LoRA safetensors file: {e}") + 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(Path(lora_file), is_lora=True) + LoraUtil._apply_transformer(transformer, lora_transformer, lora_scale) + except Exception as e: + log.error(f"Error loading the LoRA safetensors file: {e}") @staticmethod - def _apply_transformer(transformer, lora_transformer, lora_scale): + def _apply_transformer(transformer: dict, lora_transformer: dict, lora_scale: float) -> None: lora_weights = tree_flatten(lora_transformer) visited = {} @@ -81,22 +70,3 @@ class LoraUtil: transWeight = target['weight'] weight = transWeight + lora_scale * (lora_b @ lora_a) target['weight'] = weight - - @staticmethod - def _lora_transformer(lora_file: Path) -> (dict, int): - quantization_level = safe_open(lora_file, framework="pt").metadata().get("quantization_level") - weights = list(mx.load(str(lora_file)).items()) - weights = [WeightUtil.reshape_weights(k, v) for k, v in weights] - weights = WeightUtil.flatten(weights) - unflatten = tree_unflatten(weights) - for block in unflatten["transformer"]["transformer_blocks"]: - block["ff"] = { - "linear1": block["ff"]["net"][0]["proj"], - "linear2": block["ff"]["net"][2] - } - if block.get("ff_context") is not None: - block["ff_context"] = { - "linear1": block["ff_context"]["net"][0]["proj"], - "linear2": block["ff_context"]["net"][2] - } - return unflatten, quantization_level diff --git a/src/flux_1/weights/weight_handler.py b/src/flux_1/weights/weight_handler.py index 0132746..21829ee 100644 --- a/src/flux_1/weights/weight_handler.py +++ b/src/flux_1/weights/weight_handler.py @@ -14,31 +14,26 @@ class WeightHandler: self, repo_id: str | None = None, local_path: str | None = None, - lora_files=None, - lora_scales=None + lora_path: str | None = None, + lora_scale: float | None = None, ): - if lora_files is None: - lora_files = [] - if lora_scales is None: - lora_scales = [1.0] - root_path = Path(local_path) if local_path else WeightHandler._download_or_get_cached_weights(repo_id) - self.clip_encoder, _ = WeightHandler._clip_encoder(root_path=root_path) - self.t5_encoder, _ = WeightHandler._t5_encoder(root_path=root_path) - self.vae, _ = WeightHandler._vae(root_path=root_path) - self.transformer, self.quantization_level = WeightHandler._transformer(root_path=root_path) + self.clip_encoder, _ = WeightHandler.load_clip_encoder(root_path=root_path) + self.t5_encoder, _ = WeightHandler.load_t5_encoder(root_path=root_path) + self.vae, _ = WeightHandler.load_vae(root_path=root_path) + self.transformer, self.quantization_level = WeightHandler.load_transformer(root_path=root_path) - # Optionally apply LoRA weights - LoraUtil.apply_lora(self.transformer, lora_files, lora_scales) + if lora_path: + LoraUtil.apply_lora(self.transformer, lora_path, lora_scale) @staticmethod - def _clip_encoder(root_path: Path) -> (dict, int): + def load_clip_encoder(root_path: Path) -> (dict, int): weights, quantization_level = WeightHandler._get_weights("text_encoder", root_path) return weights, quantization_level @staticmethod - def _t5_encoder(root_path: Path) -> (dict, int): + def load_t5_encoder(root_path: Path) -> (dict, int): weights, quantization_level = WeightHandler._get_weights("text_encoder_2", root_path) # Quantized weights (i.e. ones exported from this project) don't need any post-processing. @@ -65,28 +60,34 @@ class WeightHandler: return weights, quantization_level @staticmethod - def _transformer(root_path: Path) -> (dict, int): + def load_transformer(root_path: Path, is_lora: bool = False) -> (dict, int): weights, quantization_level = WeightHandler._get_weights("transformer", root_path) + if is_lora: + if 'transformer' not in weights: + raise Exception("The key `transformer` is missing in the LoRA safetensors file. Please ensure that the file is correctly formatted and contains the expected keys.") + weights = weights["transformer"] + # Quantized weights (i.e. ones exported from this project) don't need any post-processing. if quantization_level is not None: return weights, quantization_level # Reshape and process the huggingface weights - for block in weights["transformer_blocks"]: - block["ff"] = { - "linear1": block["ff"]["net"][0]["proj"], - "linear2": block["ff"]["net"][2] - } - if block.get("ff_context") is not None: - block["ff_context"] = { - "linear1": block["ff_context"]["net"][0]["proj"], - "linear2": block["ff_context"]["net"][2] + if "transformer_blocks" in weights: + for block in weights["transformer_blocks"]: + block["ff"] = { + "linear1": block["ff"]["net"][0]["proj"], + "linear2": block["ff"]["net"][2] } + if block.get("ff_context") is not None: + block["ff_context"] = { + "linear1": block["ff_context"]["net"][0]["proj"], + "linear2": block["ff_context"]["net"][2] + } return weights, quantization_level @staticmethod - def _vae(root_path: Path) -> (dict, int): + def load_vae(root_path: Path) -> (dict, int): weights, quantization_level = WeightHandler._get_weights("vae", root_path) # Quantized weights (i.e. ones exported from this project) don't need any post-processing.