wehub-resource-sync
131 lines
5.2 KiB
Python
131 lines
5.2 KiB
Python
#!/usr/bin/env python3
|
|
# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES
|
|
#
|
|
# Licensed under the Apache License, Version 2.0 (the "License").
|
|
# You may not use this file except in compliance with the License.
|
|
# To view a copy of this license, visit http://www.apache.org/licenses/LICENSE-2.0
|
|
#
|
|
# No warranties are given. The work is provided "AS IS", without warranty of any kind, express or implied.
|
|
#
|
|
# SPDX-License-Identifier: Apache-2.0
|
|
"""Merge a LongLive generator checkpoint with LoRA weights for simple inference."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import argparse
|
|
import os
|
|
import sys
|
|
from pathlib import Path
|
|
|
|
PROJECT_ROOT = Path(__file__).resolve().parents[1]
|
|
if str(PROJECT_ROOT) not in sys.path:
|
|
sys.path.insert(0, str(PROJECT_ROOT))
|
|
|
|
|
|
def _torch_load(path: str):
|
|
import torch
|
|
|
|
try:
|
|
return torch.load(path, map_location="cpu", weights_only=False)
|
|
except TypeError:
|
|
return torch.load(path, map_location="cpu")
|
|
|
|
|
|
def _load_lora_state(path: str):
|
|
checkpoint = _torch_load(path)
|
|
if isinstance(checkpoint, dict) and "generator_lora" in checkpoint:
|
|
return checkpoint["generator_lora"]
|
|
return checkpoint
|
|
|
|
|
|
def parse_args() -> argparse.Namespace:
|
|
parser = argparse.ArgumentParser(description=__doc__)
|
|
parser.add_argument("--config_path", required=True, help="Inference yaml containing model, checkpoint, and adapter settings.")
|
|
parser.add_argument("--output_path", required=True, help="Path to save the merged generator checkpoint.")
|
|
parser.add_argument("--generator_ckpt", default=None, help="Override checkpoints.generator_ckpt from the yaml.")
|
|
parser.add_argument("--lora_ckpt", default=None, help="Override checkpoints.lora_ckpt from the yaml.")
|
|
parser.add_argument("--device", default="cuda:0", help="Device used for merging, e.g. cuda:0 or cpu.")
|
|
parser.add_argument("--dtype", choices=("bf16", "fp32"), default="bf16", help="Save merged weights in this dtype.")
|
|
return parser.parse_args()
|
|
|
|
|
|
def main() -> None:
|
|
args = parse_args()
|
|
|
|
import torch
|
|
from omegaconf import OmegaConf
|
|
|
|
from utils.config import normalize_config
|
|
from utils.inference_utils import load_generator_checkpoint
|
|
from utils.lora_utils import configure_lora_for_model
|
|
from utils.nvfp4_checkpoint import cpu_state_dict
|
|
from utils.wan_5b_wrapper import WanDiffusionWrapper
|
|
|
|
config = normalize_config(OmegaConf.load(args.config_path))
|
|
generator_ckpt = args.generator_ckpt or getattr(config, "generator_ckpt", None)
|
|
lora_ckpt = args.lora_ckpt or getattr(config, "lora_ckpt", None)
|
|
if not generator_ckpt:
|
|
raise ValueError("Missing generator checkpoint. Set checkpoints.generator_ckpt or pass --generator_ckpt.")
|
|
if not lora_ckpt:
|
|
raise ValueError("Missing LoRA checkpoint. Set checkpoints.lora_ckpt or pass --lora_ckpt.")
|
|
if not getattr(config, "adapter", None):
|
|
raise ValueError("Missing adapter config. The merge script needs the LoRA rank/alpha/dropout settings.")
|
|
|
|
device = torch.device(args.device)
|
|
dtype = torch.bfloat16 if args.dtype == "bf16" else torch.float32
|
|
|
|
print(f"Building generator: {config.model_kwargs}")
|
|
generator = WanDiffusionWrapper(**getattr(config, "model_kwargs", {}), is_causal=True)
|
|
generator.eval().requires_grad_(False)
|
|
|
|
print(f"Loading generator checkpoint: {generator_ckpt}")
|
|
incompatible = load_generator_checkpoint(
|
|
generator,
|
|
generator_ckpt,
|
|
use_ema=bool(getattr(config, "use_ema", False)),
|
|
)
|
|
missing = getattr(incompatible, "missing_keys", [])
|
|
unexpected = getattr(incompatible, "unexpected_keys", [])
|
|
if missing:
|
|
print(f"[Warning] Missing generator keys: {missing[:8]} ...")
|
|
if unexpected:
|
|
print(f"[Warning] Unexpected generator keys: {unexpected[:8]} ...")
|
|
|
|
print(f"Applying LoRA config: {config.adapter}")
|
|
generator.model = configure_lora_for_model(
|
|
generator.model,
|
|
model_name="generator",
|
|
lora_config=config.adapter,
|
|
is_main_process=True,
|
|
)
|
|
|
|
import peft
|
|
|
|
print(f"Loading LoRA checkpoint: {lora_ckpt}")
|
|
peft.set_peft_model_state_dict(generator.model, _load_lora_state(lora_ckpt)) # type: ignore[arg-type]
|
|
|
|
print(f"Merging LoRA on {device} in {dtype}...")
|
|
generator.to(device=device, dtype=dtype)
|
|
generator.model = generator.model.merge_and_unload(safe_merge=True)
|
|
generator.eval().requires_grad_(False)
|
|
|
|
output_path = Path(args.output_path)
|
|
output_path.parent.mkdir(parents=True, exist_ok=True)
|
|
checkpoint = {
|
|
"generator": cpu_state_dict(generator),
|
|
"checkpoint_format": "longlive_generator_merged_lora",
|
|
"source_generator_ckpt": str(generator_ckpt),
|
|
"source_lora_ckpt": str(lora_ckpt),
|
|
"model_name": getattr(config.model_kwargs, "model_name", None),
|
|
"dtype": str(dtype).replace("torch.", ""),
|
|
"merged_lora": True,
|
|
}
|
|
torch.save(checkpoint, output_path)
|
|
size_gib = os.path.getsize(output_path) / (1024 ** 3)
|
|
print(f"Saved merged generator to {output_path} ({size_gib:.2f} GiB).")
|
|
print("Use this file as checkpoints.generator_ckpt for inference and remove adapter/lora_ckpt from the inference config.")
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|