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
257 lines
12 KiB
Python
257 lines
12 KiB
Python
"""Model-local adapter names bound to complete audio.cpp GGUF tensor names."""
|
|
|
|
from dataclasses import dataclass
|
|
import json
|
|
from pathlib import Path
|
|
import re
|
|
import shutil
|
|
|
|
import gguf
|
|
import numpy as np
|
|
import torch
|
|
from safetensors import safe_open
|
|
|
|
from checkpoint import EXTENSIONS
|
|
|
|
|
|
FAMILIES = (
|
|
"qwen3_tts", "qwen3_asr", "voxcpm2", "voxcpm1", "fish_audio",
|
|
"cosyvoice3", "f5_tts", "outetts", "yue2", "ace_step", "vibevoice", "audio_flamingo",
|
|
)
|
|
|
|
MERGE_MODES = {"yue2": "rounded_bf16_delta", "vibevoice": "rankwise_f32"}
|
|
|
|
COMPONENTS = {
|
|
"qwen3_tts": "model_weights", "voxcpm2": "weights",
|
|
"fish_audio": "model_weights", "cosyvoice3": "llm",
|
|
"f5_tts": "transformer", "outetts": "model_weights", "audio_flamingo": "weights",
|
|
}
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
class Target:
|
|
name: str
|
|
first_row: int = 0
|
|
rows: int = 0
|
|
|
|
|
|
class Bindings:
|
|
def __init__(self, family, shapes, component=None):
|
|
self.targets = {}
|
|
component = component or COMPONENTS.get(family)
|
|
names = list(shapes)
|
|
if component:
|
|
names = [name for name in names if name.startswith(component + "/")]
|
|
if not names:
|
|
raise ValueError(f"Base GGUF has no {component!r} component")
|
|
for name in names:
|
|
logical = name.split("/", 1)[-1]
|
|
target = (Target(name),)
|
|
self.add(name, target)
|
|
self.add(logical, target)
|
|
if family == "qwen3_tts" and logical.startswith("talker."):
|
|
self.add(logical.removeprefix("talker."), target)
|
|
elif family == "qwen3_asr":
|
|
if logical.startswith("thinker."):
|
|
self.add(logical.removeprefix("thinker."), target)
|
|
elif logical.startswith("model.language_model."):
|
|
self.add("model." + logical.removeprefix("model.language_model."), target)
|
|
elif logical.startswith("model.audio_tower."):
|
|
self.add(logical.removeprefix("model."), target)
|
|
elif family == "cosyvoice3" and logical.startswith("llm.model."):
|
|
self.add(logical.removeprefix("llm.model."), target)
|
|
elif family == "vibevoice" and logical.startswith("model.language_model."):
|
|
self.add(logical.removeprefix("model.language_model."), target)
|
|
elif family == "yue2" and logical.startswith("model."):
|
|
self.add(logical.removeprefix("model."), target)
|
|
elif family == "ace_step" and logical.startswith("decoder."):
|
|
self.add(logical.removeprefix("decoder."), target)
|
|
elif family == "fish_audio":
|
|
self.add("model." + logical, target)
|
|
elif family == "f5_tts":
|
|
stripped = logical.removeprefix("ema_model.transformer.")
|
|
self.add(stripped, target)
|
|
self.add("transformer." + stripped, target)
|
|
elif family == "voxcpm1":
|
|
self.add_voxcpm1(logical, target)
|
|
|
|
if family in {"fish_audio", "yue2"}:
|
|
for name in names:
|
|
if name.endswith(".q_proj.weight"):
|
|
prefix = name.removesuffix(".q_proj.weight")
|
|
projections = [prefix + "." + part + ".weight" for part in ("q_proj", "k_proj", "v_proj")]
|
|
fused = ".wqkv" if family == "fish_audio" else ".qkv_proj"
|
|
elif family == "yue2" and name.endswith(".gate_proj.weight"):
|
|
prefix = name.removesuffix(".gate_proj.weight")
|
|
projections = [prefix + "." + part + ".weight" for part in ("gate_proj", "up_proj")]
|
|
fused = ".gate_up_proj"
|
|
else:
|
|
continue
|
|
offset = 0
|
|
targets = []
|
|
for projection in projections:
|
|
rows = shapes[projection][0]
|
|
targets.append(Target(projection, offset, rows))
|
|
offset += rows
|
|
logical = (prefix + fused).split("/", 1)[-1]
|
|
self.add(logical, tuple(targets))
|
|
if family == "fish_audio":
|
|
self.add("model." + logical, tuple(targets))
|
|
else:
|
|
self.add(logical.removeprefix("model."), tuple(targets))
|
|
|
|
def add(self, name, targets):
|
|
name = name.removesuffix(".weight")
|
|
for alias in {name, name.replace(".", "_")}:
|
|
if alias in self.targets and self.targets[alias] != targets:
|
|
raise ValueError(f"Ambiguous adapter binding {alias!r}; select --component")
|
|
self.targets[alias] = targets
|
|
|
|
def resolve(self, name):
|
|
for prefix in ("base_model.model.", "text_encoders.", "lora_unet_", "lora_transformer_", "lycoris_"):
|
|
if name.startswith(prefix):
|
|
name = name.removeprefix(prefix)
|
|
break
|
|
name = name.replace(".modules_to_save.default", "")
|
|
key = name.removesuffix(".weight")
|
|
if key not in self.targets:
|
|
raise ValueError(f"Adapter target has no model binding: {name}")
|
|
return self.targets[key]
|
|
|
|
def add_voxcpm1(self, name, target):
|
|
match = re.fullmatch(
|
|
r"(?:(residual_lm|locenc|locdit)\.)?blk\.(\d+)\."
|
|
r"(attn_q|attn_k|attn_v|attn_output|attn_norm|ffn_norm|ffn_gate|ffn_up|ffn_down)\.weight", name)
|
|
if match:
|
|
prefix = {None: "base_lm", "residual_lm": "residual_lm", "locenc": "feat_encoder.encoder",
|
|
"locdit": "feat_decoder.estimator.decoder"}[match[1]]
|
|
projection = {
|
|
"attn_q": "self_attn.q_proj", "attn_k": "self_attn.k_proj", "attn_v": "self_attn.v_proj",
|
|
"attn_output": "self_attn.o_proj", "attn_norm": "input_layernorm",
|
|
"ffn_norm": "post_attention_layernorm", "ffn_gate": "mlp.gate_proj",
|
|
"ffn_up": "mlp.up_proj", "ffn_down": "mlp.down_proj",
|
|
}[match[3]]
|
|
self.add(f"{prefix}.layers.{match[2]}.{projection}", target)
|
|
return
|
|
replacements = {
|
|
"token_embd": "base_lm.embed_tokens", "output_norm": "base_lm.norm",
|
|
"residual_lm.output_norm": "residual_lm.norm", "locenc.output_norm": "feat_encoder.encoder.norm",
|
|
"locdit.output_norm": "feat_decoder.estimator.decoder.norm",
|
|
"proj.enc_to_lm": "enc_to_lm_proj", "proj.lm_to_dit": "lm_to_dit_proj",
|
|
"proj.res_to_dit": "res_to_dit_proj", "stop.stop_proj": "stop_proj", "stop.stop_head": "stop_head",
|
|
}
|
|
for suffix in (".weight", ".bias"):
|
|
if not name.endswith(suffix):
|
|
continue
|
|
module = name.removesuffix(suffix)
|
|
if module in replacements:
|
|
self.add(replacements[module] + suffix, target)
|
|
for internal, external in (("locenc.", "feat_encoder."), ("locdit.", "feat_decoder.estimator."),
|
|
("fsq.", "fsq_layer.")):
|
|
if module.startswith(internal):
|
|
self.add(external + module.removeprefix(internal) + suffix, target)
|
|
|
|
|
|
class GGUFReader(gguf.GGUFReader):
|
|
def _get_field_parts(self, offset, raw_type):
|
|
# Embedded sidecars are byte arrays; one NumPy object per byte is costly.
|
|
if raw_type == gguf.GGUFValueType.ARRAY:
|
|
item_type = self._get(offset, np.uint32)
|
|
if item_type[0] == gguf.GGUFValueType.UINT8:
|
|
count = self._get(offset + 4, np.uint64)
|
|
values = self._get(offset + 12, np.uint8, int(count[0]))
|
|
return 12 + values.nbytes, [item_type, count, values], [2], [
|
|
gguf.GGUFValueType.ARRAY, gguf.GGUFValueType.UINT8]
|
|
return super()._get_field_parts(offset, raw_type)
|
|
|
|
|
|
class BaseGGUF:
|
|
def __init__(self, path):
|
|
self.reader = GGUFReader(path)
|
|
fields = self.reader.fields
|
|
if "audiocpp.tensor_names" in fields:
|
|
names = fields["audiocpp.tensor_names"].contents()
|
|
ranks = fields["audiocpp.tensor_ranks"].contents()
|
|
dimensions = fields["audiocpp.tensor_shapes"].contents()
|
|
self.shapes = {}
|
|
offset = 0
|
|
for name, rank in zip(names, ranks, strict=True):
|
|
self.shapes[name] = tuple(dimensions[offset:offset + rank])
|
|
offset += rank
|
|
if offset != len(dimensions):
|
|
raise ValueError("Base GGUF logical shape metadata is inconsistent")
|
|
else:
|
|
names = [tensor.name for tensor in self.reader.tensors]
|
|
self.shapes = {tensor.name: tuple(reversed(tensor.shape.tolist())) for tensor in self.reader.tensors}
|
|
self.tensors = dict(zip(names, self.reader.tensors, strict=True))
|
|
|
|
def read_f32(self, name):
|
|
tensor = self.tensors[name]
|
|
values = gguf.quants.dequantize(tensor.data, tensor.tensor_type)
|
|
return torch.from_numpy(np.asarray(values, dtype=np.float32).reshape(self.shapes[name]).copy())
|
|
|
|
def extract_sidecars(self, destination):
|
|
fields = self.reader.fields
|
|
names = fields["audiocpp.embedded_files.names"].contents()
|
|
if "audiocpp.embedded_files.data" in fields:
|
|
field = fields["audiocpp.embedded_files.data"]
|
|
first = field.parts[field.data[0]]
|
|
offset = first.__array_interface__["data"][0] - self.reader.data.__array_interface__["data"][0]
|
|
data = self.reader.data[offset:offset + int(field.parts[4][0])]
|
|
offsets = fields["audiocpp.embedded_files.offsets"].contents()
|
|
contents = (memoryview(data)[offsets[i]:offsets[i + 1]] for i in range(len(names)))
|
|
else:
|
|
contents = (value.encode("utf-8") for value in fields["audiocpp.embedded_files.contents"].contents())
|
|
for name, content in zip(names, contents, strict=True):
|
|
target = destination / name
|
|
if not target.resolve().is_relative_to(destination.resolve()):
|
|
raise ValueError(f"Unsafe embedded file name: {name}")
|
|
target.parent.mkdir(parents=True, exist_ok=True)
|
|
target.write_bytes(content)
|
|
|
|
|
|
class BaseSources:
|
|
def __init__(self, inputs, root):
|
|
self.root = root
|
|
self.shapes = {}
|
|
self.tensors = {}
|
|
for value in inputs:
|
|
prefix, separator, filename = value.partition("=")
|
|
path = Path(filename if separator else prefix)
|
|
prefix = prefix if separator else ""
|
|
if path.name.endswith(".safetensors.index.json"):
|
|
mapping = json.loads(path.read_text())["weight_map"]
|
|
shards = {name: path.parent / shard for name, shard in mapping.items()}
|
|
elif path.suffix == ".safetensors":
|
|
with safe_open(path, framework="pt", device="cpu") as source:
|
|
shards = {name: path for name in source.keys()}
|
|
else:
|
|
raise ValueError(f"Tensor inputs must be safetensors or a shard index: {path}")
|
|
for shard in sorted(set(shards.values())):
|
|
with safe_open(shard, framework="pt", device="cpu") as source:
|
|
for name, file in shards.items():
|
|
if file != shard:
|
|
continue
|
|
logical = prefix + "/" + name if prefix else name
|
|
if logical in self.shapes:
|
|
raise ValueError(f"Duplicate input tensor: {logical}")
|
|
self.shapes[logical] = tuple(source.get_slice(name).get_shape())
|
|
self.tensors[logical] = (shard, name)
|
|
|
|
def read_f32(self, name):
|
|
path, key = self.tensors[name]
|
|
with safe_open(path, framework="pt", device="cpu") as source:
|
|
return source.get_tensor(key).float().clone()
|
|
|
|
def extract_sidecars(self, destination):
|
|
for path in self.root.rglob("*"):
|
|
relative = path.relative_to(self.root)
|
|
if any(part in {".git", ".cache"} for part in relative.parts):
|
|
continue
|
|
if not path.is_file() or path.suffix.lower() in (*EXTENSIONS, ".gguf"):
|
|
continue
|
|
if path.stat().st_size > 64 * 1024 * 1024:
|
|
continue
|
|
target = destination / relative
|
|
target.parent.mkdir(parents=True, exist_ok=True)
|
|
shutil.copyfile(path, target)
|