|
import shutil |
|
from copy import deepcopy |
|
from pathlib import Path |
|
|
|
import click |
|
import hydra |
|
import torch |
|
from hydra import compose, initialize |
|
from hydra.utils import instantiate |
|
from loguru import logger |
|
|
|
from fish_speech.models.text2semantic.llama import BaseTransformer |
|
from fish_speech.models.text2semantic.lora import get_merged_state_dict |
|
|
|
|
|
@click.command() |
|
@click.option("--lora-config", type=str, default="r_8_alpha_16") |
|
@click.option("--base-weight", type=str, default="checkpoints/fish-speech-1.4") |
|
@click.option("--lora-weight", type=str, required=True) |
|
@click.option("--output", type=str, required=True) |
|
def merge(lora_config, base_weight, lora_weight, output): |
|
output = Path(output) |
|
logger.info( |
|
f"Merging {base_weight} and {lora_weight} into {output} with {lora_config}" |
|
) |
|
|
|
with initialize(version_base="1.3", config_path="../../fish_speech/configs/lora"): |
|
cfg = compose(config_name=lora_config) |
|
|
|
lora_config = instantiate(cfg) |
|
logger.info(f"Loaded lora model with config {lora_config}") |
|
|
|
llama_model = BaseTransformer.from_pretrained( |
|
path=base_weight, |
|
load_weights=True, |
|
lora_config=lora_config, |
|
) |
|
logger.info(f"Loaded llama model") |
|
|
|
llama_state_dict = llama_model.state_dict() |
|
llama_state_dict = {k: v for k, v in llama_state_dict.items() if "lora" not in k} |
|
llama_state_dict_copy = deepcopy(llama_state_dict) |
|
lora_state_dict = torch.load(lora_weight, map_location="cpu") |
|
|
|
if "state_dict" in llama_state_dict: |
|
llama_state_dict = llama_state_dict["state_dict"] |
|
|
|
if "state_dict" in lora_state_dict: |
|
lora_state_dict = lora_state_dict["state_dict"] |
|
|
|
|
|
if any(k.startswith("model.") for k in llama_state_dict.keys()): |
|
llama_state_dict = { |
|
k.replace("model.", ""): v |
|
for k, v in llama_state_dict.items() |
|
if k.startswith("model.") |
|
} |
|
if any(k.startswith("model.") for k in lora_state_dict.keys()): |
|
lora_state_dict = { |
|
k.replace("model.", ""): v |
|
for k, v in lora_state_dict.items() |
|
if k.startswith("model.") |
|
} |
|
|
|
logger.info(f"Found {len(llama_state_dict)} keys in llama model") |
|
logger.info(f"Found {len(lora_state_dict)} keys in lora model") |
|
|
|
merged_state_dict = llama_state_dict | lora_state_dict |
|
llama_model.load_state_dict(merged_state_dict, strict=True) |
|
logger.info(f"Merged model loaded") |
|
|
|
|
|
llama_model.eval() |
|
llama_model.save_pretrained(output, drop_lora=True) |
|
logger.info(f"Saved merged model to {output}, validating") |
|
|
|
new_state_dict = torch.load(output / "model.pth", map_location="cpu") |
|
original_keys = set(llama_state_dict_copy.keys()) |
|
merged_keys = set(new_state_dict.keys()) |
|
|
|
assert original_keys == merged_keys, "Keys should be same" |
|
|
|
for key in original_keys: |
|
diff_l1 = (new_state_dict[key] - llama_state_dict_copy[key]).abs().sum().item() |
|
if diff_l1 != 0: |
|
break |
|
else: |
|
logger.error("Merged model is same as the original model") |
|
exit(1) |
|
|
|
logger.info("Merged model is different from the original model, check passed") |
|
|
|
|
|
if __name__ == "__main__": |
|
merge() |
|
|
