123 lines
3.7 KiB
Python
123 lines
3.7 KiB
Python
"""LoRA weight loading utilities."""
|
|
|
|
import re
|
|
from pathlib import Path
|
|
from typing import Dict, List, Optional
|
|
|
|
import mlx.core as mx
|
|
|
|
from mlx_video.lora.types import LoRAConfig, LoRAWeights
|
|
|
|
|
|
def load_lora_weights(lora_path: Path) -> Dict[str, LoRAWeights]:
|
|
"""Load LoRA weights from a safetensors file.
|
|
|
|
Supports both key conventions:
|
|
- {module_name}.lora_A.weight / {module_name}.lora_B.weight
|
|
- {module_name}.lora_down.weight / {module_name}.lora_up.weight
|
|
|
|
Args:
|
|
lora_path: Path to the LoRA safetensors file
|
|
|
|
Returns:
|
|
Dictionary mapping module names to LoRAWeights objects
|
|
|
|
Raises:
|
|
FileNotFoundError: If the LoRA file doesn't exist
|
|
ValueError: If the LoRA file format is invalid
|
|
"""
|
|
if not lora_path.exists():
|
|
raise FileNotFoundError(f"LoRA file not found: {lora_path}")
|
|
|
|
all_weights = mx.load(str(lora_path))
|
|
|
|
# Group weights by module name, handling both naming conventions
|
|
lora_weights = {}
|
|
module_names = set()
|
|
|
|
for key in all_weights.keys():
|
|
# Format 1: {module}.lora_A.weight / {module}.lora_B.weight
|
|
match = re.match(r"(.+)\.lora_([AB])\.weight$", key)
|
|
if match:
|
|
module_names.add(match.group(1))
|
|
continue
|
|
# Format 2: {module}.lora_down.weight / {module}.lora_up.weight
|
|
match = re.match(r"(.+)\.lora_(down|up)\.weight$", key)
|
|
if match:
|
|
module_names.add(match.group(1))
|
|
|
|
for module_name in module_names:
|
|
# Try both key conventions
|
|
key_a = f"{module_name}.lora_A.weight"
|
|
key_b = f"{module_name}.lora_B.weight"
|
|
if key_a not in all_weights or key_b not in all_weights:
|
|
key_a = f"{module_name}.lora_down.weight"
|
|
key_b = f"{module_name}.lora_up.weight"
|
|
if key_a not in all_weights or key_b not in all_weights:
|
|
continue
|
|
|
|
lora_a = all_weights[key_a]
|
|
lora_b = all_weights[key_b]
|
|
|
|
if lora_a.ndim != 2 or lora_b.ndim != 2:
|
|
raise ValueError(
|
|
f"Invalid LoRA shape for {module_name}: "
|
|
f"lora_A={lora_a.shape}, lora_B={lora_b.shape}"
|
|
)
|
|
|
|
rank = lora_a.shape[0]
|
|
if lora_b.shape[1] != rank:
|
|
raise ValueError(
|
|
f"LoRA rank mismatch for {module_name}: "
|
|
f"lora_A rank={rank}, lora_B rank={lora_b.shape[1]}"
|
|
)
|
|
|
|
# Check for per-module alpha stored as a scalar tensor
|
|
alpha_key = f"{module_name}.alpha"
|
|
if alpha_key in all_weights:
|
|
alpha = float(all_weights[alpha_key].item())
|
|
else:
|
|
alpha = float(rank)
|
|
|
|
lora_weights[module_name] = LoRAWeights(
|
|
lora_A=lora_a,
|
|
lora_B=lora_b,
|
|
rank=rank,
|
|
alpha=alpha,
|
|
module_name=module_name,
|
|
)
|
|
|
|
if not lora_weights:
|
|
raise ValueError(f"No valid LoRA weights found in {lora_path}")
|
|
|
|
return lora_weights
|
|
|
|
|
|
def load_multiple_loras(
|
|
configs: List[LoRAConfig],
|
|
) -> Dict[str, List[tuple]]:
|
|
"""Load multiple LoRA configurations.
|
|
|
|
Args:
|
|
configs: List of LoRAConfig objects
|
|
|
|
Returns:
|
|
Dictionary mapping module names to lists of (LoRAWeights, strength) tuples.
|
|
"""
|
|
module_to_loras: Dict[str, list] = {}
|
|
|
|
for config in configs:
|
|
lora_weights = load_lora_weights(config.path)
|
|
|
|
for module_name, weights in lora_weights.items():
|
|
if config.target_modules is not None:
|
|
if module_name not in config.target_modules:
|
|
continue
|
|
|
|
if module_name not in module_to_loras:
|
|
module_to_loras[module_name] = []
|
|
|
|
module_to_loras[module_name].append((weights, config.strength))
|
|
|
|
return module_to_loras
|