88 lines
4.0 KiB
Python
88 lines
4.0 KiB
Python
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 mflux.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
|