From 7ea01875c76c4bfcf275a75038365abfb23eaa64 Mon Sep 17 00:00:00 2001 From: Alok-Ranjan23 Date: Fri, 24 Jul 2026 12:29:06 +0000 Subject: [PATCH 1/3] feat: add isolated LAM-A2E C API, GGUF tools, and CPU parity tests Expose lam_a2e_* beside stable-diffusion, convert/remap scripts for lam-audio2exp GGUF, and smoke/parity gates vs lipsync-ggml fixtures. --- CMakeLists.txt | 21 +- include/lam-a2e.h | 47 ++ scripts/convert_lam_a2e_to_gguf.py | 207 +++++++++ scripts/dump_lam_a2e_frontend_reference.py | 65 +++ scripts/dump_lam_a2e_stages.py | 138 ++++++ scripts/remap_lam_gguf.py | 174 +++++++ src/lam-a2e.cpp | 120 +++++ src/lam_audio2expression.cpp | 510 +++++++++++++++++++++ src/lam_audio2expression.hpp | 70 +++ src/lam_wav2vec_frontend.hpp | 348 ++++++++++++++ tests/lam_a2e_frontend_smoke.cpp | 90 ++++ tests/lam_a2e_parity.cpp | 111 +++++ tests/lam_a2e_smoke.cpp | 51 +++ 13 files changed, 1951 insertions(+), 1 deletion(-) create mode 100644 include/lam-a2e.h create mode 100644 scripts/convert_lam_a2e_to_gguf.py create mode 100644 scripts/dump_lam_a2e_frontend_reference.py create mode 100644 scripts/dump_lam_a2e_stages.py create mode 100644 scripts/remap_lam_gguf.py create mode 100644 src/lam-a2e.cpp create mode 100644 src/lam_audio2expression.cpp create mode 100644 src/lam_audio2expression.hpp create mode 100644 src/lam_wav2vec_frontend.hpp create mode 100644 tests/lam_a2e_frontend_smoke.cpp create mode 100644 tests/lam_a2e_parity.cpp create mode 100644 tests/lam_a2e_smoke.cpp diff --git a/CMakeLists.txt b/CMakeLists.txt index 242b84292..22e0c257f 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -273,6 +273,8 @@ else() add_library(${SD_LIB} STATIC ${SD_LIB_SOURCES}) endif() +option(SD_BUILD_LAM_A2E_SMOKE "Build the LAM-A2E GGUF loader smoke executable" OFF) + if(APPLE) sd_set_macos_rpaths(${SD_LIB}) endif() @@ -335,12 +337,29 @@ target_include_directories(${SD_LIB} PUBLIC . src include) target_include_directories(${SD_LIB} PUBLIC . thirdparty) target_compile_features(${SD_LIB} PUBLIC c_std_11 cxx_std_17) +if(SD_BUILD_LAM_A2E_SMOKE) + add_executable(lam-a2e-smoke tests/lam_a2e_smoke.cpp) + target_include_directories(lam-a2e-smoke PRIVATE include) + target_link_libraries(lam-a2e-smoke PRIVATE ${SD_LIB}) + target_compile_features(lam-a2e-smoke PRIVATE cxx_std_17) + + add_executable(lam-a2e-frontend-smoke tests/lam_a2e_frontend_smoke.cpp) + target_include_directories(lam-a2e-frontend-smoke PRIVATE include src) + target_link_libraries(lam-a2e-frontend-smoke PRIVATE ${SD_LIB}) + target_compile_features(lam-a2e-frontend-smoke PRIVATE cxx_std_17) + + add_executable(lam-a2e-parity tests/lam_a2e_parity.cpp) + target_include_directories(lam-a2e-parity PRIVATE include src) + target_link_libraries(lam-a2e-parity PRIVATE ${SD_LIB}) + target_compile_features(lam-a2e-parity PRIVATE cxx_std_17) +endif() + if (SD_BUILD_EXAMPLES) add_subdirectory(examples) endif() -set(SD_PUBLIC_HEADERS include/stable-diffusion.h) +set(SD_PUBLIC_HEADERS include/stable-diffusion.h include/lam-a2e.h) set_target_properties(${SD_LIB} PROPERTIES PUBLIC_HEADER "${SD_PUBLIC_HEADERS}") install(TARGETS ${SD_LIB} LIBRARY PUBLIC_HEADER) diff --git a/include/lam-a2e.h b/include/lam-a2e.h new file mode 100644 index 000000000..08ca6bb1b --- /dev/null +++ b/include/lam-a2e.h @@ -0,0 +1,47 @@ +#pragma once + +#include +#include + +#ifdef __cplusplus +extern "C" { +#endif + +typedef struct lam_a2e_context lam_a2e_context; + +typedef enum lam_a2e_status { + LAM_A2E_STATUS_OK = 0, + LAM_A2E_STATUS_INVALID_ARGUMENT = 1, + LAM_A2E_STATUS_MODEL_LOAD_FAILED = 2, + LAM_A2E_STATUS_NOT_IMPLEMENTED = 3, +} lam_a2e_status; + +typedef struct lam_a2e_params { + const char * model_path; + int32_t identity_index; + int32_t n_threads; + bool use_gpu; +} lam_a2e_params; + +typedef struct lam_a2e_frame { + int64_t timestamp_us; + float arkit_52[52]; +} lam_a2e_frame; + +lam_a2e_context * lam_a2e_create(const lam_a2e_params * params); +void lam_a2e_free(lam_a2e_context * ctx); + +lam_a2e_status lam_a2e_process_pcm_f32( + lam_a2e_context * ctx, + const float * pcm, + int64_t pcm_sample_count, + int32_t sample_rate, + lam_a2e_frame ** frames, + int32_t * frame_count); + +void lam_a2e_free_frames(lam_a2e_frame * frames); +const char * lam_a2e_get_last_error(const lam_a2e_context * ctx); + +#ifdef __cplusplus +} +#endif diff --git a/scripts/convert_lam_a2e_to_gguf.py b/scripts/convert_lam_a2e_to_gguf.py new file mode 100644 index 000000000..7102c4c3e --- /dev/null +++ b/scripts/convert_lam_a2e_to_gguf.py @@ -0,0 +1,207 @@ +#!/usr/bin/env python3 +"""LAM Audio2Expression PyTorch checkpoint → GGUF converter. + +Reads the upstream ``lam_audio2exp_streaming.tar`` checkpoint (Apache-2.0, +https://github.com/aigc3d/LAM_Audio2Expression) and emits a single GGUF +consumed by the isolated LAM-A2E target in qvac-ext-stable-diffusion.cpp. + +Layout notes: + +- ``general.architecture = "lam-audio2exp"`` — selected by the addon's + model factory. +- PyTorch conv1d weights are (C_out, C_in, K) row-major, which lands as + ggml ne = [K, C_in, C_out] on a straight copy — exactly what + ``ggml_conv_1d`` expects. No transposes are performed anywhere; the + converter is a 1:1 map of the state dict. +- The positional-conv weight-norm (weight_g/weight_v, dim=2) is folded + into a plain conv weight at conversion time. +- Inference-unused tensors (``lm_head``, ``identity_encoder.grus``, + ``masked_spec_embed``) are dropped. +- ``--dtype f32`` (default) keeps everything float32; ``--dtype f16`` + stores matmul/conv weights as f16 and keeps norms + biases f32. + +Usage: + python3 scripts/convert_lam_a2e_to_gguf.py \ + --checkpoint pretrained_models/lam_audio2exp_streaming.tar \ + --out lam-audio2exp-f32.gguf --dtype f32 +""" + +from __future__ import annotations + +import argparse +import os + +import numpy as np +import torch + +from gguf import GGUFWriter + +ARCH = "lam-audio2exp" + +ARKIT_BLENDSHAPES = [ + "browDownLeft", "browDownRight", "browInnerUp", "browOuterUpLeft", + "browOuterUpRight", "cheekPuff", "cheekSquintLeft", "cheekSquintRight", + "eyeBlinkLeft", "eyeBlinkRight", "eyeLookDownLeft", "eyeLookDownRight", + "eyeLookInLeft", "eyeLookInRight", "eyeLookOutLeft", "eyeLookOutRight", + "eyeLookUpLeft", "eyeLookUpRight", "eyeSquintLeft", "eyeSquintRight", + "eyeWideLeft", "eyeWideRight", "jawForward", "jawLeft", "jawOpen", + "jawRight", "mouthClose", "mouthDimpleLeft", "mouthDimpleRight", + "mouthFrownLeft", "mouthFrownRight", "mouthFunnel", "mouthLeft", + "mouthLowerDownLeft", "mouthLowerDownRight", "mouthPressLeft", + "mouthPressRight", "mouthPucker", "mouthRight", "mouthRollLower", + "mouthRollUpper", "mouthShrugLower", "mouthShrugUpper", "mouthSmileLeft", + "mouthSmileRight", "mouthStretchLeft", "mouthStretchRight", + "mouthUpperUpLeft", "mouthUpperUpRight", "noseSneerLeft", + "noseSneerRight", "tongueOut", +] + +# wav2vec2-base feature extractor geometry (configs/wav2vec2_config.json) +FE_KERNELS = [10, 3, 3, 3, 3, 2, 2] +FE_STRIDES = [5, 2, 2, 2, 2, 2, 2] + +SKIP_PREFIXES = ( + "audio_encoder.lm_head.", + "identity_encoder.grus.", +) +SKIP_KEYS = ("audio_encoder.masked_spec_embed",) + + +def load_state(checkpoint_path): + ckpt = torch.load(checkpoint_path, map_location="cpu", weights_only=True) + state = {} + for key, value in ckpt["state_dict"].items(): + if key.startswith("module."): + key = key[7:] + if key.startswith("backbone."): + key = key[9:] + if key in SKIP_KEYS or key.startswith(SKIP_PREFIXES): + continue + state[key] = value + return state + + +def fold_pos_conv_weight_norm(state): + """weight = g * v / ||v||, norm over dims (0,1) per kernel position (dim=2).""" + g = state.pop("audio_encoder.encoder.pos_conv_embed.conv.weight_g") # (1,1,128) + v = state.pop("audio_encoder.encoder.pos_conv_embed.conv.weight_v") # (768,48,128) + norm = v.norm(p=2, dim=(0, 1), keepdim=True) + state["audio_encoder.encoder.pos_conv_embed.conv.weight"] = g * v / norm + return state + + +def build_name_map(): + """checkpoint key -> (gguf name, is_matmul_weight).""" + m = {} + for i in range(7): + m[f"audio_encoder.feature_extractor.conv_layers.{i}.conv.weight"] = (f"fe.conv{i}.weight", True) + m["audio_encoder.feature_extractor.conv_layers.0.layer_norm.weight"] = ("fe.gn.weight", False) + m["audio_encoder.feature_extractor.conv_layers.0.layer_norm.bias"] = ("fe.gn.bias", False) + + m["audio_encoder.feature_projection.layer_norm.weight"] = ("fp.ln.weight", False) + m["audio_encoder.feature_projection.layer_norm.bias"] = ("fp.ln.bias", False) + m["audio_encoder.feature_projection.projection.weight"] = ("fp.proj.weight", True) + m["audio_encoder.feature_projection.projection.bias"] = ("fp.proj.bias", False) + + m["audio_encoder.encoder.pos_conv_embed.conv.weight"] = ("enc.pos_conv.weight", True) + m["audio_encoder.encoder.pos_conv_embed.conv.bias"] = ("enc.pos_conv.bias", False) + m["audio_encoder.encoder.layer_norm.weight"] = ("enc.ln.weight", False) + m["audio_encoder.encoder.layer_norm.bias"] = ("enc.ln.bias", False) + + for i in range(12): + src = f"audio_encoder.encoder.layers.{i}" + dst = f"enc.blk{i}" + for proj in ("q", "k", "v"): + m[f"{src}.attention.{proj}_proj.weight"] = (f"{dst}.attn_{proj}.weight", True) + m[f"{src}.attention.{proj}_proj.bias"] = (f"{dst}.attn_{proj}.bias", False) + m[f"{src}.attention.out_proj.weight"] = (f"{dst}.attn_o.weight", True) + m[f"{src}.attention.out_proj.bias"] = (f"{dst}.attn_o.bias", False) + m[f"{src}.layer_norm.weight"] = (f"{dst}.ln1.weight", False) + m[f"{src}.layer_norm.bias"] = (f"{dst}.ln1.bias", False) + m[f"{src}.feed_forward.intermediate_dense.weight"] = (f"{dst}.ffn_up.weight", True) + m[f"{src}.feed_forward.intermediate_dense.bias"] = (f"{dst}.ffn_up.bias", False) + m[f"{src}.feed_forward.output_dense.weight"] = (f"{dst}.ffn_down.weight", True) + m[f"{src}.feed_forward.output_dense.bias"] = (f"{dst}.ffn_down.bias", False) + m[f"{src}.final_layer_norm.weight"] = (f"{dst}.ln2.weight", False) + m[f"{src}.final_layer_norm.bias"] = (f"{dst}.ln2.bias", False) + + m["feature_projection.weight"] = ("head.proj.weight", True) + m["feature_projection.bias"] = ("head.proj.bias", False) + m["identity_encoder.id_mlp.weight"] = ("head.id_mlp.weight", True) + m["identity_encoder.id_mlp.bias"] = ("head.id_mlp.bias", False) + for i in range(3): + src = f"identity_encoder.first_net.conv_layers.{i}" + m[f"{src}.conv.weight"] = (f"head.first{i}.conv.weight", True) + m[f"{src}.conv.bias"] = (f"head.first{i}.conv.bias", False) + m[f"{src}.norm.weight"] = (f"head.first{i}.ln.weight", False) + m[f"{src}.norm.bias"] = (f"head.first{i}.ln.bias", False) + m["identity_encoder.first_net.conv_layers.0.residual_layer.0.weight"] = ("head.first0.res.weight", True) + m["identity_encoder.first_net.conv_layers.0.residual_layer.0.bias"] = ("head.first0.res.bias", False) + for i in range(3): + m[f"decoder.0.{i}.conv.weight"] = (f"head.dec{i}.conv.weight", True) + m[f"decoder.0.{i}.conv.bias"] = (f"head.dec{i}.conv.bias", False) + m[f"decoder.0.{i}.norm.weight"] = (f"head.dec{i}.ln.weight", False) + m[f"decoder.0.{i}.norm.bias"] = (f"head.dec{i}.ln.bias", False) + m["output_proj.weight"] = ("head.out.weight", True) + m["output_proj.bias"] = ("head.out.bias", False) + return m + + +def main(): + parser = argparse.ArgumentParser() + parser.add_argument("--checkpoint", required=True) + parser.add_argument("--out", required=True) + parser.add_argument("--dtype", choices=["f32", "f16"], default="f32") + args = parser.parse_args() + + state = fold_pos_conv_weight_norm(load_state(args.checkpoint)) + name_map = build_name_map() + + unmapped = sorted(set(state) - set(name_map)) + if unmapped: + raise RuntimeError(f"unmapped checkpoint tensors: {unmapped}") + missing = sorted(set(name_map) - set(state)) + if missing: + raise RuntimeError(f"expected checkpoint tensors not found: {missing}") + + os.makedirs(os.path.dirname(os.path.abspath(args.out)), exist_ok=True) + writer = GGUFWriter(args.out, ARCH) + writer.add_name("LAM Audio2Expression (streaming)") + writer.add_string(f"{ARCH}.dtype", args.dtype) + writer.add_uint32(f"{ARCH}.sample_rate", 16000) + writer.add_uint32(f"{ARCH}.fps", 30) + writer.add_uint32(f"{ARCH}.n_coeffs", 52) + writer.add_uint32(f"{ARCH}.n_identity", 12) + writer.add_uint32(f"{ARCH}.identity_feat_dim", 64) + writer.add_uint32(f"{ARCH}.hidden_dim", 512) + writer.add_uint32(f"{ARCH}.window_frames", 64) + writer.add_float32(f"{ARCH}.layer_norm_eps", 1e-5) + writer.add_uint32(f"{ARCH}.enc.n_layers", 12) + writer.add_uint32(f"{ARCH}.enc.n_heads", 12) + writer.add_uint32(f"{ARCH}.enc.hidden", 768) + writer.add_uint32(f"{ARCH}.enc.ffn", 3072) + writer.add_uint32(f"{ARCH}.enc.pos_conv_kernel", 128) + writer.add_uint32(f"{ARCH}.enc.pos_conv_groups", 16) + writer.add_array(f"{ARCH}.fe.kernels", FE_KERNELS) + writer.add_array(f"{ARCH}.fe.strides", FE_STRIDES) + writer.add_array(f"{ARCH}.coeff_names", ARKIT_BLENDSHAPES) + + total_bytes = 0 + for key in sorted(state, key=lambda k: name_map[k][0]): + gguf_name, is_matmul = name_map[key] + arr = state[key].detach().cpu().float().numpy() + arr = np.ascontiguousarray(arr) + if args.dtype == "f16" and is_matmul: + arr = arr.astype(np.float16) + writer.add_tensor(gguf_name, arr) + total_bytes += arr.nbytes + print(f" {gguf_name}: {list(arr.shape)} {arr.dtype}") + + writer.write_header_to_file() + writer.write_kv_data_to_file() + writer.write_tensors_to_file() + writer.close() + print(f"wrote {args.out} ({total_bytes / 1e6:.1f} MB tensor data, dtype={args.dtype})") + + +if __name__ == "__main__": + main() diff --git a/scripts/dump_lam_a2e_frontend_reference.py b/scripts/dump_lam_a2e_frontend_reference.py new file mode 100644 index 000000000..df8245885 --- /dev/null +++ b/scripts/dump_lam_a2e_frontend_reference.py @@ -0,0 +1,65 @@ +"""Create a Wav2Vec2 frontend parity fixture from trusted LAM-A2E PyTorch. + +Usage: + python scripts/dump_lam_a2e_frontend_reference.py \ + +""" + +from __future__ import annotations + +import json +import runpy +import sys +from pathlib import Path + +import librosa +import numpy as np +import torch + + +def main() -> None: + project_root = Path(sys.argv[1]).resolve() + checkpoint_path = Path(sys.argv[2]).resolve() + audio_path = Path(sys.argv[3]).resolve() + output_path = Path(sys.argv[4]).resolve() + + sys.path.insert(0, str(project_root)) + from models import build_model + + config = runpy.run_path( + str(project_root / "configs" / "lam_audio2exp_config_streaming.py") + ) + model = build_model(config["model"]) + checkpoint = torch.load(checkpoint_path, map_location="cpu", weights_only=False) + model.load_state_dict(checkpoint["state_dict"], strict=True) + model.cuda().eval() + + pcm, sample_rate = librosa.load(audio_path, sr=16000, mono=True) + pcm_tensor = torch.from_numpy(pcm).unsqueeze(0).cuda() + with torch.no_grad(): + frontend = model.backbone.audio_encoder.feature_extractor(pcm_tensor) + + output_path.parent.mkdir(parents=True, exist_ok=True) + np.savez( + output_path, + pcm=pcm.astype(np.float32), + frontend=frontend.cpu().numpy().astype(np.float32), + sample_rate=np.array(sample_rate, dtype=np.int32), + ) + output_path.with_suffix(".json").write_text( + json.dumps( + { + "sample_rate": sample_rate, + "pcm_shape": list(pcm.shape), + "frontend_shape": list(frontend.shape), + "weights": checkpoint_path.name, + "stage": "Wav2Vec2 feature_extractor", + }, + indent=2, + ) + + "\n" + ) + + +if __name__ == "__main__": + main() diff --git a/scripts/dump_lam_a2e_stages.py b/scripts/dump_lam_a2e_stages.py new file mode 100644 index 000000000..3b9978ad4 --- /dev/null +++ b/scripts/dump_lam_a2e_stages.py @@ -0,0 +1,138 @@ +"""Generate deterministic per-stage LAM-A2E parity fixtures. + +The fixture uses the first second of a trusted 16 kHz mono WAV so each +intermediate stage stays small enough to commit/store alongside test metadata. +""" + +from __future__ import annotations + +import argparse +import json +import runpy +import sys +from pathlib import Path + +import librosa +import numpy as np +import torch + + +def parse_args() -> argparse.Namespace: + parser = argparse.ArgumentParser() + parser.add_argument("project_root", type=Path) + parser.add_argument("checkpoint", type=Path) + parser.add_argument("audio", type=Path) + parser.add_argument("output", type=Path) + parser.add_argument("--seconds", type=float, default=1.0) + parser.add_argument("--identity-index", type=int, default=0) + return parser.parse_args() + + +def tensor_output(value: object) -> torch.Tensor: + if isinstance(value, torch.Tensor): + return value + if hasattr(value, "last_hidden_state"): + return value.last_hidden_state + if isinstance(value, (tuple, list)) and value and isinstance(value[0], torch.Tensor): + return value[0] + raise TypeError(f"Cannot serialize hook output of type {type(value)!r}") + + +def main() -> None: + args = parse_args() + project_root = args.project_root.resolve() + sys.path.insert(0, str(project_root)) + from models import build_model + + config = runpy.run_path( + str(project_root / "configs" / "lam_audio2exp_config_streaming.py") + ) + model = build_model(config["model"]) + checkpoint = torch.load(args.checkpoint, map_location="cpu", weights_only=False) + model.load_state_dict(checkpoint["state_dict"], strict=True) + model.cuda().eval() + + stage_outputs: dict[str, np.ndarray] = {} + + def capture(name: str): + def hook(_module: torch.nn.Module, _inputs: tuple[object, ...], output: object): + stage_outputs[name] = ( + tensor_output(output).detach().float().cpu().numpy() + ) + + return hook + + def capture_input(name: str): + def hook(_module: torch.nn.Module, inputs: tuple[object, ...]): + stage_outputs[name] = ( + tensor_output(inputs[0]).detach().float().cpu().numpy() + ) + + return hook + + backbone = model.backbone + hooks = [ + backbone.audio_encoder.feature_extractor.register_forward_hook(capture("frontend")), + backbone.audio_encoder.feature_projection.layer_norm.register_forward_pre_hook( + capture_input("interpolated_frontend") + ), + backbone.audio_encoder.feature_projection.layer_norm.register_forward_hook( + capture("feature_layer_norm") + ), + backbone.audio_encoder.feature_projection.register_forward_hook(capture("wav2vec_projection")), + backbone.audio_encoder.encoder.pos_conv_embed.register_forward_hook(capture("position")), + backbone.feature_projection.register_forward_hook(capture("lam_projection")), + backbone.identity_encoder.register_forward_hook(capture("identity")), + backbone.decoder[0].register_forward_hook(capture("decoder")), + backbone.output_proj.register_forward_hook(capture("output_projection")), + ] + for index, layer in enumerate(backbone.audio_encoder.feature_extractor.conv_layers): + hooks.append(layer.conv.register_forward_hook(capture(f"frontend_conv_{index}"))) + if index == 0: + hooks.append(layer.layer_norm.register_forward_hook(capture("frontend_group_norm_0"))) + for index, layer in enumerate(backbone.audio_encoder.encoder.layers): + hooks.append(layer.register_forward_hook(capture(f"transformer_{index:02d}"))) + + pcm, sample_rate = librosa.load(args.audio, sr=16000, mono=True) + pcm = pcm[: int(args.seconds * sample_rate)] + identity_class_count = config["model"]["backbone"]["num_identity_classes"] + identity = torch.nn.functional.one_hot( + torch.tensor(args.identity_index), + identity_class_count, + ).cuda()[None, ...] + input_dict = { + "id_idx": identity, + "input_audio_array": torch.from_numpy(pcm).unsqueeze(0).cuda(), + } + + with torch.no_grad(): + final_output = backbone(input_dict) + for hook in hooks: + hook.remove() + + args.output.parent.mkdir(parents=True, exist_ok=True) + np.savez( + args.output, + pcm=pcm.astype(np.float32), + final=final_output.detach().float().cpu().numpy(), + **stage_outputs, + ) + args.output.with_suffix(".json").write_text( + json.dumps( + { + "sample_rate": sample_rate, + "seconds": args.seconds, + "identity_index": args.identity_index, + "fps": 30, + "stages": {name: list(value.shape) for name, value in stage_outputs.items()}, + "final_shape": list(final_output.shape), + "checkpoint": args.checkpoint.name, + }, + indent=2, + ) + + "\n" + ) + + +if __name__ == "__main__": + main() diff --git a/scripts/remap_lam_gguf.py b/scripts/remap_lam_gguf.py new file mode 100644 index 000000000..6dd03f8ff --- /dev/null +++ b/scripts/remap_lam_gguf.py @@ -0,0 +1,174 @@ +#!/usr/bin/env python3 +"""Remap a raw LAM-A2E GGUF (pytorch/backbone names) into lam-audio2exp dialect.""" + +from __future__ import annotations + +import argparse +import os + +import numpy as np +from gguf import GGUFReader, GGUFWriter + +ARCH = "lam-audio2exp" + +ARKIT_BLENDSHAPES = [ + "browDownLeft", "browDownRight", "browInnerUp", "browOuterUpLeft", + "browOuterUpRight", "cheekPuff", "cheekSquintLeft", "cheekSquintRight", + "eyeBlinkLeft", "eyeBlinkRight", "eyeLookDownLeft", "eyeLookDownRight", + "eyeLookInLeft", "eyeLookInRight", "eyeLookOutLeft", "eyeLookOutRight", + "eyeLookUpLeft", "eyeLookUpRight", "eyeSquintLeft", "eyeSquintRight", + "eyeWideLeft", "eyeWideRight", "jawForward", "jawLeft", "jawOpen", + "jawRight", "mouthClose", "mouthDimpleLeft", "mouthDimpleRight", + "mouthFrownLeft", "mouthFrownRight", "mouthFunnel", "mouthLeft", + "mouthLowerDownLeft", "mouthLowerDownRight", "mouthPressLeft", + "mouthPressRight", "mouthPucker", "mouthRight", "mouthRollLower", + "mouthRollUpper", "mouthShrugLower", "mouthShrugUpper", "mouthSmileLeft", + "mouthSmileRight", "mouthStretchLeft", "mouthStretchRight", + "mouthUpperUpLeft", "mouthUpperUpRight", "noseSneerLeft", + "noseSneerRight", "tongueOut", +] + +FE_KERNELS = [10, 3, 3, 3, 3, 2, 2] +FE_STRIDES = [5, 2, 2, 2, 2, 2, 2] +SKIP_PREFIXES = ( + "audio_encoder.lm_head.", + "identity_encoder.grus.", +) +SKIP_KEYS = ("audio_encoder.masked_spec_embed",) + + +def build_name_map(): + m = {} + for i in range(7): + m[f"audio_encoder.feature_extractor.conv_layers.{i}.conv.weight"] = (f"fe.conv{i}.weight", True) + m["audio_encoder.feature_extractor.conv_layers.0.layer_norm.weight"] = ("fe.gn.weight", False) + m["audio_encoder.feature_extractor.conv_layers.0.layer_norm.bias"] = ("fe.gn.bias", False) + m["audio_encoder.feature_projection.layer_norm.weight"] = ("fp.ln.weight", False) + m["audio_encoder.feature_projection.layer_norm.bias"] = ("fp.ln.bias", False) + m["audio_encoder.feature_projection.projection.weight"] = ("fp.proj.weight", True) + m["audio_encoder.feature_projection.projection.bias"] = ("fp.proj.bias", False) + m["audio_encoder.encoder.pos_conv_embed.conv.weight"] = ("enc.pos_conv.weight", True) + m["audio_encoder.encoder.pos_conv_embed.conv.bias"] = ("enc.pos_conv.bias", False) + m["audio_encoder.encoder.layer_norm.weight"] = ("enc.ln.weight", False) + m["audio_encoder.encoder.layer_norm.bias"] = ("enc.ln.bias", False) + for i in range(12): + src = f"audio_encoder.encoder.layers.{i}" + dst = f"enc.blk{i}" + for proj in ("q", "k", "v"): + m[f"{src}.attention.{proj}_proj.weight"] = (f"{dst}.attn_{proj}.weight", True) + m[f"{src}.attention.{proj}_proj.bias"] = (f"{dst}.attn_{proj}.bias", False) + m[f"{src}.attention.out_proj.weight"] = (f"{dst}.attn_o.weight", True) + m[f"{src}.attention.out_proj.bias"] = (f"{dst}.attn_o.bias", False) + m[f"{src}.layer_norm.weight"] = (f"{dst}.ln1.weight", False) + m[f"{src}.layer_norm.bias"] = (f"{dst}.ln1.bias", False) + m[f"{src}.feed_forward.intermediate_dense.weight"] = (f"{dst}.ffn_up.weight", True) + m[f"{src}.feed_forward.intermediate_dense.bias"] = (f"{dst}.ffn_up.bias", False) + m[f"{src}.feed_forward.output_dense.weight"] = (f"{dst}.ffn_down.weight", True) + m[f"{src}.feed_forward.output_dense.bias"] = (f"{dst}.ffn_down.bias", False) + m[f"{src}.final_layer_norm.weight"] = (f"{dst}.ln2.weight", False) + m[f"{src}.final_layer_norm.bias"] = (f"{dst}.ln2.bias", False) + m["feature_projection.weight"] = ("head.proj.weight", True) + m["feature_projection.bias"] = ("head.proj.bias", False) + m["identity_encoder.id_mlp.weight"] = ("head.id_mlp.weight", True) + m["identity_encoder.id_mlp.bias"] = ("head.id_mlp.bias", False) + for i in range(3): + src = f"identity_encoder.first_net.conv_layers.{i}" + m[f"{src}.conv.weight"] = (f"head.first{i}.conv.weight", True) + m[f"{src}.conv.bias"] = (f"head.first{i}.conv.bias", False) + m[f"{src}.norm.weight"] = (f"head.first{i}.ln.weight", False) + m[f"{src}.norm.bias"] = (f"head.first{i}.ln.bias", False) + m["identity_encoder.first_net.conv_layers.0.residual_layer.0.weight"] = ("head.first0.res.weight", True) + m["identity_encoder.first_net.conv_layers.0.residual_layer.0.bias"] = ("head.first0.res.bias", False) + for i in range(3): + m[f"decoder.0.{i}.conv.weight"] = (f"head.dec{i}.conv.weight", True) + m[f"decoder.0.{i}.conv.bias"] = (f"head.dec{i}.conv.bias", False) + m[f"decoder.0.{i}.norm.weight"] = (f"head.dec{i}.ln.weight", False) + m[f"decoder.0.{i}.norm.bias"] = (f"head.dec{i}.ln.bias", False) + m["output_proj.weight"] = ("head.out.weight", True) + m["output_proj.bias"] = ("head.out.bias", False) + return m + + +def strip_prefix(name: str) -> str: + if name.startswith("module."): + name = name[7:] + if name.startswith("backbone."): + name = name[9:] + return name + + +def main() -> None: + parser = argparse.ArgumentParser() + parser.add_argument("--in", dest="inp", required=True) + parser.add_argument("--out", required=True) + parser.add_argument("--dtype", choices=["f32", "f16"], default="f32") + args = parser.parse_args() + + reader = GGUFReader(args.inp) + state = {} + for tensor in reader.tensors: + key = strip_prefix(tensor.name) + if key in SKIP_KEYS or key.startswith(SKIP_PREFIXES): + continue + state[key] = np.array(tensor.data, dtype=np.float32, copy=True) + + g = state.pop("audio_encoder.encoder.pos_conv_embed.conv.weight_g") + v = state.pop("audio_encoder.encoder.pos_conv_embed.conv.weight_v") + if g.shape != (1, 1, 128) or v.shape != (768, 48, 128): + raise RuntimeError(f"unexpected pos-conv shapes g={g.shape} v={v.shape}") + norm = np.linalg.norm(v, axis=(0, 1), keepdims=True) + folded = (g * v / np.maximum(norm, 1e-12)).astype(np.float32, copy=True) + state["audio_encoder.encoder.pos_conv_embed.conv.weight"] = folded + + name_map = build_name_map() + unmapped = sorted(set(state) - set(name_map)) + if unmapped: + raise RuntimeError(f"unmapped tensors ({len(unmapped)}): {unmapped[:30]}") + missing = sorted(set(name_map) - set(state)) + if missing: + raise RuntimeError(f"missing tensors: {missing}") + + out_dir = os.path.dirname(os.path.abspath(args.out)) + if out_dir: + os.makedirs(out_dir, exist_ok=True) + + writer = GGUFWriter(args.out, ARCH) + writer.add_name("LAM Audio2Expression (streaming)") + writer.add_string(f"{ARCH}.dtype", args.dtype) + writer.add_uint32(f"{ARCH}.sample_rate", 16000) + writer.add_uint32(f"{ARCH}.fps", 30) + writer.add_uint32(f"{ARCH}.n_coeffs", 52) + writer.add_uint32(f"{ARCH}.n_identity", 12) + writer.add_uint32(f"{ARCH}.identity_feat_dim", 64) + writer.add_uint32(f"{ARCH}.hidden_dim", 512) + writer.add_uint32(f"{ARCH}.window_frames", 64) + writer.add_float32(f"{ARCH}.layer_norm_eps", 1e-5) + writer.add_uint32(f"{ARCH}.enc.n_layers", 12) + writer.add_uint32(f"{ARCH}.enc.n_heads", 12) + writer.add_uint32(f"{ARCH}.enc.hidden", 768) + writer.add_uint32(f"{ARCH}.enc.ffn", 3072) + writer.add_uint32(f"{ARCH}.enc.pos_conv_kernel", 128) + writer.add_uint32(f"{ARCH}.enc.pos_conv_groups", 16) + writer.add_array(f"{ARCH}.fe.kernels", FE_KERNELS) + writer.add_array(f"{ARCH}.fe.strides", FE_STRIDES) + writer.add_array(f"{ARCH}.coeff_names", ARKIT_BLENDSHAPES) + + total = 0 + for key in sorted(state, key=lambda k: name_map[k][0]): + gguf_name, is_matmul = name_map[key] + arr = np.ascontiguousarray(state[key]) + if args.dtype == "f16" and is_matmul: + arr = arr.astype(np.float16) + writer.add_tensor(gguf_name, arr) + total += arr.nbytes + print(f" {gguf_name}: {list(arr.shape)} {arr.dtype}") + + writer.write_header_to_file() + writer.write_kv_data_to_file() + writer.write_tensors_to_file() + writer.close() + print(f"wrote {args.out} ({total / 1e6:.1f} MB)") + + +if __name__ == "__main__": + main() diff --git a/src/lam-a2e.cpp b/src/lam-a2e.cpp new file mode 100644 index 000000000..d4fe08b04 --- /dev/null +++ b/src/lam-a2e.cpp @@ -0,0 +1,120 @@ +#include "lam-a2e.h" + +#include +#include +#include +#include +#include + +#include "lam_audio2expression.hpp" + +struct lam_a2e_context { + std::string last_error; + std::unique_ptr model; + int32_t identity_index = 0; + int32_t n_threads = 0; + bool use_gpu = false; + bool model_loaded = false; +}; + +lam_a2e_context * lam_a2e_create(const lam_a2e_params * params) { + if (params == nullptr || params->model_path == nullptr || + std::strlen(params->model_path) == 0) { + return nullptr; + } + + auto * ctx = new lam_a2e_context(); + ctx->identity_index = params->identity_index; + ctx->n_threads = params->n_threads; + ctx->use_gpu = params->use_gpu; + + if (params->use_gpu) { + ctx->last_error = + "LAM-A2E GPU backends are not enabled until CPU parity is complete."; + return ctx; + } + + ctx->model = std::make_unique(); + if (!ctx->model->load(params->model_path, nullptr, params->n_threads)) { + ctx->last_error = ctx->model->lastError(); + ctx->model.reset(); + return ctx; + } + + ctx->model_loaded = true; + ctx->last_error.clear(); + return ctx; +} + +void lam_a2e_free(lam_a2e_context * ctx) { + delete ctx; +} + +lam_a2e_status lam_a2e_process_pcm_f32( + lam_a2e_context * ctx, + const float * pcm, + int64_t pcm_sample_count, + int32_t sample_rate, + lam_a2e_frame ** frames, + int32_t * frame_count) { + if (ctx == nullptr || pcm == nullptr || pcm_sample_count <= 0 || + frames == nullptr || frame_count == nullptr) { + return LAM_A2E_STATUS_INVALID_ARGUMENT; + } + + *frames = nullptr; + *frame_count = 0; + + if (!ctx->model_loaded || ctx->model == nullptr) { + if (ctx->last_error.empty()) { + ctx->last_error = "LAM-A2E model is not loaded."; + } + return LAM_A2E_STATUS_MODEL_LOAD_FAILED; + } + if (sample_rate != 16000) { + ctx->last_error = "LAM-A2E requires 16 kHz mono PCM input."; + return LAM_A2E_STATUS_INVALID_ARGUMENT; + } + + std::vector pcm_vec(pcm, pcm + pcm_sample_count); + std::vector coeffs; + if (!ctx->model->run(pcm_vec, static_cast(ctx->identity_index), coeffs)) { + ctx->last_error = ctx->model->lastError(); + return LAM_A2E_STATUS_INVALID_ARGUMENT; + } + + const auto& hp = ctx->model->hparams(); + const int64_t n_frames = ctx->model->frameCount(pcm_sample_count); + if (n_frames <= 0 || + coeffs.size() != static_cast(n_frames) * hp.nCoeffs) { + ctx->last_error = "LAM-A2E produced an unexpected coefficient buffer size."; + return LAM_A2E_STATUS_INVALID_ARGUMENT; + } + + auto * out = static_cast( + std::malloc(sizeof(lam_a2e_frame) * static_cast(n_frames))); + if (out == nullptr) { + ctx->last_error = "Failed to allocate LAM-A2E frame buffer."; + return LAM_A2E_STATUS_INVALID_ARGUMENT; + } + + for (int64_t i = 0; i < n_frames; ++i) { + out[i].timestamp_us = + (i * 1000000LL) / static_cast(hp.fps > 0 ? hp.fps : 30); + std::memcpy(out[i].arkit_52, coeffs.data() + i * hp.nCoeffs, + sizeof(float) * hp.nCoeffs); + } + + *frames = out; + *frame_count = static_cast(n_frames); + ctx->last_error.clear(); + return LAM_A2E_STATUS_OK; +} + +void lam_a2e_free_frames(lam_a2e_frame * frames) { + std::free(frames); +} + +const char * lam_a2e_get_last_error(const lam_a2e_context * ctx) { + return ctx == nullptr ? "LAM-A2E context is null." : ctx->last_error.c_str(); +} diff --git a/src/lam_audio2expression.cpp b/src/lam_audio2expression.cpp new file mode 100644 index 000000000..6eb84832d --- /dev/null +++ b/src/lam_audio2expression.cpp @@ -0,0 +1,510 @@ +#include "lam_audio2expression.hpp" + +#include +#include +#include + +#include "ggml-alloc.h" +#include "ggml-cpu.h" +#include "gguf.h" + +namespace { + +constexpr const char* kArch = "lam-audio2exp"; + +uint32_t +ggufGetU32Or(struct gguf_context* g, const std::string& key, uint32_t dflt) { + const int64_t idx = gguf_find_key(g, key.c_str()); + if (idx < 0 || gguf_get_kv_type(g, idx) != GGUF_TYPE_UINT32) { + return dflt; + } + return gguf_get_val_u32(g, idx); +} + +float +ggufGetF32Or(struct gguf_context* g, const std::string& key, float dflt) { + const int64_t idx = gguf_find_key(g, key.c_str()); + if (idx < 0 || gguf_get_kv_type(g, idx) != GGUF_TYPE_FLOAT32) { + return dflt; + } + return gguf_get_val_f32(g, idx); +} + +// conv1d that preserves the kernel's precision: the stock ggml_conv_1d +// always runs im2col in F16, which breaks exact-f32 parity. This builds the +// same im2col + mul_mat pair with the dst type matching the kernel. +// kernel: [k, cIn, cOut], data: [t, cIn] → [tOut, cOut]. +struct ggml_tensor* +conv1d(struct ggml_context* ctx, struct ggml_tensor* kernel, + struct ggml_tensor* data, int stride, int padding) { + struct ggml_tensor* cols = + ggml_im2col(ctx, kernel, data, stride, 0, padding, 0, 1, 0, + /*is_2D*/ false, kernel->type); // [k*cIn, tOut, 1] + struct ggml_tensor* out = ggml_mul_mat( + ctx, + ggml_reshape_2d(ctx, cols, cols->ne[0], cols->ne[1] * cols->ne[2]), + ggml_reshape_2d(ctx, kernel, kernel->ne[0] * kernel->ne[1], + kernel->ne[2])); // [tOut, cOut] + return out; +} + +// Broadcast a 1-D bias [c] over the ne0 (time) axis of a [t, c] tensor. +struct ggml_tensor* +addBiasTimeMajor(struct ggml_context* ctx, struct ggml_tensor* x, + struct ggml_tensor* bias) { + return ggml_add(ctx, x, ggml_reshape_2d(ctx, bias, 1, bias->ne[0])); +} + +// LayerNorm over ne0 (features) of a feature-major [c, t] tensor. +struct ggml_tensor* +layerNorm(struct ggml_context* ctx, struct ggml_tensor* x, + struct ggml_tensor* w, struct ggml_tensor* b, float eps) { + struct ggml_tensor* cur = ggml_norm(ctx, x, eps); + cur = ggml_mul(ctx, cur, w); + return ggml_add(ctx, cur, b); +} + +} // namespace + +LamAudio2Expression::~LamAudio2Expression() { + if (weightBuffer_ != nullptr) { + ggml_backend_buffer_free(weightBuffer_); + } + if (weightCtx_ != nullptr) { + ggml_free(weightCtx_); + } + if (ownsBackend_ && backend_ != nullptr) { + ggml_backend_free(backend_); + } +} + +struct ggml_tensor* +LamAudio2Expression::weight(const std::string& name) { + auto it = weights_.find(name); + return it == weights_.end() ? nullptr : it->second; +} + +int64_t +LamAudio2Expression::frameCount(int64_t nSamples) const { + return (nSamples * hparams_.fps + hparams_.sampleRate - 1) / + hparams_.sampleRate; +} + +int64_t +LamAudio2Expression::convOutLen(int64_t nSamples) const { + int64_t len = nSamples; + for (size_t i = 0; i < hparams_.feKernels.size(); ++i) { + len = (len - hparams_.feKernels[i]) / hparams_.feStrides[i] + 1; + } + return len; +} + +bool +LamAudio2Expression::load(const std::string& ggufPath, ggml_backend_t backend, int n_threads) { + if (backend != nullptr) { + backend_ = backend; + ownsBackend_ = false; + } else { + backend_ = ggml_backend_cpu_init(); + ownsBackend_ = true; + if (backend_ != nullptr) { + const unsigned hw = std::thread::hardware_concurrency(); + const int threads = n_threads > 0 + ? n_threads + : static_cast(hw > 2 ? hw - 2 : 1); + ggml_backend_cpu_set_n_threads(backend_, threads); + } + } + if (backend_ == nullptr) { + lastError_ = "failed to initialise backend"; + return false; + } + + struct ggml_context* metaCtx = nullptr; + struct gguf_init_params params = {/*no_alloc*/ true, /*ctx*/ &metaCtx}; + struct gguf_context* gguf = gguf_init_from_file(ggufPath.c_str(), params); + if (gguf == nullptr) { + lastError_ = "failed to open GGUF: " + ggufPath; + return false; + } + + const std::string arch = [&] { + const int64_t idx = gguf_find_key(gguf, "general.architecture"); + return idx >= 0 ? std::string(gguf_get_val_str(gguf, idx)) : std::string(); + }(); + if (arch != kArch) { + lastError_ = "unexpected architecture '" + arch + "'"; + gguf_free(gguf); + ggml_free(metaCtx); + return false; + } + + const std::string pfx = std::string(kArch) + "."; + hparams_.sampleRate = ggufGetU32Or(gguf, pfx + "sample_rate", 16000); + hparams_.fps = ggufGetU32Or(gguf, pfx + "fps", 30); + hparams_.nCoeffs = ggufGetU32Or(gguf, pfx + "n_coeffs", 52); + hparams_.nIdentity = ggufGetU32Or(gguf, pfx + "n_identity", 12); + hparams_.identityFeatDim = ggufGetU32Or(gguf, pfx + "identity_feat_dim", 64); + hparams_.hiddenDim = ggufGetU32Or(gguf, pfx + "hidden_dim", 512); + hparams_.windowFrames = ggufGetU32Or(gguf, pfx + "window_frames", 64); + hparams_.layerNormEps = ggufGetF32Or(gguf, pfx + "layer_norm_eps", 1e-5F); + hparams_.encLayers = ggufGetU32Or(gguf, pfx + "enc.n_layers", 12); + hparams_.encHeads = ggufGetU32Or(gguf, pfx + "enc.n_heads", 12); + hparams_.encHidden = ggufGetU32Or(gguf, pfx + "enc.hidden", 768); + hparams_.encFfn = ggufGetU32Or(gguf, pfx + "enc.ffn", 3072); + hparams_.posConvKernel = ggufGetU32Or(gguf, pfx + "enc.pos_conv_kernel", 128); + hparams_.posConvGroups = ggufGetU32Or(gguf, pfx + "enc.pos_conv_groups", 16); + + const int64_t namesIdx = gguf_find_key(gguf, (pfx + "coeff_names").c_str()); + if (namesIdx >= 0 && gguf_get_kv_type(gguf, namesIdx) == GGUF_TYPE_ARRAY && + gguf_get_arr_type(gguf, namesIdx) == GGUF_TYPE_STRING) { + const size_t n = gguf_get_arr_n(gguf, namesIdx); + hparams_.coeffNames.reserve(n); + for (size_t i = 0; i < n; ++i) { + hparams_.coeffNames.emplace_back(gguf_get_arr_str(gguf, namesIdx, i)); + } + } + + // Copy tensor metadata into our own context, then load the data through + // the backend buffer (works for CPU and GPU backends alike). + const int64_t nTensors = gguf_get_n_tensors(gguf); + const size_t ctxSize = + (static_cast(nTensors) + 1) * ggml_tensor_overhead(); + struct ggml_init_params wparams = {ctxSize, nullptr, /*no_alloc*/ true}; + weightCtx_ = ggml_init(wparams); + + for (struct ggml_tensor* meta = ggml_get_first_tensor(metaCtx); + meta != nullptr; meta = ggml_get_next_tensor(metaCtx, meta)) { + struct ggml_tensor* dst = ggml_dup_tensor(weightCtx_, meta); + ggml_set_name(dst, ggml_get_name(meta)); + weights_[ggml_get_name(meta)] = dst; + } + + weightBuffer_ = ggml_backend_alloc_ctx_tensors(weightCtx_, backend_); + if (weightBuffer_ == nullptr) { + lastError_ = "failed to allocate weight buffer"; + gguf_free(gguf); + ggml_free(metaCtx); + return false; + } + + FILE* file = fopen(ggufPath.c_str(), "rb"); + if (file == nullptr) { + lastError_ = "failed to reopen GGUF: " + ggufPath; + gguf_free(gguf); + ggml_free(metaCtx); + return false; + } + const size_t dataOffset = gguf_get_data_offset(gguf); + std::vector readBuf; + bool ok = true; + for (int64_t i = 0; i < nTensors; ++i) { + const char* name = gguf_get_tensor_name(gguf, i); + struct ggml_tensor* dst = weights_[name]; + const size_t offset = dataOffset + gguf_get_tensor_offset(gguf, i); + const size_t nbytes = ggml_nbytes(dst); + readBuf.resize(nbytes); + if (fseek(file, static_cast(offset), SEEK_SET) != 0 || + fread(readBuf.data(), 1, nbytes, file) != nbytes) { + lastError_ = std::string("failed to read tensor ") + name; + ok = false; + break; + } + ggml_backend_tensor_set(dst, readBuf.data(), 0, nbytes); + } + fclose(file); + gguf_free(gguf); + ggml_free(metaCtx); + return ok; +} + +bool +LamAudio2Expression::run(const std::vector& pcm, uint32_t idIdx, + std::vector& framesOut, + std::map>* taps) { + const auto& hp = hparams_; + const int64_t nSamples = static_cast(pcm.size()); + const int64_t frames = frameCount(nSamples); + const int64_t tConv = convOutLen(nSamples); + if (nSamples == 0 || frames < 2 || tConv < 2) { + lastError_ = "input too short"; + return false; + } + if (idIdx >= hp.nIdentity) { + lastError_ = "identity index out of range"; + return false; + } + + // Generous node bound: the encoder dominates (~30 nodes/layer) plus the + // grouped positional conv (~5 nodes/group). + const size_t graphNodes = 2048; + const size_t ctxSize = graphNodes * ggml_tensor_overhead() + + ggml_graph_overhead_custom(graphNodes, false) + + (1U << 16U); + struct ggml_init_params gparams = {ctxSize, nullptr, /*no_alloc*/ true}; + struct ggml_context* ctx = ggml_init(gparams); + struct ggml_cgraph* graph = ggml_new_graph_custom(ctx, graphNodes, false); + + const auto tap = [&](struct ggml_tensor* t, const char* name) { + ggml_set_name(t, name); + if (taps != nullptr) { + // Output-flag tapped tensors so the graph allocator does not reuse + // their buffers before we read them back. + ggml_set_output(t); + ggml_build_forward_expand(graph, t); + } + return t; + }; + + // ---- inputs ----------------------------------------------------------- + struct ggml_tensor* pcmIn = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, nSamples, 1); + ggml_set_name(pcmIn, "input_pcm"); + ggml_set_input(pcmIn); + + struct ggml_tensor* idIn = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, hp.nIdentity, 1); + ggml_set_name(idIn, "input_id"); + ggml_set_input(idIn); + + // 50 Hz → 30 fps linear interpolation as a [tConv, frames] matrix + // (align_corners=true), filled at set-input time below. + struct ggml_tensor* interpW = + ggml_new_tensor_2d(ctx, GGML_TYPE_F32, tConv, frames); + ggml_set_name(interpW, "input_interp"); + ggml_set_input(interpW); + + // ---- feature extractor (time-major [t, c]) ---------------------------- + struct ggml_tensor* cur = pcmIn; + for (uint32_t i = 0; i < hp.feKernels.size(); ++i) { + cur = conv1d(ctx, weight("fe.conv" + std::to_string(i) + ".weight"), cur, + hp.feStrides[i], 0); + if (i == 0) { + // GroupNorm(512, 512) == per-channel norm over time; time is ne0 here. + cur = ggml_norm(ctx, cur, hp.layerNormEps); + cur = ggml_mul(ctx, cur, + ggml_reshape_2d(ctx, weight("fe.gn.weight"), 1, + weight("fe.gn.weight")->ne[0])); + cur = addBiasTimeMajor(ctx, cur, weight("fe.gn.bias")); + } + cur = ggml_gelu_erf(ctx, cur); + } + // [tConv, 512] time-major (byte-compatible with the PyTorch (512, tConv) dump) + cur = tap(cur, "fe_out"); + + // ---- 50→30 fps interpolation ------------------------------------------ + cur = ggml_mul_mat(ctx, interpW, ggml_cont(ctx, cur)); // [frames, 512] + tap(cur, "interp_out"); + + // ---- wav2vec2 feature projection (feature-major [c, t]) --------------- + cur = ggml_cont(ctx, ggml_transpose(ctx, cur)); // [512, frames] + cur = layerNorm(ctx, cur, weight("fp.ln.weight"), weight("fp.ln.bias"), + hp.layerNormEps); + cur = ggml_mul_mat(ctx, weight("fp.proj.weight"), cur); // [768, frames] + cur = ggml_add(ctx, cur, weight("fp.proj.bias")); + tap(cur, "fp_out"); + + // ---- positional conv embedding (grouped conv, time-major) ------------- + { + struct ggml_tensor* xt = ggml_cont(ctx, ggml_transpose(ctx, cur)); // [t, 768] + struct ggml_tensor* kern = weight("enc.pos_conv.weight"); // [128, 48, 768] + const int64_t chPerGroup = hp.encHidden / hp.posConvGroups; // 48 + struct ggml_tensor* pos = nullptr; + for (uint32_t g = 0; g < hp.posConvGroups; ++g) { + struct ggml_tensor* xg = ggml_cont( + ctx, ggml_view_2d(ctx, xt, xt->ne[0], chPerGroup, xt->nb[1], + g * chPerGroup * xt->nb[1])); + struct ggml_tensor* kg = ggml_cont( + ctx, ggml_view_3d(ctx, kern, kern->ne[0], kern->ne[1], chPerGroup, + kern->nb[1], kern->nb[2], + g * chPerGroup * kern->nb[2])); + struct ggml_tensor* cg = + conv1d(ctx, kg, xg, 1, static_cast(hp.posConvKernel / 2)); + pos = pos == nullptr ? cg : ggml_concat(ctx, pos, cg, 1); + } + // Even kernel + pad k/2 yields t+1 outputs; drop the trailing one + // (Wav2Vec2SamePadLayer), then bias + GELU. + pos = ggml_view_2d(ctx, pos, frames, hp.encHidden, pos->nb[1], 0); + pos = addBiasTimeMajor(ctx, pos, weight("enc.pos_conv.bias")); + pos = ggml_gelu_erf(ctx, pos); + pos = ggml_cont(ctx, ggml_transpose(ctx, pos)); // [768, t] + tap(pos, "pos_conv_out"); + cur = ggml_add(ctx, cur, pos); + } + cur = layerNorm(ctx, cur, weight("enc.ln.weight"), weight("enc.ln.bias"), + hp.layerNormEps); + tap(cur, "enc_pre_ln"); + + // ---- transformer encoder (post-norm) ----------------------------------- + const int64_t headDim = hp.encHidden / hp.encHeads; + const float attnScale = 1.0F / std::sqrt(static_cast(headDim)); + for (uint32_t il = 0; il < hp.encLayers; ++il) { + const std::string blk = "enc.blk" + std::to_string(il) + "."; + struct ggml_tensor* residual = cur; + + struct ggml_tensor* q = + ggml_add(ctx, ggml_mul_mat(ctx, weight(blk + "attn_q.weight"), cur), + weight(blk + "attn_q.bias")); + q = ggml_scale(ctx, q, attnScale); + struct ggml_tensor* k = + ggml_add(ctx, ggml_mul_mat(ctx, weight(blk + "attn_k.weight"), cur), + weight(blk + "attn_k.bias")); + struct ggml_tensor* v = + ggml_add(ctx, ggml_mul_mat(ctx, weight(blk + "attn_v.weight"), cur), + weight(blk + "attn_v.bias")); + + q = ggml_cont(ctx, ggml_permute( + ctx, ggml_reshape_3d(ctx, q, headDim, hp.encHeads, frames), 0, 2, 1, 3)); + k = ggml_cont(ctx, ggml_permute( + ctx, ggml_reshape_3d(ctx, k, headDim, hp.encHeads, frames), 0, 2, 1, 3)); + v = ggml_cont(ctx, ggml_permute( + ctx, ggml_reshape_3d(ctx, v, headDim, hp.encHeads, frames), 1, 2, 0, 3)); + + struct ggml_tensor* kq = ggml_soft_max(ctx, ggml_mul_mat(ctx, k, q)); + struct ggml_tensor* kqv = ggml_mul_mat(ctx, v, kq); // [headDim, t, heads] + kqv = ggml_cont(ctx, ggml_permute(ctx, kqv, 0, 2, 1, 3)); + cur = ggml_reshape_2d(ctx, kqv, hp.encHidden, frames); + cur = ggml_add(ctx, ggml_mul_mat(ctx, weight(blk + "attn_o.weight"), cur), + weight(blk + "attn_o.bias")); + + cur = ggml_add(ctx, cur, residual); + cur = layerNorm(ctx, cur, weight(blk + "ln1.weight"), + weight(blk + "ln1.bias"), hp.layerNormEps); + + struct ggml_tensor* ffn = + ggml_add(ctx, ggml_mul_mat(ctx, weight(blk + "ffn_up.weight"), cur), + weight(blk + "ffn_up.bias")); + ffn = ggml_gelu_erf(ctx, ffn); + ffn = ggml_add(ctx, ggml_mul_mat(ctx, weight(blk + "ffn_down.weight"), ffn), + weight(blk + "ffn_down.bias")); + cur = ggml_add(ctx, cur, ffn); + cur = layerNorm(ctx, cur, weight(blk + "ln2.weight"), + weight(blk + "ln2.bias"), hp.layerNormEps); + tap(cur, ("enc_layer_" + std::to_string(il)).c_str()); + } + + // ---- LAM head ----------------------------------------------------------- + cur = ggml_add(ctx, ggml_mul_mat(ctx, weight("head.proj.weight"), cur), + weight("head.proj.bias")); // [512, t] + tap(cur, "lam_proj"); + + // identity one-hot → 64-dim feature, broadcast over time + struct ggml_tensor* idFeat = ggml_mul_mat( + ctx, + ggml_reshape_2d(ctx, weight("head.id_mlp.weight"), hp.nIdentity, + hp.identityFeatDim), + idIn); // [64, 1] + idFeat = ggml_add(ctx, idFeat, + ggml_reshape_2d(ctx, weight("head.id_mlp.bias"), + hp.identityFeatDim, 1)); + idFeat = ggml_repeat( + ctx, idFeat, + ggml_new_tensor_2d(ctx, GGML_TYPE_F32, hp.identityFeatDim, frames)); + + // features live on ne0 in this layout → concat on dim 0: [576, t] + cur = ggml_concat(ctx, cur, idFeat, 0); + + // ConvNormRelu (k=3, s=1, p=1) with LayerNorm over channels. + // residualMode: 0 none, 1 identity, 2 conv (head.first0.res). + const auto convNormRelu = [&](struct ggml_tensor* x, const std::string& name, + int residualMode) { + struct ggml_tensor* xt = ggml_cont(ctx, ggml_transpose(ctx, x)); // [t, c] + struct ggml_tensor* out = + conv1d(ctx, weight(name + ".conv.weight"), xt, 1, 1); + out = addBiasTimeMajor(ctx, out, weight(name + ".conv.bias")); + out = ggml_cont(ctx, ggml_transpose(ctx, out)); // [c, t] + out = layerNorm(ctx, out, weight(name + ".ln.weight"), + weight(name + ".ln.bias"), hp.layerNormEps); + if (residualMode == 1) { + out = ggml_add(ctx, out, x); + } else if (residualMode == 2) { + struct ggml_tensor* res = + conv1d(ctx, weight("head.first0.res.weight"), xt, 1, 1); + res = addBiasTimeMajor(ctx, res, weight("head.first0.res.bias")); + out = ggml_add(ctx, out, ggml_cont(ctx, ggml_transpose(ctx, res))); + } + return ggml_relu(ctx, out); + }; + + cur = tap(convNormRelu(cur, "head.first0", 2), "ident_cnr_0"); + cur = tap(convNormRelu(cur, "head.first1", 1), "ident_cnr_1"); + cur = tap(convNormRelu(cur, "head.first2", 1), "ident_cnr_2"); + cur = tap(convNormRelu(cur, "head.dec0", 0), "dec_cnr_0"); + cur = tap(convNormRelu(cur, "head.dec1", 0), "dec_cnr_1"); + cur = tap(convNormRelu(cur, "head.dec2", 0), "dec_cnr_2"); + + cur = ggml_add(ctx, ggml_mul_mat(ctx, weight("head.out.weight"), cur), + weight("head.out.bias")); // [52, t] + tap(cur, "expr_logits"); + struct ggml_tensor* expr = ggml_sigmoid(ctx, cur); + ggml_set_name(expr, "expr"); + ggml_set_output(expr); + ggml_build_forward_expand(graph, expr); + + // ---- allocate, set inputs, compute ------------------------------------- + ggml_gallocr_t alloc = + ggml_gallocr_new(ggml_backend_get_default_buffer_type(backend_)); + if (!ggml_gallocr_alloc_graph(alloc, graph)) { + lastError_ = "failed to allocate compute graph"; + ggml_gallocr_free(alloc); + ggml_free(ctx); + return false; + } + + ggml_backend_tensor_set(pcmIn, pcm.data(), 0, pcm.size() * sizeof(float)); + + std::vector idOnehot(hp.nIdentity, 0.0F); + idOnehot[idIdx] = 1.0F; + ggml_backend_tensor_set(idIn, idOnehot.data(), 0, + idOnehot.size() * sizeof(float)); + + // align_corners=true linear interpolation weights + std::vector interp(static_cast(tConv) * frames, 0.0F); + const double scale = + static_cast(tConv - 1) / static_cast(frames - 1); + for (int64_t j = 0; j < frames; ++j) { + const double pos = static_cast(j) * scale; + const auto i0 = static_cast(pos); + const int64_t i1 = i0 + 1 < tConv ? i0 + 1 : tConv - 1; + const auto w = static_cast(pos - static_cast(i0)); + interp[j * tConv + i0] += 1.0F - w; + interp[j * tConv + i1] += w; + } + ggml_backend_tensor_set(interpW, interp.data(), 0, + interp.size() * sizeof(float)); + + const ggml_status status = ggml_backend_graph_compute(backend_, graph); + if (status != GGML_STATUS_SUCCESS) { + lastError_ = "graph compute failed"; + ggml_gallocr_free(alloc); + ggml_free(ctx); + return false; + } + + const auto readTensor = [&](struct ggml_tensor* t) { + std::vector out(ggml_nelements(t)); + ggml_backend_tensor_get(t, out.data(), 0, out.size() * sizeof(float)); + return out; + }; + + // expr is [52, frames] with ne0 contiguous → already frame-major. + framesOut = readTensor(expr); + + if (taps != nullptr) { + for (int i = 0; i < ggml_graph_n_nodes(graph); ++i) { + struct ggml_tensor* node = ggml_graph_node(graph, i); + const std::string name = ggml_get_name(node); + if (name.rfind("fe_out", 0) == 0 || name.rfind("interp_out", 0) == 0 || + name.rfind("fp_out", 0) == 0 || name.rfind("pos_conv_out", 0) == 0 || + name.rfind("enc_pre_ln", 0) == 0 || + name.rfind("enc_layer_", 0) == 0 || name.rfind("lam_proj", 0) == 0 || + name.rfind("ident_cnr_", 0) == 0 || name.rfind("dec_cnr_", 0) == 0 || + name.rfind("expr_logits", 0) == 0 || name == "expr") { + (*taps)[name] = readTensor(node); + } + } + } + + ggml_gallocr_free(alloc); + ggml_free(ctx); + return true; +} + diff --git a/src/lam_audio2expression.hpp b/src/lam_audio2expression.hpp new file mode 100644 index 000000000..93efe7520 --- /dev/null +++ b/src/lam_audio2expression.hpp @@ -0,0 +1,70 @@ +#pragma once + +// LAM Audio2Expression on ggml — audio (16 kHz f32 PCM) → ARKit-52 +// blendshape coefficients at 30 fps. +// +// Ported from packages/lipsync-ggml into qvac-ext-stable-diffusion.cpp as an +// isolated inference target (no sd_ctx_t / diffusion lifecycle). + +#include +#include +#include +#include + +#include "ggml-backend.h" +#include "ggml.h" + +struct LamHParams { + uint32_t sampleRate = 16000; + uint32_t fps = 30; + uint32_t nCoeffs = 52; + uint32_t nIdentity = 12; + uint32_t identityFeatDim = 64; + uint32_t hiddenDim = 512; + uint32_t windowFrames = 64; + float layerNormEps = 1e-5F; + uint32_t encLayers = 12; + uint32_t encHeads = 12; + uint32_t encHidden = 768; + uint32_t encFfn = 3072; + uint32_t posConvKernel = 128; + uint32_t posConvGroups = 16; + std::vector feKernels{10, 3, 3, 3, 3, 2, 2}; + std::vector feStrides{5, 2, 2, 2, 2, 2, 2}; + std::vector coeffNames; +}; + +class LamAudio2Expression { +public: + LamAudio2Expression() = default; + ~LamAudio2Expression(); + + LamAudio2Expression(const LamAudio2Expression&) = delete; + LamAudio2Expression& operator=(const LamAudio2Expression&) = delete; + LamAudio2Expression(LamAudio2Expression&&) = delete; + LamAudio2Expression& operator=(LamAudio2Expression&&) = delete; + + bool load(const std::string& ggufPath, ggml_backend_t backend = nullptr, int n_threads = 0); + + [[nodiscard]] int64_t frameCount(int64_t nSamples) const; + [[nodiscard]] int64_t convOutLen(int64_t nSamples) const; + + bool run(const std::vector& pcm, uint32_t idIdx, + std::vector& framesOut, + std::map>* taps = nullptr); + + [[nodiscard]] const LamHParams& hparams() const { return hparams_; } + [[nodiscard]] const std::string& lastError() const { return lastError_; } + +private: + struct ggml_tensor* weight(const std::string& name); + + LamHParams hparams_; + std::string lastError_; + + ggml_backend_t backend_ = nullptr; + bool ownsBackend_ = false; + ggml_backend_buffer_t weightBuffer_ = nullptr; + struct ggml_context* weightCtx_ = nullptr; + std::map weights_; +}; diff --git a/src/lam_wav2vec_frontend.hpp b/src/lam_wav2vec_frontend.hpp new file mode 100644 index 000000000..dfe64887e --- /dev/null +++ b/src/lam_wav2vec_frontend.hpp @@ -0,0 +1,348 @@ +#pragma once + +#include +#include +#include +#include +#include + +#include "ggml_extend.hpp" +#include "model.h" + +class LamWav2VecFrontend final : public GGMLRunner { +public: + static constexpr int kLayers = 7; + static constexpr int kChannels = 512; + + LamWav2VecFrontend(ggml_backend_t backend, + ggml_backend_t params_backend, + const String2TensorStorage& tensor_storage_map) + : GGMLRunner(backend, params_backend) { + static constexpr std::array kernels = {10, 3, 3, 3, 3, 2, 2}; + static constexpr std::array input_channels = {1, 512, 512, 512, 512, 512, 512}; + + for (int i = 0; i < kLayers; ++i) { + const std::string name = + "backbone.audio_encoder.feature_extractor.conv_layers." + + std::to_string(i) + ".conv.weight"; + const auto found = tensor_storage_map.find(name); + const ggml_type type = + found == tensor_storage_map.end() ? GGML_TYPE_F32 : found->second.type; + conv_weights_[i] = ggml_new_tensor_3d( + params_ctx, type, kernels[i], input_channels[i], kChannels); + ggml_set_name(conv_weights_[i], name.c_str()); + } + + const auto norm_weight_name = + "backbone.audio_encoder.feature_extractor.conv_layers.0.layer_norm.weight"; + const auto norm_bias_name = + "backbone.audio_encoder.feature_extractor.conv_layers.0.layer_norm.bias"; + norm_weight_ = ggml_new_tensor_1d(params_ctx, GGML_TYPE_F32, kChannels); + norm_bias_ = ggml_new_tensor_1d(params_ctx, GGML_TYPE_F32, kChannels); + ggml_set_name(norm_weight_, norm_weight_name); + ggml_set_name(norm_bias_, norm_bias_name); + + feature_norm_weight_ = ggml_new_tensor_1d(params_ctx, GGML_TYPE_F32, kChannels); + feature_norm_bias_ = ggml_new_tensor_1d(params_ctx, GGML_TYPE_F32, kChannels); + feature_projection_weight_ = + ggml_new_tensor_2d(params_ctx, GGML_TYPE_F32, kChannels, 768); + feature_projection_bias_ = ggml_new_tensor_1d(params_ctx, GGML_TYPE_F32, 768); + ggml_set_name(feature_norm_weight_, + "backbone.audio_encoder.feature_projection.layer_norm.weight"); + ggml_set_name(feature_norm_bias_, + "backbone.audio_encoder.feature_projection.layer_norm.bias"); + ggml_set_name(feature_projection_weight_, + "backbone.audio_encoder.feature_projection.projection.weight"); + ggml_set_name(feature_projection_bias_, + "backbone.audio_encoder.feature_projection.projection.bias"); + + pos_weight_g_ = ggml_new_tensor_3d(params_ctx, GGML_TYPE_F32, 128, 1, 1); + pos_weight_v_ = ggml_new_tensor_3d(params_ctx, GGML_TYPE_F32, 128, 48, 768); + pos_bias_ = ggml_new_tensor_1d(params_ctx, GGML_TYPE_F32, 768); + ggml_set_name(pos_weight_g_, + "backbone.audio_encoder.encoder.pos_conv_embed.conv.weight_g"); + ggml_set_name(pos_weight_v_, + "backbone.audio_encoder.encoder.pos_conv_embed.conv.weight_v"); + ggml_set_name(pos_bias_, + "backbone.audio_encoder.encoder.pos_conv_embed.conv.bias"); + } + + std::string get_desc() override { + return "lam-a2e-wav2vec-frontend"; + } + + void get_param_tensors(std::map& tensors) { + for (int i = 0; i < kLayers; ++i) { + tensors[ggml_get_name(conv_weights_[i])] = conv_weights_[i]; + } + tensors[ggml_get_name(norm_weight_)] = norm_weight_; + tensors[ggml_get_name(norm_bias_)] = norm_bias_; + tensors[ggml_get_name(feature_norm_weight_)] = feature_norm_weight_; + tensors[ggml_get_name(feature_norm_bias_)] = feature_norm_bias_; + tensors[ggml_get_name(feature_projection_weight_)] = feature_projection_weight_; + tensors[ggml_get_name(feature_projection_bias_)] = feature_projection_bias_; + tensors[ggml_get_name(pos_weight_g_)] = pos_weight_g_; + tensors[ggml_get_name(pos_weight_v_)] = pos_weight_v_; + tensors[ggml_get_name(pos_bias_)] = pos_bias_; + } + + bool load(ModelLoader& loader, int n_threads) { + if (!alloc_params_buffer()) { + return false; + } + std::map tensors; + get_param_tensors(tensors); + return loader.load_tensors(tensors, {}, n_threads, false); + } + + ggml_tensor* forward_frontend(ggml_tensor* x) { + static constexpr std::array strides = {5, 2, 2, 2, 2, 2, 2}; + for (int i = 0; i < kLayers; ++i) { + x = ggml_conv_1d(compute_ctx, conv_weights_[i], x, strides[i], 0, 1); + if (i == 0) { + const int64_t frames = x->ne[0]; + const int64_t channels = x->ne[1]; + const int64_t batch = x->ne[2]; + // GGML GroupNorm treats ne[2] as channels. Convert from the + // frontend's [time, channels, batch] layout to + // [1, time, channels, batch] for PyTorch-equivalent + // GroupNorm(channels, channels), then restore the layout. + auto* grouped = + ggml_reshape_4d(compute_ctx, x, 1, frames, channels, batch); + grouped = ggml_group_norm(compute_ctx, grouped, kChannels, 1e-5f); + x = ggml_reshape_3d(compute_ctx, grouped, frames, channels, batch); + auto* norm_weight = ggml_reshape_3d( + compute_ctx, norm_weight_, 1, kChannels, 1); + auto* norm_bias = ggml_reshape_3d( + compute_ctx, norm_bias_, 1, kChannels, 1); + x = ggml_add( + compute_ctx, + ggml_mul(compute_ctx, x, ggml_repeat(compute_ctx, norm_weight, x)), + ggml_repeat(compute_ctx, norm_bias, x)); + } + // HuggingFace Wav2Vec2 `hidden_act="gelu"` uses the exact + // error-function form, not the tanh approximation. + x = ggml_gelu_erf(compute_ctx, x); + } + return x; + } + + ggml_tensor* interpolate_frontend(ggml_tensor* x) { + // PyTorch F.interpolate(..., mode="linear", align_corners=True). + // A precomputed interpolation matrix avoids backend-specific image + // interpolation semantics and preserves the exact audio time mapping. + const int64_t input_frames = x->ne[0]; + constexpr int64_t output_frames = 30; + interpolation_weights_.assign( + static_cast(input_frames * output_frames), 0.0f); + + for (int64_t output_index = 0; output_index < output_frames; ++output_index) { + const float source = + static_cast(output_index) * static_cast(input_frames - 1) / + static_cast(output_frames - 1); + const int64_t lower = static_cast(std::floor(source)); + const int64_t upper = std::min(lower + 1, input_frames - 1); + const float upper_weight = source - static_cast(lower); + const size_t column_offset = static_cast(input_frames * output_index); + interpolation_weights_[column_offset + static_cast(lower)] += + 1.0f - upper_weight; + interpolation_weights_[column_offset + static_cast(upper)] += + upper_weight; + } + + auto* weights = ggml_new_tensor_2d( + compute_ctx, GGML_TYPE_F32, input_frames, output_frames); + set_backend_tensor_data(weights, interpolation_weights_.data()); + return ggml_mul_mat(compute_ctx, weights, x); + } + + ggml_tensor* normalize_features(ggml_tensor* x) { + // PyTorch receives [batch, time, channel] and normalizes channel. + // GGML normalizes ne0, so transpose to [channel, time, batch] first. + x = ggml_cont(compute_ctx, ggml_transpose(compute_ctx, x)); + x = ggml_norm(compute_ctx, x, 1e-5f); + auto* norm_weight = ggml_reshape_4d(compute_ctx, feature_norm_weight_, kChannels, 1, 1, 1); + auto* norm_bias = ggml_reshape_4d(compute_ctx, feature_norm_bias_, kChannels, 1, 1, 1); + x = ggml_add( + compute_ctx, + ggml_mul(compute_ctx, x, ggml_repeat(compute_ctx, norm_weight, x)), + ggml_repeat(compute_ctx, norm_bias, x)); + return ggml_cont(compute_ctx, x); + } + + ggml_tensor* forward_projection(ggml_tensor* x) { + x = interpolate_frontend(x); + x = normalize_features(x); + return project_normalized(x); + } + + ggml_tensor* project_normalized(ggml_tensor* x) { + return ggml_ext_linear( + compute_ctx, + x, + feature_projection_weight_, + feature_projection_bias_, + false, + 1.0f); + } + + ggml_cgraph* build_graph(const sd::Tensor& pcm, int stage) { + ggml_cgraph* graph = ggml_new_graph(compute_ctx); + ggml_tensor* x = forward_frontend(make_input(pcm)); + if (stage == 1) { + x = interpolate_frontend(x); + } else if (stage == 2) { + x = normalize_features(interpolate_frontend(x)); + } else if (stage >= 3) { + x = forward_projection(x); + } + + ggml_build_forward_expand(graph, x); + return graph; + } + + sd::Tensor compute_frontend(const sd::Tensor& pcm, int n_threads) { + auto build = [&]() { return build_graph(pcm, 0); }; + return restore_trailing_singleton_dims( + GGMLRunner::compute(build, n_threads, true), 3); + } + + sd::Tensor compute_projection(const sd::Tensor& pcm, int n_threads) { + const auto normalized = compute_normalized_cpu_reference(pcm, n_threads); + auto build = [&]() { + ggml_cgraph* graph = ggml_new_graph(compute_ctx); + ggml_tensor* x = project_normalized(make_input(normalized)); + ggml_build_forward_expand(graph, x); + return graph; + }; + return restore_trailing_singleton_dims( + GGMLRunner::compute(build, n_threads, true), 3); + } + + sd::Tensor compute_interpolated(const sd::Tensor& pcm, int n_threads) { + auto build = [&]() { return build_graph(pcm, 1); }; + return restore_trailing_singleton_dims( + GGMLRunner::compute(build, n_threads, true), 4); + } + + sd::Tensor compute_normalized(const sd::Tensor& pcm, int n_threads) { + auto build = [&]() { return build_graph(pcm, 2); }; + return restore_trailing_singleton_dims( + GGMLRunner::compute(build, n_threads, true), 4); + } + + sd::Tensor compute_normalized_cpu_reference( + const sd::Tensor& pcm, + int n_threads) { + const auto frontend = compute_frontend(pcm, n_threads); + const int64_t input_frames = frontend.shape()[0]; + constexpr int64_t output_frames = 30; + sd::Tensor normalized({kChannels, output_frames, 1}); + + const auto* gamma = static_cast(ggml_get_data(feature_norm_weight_)); + const auto* beta = static_cast(ggml_get_data(feature_norm_bias_)); + for (int64_t output_index = 0; output_index < output_frames; ++output_index) { + const float source = + static_cast(output_index) * static_cast(input_frames - 1) / + static_cast(output_frames - 1); + const int64_t lower = static_cast(std::floor(source)); + const int64_t upper = std::min(lower + 1, input_frames - 1); + const float upper_weight = source - static_cast(lower); + + float mean = 0.0f; + for (int64_t channel = 0; channel < kChannels; ++channel) { + const float value = + frontend.values()[lower + input_frames * channel] * (1.0f - upper_weight) + + frontend.values()[upper + input_frames * channel] * upper_weight; + normalized.values()[channel + kChannels * output_index] = value; + mean += value; + } + mean /= kChannels; + + float variance = 0.0f; + for (int64_t channel = 0; channel < kChannels; ++channel) { + const float delta = + normalized.values()[channel + kChannels * output_index] - mean; + variance += delta * delta; + } + const float inverse_std = 1.0f / std::sqrt(variance / kChannels + 1e-5f); + for (int64_t channel = 0; channel < kChannels; ++channel) { + const float value = + normalized.values()[channel + kChannels * output_index]; + normalized.values()[channel + kChannels * output_index] = + (value - mean) * inverse_std * gamma[channel] + beta[channel]; + } + } + return normalized; + } + + sd::Tensor compute_position_cpu_reference( + const sd::Tensor& pcm, + int n_threads) { + const auto projected = compute_projection(pcm, n_threads); + constexpr int64_t channels = 768; + constexpr int64_t groups = 16; + constexpr int64_t channels_per_group = channels / groups; + constexpr int64_t kernel = 128; + constexpr int64_t padding = kernel / 2; + constexpr int64_t frames = 30; + + const auto* weight_g = static_cast(ggml_get_data(pos_weight_g_)); + const auto* weight_v = static_cast(ggml_get_data(pos_weight_v_)); + const auto* bias = static_cast(ggml_get_data(pos_bias_)); + sd::Tensor output({channels, frames, 1}); + + std::array norms{}; + for (int64_t kernel_index = 0; kernel_index < kernel; ++kernel_index) { + double sum = 0.0; + for (int64_t output_channel = 0; output_channel < channels; ++output_channel) { + for (int64_t input_channel = 0; input_channel < channels_per_group; ++input_channel) { + const size_t offset = static_cast( + kernel_index + kernel * (input_channel + channels_per_group * output_channel)); + sum += static_cast(weight_v[offset]) * weight_v[offset]; + } + } + norms[kernel_index] = static_cast(std::sqrt(sum)); + } + + for (int64_t output_channel = 0; output_channel < channels; ++output_channel) { + const int64_t group = output_channel / channels_per_group; + for (int64_t frame = 0; frame < frames; ++frame) { + float value = bias[output_channel]; + for (int64_t input_channel = 0; input_channel < channels_per_group; ++input_channel) { + const int64_t source_channel = group * channels_per_group + input_channel; + for (int64_t kernel_index = 0; kernel_index < kernel; ++kernel_index) { + const int64_t source_frame = frame + kernel_index - padding; + if (source_frame < 0 || source_frame >= frames) { + continue; + } + const size_t weight_offset = static_cast( + kernel_index + kernel * (input_channel + channels_per_group * output_channel)); + const float weight = + weight_v[weight_offset] * weight_g[kernel_index] / norms[kernel_index]; + value += weight * + projected.values()[source_channel + channels * source_frame]; + } + } + // Wav2Vec2SamePadLayer drops the final padded frame. This + // loop emits the retained 30 frames directly. + output.values()[output_channel + channels * frame] = + 0.5f * value * (1.0f + std::erf(value * 0.70710678118f)); + } + } + return output; + } + +private: + std::array conv_weights_{}; + ggml_tensor* norm_weight_ = nullptr; + ggml_tensor* norm_bias_ = nullptr; + ggml_tensor* feature_norm_weight_ = nullptr; + ggml_tensor* feature_norm_bias_ = nullptr; + ggml_tensor* feature_projection_weight_ = nullptr; + ggml_tensor* feature_projection_bias_ = nullptr; + ggml_tensor* pos_weight_g_ = nullptr; + ggml_tensor* pos_weight_v_ = nullptr; + ggml_tensor* pos_bias_ = nullptr; + std::vector interpolation_weights_; +}; diff --git a/tests/lam_a2e_frontend_smoke.cpp b/tests/lam_a2e_frontend_smoke.cpp new file mode 100644 index 000000000..257f2a400 --- /dev/null +++ b/tests/lam_a2e_frontend_smoke.cpp @@ -0,0 +1,90 @@ +#include +#include +#include +#include + +#include "ggml-cpu.h" +#include "lam_wav2vec_frontend.hpp" +#include "model.h" + +namespace { + +bool read_f32_file(const char* path, std::vector* values) { + std::ifstream input(path, std::ios::binary | std::ios::ate); + if (!input.is_open()) { + return false; + } + const auto bytes = input.tellg(); + if (bytes <= 0 || bytes % static_cast(sizeof(float)) != 0) { + return false; + } + values->resize(static_cast(bytes) / sizeof(float)); + input.seekg(0); + input.read(reinterpret_cast(values->data()), bytes); + return static_cast(input); +} + +bool write_f32_file(const char* path, const std::vector& values) { + std::ofstream output(path, std::ios::binary); + output.write(reinterpret_cast(values.data()), + static_cast(values.size() * sizeof(float))); + return static_cast(output); +} + +} // namespace + +int main(int argc, char** argv) { + if (argc != 4 && argc != 5) { + std::fprintf(stderr, + "usage: lam-a2e-frontend-smoke [frontend|interpolated|norm|projection|position]\n"); + return 2; + } + + std::vector pcm; + if (!read_f32_file(argv[2], &pcm)) { + std::fprintf(stderr, "failed to read PCM fixture\n"); + return 1; + } + + ModelLoader loader; + if (!loader.init_from_file(argv[1])) { + std::fprintf(stderr, "failed to read GGUF metadata\n"); + return 1; + } + loader.process_model_files(false); + + ggml_backend_t backend = ggml_backend_cpu_init(); + LamWav2VecFrontend frontend(backend, backend, loader.get_tensor_storage_map()); + if (!frontend.load(loader, 4)) { + std::fprintf(stderr, "failed to load frontend parameters\n"); + ggml_backend_free(backend); + return 1; + } + + const int64_t sample_count = static_cast(pcm.size()); + sd::Tensor input({sample_count, 1, 1}, std::move(pcm)); + const std::string stage = argc == 5 ? argv[4] : "frontend"; + const auto output = stage == "projection" + ? frontend.compute_projection(input, 4) + : stage == "position" + ? frontend.compute_position_cpu_reference(input, 4) + : stage == "norm" + ? frontend.compute_normalized(input, 4) + : stage == "interpolated" + ? frontend.compute_interpolated(input, 4) + : frontend.compute_frontend(input, 4); + if (!write_f32_file(argv[3], output.values())) { + std::fprintf(stderr, "failed to write frontend output\n"); + ggml_backend_free(backend); + return 1; + } + + const auto& shape = output.shape(); + std::printf("frontend output shape:"); + for (const auto dim : shape) { + std::printf(" %lld", static_cast(dim)); + } + std::printf("\n"); + ggml_backend_free(backend); + return 0; +} diff --git a/tests/lam_a2e_parity.cpp b/tests/lam_a2e_parity.cpp new file mode 100644 index 000000000..67bc65677 --- /dev/null +++ b/tests/lam_a2e_parity.cpp @@ -0,0 +1,111 @@ +#include +#include +#include +#include +#include +#include +#include + +#include "lam_audio2expression.hpp" + +static bool read_f32_bin(const std::string& path, std::vector& out) { + std::ifstream in(path, std::ios::binary); + if (!in) { + return false; + } + in.seekg(0, std::ios::end); + const auto bytes = static_cast(in.tellg()); + in.seekg(0, std::ios::beg); + if (bytes % sizeof(float) != 0) { + return false; + } + out.resize(bytes / sizeof(float)); + in.read(reinterpret_cast(out.data()), static_cast(bytes)); + return static_cast(in); +} + +static float max_abs_diff(const std::vector& a, const std::vector& b) { + if (a.size() != b.size() || a.empty()) { + return INFINITY; + } + float m = 0.0F; + for (size_t i = 0; i < a.size(); ++i) { + m = std::max(m, std::fabs(a[i] - b[i])); + } + return m; +} + +int main(int argc, char** argv) { + if (argc != 4) { + std::fprintf(stderr, + "usage: lam-a2e-parity \n"); + return 2; + } + + const std::string model = argv[1]; + const std::string dir = argv[2]; + const std::string prefix = argv[3]; + const float tol = 1e-3F; + + LamAudio2Expression model_impl; + if (!model_impl.load(model)) { + std::fprintf(stderr, "load failed: %s\n", model_impl.lastError().c_str()); + return 1; + } + + std::vector pcm; + std::vector ref_expr; + if (!read_f32_bin(dir + "/" + prefix + "_input_pcm.bin", pcm) || + !read_f32_bin(dir + "/" + prefix + "_expr.bin", ref_expr)) { + std::fprintf(stderr, "failed to read fixture bins under %s\n", dir.c_str()); + return 1; + } + + std::vector out; + std::map> taps; + if (!model_impl.run(pcm, /*idIdx*/ 0, out, &taps)) { + std::fprintf(stderr, "run failed: %s\n", model_impl.lastError().c_str()); + return 1; + } + + // Critical path stages with matching layouts in the lipsync-ggml fixtures. + const char* gated[] = { + "fe_out", "fp_out", "pos_conv_out", "enc_pre_ln", + "enc_layer_0", "enc_layer_11", "expr" + }; + + bool ok = true; + for (const char* stage : gated) { + std::vector ref; + const std::string ref_path = dir + "/" + prefix + "_" + stage + ".bin"; + if (!read_f32_bin(ref_path, ref)) { + std::printf("skip missing stage %s\n", stage); + continue; + } + const std::vector* got = (std::string(stage) == "expr") ? &out : nullptr; + std::vector local; + if (got == nullptr) { + auto it = taps.find(stage); + if (it == taps.end()) { + std::printf("missing tap %s\n", stage); + ok = false; + continue; + } + local = it->second; + got = &local; + } + const float err = max_abs_diff(*got, ref); + std::printf("%s max_abs_diff=%.6g n=%zu\n", stage, err, got->size()); + if (!(err <= tol)) { + ok = false; + } + } + + if (!ok) { + std::fprintf(stderr, "parity failed\n"); + return 1; + } + + std::printf("LAM-A2E parity passed for case %s\n", prefix.c_str()); + return 0; +} diff --git a/tests/lam_a2e_smoke.cpp b/tests/lam_a2e_smoke.cpp new file mode 100644 index 000000000..ed398265d --- /dev/null +++ b/tests/lam_a2e_smoke.cpp @@ -0,0 +1,51 @@ +#include +#include +#include +#include + +#include "lam-a2e.h" + +int main(int argc, char** argv) { + if (argc != 2) { + std::fprintf(stderr, "usage: lam-a2e-smoke \n"); + return 2; + } + + lam_a2e_params params{}; + params.model_path = argv[1]; + params.identity_index = 0; + params.n_threads = 4; + params.use_gpu = false; + + lam_a2e_context* ctx = lam_a2e_create(¶ms); + if (ctx == nullptr) { + std::fprintf(stderr, "failed to create LAM-A2E context\n"); + return 1; + } + + std::array pcm{}; + for (size_t i = 0; i < pcm.size(); ++i) { + pcm[i] = 0.05F * std::sin(2.0F * 3.14159265F * 220.0F * + static_cast(i) / 16000.0F); + } + + lam_a2e_frame* frames = nullptr; + int32_t frame_count = 0; + const auto status = lam_a2e_process_pcm_f32( + ctx, pcm.data(), static_cast(pcm.size()), 16000, &frames, &frame_count); + + if (status != LAM_A2E_STATUS_OK || frames == nullptr || frame_count < 2) { + std::fprintf(stderr, "LAM-A2E smoke failed (%d): %s\n", + static_cast(status), lam_a2e_get_last_error(ctx)); + lam_a2e_free(ctx); + return 1; + } + + std::printf("LAM-A2E smoke passed: frames=%d ts0=%lld coeff0=%.6f\n", + frame_count, + static_cast(frames[0].timestamp_us), + frames[0].arkit_52[0]); + lam_a2e_free_frames(frames); + lam_a2e_free(ctx); + return 0; +} From f2daa4c51dedb300ad2c724261b7de93be1dd0b3 Mon Sep 17 00:00:00 2001 From: Alok-Ranjan23 Date: Fri, 24 Jul 2026 16:50:31 +0000 Subject: [PATCH 2/3] infra: add CI job compiling LAM-A2E smoke/parity harness The standard build jobs compile the LAM audio2expression sources via the SD library glob, but the smoke and CPU-parity test targets are only built under -DSD_BUILD_LAM_A2E_SMOKE=ON and were therefore uncovered. Add a PR-triggered ubuntu job that builds lam-a2e-smoke, lam-a2e-frontend-smoke and lam-a2e-parity so regressions in the C API / graph / test harness are caught. Executing parity is deferred until the remapped GGUF and lipsync-ggml fixtures are available in CI (QVAC-22250 / QVAC-22249). --- .github/workflows/build.yml | 35 +++++++++++++++++++++++++++++++++++ 1 file changed, 35 insertions(+) diff --git a/.github/workflows/build.yml b/.github/workflows/build.yml index df4c07d45..195e3ecbd 100644 --- a/.github/workflows/build.yml +++ b/.github/workflows/build.yml @@ -178,6 +178,41 @@ jobs: path: | sd-${{ env.BRANCH_NAME }}-${{ steps.commit.outputs.short }}-bin-${{ steps.system-info.outputs.OS_TYPE }}-${{ steps.system-info.outputs.OS_NAME }}-${{ steps.system-info.outputs.OS_VERSION }}-${{ steps.system-info.outputs.CPU_ARCH }}-vulkan.zip + ubuntu-latest-lam-a2e: + # Compile-checks the isolated LAM audio2expression C API and its smoke / + # CPU-parity harness (only built under -DSD_BUILD_LAM_A2E_SMOKE=ON, so the + # standard build jobs above don't cover the test targets). This guards + # against regressions in include/lam-a2e.h, src/lam_audio2expression.* and + # tests/lam_a2e_*.cpp. + # + # NOTE: this job only *builds* the harness. Running lam-a2e-smoke / + # lam-a2e-parity needs the remapped lam-audio2exp GGUF plus the + # lipsync-ggml reference fixtures, neither of which is available in CI yet. + # Wire the execution step in once those assets are hosted (tracked by + # QVAC-22250 / QVAC-22249). + runs-on: ubuntu-latest + + steps: + - name: Clone + id: checkout + uses: actions/checkout@v3 + with: + submodules: recursive + + - name: Dependencies + id: depends + run: | + sudo apt-get update + sudo apt-get install build-essential + + - name: Build LAM-A2E smoke + parity targets + id: cmake_build + run: | + mkdir build + cd build + cmake .. -DGGML_AVX2=ON -DSD_BUILD_EXAMPLES=OFF -DSD_BUILD_LAM_A2E_SMOKE=ON + cmake --build . --config Release --target lam-a2e-smoke lam-a2e-frontend-smoke lam-a2e-parity + build-and-push-docker-images: name: Build and push container images if: ${{ github.event_name != 'pull_request' }} From bf8adccf6493098fd99de563281c0f89174065ec Mon Sep 17 00:00:00 2001 From: Alok-Ranjan23 Date: Mon, 27 Jul 2026 15:34:31 +0000 Subject: [PATCH 3/3] refactor: move GGUF tooling to the consuming package The Python conversion and fixture-dump scripts sat in the engine repo, but nothing here runs them: the engine only ever reads the resulting .gguf. Every comparable package (ocr-ggml and friends) keeps its conversion tooling beside the code that consumes it, so these move to qvac packages/diffusion-cpp/scripts/ where the venv, the npm scripts, and the docs already live. Keeps the engine a pure C/C++ artifact with no Python surface to maintain. --- .github/workflows/build.yml | 3 + scripts/convert_lam_a2e_to_gguf.py | 207 --------------------- scripts/dump_lam_a2e_frontend_reference.py | 65 ------- scripts/dump_lam_a2e_stages.py | 138 -------------- scripts/remap_lam_gguf.py | 174 ----------------- 5 files changed, 3 insertions(+), 584 deletions(-) delete mode 100644 scripts/convert_lam_a2e_to_gguf.py delete mode 100644 scripts/dump_lam_a2e_frontend_reference.py delete mode 100644 scripts/dump_lam_a2e_stages.py delete mode 100644 scripts/remap_lam_gguf.py diff --git a/.github/workflows/build.yml b/.github/workflows/build.yml index 195e3ecbd..c73bc2306 100644 --- a/.github/workflows/build.yml +++ b/.github/workflows/build.yml @@ -190,6 +190,9 @@ jobs: # lipsync-ggml reference fixtures, neither of which is available in CI yet. # Wire the execution step in once those assets are hosted (tracked by # QVAC-22250 / QVAC-22249). + # + # The GGUF conversion and fixture-dump tooling lives with the consumer, in + # qvac packages/diffusion-cpp/scripts/ (see README-lam-a2e.md there). runs-on: ubuntu-latest steps: diff --git a/scripts/convert_lam_a2e_to_gguf.py b/scripts/convert_lam_a2e_to_gguf.py deleted file mode 100644 index 7102c4c3e..000000000 --- a/scripts/convert_lam_a2e_to_gguf.py +++ /dev/null @@ -1,207 +0,0 @@ -#!/usr/bin/env python3 -"""LAM Audio2Expression PyTorch checkpoint → GGUF converter. - -Reads the upstream ``lam_audio2exp_streaming.tar`` checkpoint (Apache-2.0, -https://github.com/aigc3d/LAM_Audio2Expression) and emits a single GGUF -consumed by the isolated LAM-A2E target in qvac-ext-stable-diffusion.cpp. - -Layout notes: - -- ``general.architecture = "lam-audio2exp"`` — selected by the addon's - model factory. -- PyTorch conv1d weights are (C_out, C_in, K) row-major, which lands as - ggml ne = [K, C_in, C_out] on a straight copy — exactly what - ``ggml_conv_1d`` expects. No transposes are performed anywhere; the - converter is a 1:1 map of the state dict. -- The positional-conv weight-norm (weight_g/weight_v, dim=2) is folded - into a plain conv weight at conversion time. -- Inference-unused tensors (``lm_head``, ``identity_encoder.grus``, - ``masked_spec_embed``) are dropped. -- ``--dtype f32`` (default) keeps everything float32; ``--dtype f16`` - stores matmul/conv weights as f16 and keeps norms + biases f32. - -Usage: - python3 scripts/convert_lam_a2e_to_gguf.py \ - --checkpoint pretrained_models/lam_audio2exp_streaming.tar \ - --out lam-audio2exp-f32.gguf --dtype f32 -""" - -from __future__ import annotations - -import argparse -import os - -import numpy as np -import torch - -from gguf import GGUFWriter - -ARCH = "lam-audio2exp" - -ARKIT_BLENDSHAPES = [ - "browDownLeft", "browDownRight", "browInnerUp", "browOuterUpLeft", - "browOuterUpRight", "cheekPuff", "cheekSquintLeft", "cheekSquintRight", - "eyeBlinkLeft", "eyeBlinkRight", "eyeLookDownLeft", "eyeLookDownRight", - "eyeLookInLeft", "eyeLookInRight", "eyeLookOutLeft", "eyeLookOutRight", - "eyeLookUpLeft", "eyeLookUpRight", "eyeSquintLeft", "eyeSquintRight", - "eyeWideLeft", "eyeWideRight", "jawForward", "jawLeft", "jawOpen", - "jawRight", "mouthClose", "mouthDimpleLeft", "mouthDimpleRight", - "mouthFrownLeft", "mouthFrownRight", "mouthFunnel", "mouthLeft", - "mouthLowerDownLeft", "mouthLowerDownRight", "mouthPressLeft", - "mouthPressRight", "mouthPucker", "mouthRight", "mouthRollLower", - "mouthRollUpper", "mouthShrugLower", "mouthShrugUpper", "mouthSmileLeft", - "mouthSmileRight", "mouthStretchLeft", "mouthStretchRight", - "mouthUpperUpLeft", "mouthUpperUpRight", "noseSneerLeft", - "noseSneerRight", "tongueOut", -] - -# wav2vec2-base feature extractor geometry (configs/wav2vec2_config.json) -FE_KERNELS = [10, 3, 3, 3, 3, 2, 2] -FE_STRIDES = [5, 2, 2, 2, 2, 2, 2] - -SKIP_PREFIXES = ( - "audio_encoder.lm_head.", - "identity_encoder.grus.", -) -SKIP_KEYS = ("audio_encoder.masked_spec_embed",) - - -def load_state(checkpoint_path): - ckpt = torch.load(checkpoint_path, map_location="cpu", weights_only=True) - state = {} - for key, value in ckpt["state_dict"].items(): - if key.startswith("module."): - key = key[7:] - if key.startswith("backbone."): - key = key[9:] - if key in SKIP_KEYS or key.startswith(SKIP_PREFIXES): - continue - state[key] = value - return state - - -def fold_pos_conv_weight_norm(state): - """weight = g * v / ||v||, norm over dims (0,1) per kernel position (dim=2).""" - g = state.pop("audio_encoder.encoder.pos_conv_embed.conv.weight_g") # (1,1,128) - v = state.pop("audio_encoder.encoder.pos_conv_embed.conv.weight_v") # (768,48,128) - norm = v.norm(p=2, dim=(0, 1), keepdim=True) - state["audio_encoder.encoder.pos_conv_embed.conv.weight"] = g * v / norm - return state - - -def build_name_map(): - """checkpoint key -> (gguf name, is_matmul_weight).""" - m = {} - for i in range(7): - m[f"audio_encoder.feature_extractor.conv_layers.{i}.conv.weight"] = (f"fe.conv{i}.weight", True) - m["audio_encoder.feature_extractor.conv_layers.0.layer_norm.weight"] = ("fe.gn.weight", False) - m["audio_encoder.feature_extractor.conv_layers.0.layer_norm.bias"] = ("fe.gn.bias", False) - - m["audio_encoder.feature_projection.layer_norm.weight"] = ("fp.ln.weight", False) - m["audio_encoder.feature_projection.layer_norm.bias"] = ("fp.ln.bias", False) - m["audio_encoder.feature_projection.projection.weight"] = ("fp.proj.weight", True) - m["audio_encoder.feature_projection.projection.bias"] = ("fp.proj.bias", False) - - m["audio_encoder.encoder.pos_conv_embed.conv.weight"] = ("enc.pos_conv.weight", True) - m["audio_encoder.encoder.pos_conv_embed.conv.bias"] = ("enc.pos_conv.bias", False) - m["audio_encoder.encoder.layer_norm.weight"] = ("enc.ln.weight", False) - m["audio_encoder.encoder.layer_norm.bias"] = ("enc.ln.bias", False) - - for i in range(12): - src = f"audio_encoder.encoder.layers.{i}" - dst = f"enc.blk{i}" - for proj in ("q", "k", "v"): - m[f"{src}.attention.{proj}_proj.weight"] = (f"{dst}.attn_{proj}.weight", True) - m[f"{src}.attention.{proj}_proj.bias"] = (f"{dst}.attn_{proj}.bias", False) - m[f"{src}.attention.out_proj.weight"] = (f"{dst}.attn_o.weight", True) - m[f"{src}.attention.out_proj.bias"] = (f"{dst}.attn_o.bias", False) - m[f"{src}.layer_norm.weight"] = (f"{dst}.ln1.weight", False) - m[f"{src}.layer_norm.bias"] = (f"{dst}.ln1.bias", False) - m[f"{src}.feed_forward.intermediate_dense.weight"] = (f"{dst}.ffn_up.weight", True) - m[f"{src}.feed_forward.intermediate_dense.bias"] = (f"{dst}.ffn_up.bias", False) - m[f"{src}.feed_forward.output_dense.weight"] = (f"{dst}.ffn_down.weight", True) - m[f"{src}.feed_forward.output_dense.bias"] = (f"{dst}.ffn_down.bias", False) - m[f"{src}.final_layer_norm.weight"] = (f"{dst}.ln2.weight", False) - m[f"{src}.final_layer_norm.bias"] = (f"{dst}.ln2.bias", False) - - m["feature_projection.weight"] = ("head.proj.weight", True) - m["feature_projection.bias"] = ("head.proj.bias", False) - m["identity_encoder.id_mlp.weight"] = ("head.id_mlp.weight", True) - m["identity_encoder.id_mlp.bias"] = ("head.id_mlp.bias", False) - for i in range(3): - src = f"identity_encoder.first_net.conv_layers.{i}" - m[f"{src}.conv.weight"] = (f"head.first{i}.conv.weight", True) - m[f"{src}.conv.bias"] = (f"head.first{i}.conv.bias", False) - m[f"{src}.norm.weight"] = (f"head.first{i}.ln.weight", False) - m[f"{src}.norm.bias"] = (f"head.first{i}.ln.bias", False) - m["identity_encoder.first_net.conv_layers.0.residual_layer.0.weight"] = ("head.first0.res.weight", True) - m["identity_encoder.first_net.conv_layers.0.residual_layer.0.bias"] = ("head.first0.res.bias", False) - for i in range(3): - m[f"decoder.0.{i}.conv.weight"] = (f"head.dec{i}.conv.weight", True) - m[f"decoder.0.{i}.conv.bias"] = (f"head.dec{i}.conv.bias", False) - m[f"decoder.0.{i}.norm.weight"] = (f"head.dec{i}.ln.weight", False) - m[f"decoder.0.{i}.norm.bias"] = (f"head.dec{i}.ln.bias", False) - m["output_proj.weight"] = ("head.out.weight", True) - m["output_proj.bias"] = ("head.out.bias", False) - return m - - -def main(): - parser = argparse.ArgumentParser() - parser.add_argument("--checkpoint", required=True) - parser.add_argument("--out", required=True) - parser.add_argument("--dtype", choices=["f32", "f16"], default="f32") - args = parser.parse_args() - - state = fold_pos_conv_weight_norm(load_state(args.checkpoint)) - name_map = build_name_map() - - unmapped = sorted(set(state) - set(name_map)) - if unmapped: - raise RuntimeError(f"unmapped checkpoint tensors: {unmapped}") - missing = sorted(set(name_map) - set(state)) - if missing: - raise RuntimeError(f"expected checkpoint tensors not found: {missing}") - - os.makedirs(os.path.dirname(os.path.abspath(args.out)), exist_ok=True) - writer = GGUFWriter(args.out, ARCH) - writer.add_name("LAM Audio2Expression (streaming)") - writer.add_string(f"{ARCH}.dtype", args.dtype) - writer.add_uint32(f"{ARCH}.sample_rate", 16000) - writer.add_uint32(f"{ARCH}.fps", 30) - writer.add_uint32(f"{ARCH}.n_coeffs", 52) - writer.add_uint32(f"{ARCH}.n_identity", 12) - writer.add_uint32(f"{ARCH}.identity_feat_dim", 64) - writer.add_uint32(f"{ARCH}.hidden_dim", 512) - writer.add_uint32(f"{ARCH}.window_frames", 64) - writer.add_float32(f"{ARCH}.layer_norm_eps", 1e-5) - writer.add_uint32(f"{ARCH}.enc.n_layers", 12) - writer.add_uint32(f"{ARCH}.enc.n_heads", 12) - writer.add_uint32(f"{ARCH}.enc.hidden", 768) - writer.add_uint32(f"{ARCH}.enc.ffn", 3072) - writer.add_uint32(f"{ARCH}.enc.pos_conv_kernel", 128) - writer.add_uint32(f"{ARCH}.enc.pos_conv_groups", 16) - writer.add_array(f"{ARCH}.fe.kernels", FE_KERNELS) - writer.add_array(f"{ARCH}.fe.strides", FE_STRIDES) - writer.add_array(f"{ARCH}.coeff_names", ARKIT_BLENDSHAPES) - - total_bytes = 0 - for key in sorted(state, key=lambda k: name_map[k][0]): - gguf_name, is_matmul = name_map[key] - arr = state[key].detach().cpu().float().numpy() - arr = np.ascontiguousarray(arr) - if args.dtype == "f16" and is_matmul: - arr = arr.astype(np.float16) - writer.add_tensor(gguf_name, arr) - total_bytes += arr.nbytes - print(f" {gguf_name}: {list(arr.shape)} {arr.dtype}") - - writer.write_header_to_file() - writer.write_kv_data_to_file() - writer.write_tensors_to_file() - writer.close() - print(f"wrote {args.out} ({total_bytes / 1e6:.1f} MB tensor data, dtype={args.dtype})") - - -if __name__ == "__main__": - main() diff --git a/scripts/dump_lam_a2e_frontend_reference.py b/scripts/dump_lam_a2e_frontend_reference.py deleted file mode 100644 index df8245885..000000000 --- a/scripts/dump_lam_a2e_frontend_reference.py +++ /dev/null @@ -1,65 +0,0 @@ -"""Create a Wav2Vec2 frontend parity fixture from trusted LAM-A2E PyTorch. - -Usage: - python scripts/dump_lam_a2e_frontend_reference.py \ - -""" - -from __future__ import annotations - -import json -import runpy -import sys -from pathlib import Path - -import librosa -import numpy as np -import torch - - -def main() -> None: - project_root = Path(sys.argv[1]).resolve() - checkpoint_path = Path(sys.argv[2]).resolve() - audio_path = Path(sys.argv[3]).resolve() - output_path = Path(sys.argv[4]).resolve() - - sys.path.insert(0, str(project_root)) - from models import build_model - - config = runpy.run_path( - str(project_root / "configs" / "lam_audio2exp_config_streaming.py") - ) - model = build_model(config["model"]) - checkpoint = torch.load(checkpoint_path, map_location="cpu", weights_only=False) - model.load_state_dict(checkpoint["state_dict"], strict=True) - model.cuda().eval() - - pcm, sample_rate = librosa.load(audio_path, sr=16000, mono=True) - pcm_tensor = torch.from_numpy(pcm).unsqueeze(0).cuda() - with torch.no_grad(): - frontend = model.backbone.audio_encoder.feature_extractor(pcm_tensor) - - output_path.parent.mkdir(parents=True, exist_ok=True) - np.savez( - output_path, - pcm=pcm.astype(np.float32), - frontend=frontend.cpu().numpy().astype(np.float32), - sample_rate=np.array(sample_rate, dtype=np.int32), - ) - output_path.with_suffix(".json").write_text( - json.dumps( - { - "sample_rate": sample_rate, - "pcm_shape": list(pcm.shape), - "frontend_shape": list(frontend.shape), - "weights": checkpoint_path.name, - "stage": "Wav2Vec2 feature_extractor", - }, - indent=2, - ) - + "\n" - ) - - -if __name__ == "__main__": - main() diff --git a/scripts/dump_lam_a2e_stages.py b/scripts/dump_lam_a2e_stages.py deleted file mode 100644 index 3b9978ad4..000000000 --- a/scripts/dump_lam_a2e_stages.py +++ /dev/null @@ -1,138 +0,0 @@ -"""Generate deterministic per-stage LAM-A2E parity fixtures. - -The fixture uses the first second of a trusted 16 kHz mono WAV so each -intermediate stage stays small enough to commit/store alongside test metadata. -""" - -from __future__ import annotations - -import argparse -import json -import runpy -import sys -from pathlib import Path - -import librosa -import numpy as np -import torch - - -def parse_args() -> argparse.Namespace: - parser = argparse.ArgumentParser() - parser.add_argument("project_root", type=Path) - parser.add_argument("checkpoint", type=Path) - parser.add_argument("audio", type=Path) - parser.add_argument("output", type=Path) - parser.add_argument("--seconds", type=float, default=1.0) - parser.add_argument("--identity-index", type=int, default=0) - return parser.parse_args() - - -def tensor_output(value: object) -> torch.Tensor: - if isinstance(value, torch.Tensor): - return value - if hasattr(value, "last_hidden_state"): - return value.last_hidden_state - if isinstance(value, (tuple, list)) and value and isinstance(value[0], torch.Tensor): - return value[0] - raise TypeError(f"Cannot serialize hook output of type {type(value)!r}") - - -def main() -> None: - args = parse_args() - project_root = args.project_root.resolve() - sys.path.insert(0, str(project_root)) - from models import build_model - - config = runpy.run_path( - str(project_root / "configs" / "lam_audio2exp_config_streaming.py") - ) - model = build_model(config["model"]) - checkpoint = torch.load(args.checkpoint, map_location="cpu", weights_only=False) - model.load_state_dict(checkpoint["state_dict"], strict=True) - model.cuda().eval() - - stage_outputs: dict[str, np.ndarray] = {} - - def capture(name: str): - def hook(_module: torch.nn.Module, _inputs: tuple[object, ...], output: object): - stage_outputs[name] = ( - tensor_output(output).detach().float().cpu().numpy() - ) - - return hook - - def capture_input(name: str): - def hook(_module: torch.nn.Module, inputs: tuple[object, ...]): - stage_outputs[name] = ( - tensor_output(inputs[0]).detach().float().cpu().numpy() - ) - - return hook - - backbone = model.backbone - hooks = [ - backbone.audio_encoder.feature_extractor.register_forward_hook(capture("frontend")), - backbone.audio_encoder.feature_projection.layer_norm.register_forward_pre_hook( - capture_input("interpolated_frontend") - ), - backbone.audio_encoder.feature_projection.layer_norm.register_forward_hook( - capture("feature_layer_norm") - ), - backbone.audio_encoder.feature_projection.register_forward_hook(capture("wav2vec_projection")), - backbone.audio_encoder.encoder.pos_conv_embed.register_forward_hook(capture("position")), - backbone.feature_projection.register_forward_hook(capture("lam_projection")), - backbone.identity_encoder.register_forward_hook(capture("identity")), - backbone.decoder[0].register_forward_hook(capture("decoder")), - backbone.output_proj.register_forward_hook(capture("output_projection")), - ] - for index, layer in enumerate(backbone.audio_encoder.feature_extractor.conv_layers): - hooks.append(layer.conv.register_forward_hook(capture(f"frontend_conv_{index}"))) - if index == 0: - hooks.append(layer.layer_norm.register_forward_hook(capture("frontend_group_norm_0"))) - for index, layer in enumerate(backbone.audio_encoder.encoder.layers): - hooks.append(layer.register_forward_hook(capture(f"transformer_{index:02d}"))) - - pcm, sample_rate = librosa.load(args.audio, sr=16000, mono=True) - pcm = pcm[: int(args.seconds * sample_rate)] - identity_class_count = config["model"]["backbone"]["num_identity_classes"] - identity = torch.nn.functional.one_hot( - torch.tensor(args.identity_index), - identity_class_count, - ).cuda()[None, ...] - input_dict = { - "id_idx": identity, - "input_audio_array": torch.from_numpy(pcm).unsqueeze(0).cuda(), - } - - with torch.no_grad(): - final_output = backbone(input_dict) - for hook in hooks: - hook.remove() - - args.output.parent.mkdir(parents=True, exist_ok=True) - np.savez( - args.output, - pcm=pcm.astype(np.float32), - final=final_output.detach().float().cpu().numpy(), - **stage_outputs, - ) - args.output.with_suffix(".json").write_text( - json.dumps( - { - "sample_rate": sample_rate, - "seconds": args.seconds, - "identity_index": args.identity_index, - "fps": 30, - "stages": {name: list(value.shape) for name, value in stage_outputs.items()}, - "final_shape": list(final_output.shape), - "checkpoint": args.checkpoint.name, - }, - indent=2, - ) - + "\n" - ) - - -if __name__ == "__main__": - main() diff --git a/scripts/remap_lam_gguf.py b/scripts/remap_lam_gguf.py deleted file mode 100644 index 6dd03f8ff..000000000 --- a/scripts/remap_lam_gguf.py +++ /dev/null @@ -1,174 +0,0 @@ -#!/usr/bin/env python3 -"""Remap a raw LAM-A2E GGUF (pytorch/backbone names) into lam-audio2exp dialect.""" - -from __future__ import annotations - -import argparse -import os - -import numpy as np -from gguf import GGUFReader, GGUFWriter - -ARCH = "lam-audio2exp" - -ARKIT_BLENDSHAPES = [ - "browDownLeft", "browDownRight", "browInnerUp", "browOuterUpLeft", - "browOuterUpRight", "cheekPuff", "cheekSquintLeft", "cheekSquintRight", - "eyeBlinkLeft", "eyeBlinkRight", "eyeLookDownLeft", "eyeLookDownRight", - "eyeLookInLeft", "eyeLookInRight", "eyeLookOutLeft", "eyeLookOutRight", - "eyeLookUpLeft", "eyeLookUpRight", "eyeSquintLeft", "eyeSquintRight", - "eyeWideLeft", "eyeWideRight", "jawForward", "jawLeft", "jawOpen", - "jawRight", "mouthClose", "mouthDimpleLeft", "mouthDimpleRight", - "mouthFrownLeft", "mouthFrownRight", "mouthFunnel", "mouthLeft", - "mouthLowerDownLeft", "mouthLowerDownRight", "mouthPressLeft", - "mouthPressRight", "mouthPucker", "mouthRight", "mouthRollLower", - "mouthRollUpper", "mouthShrugLower", "mouthShrugUpper", "mouthSmileLeft", - "mouthSmileRight", "mouthStretchLeft", "mouthStretchRight", - "mouthUpperUpLeft", "mouthUpperUpRight", "noseSneerLeft", - "noseSneerRight", "tongueOut", -] - -FE_KERNELS = [10, 3, 3, 3, 3, 2, 2] -FE_STRIDES = [5, 2, 2, 2, 2, 2, 2] -SKIP_PREFIXES = ( - "audio_encoder.lm_head.", - "identity_encoder.grus.", -) -SKIP_KEYS = ("audio_encoder.masked_spec_embed",) - - -def build_name_map(): - m = {} - for i in range(7): - m[f"audio_encoder.feature_extractor.conv_layers.{i}.conv.weight"] = (f"fe.conv{i}.weight", True) - m["audio_encoder.feature_extractor.conv_layers.0.layer_norm.weight"] = ("fe.gn.weight", False) - m["audio_encoder.feature_extractor.conv_layers.0.layer_norm.bias"] = ("fe.gn.bias", False) - m["audio_encoder.feature_projection.layer_norm.weight"] = ("fp.ln.weight", False) - m["audio_encoder.feature_projection.layer_norm.bias"] = ("fp.ln.bias", False) - m["audio_encoder.feature_projection.projection.weight"] = ("fp.proj.weight", True) - m["audio_encoder.feature_projection.projection.bias"] = ("fp.proj.bias", False) - m["audio_encoder.encoder.pos_conv_embed.conv.weight"] = ("enc.pos_conv.weight", True) - m["audio_encoder.encoder.pos_conv_embed.conv.bias"] = ("enc.pos_conv.bias", False) - m["audio_encoder.encoder.layer_norm.weight"] = ("enc.ln.weight", False) - m["audio_encoder.encoder.layer_norm.bias"] = ("enc.ln.bias", False) - for i in range(12): - src = f"audio_encoder.encoder.layers.{i}" - dst = f"enc.blk{i}" - for proj in ("q", "k", "v"): - m[f"{src}.attention.{proj}_proj.weight"] = (f"{dst}.attn_{proj}.weight", True) - m[f"{src}.attention.{proj}_proj.bias"] = (f"{dst}.attn_{proj}.bias", False) - m[f"{src}.attention.out_proj.weight"] = (f"{dst}.attn_o.weight", True) - m[f"{src}.attention.out_proj.bias"] = (f"{dst}.attn_o.bias", False) - m[f"{src}.layer_norm.weight"] = (f"{dst}.ln1.weight", False) - m[f"{src}.layer_norm.bias"] = (f"{dst}.ln1.bias", False) - m[f"{src}.feed_forward.intermediate_dense.weight"] = (f"{dst}.ffn_up.weight", True) - m[f"{src}.feed_forward.intermediate_dense.bias"] = (f"{dst}.ffn_up.bias", False) - m[f"{src}.feed_forward.output_dense.weight"] = (f"{dst}.ffn_down.weight", True) - m[f"{src}.feed_forward.output_dense.bias"] = (f"{dst}.ffn_down.bias", False) - m[f"{src}.final_layer_norm.weight"] = (f"{dst}.ln2.weight", False) - m[f"{src}.final_layer_norm.bias"] = (f"{dst}.ln2.bias", False) - m["feature_projection.weight"] = ("head.proj.weight", True) - m["feature_projection.bias"] = ("head.proj.bias", False) - m["identity_encoder.id_mlp.weight"] = ("head.id_mlp.weight", True) - m["identity_encoder.id_mlp.bias"] = ("head.id_mlp.bias", False) - for i in range(3): - src = f"identity_encoder.first_net.conv_layers.{i}" - m[f"{src}.conv.weight"] = (f"head.first{i}.conv.weight", True) - m[f"{src}.conv.bias"] = (f"head.first{i}.conv.bias", False) - m[f"{src}.norm.weight"] = (f"head.first{i}.ln.weight", False) - m[f"{src}.norm.bias"] = (f"head.first{i}.ln.bias", False) - m["identity_encoder.first_net.conv_layers.0.residual_layer.0.weight"] = ("head.first0.res.weight", True) - m["identity_encoder.first_net.conv_layers.0.residual_layer.0.bias"] = ("head.first0.res.bias", False) - for i in range(3): - m[f"decoder.0.{i}.conv.weight"] = (f"head.dec{i}.conv.weight", True) - m[f"decoder.0.{i}.conv.bias"] = (f"head.dec{i}.conv.bias", False) - m[f"decoder.0.{i}.norm.weight"] = (f"head.dec{i}.ln.weight", False) - m[f"decoder.0.{i}.norm.bias"] = (f"head.dec{i}.ln.bias", False) - m["output_proj.weight"] = ("head.out.weight", True) - m["output_proj.bias"] = ("head.out.bias", False) - return m - - -def strip_prefix(name: str) -> str: - if name.startswith("module."): - name = name[7:] - if name.startswith("backbone."): - name = name[9:] - return name - - -def main() -> None: - parser = argparse.ArgumentParser() - parser.add_argument("--in", dest="inp", required=True) - parser.add_argument("--out", required=True) - parser.add_argument("--dtype", choices=["f32", "f16"], default="f32") - args = parser.parse_args() - - reader = GGUFReader(args.inp) - state = {} - for tensor in reader.tensors: - key = strip_prefix(tensor.name) - if key in SKIP_KEYS or key.startswith(SKIP_PREFIXES): - continue - state[key] = np.array(tensor.data, dtype=np.float32, copy=True) - - g = state.pop("audio_encoder.encoder.pos_conv_embed.conv.weight_g") - v = state.pop("audio_encoder.encoder.pos_conv_embed.conv.weight_v") - if g.shape != (1, 1, 128) or v.shape != (768, 48, 128): - raise RuntimeError(f"unexpected pos-conv shapes g={g.shape} v={v.shape}") - norm = np.linalg.norm(v, axis=(0, 1), keepdims=True) - folded = (g * v / np.maximum(norm, 1e-12)).astype(np.float32, copy=True) - state["audio_encoder.encoder.pos_conv_embed.conv.weight"] = folded - - name_map = build_name_map() - unmapped = sorted(set(state) - set(name_map)) - if unmapped: - raise RuntimeError(f"unmapped tensors ({len(unmapped)}): {unmapped[:30]}") - missing = sorted(set(name_map) - set(state)) - if missing: - raise RuntimeError(f"missing tensors: {missing}") - - out_dir = os.path.dirname(os.path.abspath(args.out)) - if out_dir: - os.makedirs(out_dir, exist_ok=True) - - writer = GGUFWriter(args.out, ARCH) - writer.add_name("LAM Audio2Expression (streaming)") - writer.add_string(f"{ARCH}.dtype", args.dtype) - writer.add_uint32(f"{ARCH}.sample_rate", 16000) - writer.add_uint32(f"{ARCH}.fps", 30) - writer.add_uint32(f"{ARCH}.n_coeffs", 52) - writer.add_uint32(f"{ARCH}.n_identity", 12) - writer.add_uint32(f"{ARCH}.identity_feat_dim", 64) - writer.add_uint32(f"{ARCH}.hidden_dim", 512) - writer.add_uint32(f"{ARCH}.window_frames", 64) - writer.add_float32(f"{ARCH}.layer_norm_eps", 1e-5) - writer.add_uint32(f"{ARCH}.enc.n_layers", 12) - writer.add_uint32(f"{ARCH}.enc.n_heads", 12) - writer.add_uint32(f"{ARCH}.enc.hidden", 768) - writer.add_uint32(f"{ARCH}.enc.ffn", 3072) - writer.add_uint32(f"{ARCH}.enc.pos_conv_kernel", 128) - writer.add_uint32(f"{ARCH}.enc.pos_conv_groups", 16) - writer.add_array(f"{ARCH}.fe.kernels", FE_KERNELS) - writer.add_array(f"{ARCH}.fe.strides", FE_STRIDES) - writer.add_array(f"{ARCH}.coeff_names", ARKIT_BLENDSHAPES) - - total = 0 - for key in sorted(state, key=lambda k: name_map[k][0]): - gguf_name, is_matmul = name_map[key] - arr = np.ascontiguousarray(state[key]) - if args.dtype == "f16" and is_matmul: - arr = arr.astype(np.float16) - writer.add_tensor(gguf_name, arr) - total += arr.nbytes - print(f" {gguf_name}: {list(arr.shape)} {arr.dtype}") - - writer.write_header_to_file() - writer.write_kv_data_to_file() - writer.write_tensors_to_file() - writer.close() - print(f"wrote {args.out} ({total / 1e6:.1f} MB)") - - -if __name__ == "__main__": - main()