Merge pull request #30 from elitexp/main

Apply lora with --apply-lora <path-to-safetensors>

.gitignore

@@ -11,3 +11,4 @@
 *.png
 *.jpg
 *.pyc
+*.safetensors

README.md

@@ -83,6 +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`).
 
+- **`--lora-paths`** (optional, `[str]`, default: `None`): The paths to the [LoRA](#LoRA) weights.
+
+- **`--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
 
@@ -295,13 +298,55 @@ 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 trained [LoRA](https://huggingface.co/docs/diffusers/en/training/lora) adapters (actual training support is coming).
+
+The following example [The_Hound](https://huggingface.co/TheLastBen/The_Hound) LoRA from [@TheLastBen](https://github.com/TheLastBen): 
+
+```
+python main.py --prompt "sandor clegane" --model dev --steps 20 --seed 43 -q 8 --lora-paths "sandor_clegane_single_layer.safetensors"
+```
+
+![image](src/flux_1/assets/lora1.jpg)
+---
+
+The following example is [Flux_1_Dev_LoRA_Paper-Cutout-Style](https://huggingface.co/Norod78/Flux_1_Dev_LoRA_Paper-Cutout-Style) LoRA from [@Norod78](https://huggingface.co/Norod78):
+
+```
+python main.py --prompt "pikachu, Paper Cutout Style" --model schnell --steps 4 --seed 43 -q 8 --lora-paths "Flux_1_Dev_LoRA_Paper-Cutout-Style.safetensors"
+```
+![image](src/flux_1/assets/lora2.jpg)
+
+*Note that LoRA trained weights are typically trained with a **trigger word or phrase**. For example, in the latter case, the sentence should include the phrase **"Paper Cutout Style"**.*
+
+*Also note that the same LoRA weights can work well with both the `schnell` and `dev` models. Refer to the original LoRA repository to see what mode it was trained for.*
+
+#### Multi-LoRA
+
+Multiple LoRAs can be sent in to combine the effects of the individual adapters. The following example combines both of the above LoRAs:
+
+```
+python main.py \
+--prompt "sandor clegane in a forest, Paper Cutout Style" \ 
+--model dev \
+--steps 20 \
+--seed 43 \
+--lora-paths sandor.safetensors paper.safetensors \ 
+--lora-scales 1.0 1.0 -q 8 \
+```
+![image](src/flux_1/assets/lora3.jpg)
+
+Just to see the difference, this image displays the four cases: One of having both adapters fully active, partially active and no LoRA at all. 
+The example above also show the usage of `--lora-scales` flag. 
+
 ### Current limitations
 
 - Images are generated one by one.
 - Negative prompts not supported.
+- LoRA weights are only supported for the transformer part of the network.
 
 ### TODO
 
-- [ ] LoRA adapters
 - [ ] LoRA fine-tuning
 - [ ] Frontend support (Gradio/Streamlit/Other?)

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,22 +23,24 @@ 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('--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()
 
     if args.path and args.model is None:
         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
+        local_path=args.path,
+        lora_paths=args.lora_paths,
+        lora_scales=args.lora_scales
     )
 
     image = flux.generate_image(
-        seed=seed,
+        seed=int(time.time()) if args.seed is None else args.seed,
         prompt=args.prompt,
         config=Config(
             num_inference_steps=args.steps,

src/flux_1/flux.py

@@ -28,6 +28,8 @@ class Flux1:
             model_config: ModelConfig,
             quantize_full_weights: int | None = None,
             local_path: str | None = None,
+            lora_paths: list[str] | None = None,
+            lora_scales: list[float] | None = None,
     ):
         self.model_config = model_config
         self.quantize_full_weights = quantize_full_weights
@@ -44,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)
+        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:

src/flux_1/weights/lora_util.py

@@ -0,0 +1,87 @@
+import logging
+
+from mlx.utils import tree_flatten
+
+log = logging.getLogger(__name__)
+
+
+class LoraUtil:
+
+    @staticmethod
+    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]) -> 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")
+        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] ")
+
+        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:
+        lora_weights = tree_flatten(lora_transformer)
+        visited = {}
+
+        for key, weight in lora_weights:
+            splits = key.split(".")
+            target = transformer
+            visiting = []
+            for splitKey in splits:
+                if isinstance(target, dict) and splitKey in target:
+                    target = target[splitKey]
+                    visiting.append(splitKey)
+                elif isinstance(target, list) and len(target) > 0:
+                    if len(target) < int(splitKey):
+                        for _ in range(int(splitKey) - len(target) + 1):
+                            target.append({})
+
+                    target = target[int(splitKey)]
+                    visiting.append(splitKey)
+                else:
+                    parentKey = ".".join(visiting)
+                    if parentKey in visited and 'lora_A' in visited[parentKey] and 'lora_B' in visited[parentKey]:
+                        continue
+                    if not splitKey.startswith("lora_"):
+                        visiting.append(splitKey)
+                        parentKey = ".".join(visiting)
+                        if splitKey == "net":
+                            target['net'] = list({})
+                            target = target['net']
+                        elif splitKey == "0":
+                            target.append({})
+                            target = target[0]
+                            continue
+                        elif splitKey == "proj":
+                            target[splitKey] = weight
+                            if parentKey not in visited:
+                                visited[parentKey] = {}
+                        continue
+                    if parentKey not in visited:
+                        visited[parentKey] = {}
+                    visited[parentKey][splitKey] = weight
+                    if not 'weight' in target:
+                        raise ValueError(f"LoRA weights for layer {parentKey} cannot be loaded into the model.")
+                    if 'lora_A' in visited[parentKey] and 'lora_B' in visited[parentKey]:
+                        lora_a = visited[parentKey]['lora_A']
+                        lora_b = visited[parentKey]['lora_B']
+                        transWeight = target['weight']
+                        weight = transWeight + lora_scale * (lora_b @ lora_a)
+                        target['weight'] = weight

src/flux_1/weights/weight_handler.py

@@ -3,9 +3,9 @@ from pathlib import Path
 import mlx.core as mx
 from huggingface_hub import snapshot_download
 from mlx.utils import tree_unflatten
-from safetensors import safe_open
 
-from flux_1.config.config import Config
+from flux_1.weights.lora_util import LoraUtil
+from flux_1.weights.weight_util import WeightUtil
 
 
 class WeightHandler:
@@ -14,21 +14,26 @@ class WeightHandler:
             self,
             repo_id: str | None = None,
             local_path: str | None = None,
+            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)
 
-        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)
+
+        if lora_paths:
+            LoraUtil.apply_loras(self.transformer, lora_paths, lora_scales)
 
     @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.
@@ -55,28 +60,34 @@ class WeightHandler:
         return weights, quantization_level
 
     @staticmethod
-    def _transformer(root_path: Path) -> (dict, int):
-        weights, quantization_level = WeightHandler._get_weights("transformer", root_path)
+    def load_transformer(root_path: Path | None = None, lora_path: str | None = None) -> (dict, int):
+        weights, quantization_level = WeightHandler._get_weights("transformer", root_path, lora_path)
+
+        if lora_path:
+            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.
@@ -93,12 +104,18 @@ class WeightHandler:
         return weights, quantization_level
 
     @staticmethod
-    def _get_weights(model_name: str, root_path: Path) -> (dict, int):
+    def _get_weights(model_name: str, root_path: Path | None = None, lora_path: str | None = None) -> (dict, int):
         weights = []
         quantization_level = None
-        for file in sorted(root_path.glob(model_name + "/*.safetensors")):
-            quantization_level = safe_open(file, framework="pt").metadata().get("quantization_level")
-            weight = list(mx.load(str(file)).items())
+
+        if root_path is not None:
+            for file in sorted(root_path.glob(model_name + "/*.safetensors")):
+                quantization_level = mx.load(str(file), return_metadata=True)[1].get("quantization_level")
+                weight = list(mx.load(str(file)).items())
+                weights.extend(weight)
+
+        if lora_path and root_path is None:
+            weight = list(mx.load(lora_path).items())
             weights.extend(weight)
 
         # Non huggingface weights (i.e. ones exported from this project) don't need any reshaping.
@@ -106,22 +123,11 @@ class WeightHandler:
             return tree_unflatten(weights), quantization_level
 
         # Huggingface weights needs to be reshaped
-        weights = [WeightHandler._reshape_weights(k, v) for k, v in weights]
-        weights = WeightHandler._flatten(weights)
+        weights = [WeightUtil.reshape_weights(k, v) for k, v in weights]
+        weights = WeightUtil.flatten(weights)
         unflatten = tree_unflatten(weights)
         return unflatten, quantization_level
 
-    @staticmethod
-    def _flatten(params):
-        return [(k, v) for p in params for (k, v) in p]
-
-    @staticmethod
-    def _reshape_weights(key, value):
-        if len(value.shape) == 4:
-            value = value.transpose(0, 2, 3, 1)
-        value = value.reshape(-1).reshape(value.shape).astype(Config.precision)
-        return [(key, value)]
-
     @staticmethod
     def _download_or_get_cached_weights(repo_id: str) -> Path:
         return Path(

src/flux_1/weights/weight_util.py

@@ -0,0 +1,15 @@
+from flux_1.config.config import Config
+
+
+class WeightUtil:
+
+    @staticmethod
+    def flatten(params):
+        return [(k, v) for p in params for (k, v) in p]
+
+    @staticmethod
+    def reshape_weights(key, value):
+        if len(value.shape) == 4:
+            value = value.transpose(0, 2, 3, 1)
+        value = value.reshape(-1).reshape(value.shape).astype(Config.precision)
+        return [(key, value)]