Spaces:
Running
Running
Add v1.6 + v1.7 model selector (7 selectable generators); sync v1.7-capable inference bundle
Browse files- .gradio/certificate.pem +31 -0
- README.md +6 -0
- app.py +39 -12
- softchart/generate.py +54 -38
- softchart/hf.py +16 -4
- softchart/model.py +250 -158
- static/app.css +23 -0
- static/app.js +3 -1
- static/index.html +21 -0
.gradio/certificate.pem
ADDED
|
@@ -0,0 +1,31 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
-----BEGIN CERTIFICATE-----
|
| 2 |
+
MIIFazCCA1OgAwIBAgIRAIIQz7DSQONZRGPgu2OCiwAwDQYJKoZIhvcNAQELBQAw
|
| 3 |
+
TzELMAkGA1UEBhMCVVMxKTAnBgNVBAoTIEludGVybmV0IFNlY3VyaXR5IFJlc2Vh
|
| 4 |
+
cmNoIEdyb3VwMRUwEwYDVQQDEwxJU1JHIFJvb3QgWDEwHhcNMTUwNjA0MTEwNDM4
|
| 5 |
+
WhcNMzUwNjA0MTEwNDM4WjBPMQswCQYDVQQGEwJVUzEpMCcGA1UEChMgSW50ZXJu
|
| 6 |
+
ZXQgU2VjdXJpdHkgUmVzZWFyY2ggR3JvdXAxFTATBgNVBAMTDElTUkcgUm9vdCBY
|
| 7 |
+
MTCCAiIwDQYJKoZIhvcNAQEBBQADggIPADCCAgoCggIBAK3oJHP0FDfzm54rVygc
|
| 8 |
+
h77ct984kIxuPOZXoHj3dcKi/vVqbvYATyjb3miGbESTtrFj/RQSa78f0uoxmyF+
|
| 9 |
+
0TM8ukj13Xnfs7j/EvEhmkvBioZxaUpmZmyPfjxwv60pIgbz5MDmgK7iS4+3mX6U
|
| 10 |
+
A5/TR5d8mUgjU+g4rk8Kb4Mu0UlXjIB0ttov0DiNewNwIRt18jA8+o+u3dpjq+sW
|
| 11 |
+
T8KOEUt+zwvo/7V3LvSye0rgTBIlDHCNAymg4VMk7BPZ7hm/ELNKjD+Jo2FR3qyH
|
| 12 |
+
B5T0Y3HsLuJvW5iB4YlcNHlsdu87kGJ55tukmi8mxdAQ4Q7e2RCOFvu396j3x+UC
|
| 13 |
+
B5iPNgiV5+I3lg02dZ77DnKxHZu8A/lJBdiB3QW0KtZB6awBdpUKD9jf1b0SHzUv
|
| 14 |
+
KBds0pjBqAlkd25HN7rOrFleaJ1/ctaJxQZBKT5ZPt0m9STJEadao0xAH0ahmbWn
|
| 15 |
+
OlFuhjuefXKnEgV4We0+UXgVCwOPjdAvBbI+e0ocS3MFEvzG6uBQE3xDk3SzynTn
|
| 16 |
+
jh8BCNAw1FtxNrQHusEwMFxIt4I7mKZ9YIqioymCzLq9gwQbooMDQaHWBfEbwrbw
|
| 17 |
+
qHyGO0aoSCqI3Haadr8faqU9GY/rOPNk3sgrDQoo//fb4hVC1CLQJ13hef4Y53CI
|
| 18 |
+
rU7m2Ys6xt0nUW7/vGT1M0NPAgMBAAGjQjBAMA4GA1UdDwEB/wQEAwIBBjAPBgNV
|
| 19 |
+
HRMBAf8EBTADAQH/MB0GA1UdDgQWBBR5tFnme7bl5AFzgAiIyBpY9umbbjANBgkq
|
| 20 |
+
hkiG9w0BAQsFAAOCAgEAVR9YqbyyqFDQDLHYGmkgJykIrGF1XIpu+ILlaS/V9lZL
|
| 21 |
+
ubhzEFnTIZd+50xx+7LSYK05qAvqFyFWhfFQDlnrzuBZ6brJFe+GnY+EgPbk6ZGQ
|
| 22 |
+
3BebYhtF8GaV0nxvwuo77x/Py9auJ/GpsMiu/X1+mvoiBOv/2X/qkSsisRcOj/KK
|
| 23 |
+
NFtY2PwByVS5uCbMiogziUwthDyC3+6WVwW6LLv3xLfHTjuCvjHIInNzktHCgKQ5
|
| 24 |
+
ORAzI4JMPJ+GslWYHb4phowim57iaztXOoJwTdwJx4nLCgdNbOhdjsnvzqvHu7Ur
|
| 25 |
+
TkXWStAmzOVyyghqpZXjFaH3pO3JLF+l+/+sKAIuvtd7u+Nxe5AW0wdeRlN8NwdC
|
| 26 |
+
jNPElpzVmbUq4JUagEiuTDkHzsxHpFKVK7q4+63SM1N95R1NbdWhscdCb+ZAJzVc
|
| 27 |
+
oyi3B43njTOQ5yOf+1CceWxG1bQVs5ZufpsMljq4Ui0/1lvh+wjChP4kqKOJ2qxq
|
| 28 |
+
4RgqsahDYVvTH9w7jXbyLeiNdd8XM2w9U/t7y0Ff/9yi0GE44Za4rF2LN9d11TPA
|
| 29 |
+
mRGunUHBcnWEvgJBQl9nJEiU0Zsnvgc/ubhPgXRR4Xq37Z0j4r7g1SgEEzwxA57d
|
| 30 |
+
emyPxgcYxn/eR44/KJ4EBs+lVDR3veyJm+kXQ99b21/+jh5Xos1AnX5iItreGCc=
|
| 31 |
+
-----END CERTIFICATE-----
|
README.md
CHANGED
|
@@ -10,6 +10,12 @@ app_file: app.py
|
|
| 10 |
pinned: false
|
| 11 |
license: mit
|
| 12 |
models:
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 13 |
- JacobLinCool/softchart-v15
|
| 14 |
- JacobLinCool/softchart-planner
|
| 15 |
---
|
|
|
|
| 10 |
pinned: false
|
| 11 |
license: mit
|
| 12 |
models:
|
| 13 |
+
- JacobLinCool/softchart-v17
|
| 14 |
+
- JacobLinCool/softchart-v17-small
|
| 15 |
+
- JacobLinCool/softchart-v17-tiny
|
| 16 |
+
- JacobLinCool/softchart-v16
|
| 17 |
+
- JacobLinCool/softchart-v16-small
|
| 18 |
+
- JacobLinCool/softchart-v16-tiny
|
| 19 |
- JacobLinCool/softchart-v15
|
| 20 |
- JacobLinCool/softchart-planner
|
| 21 |
---
|
app.py
CHANGED
|
@@ -33,30 +33,52 @@ from softchart.vocab import FPS, HOP, N_FFT, N_MELS, SR
|
|
| 33 |
LOGGER = logging.getLogger("softchart.space")
|
| 34 |
STATIC_DIR = Path(__file__).with_name("static")
|
| 35 |
|
| 36 |
-
V15_REPO = os.environ.get("SC_V15", "JacobLinCool/softchart-v15")
|
| 37 |
PLAN_REPO = os.environ.get("SC_PLAN", "JacobLinCool/softchart-planner")
|
| 38 |
DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
|
| 39 |
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 40 |
COURSE_DENS = {"easy": 1, "normal": 2, "hard": 4, "oni": 7}
|
| 41 |
CHAR = {"don": "1", "ka": "2", "don_big": "3", "ka_big": "4",
|
| 42 |
"roll": "5", "roll_big": "6", "balloon": "7"}
|
| 43 |
SUB = 96
|
| 44 |
|
| 45 |
-
_MODELS = {}
|
|
|
|
| 46 |
|
| 47 |
|
| 48 |
-
def get_models(*, include_planner=False):
|
| 49 |
-
|
| 50 |
-
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 51 |
if (not getattr(u, "_dual", False) or u.beat is None
|
| 52 |
or not getattr(u, "_has_plan", False)):
|
| 53 |
raise RuntimeError(
|
| 54 |
-
f"{
|
| 55 |
)
|
| 56 |
-
_MODELS
|
| 57 |
-
|
| 58 |
-
|
| 59 |
-
|
|
|
|
|
|
|
|
|
|
| 60 |
|
| 61 |
|
| 62 |
def load_logmel(path):
|
|
@@ -356,6 +378,7 @@ def generate_chart(
|
|
| 356 |
temperature: float,
|
| 357 |
top_p: float,
|
| 358 |
drum_volume: float,
|
|
|
|
| 359 |
) -> dict[str, object]:
|
| 360 |
"""Stream the actual inference stages to the custom frontend."""
|
| 361 |
workdir = tempfile.mkdtemp(prefix="softchart-request-")
|
|
@@ -368,11 +391,15 @@ def generate_chart(
|
|
| 368 |
top_p = float(top_p)
|
| 369 |
drum_volume = float(drum_volume)
|
| 370 |
|
|
|
|
|
|
|
|
|
|
|
|
|
| 371 |
yield _progress(
|
| 372 |
"loading", 0.03, "Preparing",
|
| 373 |
-
"Loading SoftChart.",
|
| 374 |
)
|
| 375 |
-
models = get_models(include_planner=use_planner)
|
| 376 |
|
| 377 |
yield _progress(
|
| 378 |
"audio", 0.11, "Listening",
|
|
|
|
| 33 |
LOGGER = logging.getLogger("softchart.space")
|
| 34 |
STATIC_DIR = Path(__file__).with_name("static")
|
| 35 |
|
|
|
|
| 36 |
PLAN_REPO = os.environ.get("SC_PLAN", "JacobLinCool/softchart-planner")
|
| 37 |
DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
|
| 38 |
|
| 39 |
+
# Selectable generators: three generations x three sizes (plus the v1.5 single
|
| 40 |
+
# model). All are dual (slot + time), beat- and plan-conditioned, MIT licensed.
|
| 41 |
+
# https://huggingface.co/collections/JacobLinCool/softchart-generators
|
| 42 |
+
MODELS = {
|
| 43 |
+
"v1.7": "JacobLinCool/softchart-v17",
|
| 44 |
+
"v1.7-small": "JacobLinCool/softchart-v17-small",
|
| 45 |
+
"v1.7-tiny": "JacobLinCool/softchart-v17-tiny",
|
| 46 |
+
"v1.6": "JacobLinCool/softchart-v16",
|
| 47 |
+
"v1.6-small": "JacobLinCool/softchart-v16-small",
|
| 48 |
+
"v1.6-tiny": "JacobLinCool/softchart-v16-tiny",
|
| 49 |
+
"v1.5": "JacobLinCool/softchart-v15",
|
| 50 |
+
}
|
| 51 |
+
DEFAULT_MODEL = os.environ.get("SC_MODEL", "v1.7")
|
| 52 |
+
|
| 53 |
COURSE_DENS = {"easy": 1, "normal": 2, "hard": 4, "oni": 7}
|
| 54 |
CHAR = {"don": "1", "ka": "2", "don_big": "3", "ka_big": "4",
|
| 55 |
"roll": "5", "roll_big": "6", "balloon": "7"}
|
| 56 |
SUB = 96
|
| 57 |
|
| 58 |
+
_MODELS = {} # repo id -> {"gen","slot","beat"}, loaded lazily and cached
|
| 59 |
+
_PLANNER = {} # the planner is model-agnostic and shared across generators
|
| 60 |
|
| 61 |
|
| 62 |
+
def get_models(model_choice=DEFAULT_MODEL, *, include_planner=False):
|
| 63 |
+
repo = MODELS.get(model_choice)
|
| 64 |
+
if repo is None:
|
| 65 |
+
raise ValueError(
|
| 66 |
+
f"unknown model {model_choice!r}; choose one of {list(MODELS)}"
|
| 67 |
+
)
|
| 68 |
+
if repo not in _MODELS:
|
| 69 |
+
u = load_hf(repo, device=DEVICE)
|
| 70 |
if (not getattr(u, "_dual", False) or u.beat is None
|
| 71 |
or not getattr(u, "_has_plan", False)):
|
| 72 |
raise RuntimeError(
|
| 73 |
+
f"{repo} must provide dual generation, beat, and plan conditioning"
|
| 74 |
)
|
| 75 |
+
_MODELS[repo] = {"gen": u, "slot": u, "beat": u}
|
| 76 |
+
models = dict(_MODELS[repo])
|
| 77 |
+
if include_planner:
|
| 78 |
+
if "plan" not in _PLANNER:
|
| 79 |
+
_PLANNER["plan"] = SoftChartPlanner.from_pretrained(PLAN_REPO).to(DEVICE).eval()
|
| 80 |
+
models["plan"] = _PLANNER["plan"]
|
| 81 |
+
return models
|
| 82 |
|
| 83 |
|
| 84 |
def load_logmel(path):
|
|
|
|
| 378 |
temperature: float,
|
| 379 |
top_p: float,
|
| 380 |
drum_volume: float,
|
| 381 |
+
model_choice: str = DEFAULT_MODEL,
|
| 382 |
) -> dict[str, object]:
|
| 383 |
"""Stream the actual inference stages to the custom frontend."""
|
| 384 |
workdir = tempfile.mkdtemp(prefix="softchart-request-")
|
|
|
|
| 391 |
top_p = float(top_p)
|
| 392 |
drum_volume = float(drum_volume)
|
| 393 |
|
| 394 |
+
if model_choice not in MODELS:
|
| 395 |
+
raise ValueError(
|
| 396 |
+
f"Unknown model {model_choice!r}. Choose one of: {', '.join(MODELS)}."
|
| 397 |
+
)
|
| 398 |
yield _progress(
|
| 399 |
"loading", 0.03, "Preparing",
|
| 400 |
+
f"Loading SoftChart {model_choice}.",
|
| 401 |
)
|
| 402 |
+
models = get_models(model_choice, include_planner=use_planner)
|
| 403 |
|
| 404 |
yield _progress(
|
| 405 |
"audio", 0.11, "Listening",
|
softchart/generate.py
CHANGED
|
@@ -24,12 +24,11 @@ def _dev_type(device):
|
|
| 24 |
|
| 25 |
|
| 26 |
def _autocast(device):
|
| 27 |
-
"""bf16 autocast on CUDA
|
| 28 |
import contextlib
|
| 29 |
|
| 30 |
-
|
| 31 |
-
|
| 32 |
-
return torch.autocast(dev, dtype=torch.bfloat16)
|
| 33 |
return contextlib.nullcontext()
|
| 34 |
CTX_LEN = 12 # fixed-length previous-window context (note-type tokens)
|
| 35 |
MIN_GAP_FRAMES = 4 # ~46 ms minimum inter-note gap enforced during decoding
|
|
@@ -42,24 +41,38 @@ SPAN_MAX = {"roll": 3.5, "roll_big": 4.0, "balloon": 6.5}
|
|
| 42 |
def load_model(ckpt_path, device="cuda"):
|
| 43 |
ck = torch.load(ckpt_path, map_location=device)
|
| 44 |
sd = ck["model"]
|
| 45 |
-
a = ck.get("args")
|
| 46 |
-
if not isinstance(a, dict) or not a:
|
| 47 |
-
raise ValueError(
|
| 48 |
-
f"{ckpt_path} is missing checkpoint args; refusing to infer model "
|
| 49 |
-
"capabilities from weight shapes."
|
| 50 |
-
)
|
| 51 |
vocab_size = sd["tok_emb.weight"].shape[0]
|
| 52 |
aux = any(k.startswith("aux") for k in sd)
|
| 53 |
gctx = any(k.startswith("gsum_proj") for k in sd)
|
| 54 |
-
|
|
|
|
| 55 |
emb_w = sd["tok_emb.weight"].shape[1]
|
| 56 |
d_model = a.get("d_model", 512)
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 57 |
model = ChartModel(
|
| 58 |
-
d_model=d_model,
|
|
|
|
| 59 |
dec_layers=a.get("dec_layers", 6), ffn=a.get("ffn", 2048),
|
| 60 |
vocab_size=vocab_size, aux=aux, global_ctx=gctx,
|
| 61 |
func_time=a.get("func_time", False), in_ch=in_ch,
|
| 62 |
emb_factor=emb_w if emb_w != d_model else None,
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 63 |
).to(device)
|
| 64 |
if any(k.startswith("ptr") for k in sd):
|
| 65 |
model.enable_ptr()
|
|
@@ -68,13 +81,24 @@ def load_model(ckpt_path, device="cuda"):
|
|
| 68 |
model.to(device) # newly enabled heads default to CPU
|
| 69 |
model.load_state_dict(sd)
|
| 70 |
model.eval()
|
| 71 |
-
|
| 72 |
-
|
| 73 |
-
|
| 74 |
-
|
| 75 |
-
|
| 76 |
-
|
| 77 |
-
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 78 |
model._has_gctx = gctx
|
| 79 |
return model
|
| 80 |
|
|
@@ -156,12 +180,16 @@ def decode_windows(model, mels, prefixes, device="cuda", temperature=1.0, top_p=
|
|
| 156 |
gen_device = "cpu" if _dev_type(device) == "mps" else device
|
| 157 |
gen = torch.Generator(device=gen_device)
|
| 158 |
gen.manual_seed(seed)
|
|
|
|
| 159 |
with _autocast(device):
|
| 160 |
memory = model.encode(
|
| 161 |
-
|
| 162 |
gsum=gsum.to(device) if gsum is not None else None,
|
| 163 |
pos_bucket=pos_buckets.to(device) if pos_buckets is not None else None,
|
| 164 |
)
|
|
|
|
|
|
|
|
|
|
| 165 |
|
| 166 |
seqs = torch.tensor(prefixes, dtype=torch.long, device=device)
|
| 167 |
if use_cfg: # rows B..2B: same ctx, conditions replaced by UNK
|
|
@@ -177,20 +205,10 @@ def decode_windows(model, mels, prefixes, device="cuda", temperature=1.0, top_p=
|
|
| 177 |
# lattice state: previous inter-onset interval (frames); -1 = unknown
|
| 178 |
prev_ioi = torch.full((B,), -1, dtype=torch.long, device=device)
|
| 179 |
ratios = torch.tensor(LATTICE_RATIOS, device=device)
|
| 180 |
-
time_pos = torch.arange(WINDOW, device=device).unsqueeze(0)
|
| 181 |
-
|
| 182 |
-
if model.training or not hasattr(model, "init_decode_cache"):
|
| 183 |
-
raise RuntimeError("decode_windows requires an eval model with SDPA/KV-cache decoding")
|
| 184 |
|
| 185 |
-
|
| 186 |
-
|
| 187 |
-
|
| 188 |
-
for p in range(plen):
|
| 189 |
-
logits_all = model.decode_cached_step(seqs[:, p], cache).float()
|
| 190 |
-
|
| 191 |
-
gen_toks = torch.full((B, MAX_TGT - plen), v.pad, dtype=torch.long, device=device)
|
| 192 |
-
n_steps = 0
|
| 193 |
-
for step in range(MAX_TGT - plen):
|
| 194 |
if use_cfg:
|
| 195 |
logits = logits_all[B:] + cfg_w * (logits_all[:B] - logits_all[B:])
|
| 196 |
else:
|
|
@@ -202,6 +220,7 @@ def decode_windows(model, mels, prefixes, device="cuda", temperature=1.0, top_p=
|
|
| 202 |
mask[expect_note] = note_row
|
| 203 |
idx = (~expect_note).nonzero(as_tuple=True)[0]
|
| 204 |
if len(idx):
|
|
|
|
| 205 |
sub = torch.full((len(idx), logits.shape[1]), float("-inf"), device=device)
|
| 206 |
sub[:, v.eos] = 0.0
|
| 207 |
ok = time_pos >= (last_time[idx] + min_gap).unsqueeze(1)
|
|
@@ -245,9 +264,8 @@ def decode_windows(model, mels, prefixes, device="cuda", temperature=1.0, top_p=
|
|
| 245 |
nxt = si[torch.arange(B, device=device), pick]
|
| 246 |
|
| 247 |
nxt = torch.where(done, torch.full_like(nxt, v.pad), nxt)
|
| 248 |
-
gen_toks[:, step] = nxt
|
| 249 |
-
n_steps = step + 1
|
| 250 |
step_tok = torch.cat([nxt, nxt], dim=0) if use_cfg else nxt
|
|
|
|
| 251 |
is_time = (nxt >= v.time0) & (nxt < v.time0 + WINDOW)
|
| 252 |
new_time = nxt - v.time0
|
| 253 |
upd = is_time & (last_time >= 0)
|
|
@@ -257,12 +275,10 @@ def decode_windows(model, mels, prefixes, device="cuda", temperature=1.0, top_p=
|
|
| 257 |
done = done | (nxt == v.eos)
|
| 258 |
if done.all():
|
| 259 |
break
|
| 260 |
-
with _autocast(device):
|
| 261 |
-
logits_all = model.decode_cached_step(step_tok, cache).float()
|
| 262 |
|
| 263 |
out = []
|
| 264 |
for b in range(B):
|
| 265 |
-
toks =
|
| 266 |
events = []
|
| 267 |
cur_t = None
|
| 268 |
for t in toks:
|
|
|
|
| 24 |
|
| 25 |
|
| 26 |
def _autocast(device):
|
| 27 |
+
"""bf16 autocast on CUDA; no-op elsewhere (MPS/CPU run fp32)."""
|
| 28 |
import contextlib
|
| 29 |
|
| 30 |
+
if _dev_type(device) == "cuda":
|
| 31 |
+
return torch.autocast("cuda", dtype=torch.bfloat16)
|
|
|
|
| 32 |
return contextlib.nullcontext()
|
| 33 |
CTX_LEN = 12 # fixed-length previous-window context (note-type tokens)
|
| 34 |
MIN_GAP_FRAMES = 4 # ~46 ms minimum inter-note gap enforced during decoding
|
|
|
|
| 41 |
def load_model(ckpt_path, device="cuda"):
|
| 42 |
ck = torch.load(ckpt_path, map_location=device)
|
| 43 |
sd = ck["model"]
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 44 |
vocab_size = sd["tok_emb.weight"].shape[0]
|
| 45 |
aux = any(k.startswith("aux") for k in sd)
|
| 46 |
gctx = any(k.startswith("gsum_proj") for k in sd)
|
| 47 |
+
a = ck.get("args", {}) or {}
|
| 48 |
+
front_ch = sd["frontend.0.weight"].shape[1] # 128 (clean) or 130 (legacy slot)
|
| 49 |
emb_w = sd["tok_emb.weight"].shape[1]
|
| 50 |
d_model = a.get("d_model", 512)
|
| 51 |
+
# clean_phase (v1.6): encoder frontend is 128ch and a phase_proj is present.
|
| 52 |
+
clean_phase = bool(a.get("clean_phase", False)) or ("phase_proj.0.weight" in sd)
|
| 53 |
+
# legacy slot models fed phase to the encoder (in_ch=130); clean models keep
|
| 54 |
+
# in_ch=130 at the data level but only 128 reach the frontend.
|
| 55 |
+
in_ch = 130 if (clean_phase or front_ch == 130) else front_ch
|
| 56 |
+
# v1.7 arch flags: prefer saved args; fall back to state-dict keys so
|
| 57 |
+
# args-less checkpoints of the new arch still load
|
| 58 |
+
adapter_rank_ffn = int(a.get("adapter_rank_ffn", 0) or 0)
|
| 59 |
+
if not adapter_rank_ffn and "decoder.ffn_adapters.0.down.weight" in sd:
|
| 60 |
+
adapter_rank_ffn = sd["decoder.ffn_adapters.0.down.weight"].shape[0]
|
| 61 |
+
depth_emb = bool(a.get("depth_emb", False)) or ("decoder.depth_emb" in sd)
|
| 62 |
model = ChartModel(
|
| 63 |
+
d_model=d_model, nhead=a.get("nhead", 8),
|
| 64 |
+
enc_layers=a.get("enc_layers", 6),
|
| 65 |
dec_layers=a.get("dec_layers", 6), ffn=a.get("ffn", 2048),
|
| 66 |
vocab_size=vocab_size, aux=aux, global_ctx=gctx,
|
| 67 |
func_time=a.get("func_time", False), in_ch=in_ch,
|
| 68 |
emb_factor=emb_w if emb_w != d_model else None,
|
| 69 |
+
clean_phase=clean_phase,
|
| 70 |
+
enc_share=a.get("enc_share") or None,
|
| 71 |
+
dec_share=a.get("dec_share") or None,
|
| 72 |
+
adapter_rank=int(a.get("adapter_rank", 0) or 0),
|
| 73 |
+
unique_layernorm=bool(a.get("unique_layernorm", False)),
|
| 74 |
+
adapter_rank_ffn=adapter_rank_ffn, depth_emb=depth_emb,
|
| 75 |
+
unshare_last_dec=bool(a.get("unshare_last_dec", False)),
|
| 76 |
).to(device)
|
| 77 |
if any(k.startswith("ptr") for k in sd):
|
| 78 |
model.enable_ptr()
|
|
|
|
| 81 |
model.to(device) # newly enabled heads default to CPU
|
| 82 |
model.load_state_dict(sd)
|
| 83 |
model.eval()
|
| 84 |
+
# capability detection: prefer explicit training args (vocab-size inference
|
| 85 |
+
# wrongly added prefix slots to models that never trained with them)
|
| 86 |
+
if a:
|
| 87 |
+
model._has_ctx = bool(a.get("ctx", False))
|
| 88 |
+
model._has_sib = bool(a.get("sibling", False))
|
| 89 |
+
model._has_style = bool(a.get("style", False))
|
| 90 |
+
model._has_sync = bool(a.get("sync_token", False))
|
| 91 |
+
model._has_plan = bool(a.get("plan", False))
|
| 92 |
+
model._dual = bool(a.get("dual", False))
|
| 93 |
+
model._slot = bool(a.get("slot", False)) or model._dual
|
| 94 |
+
else: # legacy fallback
|
| 95 |
+
model._has_ctx = vocab_size > VOCAB.sep
|
| 96 |
+
model._has_sib = vocab_size > VOCAB.sib
|
| 97 |
+
model._has_style = False
|
| 98 |
+
model._has_sync = False
|
| 99 |
+
model._has_plan = False
|
| 100 |
+
model._slot = False
|
| 101 |
+
model._dual = False
|
| 102 |
model._has_gctx = gctx
|
| 103 |
return model
|
| 104 |
|
|
|
|
| 180 |
gen_device = "cpu" if _dev_type(device) == "mps" else device
|
| 181 |
gen = torch.Generator(device=gen_device)
|
| 182 |
gen.manual_seed(seed)
|
| 183 |
+
mels_dev = mels.to(device)
|
| 184 |
with _autocast(device):
|
| 185 |
memory = model.encode(
|
| 186 |
+
mels_dev,
|
| 187 |
gsum=gsum.to(device) if gsum is not None else None,
|
| 188 |
pos_bucket=pos_buckets.to(device) if pos_buckets is not None else None,
|
| 189 |
)
|
| 190 |
+
# clean_phase (v1.6): inject the grid-phase channels into the memory the
|
| 191 |
+
# DECODER reads. No-op for legacy models (in_ch already carried phase).
|
| 192 |
+
memory = model.phase_mem(memory, mels_dev)
|
| 193 |
|
| 194 |
seqs = torch.tensor(prefixes, dtype=torch.long, device=device)
|
| 195 |
if use_cfg: # rows B..2B: same ctx, conditions replaced by UNK
|
|
|
|
| 205 |
# lattice state: previous inter-onset interval (frames); -1 = unknown
|
| 206 |
prev_ioi = torch.full((B,), -1, dtype=torch.long, device=device)
|
| 207 |
ratios = torch.tensor(LATTICE_RATIOS, device=device)
|
|
|
|
|
|
|
|
|
|
|
|
|
| 208 |
|
| 209 |
+
for _ in range(MAX_TGT - plen):
|
| 210 |
+
with _autocast(device):
|
| 211 |
+
logits_all = model.decode(seqs, memory)[:, -1].float()
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 212 |
if use_cfg:
|
| 213 |
logits = logits_all[B:] + cfg_w * (logits_all[:B] - logits_all[B:])
|
| 214 |
else:
|
|
|
|
| 220 |
mask[expect_note] = note_row
|
| 221 |
idx = (~expect_note).nonzero(as_tuple=True)[0]
|
| 222 |
if len(idx):
|
| 223 |
+
time_pos = torch.arange(WINDOW, device=device).unsqueeze(0)
|
| 224 |
sub = torch.full((len(idx), logits.shape[1]), float("-inf"), device=device)
|
| 225 |
sub[:, v.eos] = 0.0
|
| 226 |
ok = time_pos >= (last_time[idx] + min_gap).unsqueeze(1)
|
|
|
|
| 264 |
nxt = si[torch.arange(B, device=device), pick]
|
| 265 |
|
| 266 |
nxt = torch.where(done, torch.full_like(nxt, v.pad), nxt)
|
|
|
|
|
|
|
| 267 |
step_tok = torch.cat([nxt, nxt], dim=0) if use_cfg else nxt
|
| 268 |
+
seqs = torch.cat([seqs, step_tok.unsqueeze(1)], dim=1)
|
| 269 |
is_time = (nxt >= v.time0) & (nxt < v.time0 + WINDOW)
|
| 270 |
new_time = nxt - v.time0
|
| 271 |
upd = is_time & (last_time >= 0)
|
|
|
|
| 275 |
done = done | (nxt == v.eos)
|
| 276 |
if done.all():
|
| 277 |
break
|
|
|
|
|
|
|
| 278 |
|
| 279 |
out = []
|
| 280 |
for b in range(B):
|
| 281 |
+
toks = seqs[b, plen:].tolist()
|
| 282 |
events = []
|
| 283 |
cur_t = None
|
| 284 |
for t in toks:
|
softchart/hf.py
CHANGED
|
@@ -32,12 +32,24 @@ class SoftChartGenerator(
|
|
| 32 |
def __init__(self, d_model=256, enc_layers=4, dec_layers=4, ffn=1024,
|
| 33 |
vocab_size=1802, aux=True, global_ctx=False, func_time=False,
|
| 34 |
ptr=False, beat=False, in_ch=None, emb_factor=None,
|
| 35 |
-
capabilities=None
|
|
|
|
|
|
|
|
|
|
| 36 |
super().__init__()
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 37 |
self.net = ChartModel(
|
| 38 |
-
d_model=d_model, enc_layers=enc_layers, dec_layers=dec_layers,
|
| 39 |
-
vocab_size=vocab_size, aux=aux, global_ctx=global_ctx,
|
| 40 |
-
ptr=ptr, beat=beat, in_ch=in_ch, emb_factor=emb_factor,
|
|
|
|
|
|
|
|
|
|
|
|
|
| 41 |
)
|
| 42 |
# capabilities: which condition axes this checkpoint was trained with,
|
| 43 |
# so inference knows which prefix tokens to emit
|
|
|
|
| 32 |
def __init__(self, d_model=256, enc_layers=4, dec_layers=4, ffn=1024,
|
| 33 |
vocab_size=1802, aux=True, global_ctx=False, func_time=False,
|
| 34 |
ptr=False, beat=False, in_ch=None, emb_factor=None,
|
| 35 |
+
capabilities=None, nhead=8, clean_phase=False,
|
| 36 |
+
enc_share=None, dec_share=None, adapter_rank=0,
|
| 37 |
+
unique_layernorm=False, adapter_rank_ffn=0, depth_emb=False,
|
| 38 |
+
unshare_last_dec=False):
|
| 39 |
super().__init__()
|
| 40 |
+
# nhead and the v1.6/v1.7 sharing flags (enc_share/dec_share/adapter_rank/
|
| 41 |
+
# unique_layernorm/adapter_rank_ffn/depth_emb/unshare_last_dec) and
|
| 42 |
+
# clean_phase must be persisted in config.json: from_pretrained rebuilds
|
| 43 |
+
# the architecture from these kwargs alone, before any weights are loaded.
|
| 44 |
+
# Defaults reproduce the pre-v1.6 unshared architecture for older repos.
|
| 45 |
self.net = ChartModel(
|
| 46 |
+
d_model=d_model, nhead=nhead, enc_layers=enc_layers, dec_layers=dec_layers,
|
| 47 |
+
ffn=ffn, vocab_size=vocab_size, aux=aux, global_ctx=global_ctx,
|
| 48 |
+
func_time=func_time, ptr=ptr, beat=beat, in_ch=in_ch, emb_factor=emb_factor,
|
| 49 |
+
clean_phase=clean_phase, enc_share=enc_share, dec_share=dec_share,
|
| 50 |
+
adapter_rank=adapter_rank, unique_layernorm=unique_layernorm,
|
| 51 |
+
adapter_rank_ffn=adapter_rank_ffn, depth_emb=depth_emb,
|
| 52 |
+
unshare_last_dec=unshare_last_dec,
|
| 53 |
)
|
| 54 |
# capabilities: which condition axes this checkpoint was trained with,
|
| 55 |
# so inference knows which prefix tokens to emit
|
softchart/model.py
CHANGED
|
@@ -4,7 +4,6 @@ import math
|
|
| 4 |
|
| 5 |
import torch
|
| 6 |
import torch.nn as nn
|
| 7 |
-
import torch.nn.functional as F
|
| 8 |
|
| 9 |
from .vocab import N_MELS, VOCAB, WINDOW
|
| 10 |
|
|
@@ -34,6 +33,164 @@ class BeatHiRes(nn.Module):
|
|
| 34 |
return self.out(h).transpose(1, 2)
|
| 35 |
|
| 36 |
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 37 |
class ChartModel(nn.Module):
|
| 38 |
def __init__(
|
| 39 |
self,
|
|
@@ -52,6 +209,14 @@ class ChartModel(nn.Module):
|
|
| 52 |
beat=False,
|
| 53 |
in_ch=None,
|
| 54 |
emb_factor=None,
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 55 |
):
|
| 56 |
super().__init__()
|
| 57 |
from .vocab import MAX_TGT
|
|
@@ -59,11 +224,17 @@ class ChartModel(nn.Module):
|
|
| 59 |
vocab_size = vocab_size or VOCAB.size
|
| 60 |
max_tgt = max_tgt or MAX_TGT
|
| 61 |
self.d_model = d_model
|
| 62 |
-
|
| 63 |
-
|
| 64 |
-
#
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 65 |
self.frontend = nn.Sequential(
|
| 66 |
-
nn.Conv1d(
|
| 67 |
nn.GELU(),
|
| 68 |
nn.Conv1d(d_model, d_model, 3, stride=2, padding=1),
|
| 69 |
nn.GELU(),
|
|
@@ -72,14 +243,51 @@ class ChartModel(nn.Module):
|
|
| 72 |
self.register_buffer("enc_pos", sinusoidal(enc_len, d_model), persistent=False)
|
| 73 |
self.register_buffer("dec_pos", sinusoidal(max_tgt, d_model), persistent=False)
|
| 74 |
|
| 75 |
-
|
| 76 |
-
|
| 77 |
-
)
|
| 78 |
-
|
| 79 |
-
|
| 80 |
-
|
| 81 |
-
|
| 82 |
-
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 83 |
|
| 84 |
# ALBERT-style factorized embeddings (v1.5): vocab->E->d. Position
|
| 85 |
# tokens dominate the table (1728/1810 rows), so E=64 cuts the
|
|
@@ -132,9 +340,13 @@ class ChartModel(nn.Module):
|
|
| 132 |
self.time_proj = None
|
| 133 |
|
| 134 |
def encode(self, mel, gsum=None, pos_bucket=None):
|
| 135 |
-
# mel: (B,
|
| 136 |
-
#
|
| 137 |
-
|
|
|
|
|
|
|
|
|
|
|
|
|
| 138 |
h = h + self.enc_pos[: h.shape[1]]
|
| 139 |
if self.gsum_proj is not None and gsum is not None:
|
| 140 |
g = self.gsum_proj(gsum) + self.seg_emb[1] + self.gpos[: gsum.shape[1]]
|
|
@@ -143,6 +355,24 @@ class ChartModel(nn.Module):
|
|
| 143 |
h = torch.cat([p, g, h], dim=1)
|
| 144 |
return self.encoder(self.dropout(h))
|
| 145 |
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 146 |
def emb_matrix(self):
|
| 147 |
"""Full embedding matrix; factorized and/or functional rows resolved."""
|
| 148 |
E = self.tok_emb.weight
|
|
@@ -153,146 +383,6 @@ class ChartModel(nn.Module):
|
|
| 153 |
E = torch.cat([E[: VOCAB.time0], t, E[VOCAB.time0 + WINDOW :]], dim=0)
|
| 154 |
return E
|
| 155 |
|
| 156 |
-
@staticmethod
|
| 157 |
-
def _check_mha_supported(mha):
|
| 158 |
-
if not getattr(mha, "_qkv_same_embed_dim", True):
|
| 159 |
-
raise NotImplementedError("cached decode requires same q/k/v embed dims")
|
| 160 |
-
if getattr(mha, "bias_k", None) is not None or getattr(mha, "bias_v", None) is not None:
|
| 161 |
-
raise NotImplementedError("cached decode does not support bias_k/bias_v")
|
| 162 |
-
if getattr(mha, "add_zero_attn", False):
|
| 163 |
-
raise NotImplementedError("cached decode does not support add_zero_attn")
|
| 164 |
-
|
| 165 |
-
@staticmethod
|
| 166 |
-
def _split_heads(x, mha):
|
| 167 |
-
B, L, D = x.shape
|
| 168 |
-
H = mha.num_heads
|
| 169 |
-
return x.view(B, L, H, D // H).transpose(1, 2).contiguous()
|
| 170 |
-
|
| 171 |
-
@staticmethod
|
| 172 |
-
def _merge_heads(x):
|
| 173 |
-
B, H, L, Dh = x.shape
|
| 174 |
-
return x.transpose(1, 2).contiguous().view(B, L, H * Dh)
|
| 175 |
-
|
| 176 |
-
@staticmethod
|
| 177 |
-
def _in_proj_qkv(mha, x):
|
| 178 |
-
ChartModel._check_mha_supported(mha)
|
| 179 |
-
D = mha.embed_dim
|
| 180 |
-
w = mha.in_proj_weight
|
| 181 |
-
b = mha.in_proj_bias
|
| 182 |
-
q = F.linear(x, w[:D], b[:D] if b is not None else None)
|
| 183 |
-
k = F.linear(x, w[D:2 * D], b[D:2 * D] if b is not None else None)
|
| 184 |
-
v = F.linear(x, w[2 * D:], b[2 * D:] if b is not None else None)
|
| 185 |
-
return q, k, v
|
| 186 |
-
|
| 187 |
-
@staticmethod
|
| 188 |
-
def _in_proj_q(mha, x):
|
| 189 |
-
ChartModel._check_mha_supported(mha)
|
| 190 |
-
D = mha.embed_dim
|
| 191 |
-
b = mha.in_proj_bias
|
| 192 |
-
return F.linear(x, mha.in_proj_weight[:D], b[:D] if b is not None else None)
|
| 193 |
-
|
| 194 |
-
@staticmethod
|
| 195 |
-
def _in_proj_kv(mha, x):
|
| 196 |
-
ChartModel._check_mha_supported(mha)
|
| 197 |
-
D = mha.embed_dim
|
| 198 |
-
w = mha.in_proj_weight
|
| 199 |
-
b = mha.in_proj_bias
|
| 200 |
-
k = F.linear(x, w[D:2 * D], b[D:2 * D] if b is not None else None)
|
| 201 |
-
v = F.linear(x, w[2 * D:], b[2 * D:] if b is not None else None)
|
| 202 |
-
return ChartModel._split_heads(k, mha), ChartModel._split_heads(v, mha)
|
| 203 |
-
|
| 204 |
-
@staticmethod
|
| 205 |
-
def _sdpa(mha, q, k, v, dropout_p=0.0):
|
| 206 |
-
q = ChartModel._split_heads(q, mha)
|
| 207 |
-
y = F.scaled_dot_product_attention(
|
| 208 |
-
q, k, v, dropout_p=dropout_p, is_causal=False)
|
| 209 |
-
y = ChartModel._merge_heads(y)
|
| 210 |
-
return mha.out_proj(y)
|
| 211 |
-
|
| 212 |
-
def init_decode_cache(self, memory, max_len=None):
|
| 213 |
-
"""Prepare an inference-only KV cache for decode_cached_step().
|
| 214 |
-
|
| 215 |
-
This adds no parameters or buffers: all projections read the existing
|
| 216 |
-
nn.TransformerDecoderLayer weights.
|
| 217 |
-
Cross-attention K/V are static for a window, so they are projected once.
|
| 218 |
-
Self-attention K/V are preallocated lazily on the first generated token.
|
| 219 |
-
"""
|
| 220 |
-
max_len = max_len or self.dec_pos.shape[0]
|
| 221 |
-
layers = []
|
| 222 |
-
for layer in self.decoder.layers:
|
| 223 |
-
mem_k, mem_v = self._in_proj_kv(layer.multihead_attn, memory)
|
| 224 |
-
layers.append({
|
| 225 |
-
"self_k": None,
|
| 226 |
-
"self_v": None,
|
| 227 |
-
"mem_k": mem_k,
|
| 228 |
-
"mem_v": mem_v,
|
| 229 |
-
})
|
| 230 |
-
return {
|
| 231 |
-
"layers": layers,
|
| 232 |
-
"pos": 0,
|
| 233 |
-
"max_len": int(max_len),
|
| 234 |
-
"E": self.emb_matrix(),
|
| 235 |
-
}
|
| 236 |
-
|
| 237 |
-
def _cached_sa_block(self, layer, x, cache_layer, pos):
|
| 238 |
-
q, k, v = self._in_proj_qkv(layer.self_attn, x)
|
| 239 |
-
k = self._split_heads(k, layer.self_attn)
|
| 240 |
-
v = self._split_heads(v, layer.self_attn)
|
| 241 |
-
if cache_layer["self_k"] is None:
|
| 242 |
-
B, H, _, Dh = k.shape
|
| 243 |
-
shape = (B, H, cache_layer["max_len"], Dh)
|
| 244 |
-
cache_layer["self_k"] = k.new_empty(shape)
|
| 245 |
-
cache_layer["self_v"] = v.new_empty(shape)
|
| 246 |
-
cache_layer["self_k"][:, :, pos:pos + 1] = k
|
| 247 |
-
cache_layer["self_v"][:, :, pos:pos + 1] = v
|
| 248 |
-
k_all = cache_layer["self_k"][:, :, :pos + 1]
|
| 249 |
-
v_all = cache_layer["self_v"][:, :, :pos + 1]
|
| 250 |
-
p = layer.self_attn.dropout if self.training else 0.0
|
| 251 |
-
y = self._sdpa(layer.self_attn, q, k_all, v_all, dropout_p=p)
|
| 252 |
-
return layer.dropout1(y)
|
| 253 |
-
|
| 254 |
-
def _cached_mha_block(self, layer, x, cache_layer):
|
| 255 |
-
q = self._in_proj_q(layer.multihead_attn, x)
|
| 256 |
-
p = layer.multihead_attn.dropout if self.training else 0.0
|
| 257 |
-
y = self._sdpa(
|
| 258 |
-
layer.multihead_attn, q, cache_layer["mem_k"], cache_layer["mem_v"],
|
| 259 |
-
dropout_p=p,
|
| 260 |
-
)
|
| 261 |
-
return layer.dropout2(y)
|
| 262 |
-
|
| 263 |
-
@staticmethod
|
| 264 |
-
def _ff_block(layer, x):
|
| 265 |
-
return layer.dropout3(layer.linear2(layer.dropout(layer.activation(layer.linear1(x)))))
|
| 266 |
-
|
| 267 |
-
def decode_cached_step(self, token, cache):
|
| 268 |
-
"""Incremental decoder step using SDPA and KV cache.
|
| 269 |
-
|
| 270 |
-
token: (B,) or (B, 1) token ids for the current position. Returns
|
| 271 |
-
next-token logits for that position, matching decode(... )[:, -1] in
|
| 272 |
-
eval mode up to backend floating-point noise.
|
| 273 |
-
"""
|
| 274 |
-
if token.dim() == 1:
|
| 275 |
-
token = token[:, None]
|
| 276 |
-
pos = int(cache["pos"])
|
| 277 |
-
if pos >= cache["max_len"] or pos >= self.dec_pos.shape[0]:
|
| 278 |
-
raise RuntimeError("decoder cache exceeded configured max length")
|
| 279 |
-
E = cache["E"]
|
| 280 |
-
h = F.embedding(token, E) * math.sqrt(self.d_model) + self.dec_pos[pos:pos + 1]
|
| 281 |
-
for layer, lc in zip(self.decoder.layers, cache["layers"]):
|
| 282 |
-
lc["max_len"] = cache["max_len"]
|
| 283 |
-
if layer.norm_first:
|
| 284 |
-
h = h + self._cached_sa_block(layer, layer.norm1(h), lc, pos)
|
| 285 |
-
h = h + self._cached_mha_block(layer, layer.norm2(h), lc)
|
| 286 |
-
h = h + self._ff_block(layer, layer.norm3(h))
|
| 287 |
-
else:
|
| 288 |
-
h = layer.norm1(h + self._cached_sa_block(layer, h, lc, pos))
|
| 289 |
-
h = layer.norm2(h + self._cached_mha_block(layer, h, lc))
|
| 290 |
-
h = layer.norm3(h + self._ff_block(layer, h))
|
| 291 |
-
if self.decoder.norm is not None:
|
| 292 |
-
h = self.decoder.norm(h)
|
| 293 |
-
cache["pos"] = pos + 1
|
| 294 |
-
return (h @ E.T).squeeze(1)
|
| 295 |
-
|
| 296 |
def decode(self, tgt_in, memory):
|
| 297 |
# tgt_in: (B, L) token ids
|
| 298 |
L = tgt_in.shape[1]
|
|
@@ -319,19 +409,21 @@ class ChartModel(nn.Module):
|
|
| 319 |
return_extras=False, in_tgt=None):
|
| 320 |
# in_tgt: optional corrupted decoder input (e.g. MASKed types for the
|
| 321 |
# skeleton->color infill curriculum); gold targets stay = tgt
|
| 322 |
-
memory = self.encode(mel, gsum=gsum, pos_bucket=pos_bucket)
|
|
|
|
| 323 |
dec_in = (in_tgt if in_tgt is not None else tgt)[:, :-1]
|
| 324 |
L = dec_in.shape[1]
|
| 325 |
E = self.emb_matrix()
|
| 326 |
import math as _m
|
| 327 |
h = nn.functional.embedding(dec_in, E) * _m.sqrt(self.d_model) + self.dec_pos[:L]
|
| 328 |
dm = nn.Transformer.generate_square_subsequent_mask(L, device=tgt.device)
|
| 329 |
-
h = self.decoder(self.dropout(h),
|
| 330 |
tgt_key_padding_mask=dec_in == VOCAB.pad, tgt_is_causal=True)
|
| 331 |
logits = h @ E.T
|
| 332 |
if not (return_aux or return_extras):
|
| 333 |
return logits
|
| 334 |
outs = [logits]
|
|
|
|
| 335 |
aux_mem = memory[:, -(WINDOW // 4):]
|
| 336 |
outs.append(self.aux(aux_mem).squeeze(-1) if self.aux is not None else None)
|
| 337 |
if return_extras:
|
|
|
|
| 4 |
|
| 5 |
import torch
|
| 6 |
import torch.nn as nn
|
|
|
|
| 7 |
|
| 8 |
from .vocab import N_MELS, VOCAB, WINDOW
|
| 9 |
|
|
|
|
| 33 |
return self.out(h).transpose(1, 2)
|
| 34 |
|
| 35 |
|
| 36 |
+
def _parse_pattern(spec, n):
|
| 37 |
+
"""Layer-sharing schedule. spec is either an int (number of UNIQUE
|
| 38 |
+
physical layers, cycled to length n) or a comma string like "E0,E1,E0,E1"
|
| 39 |
+
(labels; each distinct label is one physical layer). Returns a list of n
|
| 40 |
+
physical-layer indices, e.g. [0,1,0,1] for 2-unique/4-physical-depth."""
|
| 41 |
+
if spec is None:
|
| 42 |
+
return list(range(n)) # no sharing: every layer unique
|
| 43 |
+
if isinstance(spec, int):
|
| 44 |
+
return [k % spec for k in range(n)]
|
| 45 |
+
labels = [s.strip() for s in str(spec).split(",") if s.strip()]
|
| 46 |
+
assert len(labels) == n, f"pattern {spec} has {len(labels)} labels, need {n}"
|
| 47 |
+
order = {}
|
| 48 |
+
out = []
|
| 49 |
+
for lb in labels:
|
| 50 |
+
if lb not in order:
|
| 51 |
+
order[lb] = len(order)
|
| 52 |
+
out.append(order[lb])
|
| 53 |
+
return out
|
| 54 |
+
|
| 55 |
+
|
| 56 |
+
class Adapter(nn.Module):
|
| 57 |
+
"""Rank-r residual adapter applied after a (possibly reused) transformer
|
| 58 |
+
block: x + gate * up(gelu(down(x))). gate init 0 so the whole stack starts
|
| 59 |
+
identical to hard weight sharing; the adapter only differentiates a reused
|
| 60 |
+
physical layer as training moves the gate off zero."""
|
| 61 |
+
|
| 62 |
+
def __init__(self, d, rank):
|
| 63 |
+
super().__init__()
|
| 64 |
+
self.down = nn.Linear(d, rank, bias=False)
|
| 65 |
+
self.up = nn.Linear(rank, d, bias=False)
|
| 66 |
+
nn.init.normal_(self.down.weight, std=0.02)
|
| 67 |
+
nn.init.zeros_(self.up.weight)
|
| 68 |
+
self.gate = nn.Parameter(torch.zeros(1))
|
| 69 |
+
|
| 70 |
+
def forward(self, x):
|
| 71 |
+
return x + self.gate * self.up(nn.functional.gelu(self.down(x)))
|
| 72 |
+
|
| 73 |
+
|
| 74 |
+
class SharedEncoder(nn.Module):
|
| 75 |
+
"""Encoder whose physical layers are reused according to a sharing pattern.
|
| 76 |
+
Each POSITION in the depth keeps its own (a) rank-r adapter and (b) final
|
| 77 |
+
LayerNorm — so ΔW starts at 0 (hard sharing) but every position can drift.
|
| 78 |
+
With pattern=None this is exactly n unique layers + a norm, matching the
|
| 79 |
+
stock nn.TransformerEncoder param count (adapters/pos-norms are opt-in).
|
| 80 |
+
|
| 81 |
+
v1.7 depth-diversity options (all default-off, init == v1.6 behaviour):
|
| 82 |
+
depth_emb: per-depth learned d-vector ADDED to the block input at each
|
| 83 |
+
reuse (params = depth x d, zero-init so init == hard sharing).
|
| 84 |
+
adapter_rank_ffn > 0: SPLIT adapters — instead of one post-block adapter,
|
| 85 |
+
each depth owns a rank-`adapter_rank` adapter on the attention sublayer
|
| 86 |
+
output and a rank-`adapter_rank_ffn` adapter on the FFN sublayer output
|
| 87 |
+
(patterning lives in the FFN). Gate init 0 keeps init == hard sharing.
|
| 88 |
+
"""
|
| 89 |
+
|
| 90 |
+
def __init__(self, d, nhead, ffn, dropout, pattern, adapter_rank=0,
|
| 91 |
+
unique_layernorm=False, adapter_rank_ffn=0, depth_emb=False):
|
| 92 |
+
super().__init__()
|
| 93 |
+
# `pattern` is the resolved list of physical-layer indices per depth
|
| 94 |
+
n_phys = max(pattern) + 1
|
| 95 |
+
self.layers = nn.ModuleList([
|
| 96 |
+
nn.TransformerEncoderLayer(d, nhead, ffn, dropout, activation="gelu",
|
| 97 |
+
batch_first=True, norm_first=True)
|
| 98 |
+
for _ in range(n_phys)])
|
| 99 |
+
self.plan = pattern
|
| 100 |
+
self.split = adapter_rank_ffn > 0
|
| 101 |
+
if self.split: # v1.7: per-sublayer adapters (attention rank-a, FFN rank-f)
|
| 102 |
+
self.attn_adapters = nn.ModuleList([
|
| 103 |
+
Adapter(d, adapter_rank) if adapter_rank else nn.Identity()
|
| 104 |
+
for _ in pattern])
|
| 105 |
+
self.ffn_adapters = nn.ModuleList([
|
| 106 |
+
Adapter(d, adapter_rank_ffn) for _ in pattern])
|
| 107 |
+
else: # v1.6: one residual adapter after the whole block
|
| 108 |
+
self.adapters = nn.ModuleList([
|
| 109 |
+
Adapter(d, adapter_rank) if adapter_rank else nn.Identity()
|
| 110 |
+
for _ in pattern])
|
| 111 |
+
self.depth_emb = (nn.Parameter(torch.zeros(len(pattern), d))
|
| 112 |
+
if depth_emb else None)
|
| 113 |
+
self.pos_norms = nn.ModuleList([
|
| 114 |
+
nn.LayerNorm(d) if unique_layernorm else nn.Identity()
|
| 115 |
+
for _ in pattern])
|
| 116 |
+
self.norm = nn.LayerNorm(d)
|
| 117 |
+
|
| 118 |
+
def forward(self, x, src_key_padding_mask=None):
|
| 119 |
+
for k, phys in enumerate(self.plan):
|
| 120 |
+
if self.depth_emb is not None:
|
| 121 |
+
x = x + self.depth_emb[k]
|
| 122 |
+
layer = self.layers[phys]
|
| 123 |
+
if self.split:
|
| 124 |
+
# norm_first decomposition of nn.TransformerEncoderLayer with
|
| 125 |
+
# per-depth adapters applied to each sublayer OUTPUT (Houlsby
|
| 126 |
+
# placement); at gate=0 this is bit-equal to the stock layer
|
| 127 |
+
x = x + self.attn_adapters[k](
|
| 128 |
+
layer._sa_block(layer.norm1(x), None, src_key_padding_mask))
|
| 129 |
+
x = x + self.ffn_adapters[k](layer._ff_block(layer.norm2(x)))
|
| 130 |
+
else:
|
| 131 |
+
x = layer(x, src_key_padding_mask=src_key_padding_mask)
|
| 132 |
+
x = self.adapters[k](x)
|
| 133 |
+
x = self.pos_norms[k](x)
|
| 134 |
+
return self.norm(x)
|
| 135 |
+
|
| 136 |
+
|
| 137 |
+
class SharedDecoder(nn.Module):
|
| 138 |
+
"""Decoder counterpart of SharedEncoder (see that docstring). Split mode
|
| 139 |
+
puts rank-`adapter_rank` adapters on BOTH attention sublayers (self and
|
| 140 |
+
cross) and rank-`adapter_rank_ffn` on the FFN sublayer."""
|
| 141 |
+
|
| 142 |
+
def __init__(self, d, nhead, ffn, dropout, pattern, adapter_rank=0,
|
| 143 |
+
unique_layernorm=False, adapter_rank_ffn=0, depth_emb=False):
|
| 144 |
+
super().__init__()
|
| 145 |
+
n_phys = max(pattern) + 1
|
| 146 |
+
self.layers = nn.ModuleList([
|
| 147 |
+
nn.TransformerDecoderLayer(d, nhead, ffn, dropout, activation="gelu",
|
| 148 |
+
batch_first=True, norm_first=True)
|
| 149 |
+
for _ in range(n_phys)])
|
| 150 |
+
self.plan = pattern
|
| 151 |
+
self.split = adapter_rank_ffn > 0
|
| 152 |
+
if self.split:
|
| 153 |
+
self.sa_adapters = nn.ModuleList([
|
| 154 |
+
Adapter(d, adapter_rank) if adapter_rank else nn.Identity()
|
| 155 |
+
for _ in pattern])
|
| 156 |
+
self.ca_adapters = nn.ModuleList([
|
| 157 |
+
Adapter(d, adapter_rank) if adapter_rank else nn.Identity()
|
| 158 |
+
for _ in pattern])
|
| 159 |
+
self.ffn_adapters = nn.ModuleList([
|
| 160 |
+
Adapter(d, adapter_rank_ffn) for _ in pattern])
|
| 161 |
+
else:
|
| 162 |
+
self.adapters = nn.ModuleList([
|
| 163 |
+
Adapter(d, adapter_rank) if adapter_rank else nn.Identity()
|
| 164 |
+
for _ in pattern])
|
| 165 |
+
self.depth_emb = (nn.Parameter(torch.zeros(len(pattern), d))
|
| 166 |
+
if depth_emb else None)
|
| 167 |
+
self.pos_norms = nn.ModuleList([
|
| 168 |
+
nn.LayerNorm(d) if unique_layernorm else nn.Identity()
|
| 169 |
+
for _ in pattern])
|
| 170 |
+
self.norm = nn.LayerNorm(d)
|
| 171 |
+
|
| 172 |
+
def forward(self, tgt, memory, tgt_mask=None, tgt_key_padding_mask=None,
|
| 173 |
+
tgt_is_causal=False):
|
| 174 |
+
x = tgt
|
| 175 |
+
for k, phys in enumerate(self.plan):
|
| 176 |
+
if self.depth_emb is not None:
|
| 177 |
+
x = x + self.depth_emb[k]
|
| 178 |
+
layer = self.layers[phys]
|
| 179 |
+
if self.split:
|
| 180 |
+
x = x + self.sa_adapters[k](layer._sa_block(
|
| 181 |
+
layer.norm1(x), tgt_mask, tgt_key_padding_mask, tgt_is_causal))
|
| 182 |
+
x = x + self.ca_adapters[k](layer._mha_block(
|
| 183 |
+
layer.norm2(x), memory, None, None))
|
| 184 |
+
x = x + self.ffn_adapters[k](layer._ff_block(layer.norm3(x)))
|
| 185 |
+
else:
|
| 186 |
+
x = layer(x, memory, tgt_mask=tgt_mask,
|
| 187 |
+
tgt_key_padding_mask=tgt_key_padding_mask,
|
| 188 |
+
tgt_is_causal=tgt_is_causal)
|
| 189 |
+
x = self.adapters[k](x)
|
| 190 |
+
x = self.pos_norms[k](x)
|
| 191 |
+
return self.norm(x)
|
| 192 |
+
|
| 193 |
+
|
| 194 |
class ChartModel(nn.Module):
|
| 195 |
def __init__(
|
| 196 |
self,
|
|
|
|
| 209 |
beat=False,
|
| 210 |
in_ch=None,
|
| 211 |
emb_factor=None,
|
| 212 |
+
clean_phase=False,
|
| 213 |
+
enc_share=None,
|
| 214 |
+
dec_share=None,
|
| 215 |
+
adapter_rank=0,
|
| 216 |
+
unique_layernorm=False,
|
| 217 |
+
adapter_rank_ffn=0,
|
| 218 |
+
depth_emb=False,
|
| 219 |
+
unshare_last_dec=False,
|
| 220 |
):
|
| 221 |
super().__init__()
|
| 222 |
from .vocab import MAX_TGT
|
|
|
|
| 224 |
vocab_size = vocab_size or VOCAB.size
|
| 225 |
max_tgt = max_tgt or MAX_TGT
|
| 226 |
self.d_model = d_model
|
| 227 |
+
# clean_phase (v1.6): the encoder never sees the 2 grid-phase channels
|
| 228 |
+
# (in_ch stays N_MELS), and phase is injected into the DECODER memory
|
| 229 |
+
# only, via a small projection. This removes the condition leak that
|
| 230 |
+
# forced beat supervision to be masked on slot windows in v1.5.
|
| 231 |
+
self.clean_phase = clean_phase
|
| 232 |
+
self.in_ch = in_ch or N_MELS # legacy slot mode appends 2 grid-phase channels
|
| 233 |
+
front_ch = N_MELS if clean_phase else self.in_ch
|
| 234 |
+
|
| 235 |
+
# conv frontend: (B, front_ch, T) -> (B, T/4, d)
|
| 236 |
self.frontend = nn.Sequential(
|
| 237 |
+
nn.Conv1d(front_ch, d_model, 3, stride=2, padding=1),
|
| 238 |
nn.GELU(),
|
| 239 |
nn.Conv1d(d_model, d_model, 3, stride=2, padding=1),
|
| 240 |
nn.GELU(),
|
|
|
|
| 243 |
self.register_buffer("enc_pos", sinusoidal(enc_len, d_model), persistent=False)
|
| 244 |
self.register_buffer("dec_pos", sinusoidal(max_tgt, d_model), persistent=False)
|
| 245 |
|
| 246 |
+
# phase injection (clean_phase only): 2 grid-phase channels at frame
|
| 247 |
+
# rate -> encoder rate (T/4) -> additive conditioning on the memory the
|
| 248 |
+
# DECODER reads. The beat head reads the pre-injection (clean) memory.
|
| 249 |
+
# phase_proj input = [measure_phase, beat_phase, valid_flag]. valid=0 on
|
| 250 |
+
# time-mode windows (their raw phase is -1), giving the decoder a clean
|
| 251 |
+
# zero-phase conditioning there; slot windows pass valid=1 + real phase.
|
| 252 |
+
if clean_phase:
|
| 253 |
+
self.phase_proj = nn.Sequential(
|
| 254 |
+
nn.Conv1d(3, d_model, 3, stride=2, padding=1),
|
| 255 |
+
nn.GELU(),
|
| 256 |
+
nn.Conv1d(d_model, d_model, 3, stride=2, padding=1),
|
| 257 |
+
)
|
| 258 |
+
else:
|
| 259 |
+
self.phase_proj = None
|
| 260 |
+
|
| 261 |
+
# sharing schedule (v1.6-small/tiny): reuse physical layers per pattern.
|
| 262 |
+
# enc_share/dec_share: int (#unique layers) or "E0,E1,E0,E1" label string.
|
| 263 |
+
self.enc_share = enc_share
|
| 264 |
+
self.dec_share = dec_share
|
| 265 |
+
if (enc_share is None and dec_share is None and not adapter_rank
|
| 266 |
+
and not adapter_rank_ffn and not depth_emb and not unshare_last_dec):
|
| 267 |
+
enc_layer = nn.TransformerEncoderLayer(
|
| 268 |
+
d_model, nhead, ffn, dropout, activation="gelu",
|
| 269 |
+
batch_first=True, norm_first=True)
|
| 270 |
+
self.encoder = nn.TransformerEncoder(enc_layer, enc_layers,
|
| 271 |
+
nn.LayerNorm(d_model))
|
| 272 |
+
dec_layer = nn.TransformerDecoderLayer(
|
| 273 |
+
d_model, nhead, ffn, dropout, activation="gelu",
|
| 274 |
+
batch_first=True, norm_first=True)
|
| 275 |
+
self.decoder = nn.TransformerDecoder(dec_layer, dec_layers,
|
| 276 |
+
nn.LayerNorm(d_model))
|
| 277 |
+
else:
|
| 278 |
+
enc_plan = _parse_pattern(enc_share, enc_layers)
|
| 279 |
+
dec_plan = _parse_pattern(dec_share, dec_layers)
|
| 280 |
+
# v1.7: un-share the LAST decoder layer (pattern-shaping layer in
|
| 281 |
+
# front of the generation head gets its own physical parameters);
|
| 282 |
+
# no-op if that depth is already unique.
|
| 283 |
+
if unshare_last_dec and dec_plan.count(dec_plan[-1]) > 1:
|
| 284 |
+
dec_plan = dec_plan[:-1] + [max(dec_plan) + 1]
|
| 285 |
+
self.encoder = SharedEncoder(d_model, nhead, ffn, dropout, enc_plan,
|
| 286 |
+
adapter_rank, unique_layernorm,
|
| 287 |
+
adapter_rank_ffn, depth_emb)
|
| 288 |
+
self.decoder = SharedDecoder(d_model, nhead, ffn, dropout, dec_plan,
|
| 289 |
+
adapter_rank, unique_layernorm,
|
| 290 |
+
adapter_rank_ffn, depth_emb)
|
| 291 |
|
| 292 |
# ALBERT-style factorized embeddings (v1.5): vocab->E->d. Position
|
| 293 |
# tokens dominate the table (1728/1810 rows), so E=64 cuts the
|
|
|
|
| 340 |
self.time_proj = None
|
| 341 |
|
| 342 |
def encode(self, mel, gsum=None, pos_bucket=None):
|
| 343 |
+
# mel: (B, C, T). clean_phase: C may be N_MELS (audio only) or N_MELS+2
|
| 344 |
+
# (audio + 2 grid-phase channels); only the audio channels reach the
|
| 345 |
+
# encoder, so the returned memory is CLEAN (phase-free) and safe for the
|
| 346 |
+
# beat head. gsum: (B, G, n_mels) whole-song summary chunks;
|
| 347 |
+
# pos_bucket: (B,) window-position bucket in [0, 16).
|
| 348 |
+
audio = mel[:, :N_MELS] if self.clean_phase else mel
|
| 349 |
+
h = self.frontend(audio).transpose(1, 2) # (B, T/4, d)
|
| 350 |
h = h + self.enc_pos[: h.shape[1]]
|
| 351 |
if self.gsum_proj is not None and gsum is not None:
|
| 352 |
g = self.gsum_proj(gsum) + self.seg_emb[1] + self.gpos[: gsum.shape[1]]
|
|
|
|
| 355 |
h = torch.cat([p, g, h], dim=1)
|
| 356 |
return self.encoder(self.dropout(h))
|
| 357 |
|
| 358 |
+
def phase_mem(self, memory, mel):
|
| 359 |
+
"""Phase-conditioned copy of `memory` for the DECODER to cross-attend.
|
| 360 |
+
Adds a projection of the grid-phase channels to the audio-frame slice
|
| 361 |
+
of memory; the beat head keeps reading the clean `memory`. No-op unless
|
| 362 |
+
clean_phase is on and phase channels are present. Time-mode windows
|
| 363 |
+
(raw phase = -1) get valid=0 and zeroed phase -> ~zero conditioning."""
|
| 364 |
+
nch = N_MELS
|
| 365 |
+
if not self.clean_phase or mel.shape[1] < nch + 2:
|
| 366 |
+
return memory
|
| 367 |
+
ph = mel[:, nch:nch + 2] # (B, 2, T) in [0,1); -1 marks no-grid frames
|
| 368 |
+
valid = (ph[:, :1] >= 0).float() # 1 where a grid exists, else 0
|
| 369 |
+
ph_in = torch.cat([ph.clamp(min=0.0), valid], dim=1) # (B, 3, T)
|
| 370 |
+
p = self.phase_proj(ph_in).transpose(1, 2) # (B, T/4, d)
|
| 371 |
+
L = p.shape[1]
|
| 372 |
+
out = memory.clone()
|
| 373 |
+
out[:, -L:] = out[:, -L:] + p # audio frames are the last L memory slots
|
| 374 |
+
return out
|
| 375 |
+
|
| 376 |
def emb_matrix(self):
|
| 377 |
"""Full embedding matrix; factorized and/or functional rows resolved."""
|
| 378 |
E = self.tok_emb.weight
|
|
|
|
| 383 |
E = torch.cat([E[: VOCAB.time0], t, E[VOCAB.time0 + WINDOW :]], dim=0)
|
| 384 |
return E
|
| 385 |
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 386 |
def decode(self, tgt_in, memory):
|
| 387 |
# tgt_in: (B, L) token ids
|
| 388 |
L = tgt_in.shape[1]
|
|
|
|
| 409 |
return_extras=False, in_tgt=None):
|
| 410 |
# in_tgt: optional corrupted decoder input (e.g. MASKed types for the
|
| 411 |
# skeleton->color infill curriculum); gold targets stay = tgt
|
| 412 |
+
memory = self.encode(mel, gsum=gsum, pos_bucket=pos_bucket) # CLEAN
|
| 413 |
+
dec_mem = self.phase_mem(memory, mel) # phase-conditioned for the decoder
|
| 414 |
dec_in = (in_tgt if in_tgt is not None else tgt)[:, :-1]
|
| 415 |
L = dec_in.shape[1]
|
| 416 |
E = self.emb_matrix()
|
| 417 |
import math as _m
|
| 418 |
h = nn.functional.embedding(dec_in, E) * _m.sqrt(self.d_model) + self.dec_pos[:L]
|
| 419 |
dm = nn.Transformer.generate_square_subsequent_mask(L, device=tgt.device)
|
| 420 |
+
h = self.decoder(self.dropout(h), dec_mem, tgt_mask=dm,
|
| 421 |
tgt_key_padding_mask=dec_in == VOCAB.pad, tgt_is_causal=True)
|
| 422 |
logits = h @ E.T
|
| 423 |
if not (return_aux or return_extras):
|
| 424 |
return logits
|
| 425 |
outs = [logits]
|
| 426 |
+
# aux/ptr/beat heads read the CLEAN memory (no phase leak)
|
| 427 |
aux_mem = memory[:, -(WINDOW // 4):]
|
| 428 |
outs.append(self.aux(aux_mem).squeeze(-1) if self.aux is not None else None)
|
| 429 |
if return_extras:
|
static/app.css
CHANGED
|
@@ -594,6 +594,29 @@ audio {
|
|
| 594 |
font-weight: 750;
|
| 595 |
}
|
| 596 |
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 597 |
.range {
|
| 598 |
--range-progress: 88.89%;
|
| 599 |
width: 100%;
|
|
|
|
| 594 |
font-weight: 750;
|
| 595 |
}
|
| 596 |
|
| 597 |
+
.model-select {
|
| 598 |
+
width: 100%;
|
| 599 |
+
margin-top: 8px;
|
| 600 |
+
padding: 10px 12px;
|
| 601 |
+
color: var(--ink);
|
| 602 |
+
background: var(--surface-solid);
|
| 603 |
+
border: 1px solid var(--line);
|
| 604 |
+
border-radius: 10px;
|
| 605 |
+
font-size: 13px;
|
| 606 |
+
font-weight: 600;
|
| 607 |
+
cursor: pointer;
|
| 608 |
+
transition: border-color .18s, box-shadow .18s;
|
| 609 |
+
}
|
| 610 |
+
|
| 611 |
+
.model-select:hover {
|
| 612 |
+
border-color: var(--line-strong);
|
| 613 |
+
}
|
| 614 |
+
|
| 615 |
+
.model-select:focus-visible {
|
| 616 |
+
outline: 3px solid rgba(39, 111, 145, .35);
|
| 617 |
+
outline-offset: 1px;
|
| 618 |
+
}
|
| 619 |
+
|
| 620 |
.range {
|
| 621 |
--range-progress: 88.89%;
|
| 622 |
width: 100%;
|
static/app.js
CHANGED
|
@@ -15,6 +15,7 @@ const ui = {
|
|
| 15 |
fileMeta: $("#file-meta"),
|
| 16 |
removeFile: $("#remove-file"),
|
| 17 |
fileError: $("#file-error"),
|
|
|
|
| 18 |
level: $("#level"),
|
| 19 |
levelOutput: $("#level-output"),
|
| 20 |
bpm: $("#bpm"),
|
|
@@ -529,7 +530,8 @@ function requestValues() {
|
|
| 529 |
ui.sampling.checked,
|
| 530 |
Number(ui.temperature.value),
|
| 531 |
Number(ui.topP.value),
|
| 532 |
-
Number(ui.drumVolume.value)
|
|
|
|
| 533 |
];
|
| 534 |
}
|
| 535 |
|
|
|
|
| 15 |
fileMeta: $("#file-meta"),
|
| 16 |
removeFile: $("#remove-file"),
|
| 17 |
fileError: $("#file-error"),
|
| 18 |
+
model: $("#model"),
|
| 19 |
level: $("#level"),
|
| 20 |
levelOutput: $("#level-output"),
|
| 21 |
bpm: $("#bpm"),
|
|
|
|
| 530 |
ui.sampling.checked,
|
| 531 |
Number(ui.temperature.value),
|
| 532 |
Number(ui.topP.value),
|
| 533 |
+
Number(ui.drumVolume.value),
|
| 534 |
+
ui.model ? ui.model.value : "v1.7"
|
| 535 |
];
|
| 536 |
}
|
| 537 |
|
static/index.html
CHANGED
|
@@ -81,6 +81,27 @@
|
|
| 81 |
</div>
|
| 82 |
|
| 83 |
<div class="form-section split-rule">
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 84 |
<fieldset class="field-group">
|
| 85 |
<legend>Difficulty</legend>
|
| 86 |
<div class="course-selector">
|
|
|
|
| 81 |
</div>
|
| 82 |
|
| 83 |
<div class="form-section split-rule">
|
| 84 |
+
<div class="field-group">
|
| 85 |
+
<div class="label-row">
|
| 86 |
+
<label for="model">Model</label>
|
| 87 |
+
</div>
|
| 88 |
+
<select id="model" class="model-select">
|
| 89 |
+
<optgroup label="v1.7 — richest patterns (recommended)">
|
| 90 |
+
<option value="v1.7" selected>v1.7 · 7.9M — best quality</option>
|
| 91 |
+
<option value="v1.7-small">v1.7-small · 3.6M</option>
|
| 92 |
+
<option value="v1.7-tiny">v1.7-tiny · 1.4M — fastest</option>
|
| 93 |
+
</optgroup>
|
| 94 |
+
<optgroup label="v1.6 — scaling ladder">
|
| 95 |
+
<option value="v1.6">v1.6 · 7.9M</option>
|
| 96 |
+
<option value="v1.6-small">v1.6-small · 3.6M</option>
|
| 97 |
+
<option value="v1.6-tiny">v1.6-tiny · 1.4M</option>
|
| 98 |
+
</optgroup>
|
| 99 |
+
<optgroup label="v1.5 — single model">
|
| 100 |
+
<option value="v1.5">v1.5 · 7.9M</option>
|
| 101 |
+
</optgroup>
|
| 102 |
+
</select>
|
| 103 |
+
</div>
|
| 104 |
+
|
| 105 |
<fieldset class="field-group">
|
| 106 |
<legend>Difficulty</legend>
|
| 107 |
<div class="course-selector">
|