diff --git a/.gitignore b/.gitignore index ee9ae73..afe2251 100644 --- a/.gitignore +++ b/.gitignore @@ -11,3 +11,4 @@ *.png *.jpg *.pyc +*.safetensors diff --git a/README.md b/README.md index 09a93ad..433a4fc 100644 --- a/README.md +++ b/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?) \ No newline at end of file diff --git a/main.py b/main.py index 126e783..4660c06 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,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, diff --git a/src/flux_1/assets/lora1.jpg b/src/flux_1/assets/lora1.jpg new file mode 100644 index 0000000..46188d4 Binary files /dev/null and b/src/flux_1/assets/lora1.jpg differ diff --git a/src/flux_1/assets/lora2.jpg b/src/flux_1/assets/lora2.jpg new file mode 100644 index 0000000..b317be0 Binary files /dev/null and b/src/flux_1/assets/lora2.jpg differ diff --git a/src/flux_1/assets/lora3.jpg b/src/flux_1/assets/lora3.jpg new file mode 100644 index 0000000..f83edb6 Binary files /dev/null and b/src/flux_1/assets/lora3.jpg differ diff --git a/src/flux_1/flux.py b/src/flux_1/flux.py index 74b58a8..9a00d46 100644 --- a/src/flux_1/flux.py +++ b/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: diff --git a/src/flux_1/weights/lora_util.py b/src/flux_1/weights/lora_util.py new file mode 100644 index 0000000..afb539d --- /dev/null +++ b/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 diff --git a/src/flux_1/weights/weight_handler.py b/src/flux_1/weights/weight_handler.py index 7ef50b6..0314437 100644 --- a/src/flux_1/weights/weight_handler.py +++ b/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( diff --git a/src/flux_1/weights/weight_util.py b/src/flux_1/weights/weight_util.py new file mode 100644 index 0000000..50d8359 --- /dev/null +++ b/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)]