Fabio
c36510c21e
Some checks failed
Linux Build / Linux cpu (push) Has been cancelled
Linux Build / Linux cpu-clang (push) Has been cancelled
Linux Build / Linux vulkan (push) Has been cancelled
mac-build / macOS CPU (push) Has been cancelled
Nix Build / Nix (cpu) (push) Has been cancelled
Nix Build / Nix (vulkan) (push) Has been cancelled
windows-build / Windows CPU (push) Has been cancelled
317 lines
17 KiB
Python
317 lines
17 KiB
Python
#!/usr/bin/env python3
|
|
"""Bind a model's LoRA to its base GGUF and merge with audiocpp_gguf."""
|
|
|
|
import argparse
|
|
import json
|
|
from pathlib import Path
|
|
import re
|
|
import subprocess
|
|
import tempfile
|
|
|
|
import torch
|
|
from safetensors.torch import save_file
|
|
|
|
from bindings import BaseGGUF, BaseSources, Bindings, FAMILIES, MERGE_MODES
|
|
from checkpoint import find_checkpoint, load_checkpoint
|
|
from prepare_fish_audio_adapter import prepare as prepare_fish
|
|
|
|
|
|
SUFFIXES = (
|
|
(".lora_A.default.weight", "a"), (".lora_B.default.weight", "b"),
|
|
(".lora_A.weight", "a"), (".lora_B.weight", "b"),
|
|
(".lora_A", "a"), (".lora_B", "b"),
|
|
(".lora_down.weight", "a"), (".lora_up.weight", "b"),
|
|
(".lokr_w1", "lokr_w1"), (".lokr_w2", "lokr_w2"),
|
|
(".lokr_w1_a", "lokr_w1_a"), (".lokr_w1_b", "lokr_w1_b"),
|
|
(".lokr_w2_a", "lokr_w2_a"), (".lokr_w2_b", "lokr_w2_b"),
|
|
(".alpha", "alpha"),
|
|
)
|
|
|
|
|
|
def read_adapter(path, family, work, alpha):
|
|
directory = path if path.is_dir() else path.parent
|
|
native_config = None
|
|
native_cosy = False
|
|
if family == "cosyvoice3":
|
|
native_package = (
|
|
find_checkpoint(path, "lora_weights", required=False) is not None
|
|
if path.is_dir() else path.stem == "lora_weights")
|
|
native_cosy = (directory / "config.json").is_file() and (
|
|
native_package or not (directory / "adapter_config.json").is_file())
|
|
if path.is_dir():
|
|
peft_file = find_checkpoint(path, "adapter_model", required=False) if not native_cosy else None
|
|
if native_cosy:
|
|
path = find_checkpoint(path, "lora_weights")
|
|
elif peft_file is not None:
|
|
path = peft_file
|
|
else:
|
|
stem = "lora_weights" if family in {"voxcpm1", "voxcpm2"} else "adapter_model"
|
|
path = find_checkpoint(path, stem)
|
|
if family == "fish_audio" and alpha is not None:
|
|
prepared = work / "fish.safetensors"
|
|
prepare_fish(path, prepared, alpha)
|
|
path = prepared
|
|
if family in {"voxcpm1", "voxcpm2"}:
|
|
native_config = json.loads((directory / "lora_config.json").read_text())["lora_config"]
|
|
config = {"r": native_config["r"], "lora_alpha": native_config["alpha"]}
|
|
elif native_cosy:
|
|
native_config = json.loads((directory / "config.json").read_text())["config"]
|
|
config = {"r": native_config["lora_r"], "lora_alpha": native_config["lora_alpha"]}
|
|
else:
|
|
config_file = directory / "adapter_config.json"
|
|
config = json.loads(config_file.read_text()) if config_file.is_file() else {}
|
|
state = load_checkpoint(path)
|
|
if not state:
|
|
raise ValueError(f"Adapter checkpoint is empty: {path}")
|
|
peft_suffixes = (".lora_A.weight", ".lora_B.weight", ".lora_A.default.weight", ".lora_B.default.weight")
|
|
has_peft = any(name.endswith(peft_suffixes) for name in state)
|
|
if has_peft and native_config is None and not (directory / "adapter_config.json").is_file():
|
|
for name in state:
|
|
for suffix in peft_suffixes:
|
|
if name.endswith(suffix) and name.removesuffix(suffix) + ".alpha" not in state:
|
|
raise ValueError("PEFT adapters require adapter_config.json or per-target alpha tensors")
|
|
if not has_peft and native_config is None:
|
|
config = {}
|
|
return state, config, directory, native_config
|
|
|
|
|
|
def pattern_value(config, field, name, default):
|
|
matches = [value for pattern, value in config.get(field, {}).items()
|
|
if re.fullmatch(r"(.*\.)?(" + pattern + ")", name)]
|
|
if matches and any(value != matches[0] for value in matches):
|
|
raise ValueError(f"Ambiguous {field} for {name}")
|
|
return matches[0] if matches else default
|
|
|
|
|
|
def normalize(state, config, bindings, output, scale, rs_config, native_scaling=False, require_scaling=True):
|
|
groups = {}
|
|
peft_targets = set()
|
|
replacements = {}
|
|
for name, tensor in state.items():
|
|
for suffix, part in SUFFIXES:
|
|
if name.endswith(suffix):
|
|
target = name.removesuffix(suffix)
|
|
group = groups.setdefault(target, {})
|
|
if part in group:
|
|
raise ValueError(f"Duplicate adapter factor: {name}")
|
|
group[part] = tensor
|
|
if suffix in {".lora_A.weight", ".lora_B.weight", ".lora_A.default.weight", ".lora_B.default.weight"}:
|
|
peft_targets.add(target)
|
|
break
|
|
else:
|
|
if not name.endswith((".weight", ".bias")):
|
|
raise ValueError(f"Unsupported adapter tensor: {name}")
|
|
targets = bindings.resolve(name)
|
|
if len(targets) != 1 or targets[0].rows:
|
|
raise ValueError(f"Cannot slice a full replacement: {name}")
|
|
if targets[0].name in replacements:
|
|
raise ValueError(f"Multiple replacements target {targets[0].name}")
|
|
replacements[targets[0].name] = tensor.contiguous()
|
|
for name, factors in groups.items():
|
|
targets = bindings.resolve(name)
|
|
group_config = config if name in peft_targets or native_scaling else {}
|
|
if "a" in factors or "b" in factors:
|
|
if factors.keys() - {"a", "b", "alpha"}:
|
|
raise ValueError(f"Mixed adapter algorithms: {name}")
|
|
a, b = factors["a"], factors["b"]
|
|
if a.ndim != 2 or b.ndim != 2 or a.shape[0] != b.shape[1]:
|
|
raise ValueError(f"Invalid A/B shapes: {name}")
|
|
rank = a.shape[0]
|
|
if require_scaling and not group_config and "alpha" not in factors:
|
|
raise ValueError(f"Adapter has no scaling metadata: {name}")
|
|
if pattern_value(group_config, "rank_pattern", name.removeprefix("base_model.model."), group_config.get("r", rank)) != rank:
|
|
raise ValueError(f"Configured rank differs from factors: {name}")
|
|
coefficient = pattern_value(group_config, "alpha_pattern", name.removeprefix("base_model.model."),
|
|
group_config.get("lora_alpha", rank))
|
|
coefficient = float(factors.get("alpha", torch.tensor(coefficient)).item())
|
|
next_row = 0
|
|
for target in sorted(targets, key=lambda target: target.first_row):
|
|
rows = target.rows or b.shape[0]
|
|
if target.first_row != next_row or rows <= 0 or target.first_row + rows > b.shape[0]:
|
|
raise ValueError(f"Invalid adapter output row slices: {name}")
|
|
next_row += rows
|
|
if next_row != b.shape[0]:
|
|
raise ValueError(f"Adapter output rows were not fully bound: {name}")
|
|
for target in targets:
|
|
rows = target.rows or b.shape[0]
|
|
prefix = target.name
|
|
if any(key == prefix or key.startswith(prefix + ".") for key in output):
|
|
raise ValueError(f"Multiple updates target {prefix}")
|
|
output[prefix + ".lora_A"] = a.contiguous().clone()
|
|
output[prefix + ".lora_B"] = b[target.first_row:target.first_row + rows].contiguous().clone()
|
|
effective_alpha = coefficient * scale
|
|
if group_config.get("use_rslora", False):
|
|
# Keep rsLoRA scaling in the framework importer, not in the factors.
|
|
output[prefix + ".lora_A.weight"] = output.pop(prefix + ".lora_A")
|
|
output[prefix + ".lora_B.weight"] = output.pop(prefix + ".lora_B")
|
|
rs_config["rank_pattern"][re.escape(prefix)] = rank
|
|
rs_config["alpha_pattern"][re.escape(prefix)] = effective_alpha
|
|
else:
|
|
output[prefix + ".alpha"] = torch.tensor(effective_alpha, dtype=torch.float32)
|
|
else:
|
|
if len(targets) != 1 or targets[0].rows:
|
|
raise ValueError(f"Cannot slice LoKr factors: {name}")
|
|
prefix = targets[0].name
|
|
if any(key == prefix or key.startswith(prefix + ".") for key in output):
|
|
raise ValueError(f"Multiple updates target {prefix}")
|
|
for part, tensor in factors.items():
|
|
if part == "alpha":
|
|
output[prefix + ".alpha"] = tensor.float() * scale
|
|
else:
|
|
output[prefix + "." + part] = tensor.contiguous()
|
|
if scale != 1:
|
|
raise ValueError("Non-unit LoKr export strength is not supported yet")
|
|
for name, tensor in replacements.items():
|
|
if name in output:
|
|
raise ValueError(f"Duplicate replacement: {name}")
|
|
output[name] = tensor
|
|
|
|
|
|
def sidecars(family, directory, native_config, base, bindings, output, work):
|
|
root = None
|
|
speaker_file = find_checkpoint(directory, "speaker_embedding", required=False) if family == "qwen3_tts" else None
|
|
if speaker_file is not None:
|
|
config_file = directory / "tts_config.json"
|
|
config = json.loads(config_file.read_text())
|
|
talker = config["talker_config"]
|
|
if config["tts_model_type"] != "custom_voice" or len(talker["spk_id"]) != 1:
|
|
raise ValueError("Speaker sidecar requires a single custom_voice speaker")
|
|
target = bindings.resolve("talker.model.codec_embedding.weight")[0].name
|
|
if any(key == target or key.startswith(target + ".") for key in output):
|
|
raise ValueError("Speaker sidecar conflicts with codec embedding updates")
|
|
shape = base.shapes[target]
|
|
if shape != (talker["vocab_size"], talker["hidden_size"]):
|
|
raise ValueError("Speaker sidecar configuration differs from the base checkpoint")
|
|
speaker_id = next(iter(talker["spk_id"].values()))
|
|
if not 0 <= speaker_id < shape[0]:
|
|
raise ValueError("Speaker ID is outside the codec vocabulary")
|
|
speaker = load_checkpoint(speaker_file)["speaker_embedding"]
|
|
if tuple(speaker.shape) != (shape[1],):
|
|
raise ValueError("Speaker embedding width differs from the base checkpoint")
|
|
embedding = base.read_f32(target)
|
|
embedding[speaker_id] = speaker.float()
|
|
output[target] = embedding
|
|
root = work / "model"
|
|
root.mkdir()
|
|
base.extract_sidecars(root)
|
|
(root / "config.json").write_text(json.dumps(config, ensure_ascii=False), encoding="utf-8")
|
|
if family == "cosyvoice3" and native_config is not None:
|
|
for module in ("speech_embedding", "llm_decoder"):
|
|
state = load_checkpoint(find_checkpoint(directory, module))
|
|
for name, tensor in state.items():
|
|
target = bindings.resolve(module + "." + name)[0].name
|
|
if target in output:
|
|
raise ValueError(f"Duplicate sidecar replacement: {target}")
|
|
output[target] = tensor.contiguous()
|
|
if family == "vibevoice":
|
|
candidates = (
|
|
("diffusion_head/model", "model.prediction_head."),
|
|
("diffusion_head_full", "model.prediction_head."),
|
|
("diffusion_head/diffusion_head_full", "model.prediction_head."),
|
|
("acoustic_connector/pytorch_model", "model.acoustic_connector."),
|
|
("semantic_connector/pytorch_model", "model.semantic_connector."),
|
|
)
|
|
loaded = set()
|
|
for filename, prefix in candidates:
|
|
if prefix in loaded:
|
|
continue
|
|
logical = directory / filename
|
|
path = find_checkpoint(logical.parent, logical.name, required=False)
|
|
if path is None:
|
|
continue
|
|
for name, tensor in load_checkpoint(path).items():
|
|
target = bindings.resolve(prefix + name)[0].name
|
|
output[target] = tensor.contiguous()
|
|
loaded.add(prefix)
|
|
return root
|
|
|
|
|
|
def main():
|
|
parser = argparse.ArgumentParser(description=__doc__)
|
|
parser.add_argument("--family", choices=FAMILIES, required=True)
|
|
base_options = parser.add_mutually_exclusive_group(required=True)
|
|
base_options.add_argument("--base", type=Path, help="Unmerged base GGUF, preferably original dtype or 16-bit")
|
|
base_options.add_argument("--input", action="append", help="Original converter-compatible [namespace=]safetensors or shard index; repeat for components")
|
|
parser.add_argument("--root", type=Path, help="Original checkpoint directory containing config/tokenizer files")
|
|
parser.add_argument("--adapter", type=Path, help="Adapter file or directory; AR adapter for YuE2")
|
|
parser.add_argument("--nar-adapter", type=Path, help="Additional YuE2 NAR adapter")
|
|
parser.add_argument("--scale", type=float, default=1.0, help="Adapter strength multiplier")
|
|
parser.add_argument("--nar-scale", type=float, default=1.0)
|
|
parser.add_argument("--alpha", type=float, help="Training alpha for a Fish Lightning checkpoint")
|
|
parser.add_argument("--component", help="GGUF component namespace, e.g. dit_turbo_weights for ACE-Step")
|
|
parser.add_argument("--output", type=Path, required=True)
|
|
parser.add_argument("--type", default="orig", choices=("orig", "f32", "f16", "bf16", "q8_0", "q4_0", "q2_k", "q3_k", "q4_k", "q5_k", "q6_k"))
|
|
parser.add_argument("--keep-type", action="append", default=[], help="Converter tensor-pattern dtype override")
|
|
parser.add_argument("--converter", type=Path, default=Path(__file__).resolve().parents[2] / "build/debug/bin/audiocpp_gguf")
|
|
parser.add_argument("--model-spec", type=Path, help="Override the repository's current model spec")
|
|
parser.add_argument("--overwrite", action="store_true")
|
|
args = parser.parse_args()
|
|
if args.base and args.base.resolve() == args.output.resolve():
|
|
parser.error("--output must not replace the base GGUF")
|
|
for source in args.input or []:
|
|
path = Path(source.partition("=")[2] if "=" in source else source)
|
|
if path.resolve() == args.output.resolve():
|
|
parser.error("--output must not replace an original weight source")
|
|
if args.output.exists() and not args.overwrite:
|
|
parser.error("Output exists; pass --overwrite to replace it")
|
|
if args.nar_adapter and args.family != "yue2":
|
|
parser.error("--nar-adapter is only valid for YuE2")
|
|
if args.adapter is None and args.nar_adapter is None:
|
|
parser.error("Provide --adapter, or --nar-adapter for a NAR-only YuE2 merge")
|
|
if args.alpha is not None and args.family != "fish_audio":
|
|
parser.error("--alpha is only used to prepare Fish Speech Lightning checkpoints")
|
|
if not torch.isfinite(torch.tensor([args.scale, args.nar_scale])).all():
|
|
parser.error("Adapter strengths must be finite")
|
|
if args.base:
|
|
base = BaseGGUF(args.base)
|
|
inputs = [str(args.base)]
|
|
else:
|
|
inputs = args.input
|
|
first = Path(inputs[0].partition("=")[2] if "=" in inputs[0] else inputs[0])
|
|
base = BaseSources(inputs, args.root or first.parent)
|
|
if any(path.resolve() == args.output.resolve() for path, _ in base.tensors.values()):
|
|
parser.error("--output must not replace an original weight shard")
|
|
bindings = Bindings(args.family, base.shapes, args.component)
|
|
with tempfile.TemporaryDirectory(prefix="audiocpp-lora-") as temporary:
|
|
work = Path(temporary)
|
|
output = {}
|
|
rs_config = {"r": 1, "lora_alpha": 1, "use_rslora": True, "rank_pattern": {}, "alpha_pattern": {}}
|
|
root = args.root
|
|
for adapter, strength in ((args.adapter, args.scale), (args.nar_adapter, args.nar_scale)):
|
|
if adapter is None or (strength == 0 and args.family in {"yue2", "ace_step"}):
|
|
continue
|
|
state, config, directory, native = read_adapter(adapter, args.family, work, args.alpha)
|
|
for option in ("use_dora", "use_qalora", "fan_in_fan_out", "lora_bias"):
|
|
if config.get(option, False):
|
|
raise ValueError(f"Unsupported adapter option: {option}")
|
|
normalize(state, config, bindings, output, strength if args.family == "yue2" else 1.0,
|
|
rs_config, native_scaling=native is not None, require_scaling=args.family not in {"yue2", "ace_step"})
|
|
model_root = sidecars(args.family, directory, native, base, bindings, output, work)
|
|
if model_root is not None:
|
|
root = model_root
|
|
command = [str(args.converter), "--output", str(args.output),
|
|
"--type", args.type, "--family", args.family, "--model-spec",
|
|
str(args.model_spec or Path(__file__).resolve().parents[2] / "model_specs" / (args.family + ".json")), "--log"]
|
|
for value in inputs:
|
|
command.extend(("--input", value))
|
|
if output:
|
|
normalized = work / "adapter.safetensors"
|
|
save_file(output, normalized)
|
|
if rs_config["rank_pattern"]:
|
|
(work / "adapter_config.json").write_text(json.dumps(rs_config))
|
|
command.extend(("--weight-adapter", str(normalized), "--adapter-merge-mode",
|
|
MERGE_MODES.get(args.family, "delta_f32")))
|
|
if args.family != "yue2":
|
|
command.extend(("--adapter-strength", str(args.scale)))
|
|
if root is not None:
|
|
command.extend(("--root", str(root)))
|
|
for pattern in args.keep_type:
|
|
command.extend(("--keep-type", pattern))
|
|
if args.overwrite:
|
|
command.append("--overwrite")
|
|
subprocess.run(command, check=True)
|
|
print(f"adapter_tensors={len(output)}")
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|