135 lines
3.8 KiB
Python
135 lines
3.8 KiB
Python
"""Merge a list of LoRAs into a single checkpoint."""
|
|
|
|
import argparse
|
|
import json
|
|
import os
|
|
import sys
|
|
|
|
import torch
|
|
from safetensors.torch import load_file, save_file
|
|
from safetensors import safe_open
|
|
|
|
|
|
DTYPE_MAP = {
|
|
"float32": torch.float32,
|
|
"fp32": torch.float32,
|
|
"float16": torch.float16,
|
|
"fp16": torch.float16,
|
|
"bfloat16": torch.bfloat16,
|
|
"bf16": torch.bfloat16,
|
|
}
|
|
|
|
|
|
def log(message: str) -> None:
|
|
print(message, flush=True)
|
|
|
|
|
|
def main() -> int:
|
|
parser = argparse.ArgumentParser(
|
|
description="Merge a list of LoRAs into a single checkpoint."
|
|
)
|
|
parser.add_argument(
|
|
"--loras",
|
|
required=True,
|
|
help='JSON list of {"path": "...", "strength": 1.0} entries.',
|
|
)
|
|
parser.add_argument("--output", required=True, help="Output .safetensors path.")
|
|
parser.add_argument(
|
|
"--save_dtype",
|
|
default="bfloat16",
|
|
choices=list(DTYPE_MAP.keys()),
|
|
help="Dtype of the saved tensors (merging is always done in float32).",
|
|
)
|
|
parser.add_argument(
|
|
"--device",
|
|
default="cpu",
|
|
help="Device to merge on (cpu, cuda, cuda:1, mps). Output is always saved from CPU.",
|
|
)
|
|
args = parser.parse_args()
|
|
|
|
try:
|
|
loras = json.loads(args.loras)
|
|
except json.JSONDecodeError as e:
|
|
print(f"Failed to parse --loras JSON: {e}", file=sys.stderr, flush=True)
|
|
return 2
|
|
|
|
if not isinstance(loras, list) or len(loras) == 0:
|
|
print("--loras must be a non-empty JSON list.", file=sys.stderr, flush=True)
|
|
return 2
|
|
|
|
device = torch.device(args.device)
|
|
save_dtype = DTYPE_MAP[args.save_dtype]
|
|
|
|
log(f"Merging {len(loras)} LoRA(s) on {device}, saving as {args.save_dtype}.")
|
|
|
|
merged: dict[str, torch.Tensor] = {}
|
|
|
|
metadata = {}
|
|
|
|
for i, entry in enumerate(loras):
|
|
if not isinstance(entry, dict) or "path" not in entry:
|
|
print(
|
|
f"LoRA entry {i} must be an object with a 'path' field.",
|
|
file=sys.stderr,
|
|
flush=True,
|
|
)
|
|
return 2
|
|
|
|
path = entry["path"]
|
|
strength = float(entry.get("strength", 1.0))
|
|
|
|
if not os.path.isfile(path):
|
|
print(f"LoRA file not found: {path}", file=sys.stderr, flush=True)
|
|
return 2
|
|
|
|
log(f"[{i + 1}/{len(loras)}] Loading {path} (strength={strength})")
|
|
state_dict = load_file(path, device=str(device))
|
|
|
|
for key, tensor in state_dict.items():
|
|
scaled = tensor.to(torch.float32) * strength
|
|
if key in merged:
|
|
merged[key].add_(scaled)
|
|
else:
|
|
merged[key] = scaled
|
|
del state_dict
|
|
|
|
if i == 0:
|
|
# For the first LoRA, also copy over all non-tensor metadata (e.g. base model info)
|
|
with safe_open(path, framework="pt") as f:
|
|
metadata_to_keep = [
|
|
"version",
|
|
"format",
|
|
"ss_base_model_version",
|
|
"software",
|
|
]
|
|
orig_metadata = f.metadata()
|
|
for meta_key in metadata_to_keep:
|
|
if meta_key in orig_metadata:
|
|
metadata[meta_key] = orig_metadata[meta_key]
|
|
|
|
log(f"Casting to {args.save_dtype} and moving to CPU")
|
|
final = {k: v.to(save_dtype).cpu().contiguous() for k, v in merged.items()}
|
|
merged.clear()
|
|
|
|
log(f"Saving merged checkpoint to {args.output}")
|
|
save_file(final, args.output, metadata=metadata)
|
|
|
|
print(
|
|
json.dumps(
|
|
{
|
|
"ok": True,
|
|
"output": args.output,
|
|
"num_loras": len(loras),
|
|
"num_keys": len(final),
|
|
"save_dtype": args.save_dtype,
|
|
"device": str(device),
|
|
}
|
|
),
|
|
flush=True,
|
|
)
|
|
return 0
|
|
|
|
|
|
if __name__ == "__main__":
|
|
sys.exit(main())
|