stack_adapters.py

#!/usr/bin/env python
"""Stack a voice LoRA with an emotion LoRA and a vocal-burst LoRA, at chosen merge weights.

    python stack_adapters.py --voice k325_age3_bg1 --emotion Anger --burst sobs --outdir /tmp/x

peft can hold many adapters against one frozen base and activate several at once, but it gives
you no dose knob: an activated adapter contributes at its trained scaling (alpha / r). The dose
lives in `LoraLayer.scaling[name]`, so a stack is "activate the set, then rewrite scaling".

Two traps are handled here and both cost this project real runs -- see `set_active` below.

This file is the code from the README, verbatim and runnable. It was executed on one GH200
before publication.
"""
import argparse
import os
import time

import numpy as np
import soundfile as sf
import torch
from huggingface_hub import hf_hub_download, snapshot_download
from peft import PeftModel
from peft.tuners.lora import LoraLayer
from transformers import AutoModel, AutoProcessor

BASE = "laion/moss-tts-local-transformer-4.55b-voice-acting-v2"
CODEC = "OpenMOSS-Team/MOSS-Audio-Tokenizer-v2"
LORAS = "TTS-AGI/moss-voice-profile-loras"
REFS = "TTS-AGI/moss-voice-profile-references"
EMO = "TTS-AGI/moss-emotion-loras-v3"
BURST = "laion/vocal-burst-lora-adapters"
GEN_SR = 48000


class AdapterStack:
    """Several LoRAs on one frozen base, each at its own merge weight."""

    def __init__(self, base_model):
        self.pm = None
        self.base = base_model
        self.base_scaling = {}          # module name -> {adapter: trained scaling}

    def load(self, name, model_id, subfolder=None):
        """Attach one adapter under `name`. Names are yours; they index everything below.

        `model_id` may be a hub repo or a local directory. When a subfolder is given the repo
        is resolved to a local path first: peft 0.20.0's `subfolder=` kwarg works online but
        doubles the path under HF_HUB_OFFLINE=1 (it puts the subfolder in the filename AND
        passes it as a kwarg), which raises LocalEntryNotFoundError on an air-gapped node.
        """
        path = model_id
        if subfolder:
            path = (f"{model_id}/{subfolder}" if os.path.isdir(model_id) else
                    f"{snapshot_download(model_id, allow_patterns=[subfolder + '/adapter_*'])}"
                    f"/{subfolder}")
        if self.pm is None:
            self.pm = PeftModel.from_pretrained(self.base, path, adapter_name=name).eval()
        else:
            self.pm.load_adapter(path, adapter_name=name)
        # Re-snapshot the trained scalings: adapters do NOT all target the same modules, and a
        # module first seen by adapter B has no entry for it in a snapshot taken under A.
        self.base_scaling = {nm: dict(m.scaling) for nm, m in self.pm.named_modules()
                             if isinstance(m, LoraLayer)}
        return self

    def set_active(self, spec):
        """spec: {name: merge_weight}. Weight 0 or a missing name = off; {} = pure base model."""
        pm = self.pm
        keys = {n: float(v) for n, v in spec.items() if v}
        if not keys:
            pm.base_model.disable_adapter_layers()
            return
        pm.base_model.enable_adapter_layers()
        pm.base_model.set_adapter(list(keys))

        # ------------------------------------------------------------------ TRAP 1
        # `PeftModel.active_adapter` is a PLAIN ATTRIBUTE, not a property. peft sets it once in
        # __init__ to the first adapter ever loaded and thereafter only updates it inside
        # PeftModel.set_adapter(). Going through `pm.base_model.set_adapter()` -- which is what
        # you must do to activate SEVERAL adapters with per-adapter doses -- updates the
        # LoraModel and leaves the PeftModel attribute pinned to the first adapter forever.
        #
        # Harmless until that first adapter is deleted or replaced. Then every generate() does
        # `peft_config[self.active_adapter]` and raises KeyError on a name that is no longer
        # there, permanently. It killed four runs before it was found, and the reason it hid so
        # well is that peft's delete_adapter() DOES repair active_adapter -- but only when
        # exactly one adapter is left active, which is never true for a stack.
        #
        # Fix: point it at a member of the set you just activated.
        first = next(iter(keys))
        if first in getattr(pm, "peft_config", {}):
            pm.active_adapter = first

        # ------------------------------------------------------------------ TRAP 2
        # `a in m.scaling` is not enough to know the trained scaling. Adapters of different rank
        # target different module sets, so on a module that only adapter B reaches,
        # base_scaling[nm] exists but has no key for A. Look up BOTH, and record the trained
        # value the first time a module sees an adapter.
        for nm, m in pm.named_modules():
            if not isinstance(m, LoraLayer):
                continue
            for a, s in keys.items():
                if a not in m.scaling:
                    continue
                if a in self.base_scaling.get(nm, {}):
                    m.scaling[a] = self.base_scaling[nm][a] * s
                else:
                    self.base_scaling.setdefault(nm, {})[a] = m.scaling[a]
                    m.scaling[a] = m.scaling[a] * s


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--voice", default="k325_age3_bg1")
    ap.add_argument("--emotion", default="Anger")
    ap.add_argument("--burst", default="sobs")
    ap.add_argument("--text", default="I told you what would happen. (sobs) I told you, and "
                                      "you did it anyway.")
    ap.add_argument("--instruction", default="A warm, aged baritone breaking under anger he is "
                                             "trying and failing to contain.")
    ap.add_argument("--language", default="English")
    ap.add_argument("--outdir", default=".")
    ap.add_argument("--seed", type=int, default=0)
    ap.add_argument("--device", default="cuda")
    # Doses. Voice adapters are identity and want their full trained strength. The emotion dose
    # is per-emotion and measured (0.5 moderate; 0.5-1.9 intense). The burst optimum is 0.5, and
    # on a burst-carrying line the emotion adapter is CAPPED at half the burst dose.
    ap.add_argument("--voice-lambda", type=float, default=1.0)
    ap.add_argument("--emotion-lambda", type=float, default=1.9)
    ap.add_argument("--burst-lambda", type=float, default=0.5)
    ap.add_argument("--burst-emotion-ratio", type=float, default=0.5)
    a = ap.parse_args()

    t0 = time.time()
    proc = AutoProcessor.from_pretrained(BASE, trust_remote_code=True, codec_path=CODEC)
    proc.audio_tokenizer = proc.audio_tokenizer.to(a.device).eval()
    model = AutoModel.from_pretrained(BASE, trust_remote_code=True, dtype=torch.bfloat16,
                                      attn_implementation="sdpa").to(a.device).eval()
    print(f"[{time.time() - t0:.0f}s] base loaded")

    stack = AdapterStack(model)
    stack.load("voice", LORAS, subfolder=a.voice)
    stack.load("emotion", EMO, subfolder=a.emotion)
    stack.load("burst", BURST, subfolder=a.burst)
    print(f"[{time.time() - t0:.0f}s] adapters: {list(stack.pm.peft_config)}")
    for n, c in stack.pm.peft_config.items():
        print(f"    {n:8s} r={c.r:<3d} alpha={c.lora_alpha:<3d} "
              f"targets={len(c.target_modules)} modules")

    ref = hf_hub_download(REFS, f"pilot/{a.voice}/reference.wav", repo_type="dataset")
    net = stack.pm

    @torch.no_grad()
    def gen(tag, spec, text):
        stack.set_active(spec)
        conv = [[proc.build_user_message(text=text, instruction=a.instruction,
                                         language=a.language, reference=[ref],
                                         tokens=max(8, len(text.split())))]]
        b = proc(conv, mode="generation")
        torch.manual_seed(a.seed)
        o = net.generate(input_ids=b["input_ids"].to(a.device),
                         attention_mask=b["attention_mask"].to(a.device),
                         max_new_frames=400, do_sample=True,
                         text_temperature=0.7, text_top_k=50, text_top_p=1.0,
                         audio_temperature=1.0, audio_top_p=0.95, audio_top_k=30,
                         audio_repetition_penalty=1.1)
        m = proc.decode(o)[0]
        if not m.audio_codes_list:
            print(f"  {tag}: EMPTY DECODE")
            return
        w = m.audio_codes_list[0].cpu().float().numpy()
        w = np.ascontiguousarray(w.mean(0) if w.ndim > 1 else w)
        path = f"{a.outdir}/{tag}.wav"
        sf.write(path, w, GEN_SR)
        print(f"  {tag:26s} {spec}  ->  {path}  {len(w) / GEN_SR:.2f}s "
              f"peak {np.abs(w).max():.3f}")

    plain = a.text.replace(f"({a.burst}) ", "")
    # The emotion dose on a burst-carrying line is capped, not set: a group already below the
    # cap keeps its own smaller dose.
    capped = min(a.emotion_lambda, a.burst_lambda * a.burst_emotion_ratio)

    gen("00_base", {}, plain)
    gen("01_voice", {"voice": a.voice_lambda}, plain)
    gen("02_voice_emotion", {"voice": a.voice_lambda, "emotion": a.emotion_lambda}, plain)
    gen("03_voice_emotion_burst",
        {"voice": a.voice_lambda, "emotion": capped, "burst": a.burst_lambda}, a.text)

    # ---- TRAP 1, demonstrated ------------------------------------------------------
    # Delete the FIRST-LOADED adapter. That is the exact trigger: peft pinned
    # PeftModel.active_adapter to "voice" in __init__ and every set_active() since then went
    # through base_model.set_adapter(), which does not touch it. peft's delete_adapter() only
    # repairs active_adapter when a single adapter is left active -- never true for a stack.
    print(f"  active_adapter before delete: {stack.pm.active_adapter!r}")
    stack.pm.delete_adapter("voice")
    print(f"  active_adapter after deleting 'voice': {stack.pm.active_adapter!r}  "
          f"(peft_config now holds {sorted(stack.pm.peft_config)})")
    if stack.pm.active_adapter not in stack.pm.peft_config:
        print("  ^ DANGLING. Without the one-line repair in set_active(), the next generate() "
              "raises KeyError from inside peft_config[self.active_adapter].")
    # set_active() repairs it; generation works.
    gen("04_after_deleting_first_adapter",
        {"emotion": capped, "burst": a.burst_lambda}, a.text)
    print(f"  active_adapter after set_active: {stack.pm.active_adapter!r}")
    print(f"[{time.time() - t0:.0f}s] done")


if __name__ == "__main__":
    main()