Start supporting only single LoRA (and some other simplifications)

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('--apply-lora', type=str, default=None, 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('--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_path=args.apply_lora,
-        lora_scales=args.lora_scales
+        lora_scale=args.lora_scale
     )
 
     image = flux.generate_image(

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_path: str | None = None,
-            lora_scales: [float] = [1.0]
+            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:

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):
+    def apply_lora(transformer: dict, lora_file: str, lora_scale: float) -> None:
-        if lora_files:
+        if lora_scale < 0.0 or lora_scale > 1.0:
-            if len(lora_files) < len(lora_scales):
+            raise Exception(f"Invalid scale {lora_scale} provided for {lora_file}. Valid Range [0.0 - 1.0] ")
-                lora_scales = lora_scales[0:len(lora_files)]
+        try:
-            if len(lora_scales) < len(lora_files):
+            from flux_1.weights.weight_handler import WeightHandler
-                lora_scales = lora_scales + (len(lora_files) - len(lora_scales)) * [1.0]
+            lora_transformer, _ = WeightHandler.load_transformer(Path(lora_file), is_lora=True)
-            for lora_file, lora_scale in zip(lora_files, lora_scales):
+            LoraUtil._apply_transformer(transformer, lora_transformer, lora_scale)
-                if lora_scale < 0.0 or lora_scale > 1.0:
+        except Exception as e:
-                    raise Exception(f"Invalid scale {lora_scale} provided for {lora_file}. Valid Range [0.0-1.0] ")
+            log.error(f"Error loading the LoRA safetensors file: {e}")
-                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}")
 
     @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

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_path: str | None = None,
-            lora_scales=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.clip_encoder, _ = WeightHandler.load_clip_encoder(root_path=root_path)
-        self.t5_encoder, _ = WeightHandler._t5_encoder(root_path=root_path)
+        self.t5_encoder, _ = WeightHandler.load_t5_encoder(root_path=root_path)
-        self.vae, _ = WeightHandler._vae(root_path=root_path)
+        self.vae, _ = WeightHandler.load_vae(root_path=root_path)
-        self.transformer, self.quantization_level = WeightHandler._transformer(root_path=root_path)
+        self.transformer, self.quantization_level = WeightHandler.load_transformer(root_path=root_path)
 
-        # Optionally apply LoRA weights
+        if lora_path:
-        LoraUtil.apply_lora(self.transformer, lora_files, lora_scales)
+            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"]:
+        if "transformer_blocks" in weights:
-            block["ff"] = {
+            for block in weights["transformer_blocks"]:
-                "linear1": block["ff"]["net"][0]["proj"],
+                block["ff"] = {
-                "linear2": block["ff"]["net"][2]
+                    "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 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.