|
import argparse
|
|
import os
|
|
from safetensors.torch import load_file, save_file
|
|
import toml
|
|
import re
|
|
from safetensors import safe_open
|
|
import math
|
|
|
|
def parse_key(key):
|
|
match = re.match(r"lora_unet_(input|output|up|down)_blocks_(\d+(?:_\d+)?)_(.+)\.(?:alpha|lora_(?:down|up)\.weight)", key)
|
|
if match:
|
|
return "unet", match.group(1) + "_blocks", match.group(2), match.group(3)
|
|
|
|
match = re.match(r"lora_unet_(mid_block)_(resnets|attentions)_(\d+)_(.+)\.(?:alpha|lora_(?:down|up)\.weight)", key)
|
|
if match:
|
|
return "unet", match.group(1), f"{match.group(2)}_{match.group(3)}", match.group(4)
|
|
|
|
match = re.match(r"lora_unet_(middle_block)_(\d+)_(.+)\.(?:alpha|lora_(?:down|up)\.weight)", key)
|
|
if match:
|
|
return "unet", match.group(1), match.group(2), match.group(3)
|
|
|
|
match = re.match(r"lora_te\d+_text_model_encoder_(.+)\.(?:alpha|lora_(?:down|up)\.weight)", key)
|
|
if match:
|
|
return "text_encoder", "encoder_layers", match.group(1).split("_")[0], "_".join(match.group(1).split("_")[1:])
|
|
|
|
return None, None, None, None
|
|
|
|
def extract_lora_hierarchy(lora_tensors, mode="extract"):
|
|
lora_hierarchy = {}
|
|
lora_key_groups = {"unet": {}, "text_encoder": {}} if mode == "adjust" else None
|
|
|
|
for key in lora_tensors:
|
|
if key.startswith("lora_unet_"):
|
|
model_type, block_type, block_num, layer_key = parse_key(key)
|
|
|
|
if model_type and block_type and layer_key:
|
|
parts = layer_key.split("_")
|
|
if "transformer_blocks" in layer_key:
|
|
grouped_key = "_".join(parts[:3] + [parts[3] if len(parts) > 5 else ""])
|
|
elif "attentions" in layer_key:
|
|
grouped_key = "_".join(parts[:3] + [parts[3] if len(parts) > 5 else ""])
|
|
elif "resnets" in layer_key:
|
|
grouped_key = "_".join(parts[:3])
|
|
else:
|
|
grouped_key = layer_key
|
|
|
|
if model_type not in lora_hierarchy:
|
|
lora_hierarchy[model_type] = {}
|
|
if block_type not in lora_hierarchy[model_type]:
|
|
lora_hierarchy[model_type][block_type] = {}
|
|
if block_num not in lora_hierarchy[model_type][block_type]:
|
|
lora_hierarchy[model_type][block_type][block_num] = {}
|
|
lora_hierarchy[model_type][block_type][block_num][grouped_key] = 1.0
|
|
|
|
if mode == "adjust":
|
|
group_key = f"..unet_{block_type}_{block_num}_{grouped_key}"
|
|
if group_key not in lora_key_groups["unet"]:
|
|
lora_key_groups["unet"][group_key] = []
|
|
lora_key_groups["unet"][group_key].append(key)
|
|
|
|
elif key.startswith("lora_te"):
|
|
match = re.match(r"(lora_te\d+)_text_model_encoder_layers_(\d+)_(.+)\.(?:alpha|lora_(?:down|up)\.weight)", key)
|
|
if match:
|
|
model_section = match.group(1)
|
|
block_type = "encoder"
|
|
block_num = match.group(2)
|
|
layer_key = match.group(3)
|
|
|
|
grouped_key = f"layers_{block_num}__{layer_key}"
|
|
|
|
if model_section not in lora_hierarchy:
|
|
lora_hierarchy[model_section] = {}
|
|
if block_type not in lora_hierarchy[model_section]:
|
|
lora_hierarchy[model_section][block_type] = {}
|
|
lora_hierarchy[model_section][block_type][grouped_key] = 1.0
|
|
|
|
if mode == "adjust":
|
|
group_key = f"..{model_section}_{block_num}_{layer_key}"
|
|
lora_key_groups["text_encoder"][group_key] = [key]
|
|
|
|
return lora_hierarchy if mode == "extract" else lora_key_groups
|
|
|
|
|
|
def adjust_lora_weights(lora_path, toml_path, output_path, multiplier=1.0, remove_zero_weight_keys=True):
|
|
try:
|
|
lora_tensors = load_file(lora_path)
|
|
with safe_open(lora_path, framework="pt") as f:
|
|
metadata = f.metadata()
|
|
except Exception as e:
|
|
raise Exception(f"Error loading LoRA model: {e}")
|
|
|
|
try:
|
|
with open(toml_path, "r") as f:
|
|
lora_config = toml.load(f)
|
|
except Exception as e:
|
|
raise Exception(f"Error loading TOML file: {e}")
|
|
|
|
|
|
lora_key_groups = extract_lora_hierarchy(lora_tensors, mode="adjust")
|
|
adjusted_tensors = {}
|
|
|
|
for model_section, model_config in lora_config.items():
|
|
if model_section.startswith("lora_te"):
|
|
for block_type, layers in model_config.items():
|
|
for layer_key, weight in layers.items():
|
|
block_num, layer_name = layer_key.replace("layers_", "").split("__")
|
|
group_key = f"..{model_section}_{block_num}_{layer_name}"
|
|
if group_key in lora_key_groups["text_encoder"]:
|
|
final_weight = weight * multiplier
|
|
if not remove_zero_weight_keys or final_weight != 0.0:
|
|
for target_key in lora_key_groups["text_encoder"][group_key]:
|
|
if target_key.endswith(".alpha"):
|
|
final_weight = weight * multiplier
|
|
if not remove_zero_weight_keys or final_weight != 0.0:
|
|
adjusted_tensors[target_key] = lora_tensors[target_key]
|
|
else:
|
|
final_weight = weight * multiplier
|
|
if not remove_zero_weight_keys or final_weight != 0.0:
|
|
adjusted_tensors[target_key] = lora_tensors[target_key] * math.sqrt(final_weight)
|
|
|
|
else:
|
|
for block_type, block_nums in model_config.items():
|
|
for block_num, layer_keys in block_nums.items():
|
|
for grouped_key, weight in layer_keys.items():
|
|
group_key = f"..unet_{block_type}_{block_num}_{grouped_key}"
|
|
if group_key in lora_key_groups["unet"]:
|
|
final_weight = weight * multiplier
|
|
if not remove_zero_weight_keys or final_weight != 0.0:
|
|
for target_key in lora_key_groups["unet"][group_key]:
|
|
if target_key.endswith(".alpha"):
|
|
final_weight = weight * multiplier
|
|
if not remove_zero_weight_keys or final_weight != 0.0:
|
|
adjusted_tensors[target_key] = lora_tensors[target_key]
|
|
else:
|
|
final_weight = weight * multiplier
|
|
if not remove_zero_weight_keys or final_weight != 0.0:
|
|
adjusted_tensors[target_key] = lora_tensors[target_key] * math.sqrt(final_weight)
|
|
|
|
|
|
|
|
try:
|
|
save_file(adjusted_tensors, output_path, metadata)
|
|
except Exception as e:
|
|
raise Exception(f"Error saving adjusted model: {e}")
|
|
|
|
|
|
def write_toml(lora_hierarchy, output_path):
|
|
try:
|
|
with open(output_path, "w") as f:
|
|
toml.dump(lora_hierarchy, f)
|
|
except Exception as e:
|
|
raise Exception(f"Error writing TOML file: {e}")
|
|
|
|
|
|
def main():
|
|
parser = argparse.ArgumentParser(description="Extract or adjust LoRA weights based on a TOML config.")
|
|
subparsers = parser.add_subparsers(dest="mode", help="Choose mode: 'extract' or 'adjust'")
|
|
|
|
|
|
parser_extract = subparsers.add_parser("extract", help="Extract LoRA hierarchy to a TOML file")
|
|
parser_extract.add_argument("--lora_path", required=True, help="Path to the LoRA safetensors file")
|
|
parser_extract.add_argument("--output_path", required=True, help="Path to the output TOML file")
|
|
|
|
|
|
parser_adjust = subparsers.add_parser("adjust", help="Adjust LoRA weights based on a TOML config.")
|
|
parser_adjust.add_argument("--lora_path", required=True, help="Path to the LoRA safetensors file")
|
|
parser_adjust.add_argument("--toml_path", required=True, help="Path to the TOML config file")
|
|
parser_adjust.add_argument("--output_path", required=True, help="Path to the output safetensors file")
|
|
parser_adjust.add_argument("--multiplier", type=float, default=1.0, help="Global multiplier for the LoRA weights")
|
|
parser_adjust.add_argument("--remove_zero_weight_keys", action="store_true",
|
|
help="Remove keys with resulting weight of 0. Useful for reducing file size.")
|
|
|
|
args = parser.parse_args()
|
|
|
|
try:
|
|
if args.mode == "extract":
|
|
lora_tensors = load_file(args.lora_path)
|
|
lora_hierarchy = extract_lora_hierarchy(lora_tensors)
|
|
write_toml(lora_hierarchy, args.output_path)
|
|
print(f"Successfully extracted LoRA hierarchy to {args.output_path}")
|
|
|
|
elif args.mode == "adjust":
|
|
adjust_lora_weights(args.lora_path, args.toml_path, args.output_path, args.multiplier, args.remove_zero_weight_keys)
|
|
print(f"Successfully adjusted LoRA weights and saved to {args.output_path}")
|
|
|
|
else:
|
|
parser.print_help()
|
|
|
|
except Exception as e:
|
|
print(f"An error occurred: {e}")
|
|
|
|
if __name__ == "__main__":
|
|
main() |
