diff --git a/src/flux_1/weights/lora_util.py b/src/flux_1/weights/lora_util.py new file mode 100644 index 0000000..fde465f --- /dev/null +++ b/src/flux_1/weights/lora_util.py @@ -0,0 +1,102 @@ +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_handler import WeightHandler + +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}") + + @staticmethod + def _apply_transformer(transformer, lora_transformer, lora_scale): + 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 + + @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 = [WeightHandler.reshape_weights(k, v) for k, v in weights] + weights = WeightHandler.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 diff --git a/src/flux_1/weights/weight_handler.py b/src/flux_1/weights/weight_handler.py index 41a326f..b184d9a 100644 --- a/src/flux_1/weights/weight_handler.py +++ b/src/flux_1/weights/weight_handler.py @@ -1,15 +1,12 @@ -import logging from pathlib import Path import mlx.core as mx from huggingface_hub import snapshot_download -from mlx.utils import tree_flatten from mlx.utils import tree_unflatten from safetensors import safe_open from flux_1.config.config import Config - -log = logging.getLogger(__name__) +from flux_1.weights.lora_util import LoraUtil class WeightHandler: @@ -25,99 +22,16 @@ class WeightHandler: 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) - 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, _ = 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.") - WeightHandler._apply_transformer(self.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): - 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 - - @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 = [WeightHandler._reshape_weights(k, v) for k, v in weights] - weights = WeightHandler._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 + # Optionally apply LoRA weights + LoraUtil.apply_lora(self.transformer, lora_files, lora_scales) @staticmethod def _clip_encoder(root_path: Path) -> (dict, int): @@ -203,17 +117,17 @@ 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 = [WeightHandler.reshape_weights(k, v) for k, v in weights] + weights = WeightHandler.flatten(weights) unflatten = tree_unflatten(weights) return unflatten, quantization_level @staticmethod - def _flatten(params): + def flatten(params): return [(k, v) for p in params for (k, v) in p] @staticmethod - def _reshape_weights(key, value): + 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)