Start supporting only single LoRA (and some other simplifications)
This commit is contained in:
filipstrand
parent
6d35db5e9a
commit
a3567bddd7
10
main.py
10
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(
|
||||
|
||||
@ -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:
|
||||
|
||||
@ -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
|
||||
|
||||
@ -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.
|
||||
|
||||
Loading…
Reference in New Issue
Block a user