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
77 lines
3.2 KiB
Python
77 lines
3.2 KiB
Python
#!/usr/bin/env python3
|
|
"""Export Fish Speech Lightning LoRA factors as runtime-ready safetensors."""
|
|
|
|
import argparse
|
|
import collections
|
|
from pathlib import Path
|
|
import typing
|
|
|
|
import torch
|
|
from safetensors.torch import save_file
|
|
|
|
from checkpoint import load_checkpoint
|
|
|
|
|
|
def prepare(input_path: Path, output_path: Path, alpha: float) -> None:
|
|
# Lightning stores scheduler settings as OmegaConf objects, not just tensors.
|
|
from omegaconf import ListConfig
|
|
from omegaconf.base import ContainerMetadata, Metadata
|
|
from omegaconf.nodes import AnyNode
|
|
|
|
allowed = [
|
|
dict, list, int, collections.defaultdict, typing.Any,
|
|
ListConfig, ContainerMetadata, Metadata, AnyNode,
|
|
]
|
|
with torch.serialization.safe_globals(allowed):
|
|
state = load_checkpoint(input_path)
|
|
if (input_path.parent / "adapter_config.json").is_file() or any(name.endswith(".alpha") for name in state):
|
|
# Already-prepared adapters retain their own scaling metadata.
|
|
output_path.parent.mkdir(parents=True, exist_ok=True)
|
|
save_file(state, output_path, metadata={"format": "pt", "family": "fish_audio"})
|
|
return
|
|
tensors = {}
|
|
targets = set()
|
|
for name, value in state.items():
|
|
if not isinstance(value, torch.Tensor) or not name.startswith("model."):
|
|
raise ValueError(f"Unexpected Fish adapter entry: {name}")
|
|
name = name.removeprefix("model.")
|
|
if not name.endswith((".lora_A", ".lora_B")):
|
|
raise ValueError(f"Expected LoRA factor, got {name}")
|
|
target, factor = name.rsplit(".", 1)
|
|
if value.ndim != 2:
|
|
raise ValueError(f"Expected matrix factor: {name}")
|
|
if target in {"embeddings", "codebook_embeddings", "fast_embeddings"}:
|
|
# loralib.Embedding adds (B @ A).T to its [vocabulary, width] weight.
|
|
factor = "lora_B" if factor == "lora_A" else "lora_A"
|
|
value = value.T
|
|
tensors[f"{target}.{factor}"] = value.contiguous()
|
|
targets.add(target)
|
|
if not targets:
|
|
raise ValueError("Checkpoint contains no LoRA targets")
|
|
for target in targets:
|
|
a = tensors[f"{target}.lora_A"]
|
|
b = tensors[f"{target}.lora_B"]
|
|
if a.shape[0] != b.shape[1] or a.shape[0] <= 0:
|
|
raise ValueError(f"Inconsistent LoRA rank: {target}")
|
|
tensors[f"{target}.alpha"] = torch.tensor(alpha, dtype=torch.float32)
|
|
output_path.parent.mkdir(parents=True, exist_ok=True)
|
|
save_file(tensors, output_path, metadata={"format": "pt", "family": "fish_audio"})
|
|
print(f"adapter={output_path}")
|
|
print(f"targets={len(targets)}")
|
|
|
|
|
|
def main() -> None:
|
|
parser = argparse.ArgumentParser(description=__doc__)
|
|
parser.add_argument("--input", type=Path, required=True, help="Fish Speech Lightning .ckpt")
|
|
parser.add_argument("--output", type=Path, required=True, help="Output adapter .safetensors")
|
|
parser.add_argument("--alpha", type=float, required=True, help="Training LoRA alpha, not alpha/r")
|
|
args = parser.parse_args()
|
|
if args.output.exists():
|
|
parser.error(f"Output already exists: {args.output}")
|
|
if not torch.isfinite(torch.tensor(args.alpha)):
|
|
parser.error("--alpha must be finite")
|
|
prepare(args.input, args.output, args.alpha)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|