Spaces:
Sleeping
Sleeping
File size: 4,611 Bytes
2e1dc7f | 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 | from functools import partial
from typing import Dict
# from models.musicgen_lm import MusicgenLm
def transfer(src, dest, srcname, destname):
dest[destname].copy_(src[srcname])
def transfer_weights(af_dict: Dict, lag_dict: Dict, n_layers: int,
has_xatt: bool) -> Dict:
# af_dict = audiocraft_lm.state_dict()
# lag_dict = lag_lm.state_dict()
t = partial(transfer, af_dict, lag_dict)
hop = 3 if has_xatt else 2
# transfer conditioner
for s in ("weight", "bias"):
lag_dict[f"conditioner.output_proj.{s}"] = af_dict[
f"condition_provider.conditioners."
f"description.output_proj.{s}"].clone()
# t(
# f"condition_provider.conditioners.description.output_proj.{s}",
# f"conditioner.output_proj.{s}",
# )
# transfer embeddings
for i in range(4):
t(f"emb.{i}.weight", f"decoder.token_emb.emb.{i}.weight")
# transfer decoder layers
for layeridx in range(n_layers):
lamsi_prefix = "decoder.attn_layers.layers"
af_prefix = f"transformer.layers.{layeridx}"
# ---- LAMSI SUBLAYER 0: self attention ----
lamsi_prefix_layer = f"{lamsi_prefix}.{layeridx * hop}"
# att norm
for s in ("weight", "bias"):
t(f"{af_prefix}.norm1.{s}", f"{lamsi_prefix_layer}.0.0.{s}")
# attention q, k, v, out
af_qkv = af_dict[f"{af_prefix}.self_attn.in_proj_weight"]
assert (af_qkv.shape[0] % 3) == 0
dim = af_qkv.shape[0] // 3
att_q = af_qkv[:dim]
att_k = af_qkv[dim:dim * 2]
att_v = af_qkv[dim * 2:]
lag_dict[f"{lamsi_prefix_layer}.1.to_q.weight"].copy_(att_q)
lag_dict[f"{lamsi_prefix_layer}.1.to_k.weight"].copy_(att_k)
lag_dict[f"{lamsi_prefix_layer}.1.to_v.weight"].copy_(att_v)
t(
f"{af_prefix}.self_attn.out_proj.weight",
f"{lamsi_prefix_layer}.1.to_out.weight",
)
# ---- LAMSI SUBLAYER 1: cross attention ----
if has_xatt:
lamsi_prefix_layer = f"{lamsi_prefix}.{layeridx * hop + 1}"
# xatt norm
for s in ("weight", "bias"):
t(f"{af_prefix}.norm_cross.{s}",
f"{lamsi_prefix_layer}.0.0.{s}")
# xatt q, k, v, out
af_xqkv = af_dict[f"{af_prefix}.cross_attention.in_proj_weight"]
assert (af_xqkv.shape[0] % 3) == 0
dim = af_xqkv.shape[0] // 3
xatt_q = af_xqkv[:dim]
xatt_k = af_xqkv[dim:dim * 2]
xatt_v = af_xqkv[dim * 2:]
lag_dict[f"{lamsi_prefix_layer}.1.to_q.weight"].copy_(xatt_q)
lag_dict[f"{lamsi_prefix_layer}.1.to_k.weight"].copy_(xatt_k)
lag_dict[f"{lamsi_prefix_layer}.1.to_v.weight"].copy_(xatt_v)
t(
f"{af_prefix}.cross_attention.out_proj.weight",
f"{lamsi_prefix_layer}.1.to_out.weight",
)
# ---- LAMSI SUBLAYER 2: feed forward ----
lamsi_prefix_layer = f"{lamsi_prefix}.{layeridx * hop + (hop - 1)}"
# ff norm
for s in ("weight", "bias"):
t(f"{af_prefix}.norm2.{s}", f"{lamsi_prefix_layer}.0.0.{s}")
# ff layers
t(f"{af_prefix}.linear1.weight",
f"{lamsi_prefix_layer}.1.ff.0.0.weight")
t(f"{af_prefix}.linear2.weight", f"{lamsi_prefix_layer}.1.ff.2.weight")
# transfer final norm
for s in ("weight", "bias"):
t(f"out_norm.{s}", f"decoder.attn_layers.final_norm.{s}")
# transfer output linears
for i in range(4):
t(f"linears.{i}.weight", f"decoder.to_logits.linears.{i}.weight")
# lag_lm.load_state_dict(lag_dict)
return lag_dict
if __name__ == "__main__":
import torch
import config as cfg
from models.musicgen_lm import MusicgenLm
from hyperparameters import PretrainedMelodyLmParams
lag_lm = MusicgenLm(PretrainedMelodyLmParams())
lag_state_dict = lag_lm.state_dict()
af_state_dict = torch.load(cfg.weights_dir() / "audiocraft-melody.pt")
print(
sum(p.numel()
for k, p in af_state_dict.items()
if "conditioner" not in k))
print(
sum(p.numel()
for k, p in lag_state_dict.items()
if "conditioner" not in k))
print(list((k, p.numel()) for k, p in af_state_dict.items()))
print(list((k, p.numel()) for k, p in lag_state_dict.items()))
new_lag_dict = transfer_weights(af_state_dict, lag_state_dict, 48, False)
torch.save(new_lag_dict, cfg.weights_dir() / "lm-melody-weights.pt")
|