From 8d8812675ef324e7d549b566510bfef036d8c168 Mon Sep 17 00:00:00 2001 From: Sujip Maharjan Date: Tue, 3 Sep 2024 19:36:15 +0545 Subject: [PATCH] feat: --apply-lora a.safetensors b.safetensors c.safetensor --lora-scales 0.5 1.0 0.5 --- README.md | 6 +++++ main.py | 9 ++++---- src/flux_1/flux.py | 9 ++++---- src/flux_1/weights/weight_handler.py | 34 +++++++++++++++++----------- 4 files changed, 35 insertions(+), 23 deletions(-) diff --git a/README.md b/README.md index 09a93ad..d47a23a 100644 --- a/README.md +++ b/README.md @@ -83,6 +83,12 @@ 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-scales`** (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. Or make a new separate script like the following diff --git a/main.py b/main.py index a27af81..0f95784 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=[], 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.') args = parser.parse_args() @@ -33,13 +33,12 @@ def main(): parser.error("--model must be specified when using --path") seed = int(time.time()) if args.seed is None else args.seed - flux = Flux1( 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_files=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 298be09..d76a0e4 100644 --- a/src/flux_1/flux.py +++ b/src/flux_1/flux.py @@ -29,12 +29,12 @@ 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_files: [str] =[], + lora_scales: [float] = [1.0] ): self.model_config = model_config self.quantize_full_weights = quantize_full_weights - self.lora_path = lora_path + self.lora_files = lora_files # Load and initialize the tokenizers from disk, huggingface cache, or download from huggingface tokenizers = TokenizerHandler(model_config.model_name, self.model_config.max_sequence_length, local_path) @@ -47,8 +47,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_files=lora_files,lora_scales=lora_scales) # Set the loaded weights if they are not quantized if weights.quantization_level is None: self._set_model_weights(weights) diff --git a/src/flux_1/weights/weight_handler.py b/src/flux_1/weights/weight_handler.py index 6d312cd..863559e 100644 --- a/src/flux_1/weights/weight_handler.py +++ b/src/flux_1/weights/weight_handler.py @@ -17,8 +17,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_files: [str] =[], + lora_scales: [float] = [1.0] ): root_path = Path(local_path) if local_path else WeightHandler._download_or_get_cached_weights(repo_id) @@ -26,14 +26,22 @@ class WeightHandler: 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) - if(lora_path is not None): - try: - self.lora_transformer,self.lora_quantization_level= WeightHandler._lora_transformer(lora_path=lora_path) - if 'transformer' not in self.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.") - self._apply_transformer(self.transformer,self.lora_transformer['transformer'],lora_scale) - except Exception as e: - log.error(f"Error loading the LoRA safetensors file: {e}") + if(lora_files): + if(len(lora_files)< len(lora_scales)): + lora_scales=lora_scales[0:len(lora_files)] + if(len(lora_scales)1.0): + raise Exception(f"Invalid scale {lora_scale} provided for {lora_file}. Valid Range [0.0-1.0] ") + + try: + lora_transformer,_ = WeightHandler._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.") + self._apply_transformer(self.transformer,lora_transformer['transformer'],lora_scale) + except Exception as e: + log.error(f"Error loading the LoRA safetensors file: {e}") def _apply_transformer(self,transformer,lora_transformer,lora_scale): @@ -91,9 +99,9 @@ class WeightHandler: @staticmethod - def _lora_transformer(lora_path: Path) -> (dict, int): - quantization_level = safe_open(lora_path, framework="pt").metadata().get("quantization_level") - weights = list(mx.load(str(lora_path)).items()) + 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 = [WeightHandler._reshape_weights(k, v) for k, v in weights] weights = WeightHandler._flatten(weights) unflatten = tree_unflatten(weights)