File size: 10,676 Bytes
5fe566e | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 | #!/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()
|