334 lines
12 KiB
Python
334 lines
12 KiB
Python
"""Merge LoRA adapter into base model weights.
|
|
|
|
Usage:
|
|
python merge_lora_updated.py \
|
|
--base Wan-AI/Wan2.2-TI2V-5B-Diffusers \
|
|
--adapter adapter.safetensors \
|
|
--ft FastVideo/FastWan2.2-TI2V-5B-FullAttn-Diffusers \
|
|
--output ./merged_model
|
|
"""
|
|
from __future__ import annotations
|
|
|
|
import argparse
|
|
import os
|
|
import sys
|
|
import shutil
|
|
import logging
|
|
from pathlib import Path
|
|
from collections import defaultdict
|
|
|
|
os.environ.setdefault("MASTER_ADDR", "127.0.0.1")
|
|
os.environ.setdefault("MASTER_PORT", "29500")
|
|
os.environ.setdefault("WORLD_SIZE", "1")
|
|
os.environ.setdefault("RANK", "0")
|
|
os.environ.setdefault("LOCAL_RANK", "0")
|
|
|
|
_FASTVIDEO_PATH = os.path.abspath(os.path.join(os.path.dirname(__file__), "..", "fastvideo_pr", "FastVideo"))
|
|
if _FASTVIDEO_PATH not in sys.path:
|
|
sys.path.insert(0, _FASTVIDEO_PATH)
|
|
|
|
import torch
|
|
from safetensors.torch import load_file, save_file
|
|
|
|
from extract_lora import load_transformer_state_dict_from_model, get_pipeline_class_for_model
|
|
from fastvideo.training.training_utils import custom_to_hf_state_dict
|
|
from fastvideo.models.loader.utils import get_param_names_mapping, hf_to_custom_state_dict
|
|
|
|
LOG = logging.getLogger(__name__)
|
|
|
|
|
|
def configure_logging(level: str = "INFO"):
|
|
handler = logging.StreamHandler()
|
|
handler.setFormatter(logging.Formatter("%(asctime)s %(levelname)s %(message)s", datefmt="%Y-%m-%d %H:%M:%S"))
|
|
LOG.addHandler(handler)
|
|
LOG.setLevel(level)
|
|
|
|
|
|
def fix_adapter_naming(adapter: dict) -> dict:
|
|
fixed = {}
|
|
renamed = 0
|
|
|
|
for key, tensor in adapter.items():
|
|
new_key = key
|
|
|
|
if ".lora_A.weight" in key:
|
|
base = key.replace(".lora_A.weight", "")
|
|
if not base.endswith(".weight"):
|
|
new_key = base + ".weight.lora_A.weight"
|
|
renamed += 1
|
|
elif ".lora_B.weight" in key:
|
|
base = key.replace(".lora_B.weight", "")
|
|
if not base.endswith(".weight"):
|
|
new_key = base + ".weight.lora_B.weight"
|
|
renamed += 1
|
|
elif ".lora_rank" in key:
|
|
base = key.replace(".lora_rank", "")
|
|
if not base.endswith(".weight"):
|
|
new_key = base + ".weight.lora_rank"
|
|
renamed += 1
|
|
elif ".lora_alpha" in key:
|
|
base = key.replace(".lora_alpha", "")
|
|
if not base.endswith(".weight"):
|
|
new_key = base + ".weight.lora_alpha"
|
|
renamed += 1
|
|
|
|
fixed[new_key] = tensor
|
|
|
|
if renamed > 0:
|
|
LOG.info(f"Fixed {renamed} adapter key names")
|
|
|
|
return fixed
|
|
|
|
|
|
def load_adapter(adapter_path: str) -> dict:
|
|
abs_path = os.path.abspath(adapter_path)
|
|
if not os.path.exists(abs_path):
|
|
raise FileNotFoundError(f"Adapter file not found: {abs_path}")
|
|
if not abs_path.endswith('.safetensors'):
|
|
raise ValueError(f"Adapter must be .safetensors: {abs_path}")
|
|
|
|
LOG.info(f"Loading adapter: {abs_path}")
|
|
adapter = load_file(abs_path)
|
|
file_size_mb = os.path.getsize(abs_path) / (1024 * 1024)
|
|
LOG.info(f"Loaded {len(adapter)} tensors ({file_size_mb:.1f} MB)")
|
|
|
|
return fix_adapter_naming(adapter)
|
|
|
|
|
|
def group_adapter_keys(adapter: dict) -> dict:
|
|
grouped = defaultdict(dict)
|
|
|
|
for key, tensor in adapter.items():
|
|
if key.endswith(".lora_A.weight"):
|
|
grouped[key.replace(".lora_A.weight", "")]["A"] = tensor
|
|
elif key.endswith(".lora_B.weight"):
|
|
grouped[key.replace(".lora_B.weight", "")]["B"] = tensor
|
|
elif key.endswith(".lora_rank"):
|
|
grouped[key.replace(".lora_rank", "")]["rank"] = tensor
|
|
elif key.endswith(".lora_alpha"):
|
|
grouped[key.replace(".lora_alpha", "")]["alpha"] = tensor
|
|
|
|
LOG.info(f"Grouped {len(grouped)} LoRA layers")
|
|
return grouped
|
|
|
|
|
|
def get_reverse_param_mapping(base_model_path: str):
|
|
LOG.info("Loading base model for parameter mapping")
|
|
|
|
pipeline_cls = get_pipeline_class_for_model(base_model_path)
|
|
pipeline = pipeline_cls.from_pretrained(
|
|
base_model_path,
|
|
num_gpus=1,
|
|
inference_mode=True,
|
|
dit_cpu_offload=True,
|
|
vae_cpu_offload=True,
|
|
text_encoder_cpu_offload=True,
|
|
pin_cpu_memory=True,
|
|
)
|
|
|
|
transformer = None
|
|
if hasattr(pipeline, "transformer"):
|
|
transformer = pipeline.transformer
|
|
elif hasattr(pipeline, "modules") and isinstance(pipeline.modules, dict):
|
|
if "transformer" in pipeline.modules:
|
|
transformer = pipeline.modules["transformer"]
|
|
|
|
if transformer is None:
|
|
raise RuntimeError("Could not find transformer in pipeline")
|
|
|
|
if hasattr(transformer, "reverse_param_names_mapping"):
|
|
reverse_mapping = transformer.reverse_param_names_mapping
|
|
elif hasattr(transformer, "config") and hasattr(transformer.config, "arch_config"):
|
|
arch_config = transformer.config.arch_config
|
|
if hasattr(arch_config, "reverse_param_names_mapping"):
|
|
reverse_mapping = arch_config.reverse_param_names_mapping
|
|
else:
|
|
param_mapping = arch_config.param_names_mapping
|
|
param_names_mapping_fn = get_param_names_mapping(param_mapping)
|
|
|
|
from diffusers import DiffusionPipeline
|
|
from huggingface_hub import snapshot_download
|
|
|
|
if not os.path.exists(base_model_path) or not os.path.isdir(base_model_path):
|
|
model_path = snapshot_download(repo_id=base_model_path, ignore_patterns=["*.onnx", "*.msgpack"])
|
|
else:
|
|
model_path = base_model_path
|
|
|
|
hf_pipeline = DiffusionPipeline.from_pretrained(model_path, torch_dtype=torch.float32)
|
|
hf_transformer = hf_pipeline.transformer
|
|
hf_sd = hf_transformer.state_dict()
|
|
|
|
_, reverse_mapping = hf_to_custom_state_dict(hf_sd, param_names_mapping_fn)
|
|
|
|
del hf_pipeline
|
|
del hf_transformer
|
|
torch.cuda.empty_cache()
|
|
else:
|
|
raise RuntimeError("Could not find reverse_param_names_mapping in transformer or config")
|
|
|
|
del pipeline
|
|
del transformer
|
|
torch.cuda.empty_cache()
|
|
|
|
return reverse_mapping
|
|
|
|
|
|
def merge_lora_into_base(base_sd: dict, adapter: dict) -> dict:
|
|
LOG.info("Merging LoRA into base weights")
|
|
|
|
adapter_layers = group_adapter_keys(adapter)
|
|
merged_sd = dict(base_sd)
|
|
|
|
merged_count = 0
|
|
skipped_count = 0
|
|
|
|
for base_name, parts in adapter_layers.items():
|
|
weight_key = base_name if base_name.endswith(".weight") else base_name + ".weight"
|
|
|
|
if weight_key not in base_sd:
|
|
skipped_count += 1
|
|
continue
|
|
|
|
if "A" not in parts or "B" not in parts:
|
|
skipped_count += 1
|
|
continue
|
|
|
|
lora_A = parts["A"].to(torch.float32)
|
|
lora_B = parts["B"].to(torch.float32)
|
|
base_weight = base_sd[weight_key].to(torch.float32)
|
|
|
|
out_dim, in_dim = base_weight.shape
|
|
if lora_B.shape[0] != out_dim or lora_A.shape[1] != in_dim or lora_B.shape[1] != lora_A.shape[0]:
|
|
skipped_count += 1
|
|
continue
|
|
|
|
delta = lora_B @ lora_A
|
|
|
|
rank = int(parts.get("rank", torch.tensor([lora_A.shape[0]])).item()) if "rank" in parts else lora_A.shape[0]
|
|
alpha = float(parts.get("alpha", torch.tensor([rank])).item()) if "alpha" in parts else float(rank)
|
|
|
|
if rank != 0 and alpha != rank:
|
|
delta = delta * (alpha / float(rank))
|
|
|
|
merged_weight = base_weight + delta
|
|
merged_sd[weight_key] = merged_weight.to(base_sd[weight_key].dtype)
|
|
merged_count += 1
|
|
|
|
LOG.info(f"Merged {merged_count} layers, skipped {skipped_count}")
|
|
return merged_sd
|
|
|
|
|
|
def save_merged_model(merged_sd: dict, base_model_path: str, ft_model_path: str, output_dir: str,
|
|
reverse_mapping: dict):
|
|
output_path = Path(output_dir)
|
|
output_path.mkdir(parents=True, exist_ok=True)
|
|
|
|
transformer_dir = output_path / "transformer"
|
|
transformer_dir.mkdir(parents=True, exist_ok=True)
|
|
|
|
LOG.info("Converting to HuggingFace format")
|
|
hf_merged_sd = custom_to_hf_state_dict(merged_sd, reverse_mapping)
|
|
LOG.info(f"Converted {len(hf_merged_sd)} parameters")
|
|
|
|
base_path = Path(base_model_path)
|
|
if not base_path.exists() or not base_path.is_dir():
|
|
from huggingface_hub import snapshot_download
|
|
base_path = Path(snapshot_download(repo_id=base_model_path, ignore_patterns=["*.onnx", "*.msgpack"]))
|
|
|
|
LOG.info("Copying model components")
|
|
for component in ["scheduler", "text_encoder", "tokenizer", "vae"]:
|
|
src = base_path / component
|
|
if src.exists():
|
|
dst = output_path / component
|
|
if dst.exists():
|
|
shutil.rmtree(dst)
|
|
shutil.copytree(src, dst)
|
|
|
|
LOG.info("Copying finetuned model config")
|
|
ft_path = Path(ft_model_path)
|
|
if not ft_path.exists() or not ft_path.is_dir():
|
|
from huggingface_hub import snapshot_download
|
|
ft_path = Path(
|
|
snapshot_download(repo_id=ft_model_path,
|
|
allow_patterns=["model_index.json"],
|
|
ignore_patterns=["*.onnx", "*.msgpack"]))
|
|
|
|
ft_index = ft_path / "model_index.json"
|
|
if ft_index.exists():
|
|
shutil.copy2(ft_index, output_path / "model_index.json")
|
|
else:
|
|
LOG.warning("Finetuned model_index.json not found, using base")
|
|
src_index = base_path / "model_index.json"
|
|
if src_index.exists():
|
|
shutil.copy2(src_index, output_path / "model_index.json")
|
|
|
|
weight_path = transformer_dir / "diffusion_pytorch_model.safetensors"
|
|
LOG.info(f"Saving merged weights to {weight_path}")
|
|
to_save_hf = {k: v.detach().cpu() for k, v in hf_merged_sd.items()}
|
|
save_file(to_save_hf, str(weight_path))
|
|
|
|
config_src = base_path / "transformer" / "config.json"
|
|
if config_src.exists():
|
|
shutil.copy2(config_src, transformer_dir / "config.json")
|
|
|
|
file_size_mb = weight_path.stat().st_size / (1024 * 1024)
|
|
LOG.info(f"Saved to {output_dir} ({file_size_mb:.0f} MB, {len(hf_merged_sd)} params)")
|
|
|
|
|
|
def merge_lora(
|
|
base: str,
|
|
adapter: str,
|
|
ft: str,
|
|
output: str,
|
|
log_level: str = "INFO",
|
|
) -> None:
|
|
"""Merge LoRA adapter into base model.
|
|
|
|
Args:
|
|
base: Base model ID or path
|
|
adapter: LoRA adapter .safetensors file
|
|
ft: Finetuned model ID (for config)
|
|
output: Output directory
|
|
log_level: Logging level
|
|
"""
|
|
configure_logging(log_level)
|
|
|
|
LOG.info(f"Base: {base}")
|
|
LOG.info(f"Adapter: {adapter}")
|
|
LOG.info(f"Output: {output}")
|
|
|
|
reverse_mapping = get_reverse_param_mapping(base)
|
|
|
|
LOG.info(f"Loading base model: {base}")
|
|
base_sd = load_transformer_state_dict_from_model(base)
|
|
LOG.info(f"Loaded { len(base_sd)} parameters")
|
|
|
|
adapter_sd = load_adapter(adapter)
|
|
merged_sd = merge_lora_into_base(base_sd, adapter_sd)
|
|
|
|
save_merged_model(merged_sd, base, ft, output, reverse_mapping)
|
|
LOG.info("Merge complete")
|
|
|
|
|
|
def main():
|
|
"""CLI wrapper for merge_lora."""
|
|
parser = argparse.ArgumentParser(description="Merge LoRA adapter into base model")
|
|
parser.add_argument("--base", required=True, help="Base model ID or path")
|
|
parser.add_argument("--adapter", required=True, help="LoRA adapter .safetensors file")
|
|
parser.add_argument("--ft", required=True, help="Finetuned model ID (for config)")
|
|
parser.add_argument("--output", required=True, help="Output directory")
|
|
parser.add_argument("--log-level", default="INFO", help="Logging level")
|
|
args = parser.parse_args()
|
|
|
|
merge_lora(
|
|
base=args.base,
|
|
adapter=args.adapter,
|
|
ft=args.ft,
|
|
output=args.output,
|
|
log_level=args.log_level,
|
|
)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|