Load sage sheet and use it on training with 50% dropout so it generates sheet music and handles without it.
This commit is contained in:
Jaret Burkett committed
1 parent
9ea4348d1e
commit
f9b23b40bb
14 files changed
+1224
-110
No files matched your search
@@ -192,7 +192,8 @@ class YuE2AR(YuE2Expert):
|
||||
self.model.lm_head = nn.Linear(cfg.hidden_size, cfg.vocab_size, bias=False)
|
||||
|
||||
def embed(self, ids: torch.Tensor):
|
||||
return self.model.embed_tokens(ids.to(self.model.embed_tokens.weight.device))
|
||||
# never touch embed_tokens.weight: the int8 checkpoint's table dequantizes the whole vocab on that property
|
||||
return self.model.embed_tokens(ids.to(self.device))
|
||||
|
||||
def prefill(self, inputs_embeds: torch.Tensor, return_hidden: bool = False):
|
||||
"""Causal pass over [B, L, H]. Returns per-layer post-rope (k, v) as
|
||||
@@ -355,17 +356,49 @@ class YuE2Model(nn.Module, OstrisModelMixin):
|
||||
|
||||
@classmethod
|
||||
def load_from_state_dict(cls, sd: dict, dtype=torch.bfloat16):
|
||||
"""``sd`` is the Comfy-Org all-in-one checkpoint (raw keys)."""
|
||||
"""``sd`` is a Comfy-Org all-in-one checkpoint (raw keys): bf16, or the int8 convrot repack whose
|
||||
``comfy_quant`` markers attach the shipped quantization directly (no requantization)."""
|
||||
from toolkit.util.comfy_quant_import import Int8Embedding, import_comfy_quantized_layers
|
||||
from toolkit.util.ostris_quant import OstrisLinear
|
||||
|
||||
model = cls()
|
||||
nar_sd = {k[len("model.diffusion_model.") :]: v for k, v in sd.items() if k.startswith("model.diffusion_model.")}
|
||||
ar_sd = {k[len("text_encoders.") :]: v for k, v in sd.items() if k.startswith("text_encoders.") and k != "text_encoders.yue2_tokenizer_json"}
|
||||
prequantized = any(k.endswith(".comfy_quant") for k in sd)
|
||||
for name, module, part in (("NAR", model.nar, nar_sd), ("AR", model.ar, ar_sd)):
|
||||
missing, unexpected = module.load_state_dict({k: v.to(dtype) for k, v in part.items()}, strict=False)
|
||||
allowed_missing = set()
|
||||
if prequantized:
|
||||
quantized_paths = {k[: -len(".comfy_quant")] for k in part if k.endswith(".comfy_quant")}
|
||||
part, n = import_comfy_quantized_layers(module, part, orig_dtype=dtype)
|
||||
for mod_name, m in module.named_modules():
|
||||
if isinstance(m, OstrisLinear):
|
||||
allowed_missing.add(f"{mod_name}.weight")
|
||||
if m.bias is not None:
|
||||
allowed_missing.add(f"{mod_name}.bias")
|
||||
m.bias.data = m.bias.data.to(dtype)
|
||||
elif isinstance(m, (Int8Embedding, nn.Embedding)) and mod_name in quantized_paths:
|
||||
allowed_missing.add(f"{mod_name}.weight")
|
||||
part = {k: (v.to(dtype) if v.is_floating_point() else v) for k, v in part.items()}
|
||||
missing, unexpected = module.load_state_dict(part, strict=False)
|
||||
missing = [k for k in missing if k not in allowed_missing]
|
||||
if missing:
|
||||
raise ValueError(f"YuE2 {name} missing keys: {missing[:5]} (+{max(0, len(missing) - 5)})")
|
||||
if unexpected:
|
||||
print(f" YuE2 {name} unexpected: {len(unexpected)} (first 3: {unexpected[:3]})")
|
||||
return model.to(dtype)
|
||||
if not prequantized:
|
||||
return model.to(dtype)
|
||||
# cast the remaining float tensors (norms, embeddings, buffers) without touching the quantized layers
|
||||
for m in model.modules():
|
||||
if isinstance(m, (OstrisLinear, Int8Embedding)):
|
||||
continue
|
||||
for pname, prm in list(m.named_parameters(recurse=False)):
|
||||
if prm.is_floating_point():
|
||||
prm.data = prm.data.to(dtype)
|
||||
for bname, buf in list(m.named_buffers(recurse=False)):
|
||||
if buf is not None and buf.is_floating_point():
|
||||
setattr(m, bname, buf.to(dtype))
|
||||
model.aitk_is_quantized = True # aitk_post_load keeps the shipped convrot8 layers when convrot8 is requested
|
||||
return model
|
||||
|
||||
|
||||
@torch.no_grad()
|
||||
@@ -376,6 +409,8 @@ def merge_nar_lora(model: YuE2Model, ckpt_path: str, scale: float = 1.0):
|
||||
ck = torch.load(ckpt_path, map_location="cpu", weights_only=True)
|
||||
except Exception:
|
||||
ck = torch.load(ckpt_path, map_location="cpu", weights_only=False)
|
||||
if getattr(model, "aitk_is_quantized", False):
|
||||
raise ValueError("merge_nar_lora needs the bf16 checkpoint; the int8 convrot repack cannot take merged weights")
|
||||
tensors = iter(ck["lora"])
|
||||
cfg = model.cfg
|
||||
inner = cfg.num_attention_heads * cfg.head_dim
|
||||
|
||||
@@ -11,11 +11,13 @@ from typing import Callable, List, Optional, Tuple
|
||||
import torch
|
||||
|
||||
from .model import (
|
||||
ABC_END,
|
||||
CODEC_OFFSET,
|
||||
CODEC_SIZE,
|
||||
CONTEXT,
|
||||
MUSIC_END,
|
||||
EOD,
|
||||
KVCache,
|
||||
MUSIC_END,
|
||||
YuE2Model,
|
||||
rope_cos_sin,
|
||||
)
|
||||
@@ -40,14 +42,22 @@ def sample_logits(
|
||||
penalty_window: int,
|
||||
min_tokens: int,
|
||||
generator: torch.Generator,
|
||||
phase: str = "semantic",
|
||||
legacy_off: bool = False,
|
||||
) -> int:
|
||||
scores = logits.float().clone()
|
||||
"""ComfyUI's ``distribution`` for one row: ``abc`` allows the text vocabulary + ABC_END,
|
||||
``semantic`` the codec ids + MUSIC_END. ``legacy_off`` = cot "off" (bf16 scores, keep top 3)."""
|
||||
scores = logits.clone() if legacy_off else logits.float().clone()
|
||||
end = ABC_END if phase == "abc" else MUSIC_END
|
||||
allowed = torch.full_like(scores, -torch.inf)
|
||||
allowed[CODEC_OFFSET : CODEC_OFFSET + CODEC_SIZE] = 0
|
||||
allowed[MUSIC_END] = 0
|
||||
if phase == "abc":
|
||||
allowed[:EOD] = 0
|
||||
else:
|
||||
allowed[CODEC_OFFSET : CODEC_OFFSET + CODEC_SIZE] = 0
|
||||
allowed[end] = 0
|
||||
scores += allowed
|
||||
if step < min_tokens:
|
||||
scores[MUSIC_END] = -torch.inf
|
||||
scores[end] = -torch.inf
|
||||
if repetition_penalty != 1.0 and history:
|
||||
recent = torch.tensor(history[-penalty_window:], dtype=torch.long, device=scores.device)
|
||||
counts = torch.zeros_like(scores).scatter_add_(0, recent, torch.ones_like(recent, dtype=scores.dtype))
|
||||
@@ -62,8 +72,7 @@ def sample_logits(
|
||||
values, indices = scores.sort(descending=True)
|
||||
probs = values.softmax(-1)
|
||||
removed = probs.cumsum(-1) - probs > top_p
|
||||
# released "off" mode keeps the top 3 candidates regardless of mass
|
||||
removed[:3] = False
|
||||
removed[: 3 if legacy_off else 1] = False
|
||||
values.masked_fill_(removed, -torch.inf)
|
||||
scores = values.scatter(-1, indices, values)
|
||||
probs = scores.softmax(-1)
|
||||
@@ -89,6 +98,10 @@ class ARDecodeGraph:
|
||||
self.token = torch.zeros((1, 1), device=device, dtype=torch.long)
|
||||
self.pos = torch.zeros((1, 1), device=device, dtype=torch.long)
|
||||
self.index = torch.arange(capacity, device=device)
|
||||
# the graph covers the layer stack only; embedding lookup and lm_head run outside it
|
||||
# (the int8 repack's embedding/lm_head kernels invalidate stream capture)
|
||||
self.x_in = torch.zeros((1, 1, cfg.hidden_size), device=device, dtype=dtype)
|
||||
self.h_out = torch.zeros((1, cfg.hidden_size), device=device, dtype=dtype)
|
||||
self.logits = torch.empty((1, cfg.vocab_size), device=device, dtype=dtype)
|
||||
self.graph = None
|
||||
|
||||
@@ -98,9 +111,9 @@ class ARDecodeGraph:
|
||||
self.v[i][:, :, :length] = v
|
||||
self.pos.fill_(length)
|
||||
|
||||
def _step(self):
|
||||
def _layers(self):
|
||||
ar = self.ar
|
||||
x = ar.model.embed_tokens(self.token)
|
||||
x = self.x_in
|
||||
cos, sin = rope_cos_sin(self.pos, self.head_dim, ar.cfg.rope_theta)
|
||||
mask = (self.index <= self.pos)[None, None] # [1, 1, 1, capacity]
|
||||
scale = self.head_dim ** -0.5
|
||||
@@ -115,25 +128,31 @@ class ARDecodeGraph:
|
||||
out = torch.matmul(scores, self.v[i]).reshape(1, 1, -1)
|
||||
x = x + layer.self_attn.o_proj(out)
|
||||
x = x + layer.mlp(layer.post_attention_layernorm(x))
|
||||
self.logits.copy_(ar.model.lm_head(ar.model.norm(x[:, -1])))
|
||||
self.h_out.copy_(ar.model.norm(x[:, -1]))
|
||||
|
||||
def _step(self):
|
||||
ar = self.ar
|
||||
self.x_in.copy_(ar.model.embed_tokens(self.token).to(self.x_in.dtype))
|
||||
if self.graph is not None:
|
||||
self.graph.replay()
|
||||
else:
|
||||
self._layers()
|
||||
self.logits.copy_(ar.model.lm_head(self.h_out))
|
||||
|
||||
def capture(self):
|
||||
stream = torch.cuda.Stream()
|
||||
stream.wait_stream(torch.cuda.current_stream())
|
||||
with torch.cuda.stream(stream):
|
||||
for _ in range(2):
|
||||
self._step()
|
||||
self._layers()
|
||||
torch.cuda.current_stream().wait_stream(stream)
|
||||
self.graph = torch.cuda.CUDAGraph()
|
||||
with torch.cuda.graph(self.graph):
|
||||
self._step()
|
||||
self._layers()
|
||||
|
||||
def step(self, token: int):
|
||||
self.token.fill_(token)
|
||||
if self.graph is not None:
|
||||
self.graph.replay()
|
||||
else:
|
||||
self._step()
|
||||
self._step()
|
||||
self.pos.add_(1)
|
||||
return self.logits[0]
|
||||
|
||||
@@ -146,19 +165,22 @@ class YuE2Pipeline:
|
||||
self.use_cuda_graphs = True
|
||||
|
||||
@torch.no_grad()
|
||||
def generate_codec_tokens(
|
||||
def generate_tokens(
|
||||
self,
|
||||
prefix_embeds: torch.Tensor, # [1, L, H]
|
||||
phase: str,
|
||||
max_tokens: int,
|
||||
seed: int,
|
||||
temperature: float = 1.0,
|
||||
top_p: float = 0.95,
|
||||
top_k: int = 100,
|
||||
repetition_penalty: float = 1.2,
|
||||
penalty_window: int = 50,
|
||||
min_tokens: int = 200,
|
||||
temperature: float,
|
||||
top_p: float,
|
||||
top_k: int,
|
||||
repetition_penalty: float,
|
||||
penalty_window: int,
|
||||
min_tokens: int,
|
||||
progress: Optional[Callable[[int, int], None]] = None,
|
||||
legacy_off: bool = False,
|
||||
) -> List[int]:
|
||||
"""Sample AR tokens after ``prefix_embeds`` until the phase's end token (raw vocab ids)."""
|
||||
ar = self.model.ar
|
||||
device = prefix_embeds.device
|
||||
prefix_len = prefix_embeds.shape[1]
|
||||
@@ -185,11 +207,13 @@ class YuE2Pipeline:
|
||||
logits = ar.prefill_into_cache(prefix_embeds.to(ar.dtype), cache)[0]
|
||||
history: List[int] = []
|
||||
min_tokens = min(min_tokens, max_tokens)
|
||||
end = ABC_END if phase == "abc" else MUSIC_END
|
||||
for step in range(max_tokens):
|
||||
token = sample_logits(
|
||||
logits, history, step, temperature, top_p, top_k, repetition_penalty, penalty_window, min_tokens, generator
|
||||
logits, history, step, temperature, top_p, top_k, repetition_penalty, penalty_window, min_tokens, generator,
|
||||
phase=phase, legacy_off=legacy_off,
|
||||
)
|
||||
if token == MUSIC_END:
|
||||
if token == end:
|
||||
break
|
||||
history.append(token)
|
||||
if progress is not None:
|
||||
@@ -203,7 +227,46 @@ class YuE2Pipeline:
|
||||
if graph is None:
|
||||
del cache
|
||||
del graph
|
||||
return [t - CODEC_OFFSET for t in history]
|
||||
return history
|
||||
|
||||
def generate_codec_tokens(
|
||||
self,
|
||||
prefix_embeds: torch.Tensor,
|
||||
max_tokens: int,
|
||||
seed: int,
|
||||
temperature: float = 1.0,
|
||||
top_p: float = 0.95,
|
||||
top_k: int = 100,
|
||||
repetition_penalty: float = 1.2,
|
||||
penalty_window: int = 50,
|
||||
min_tokens: int = 200,
|
||||
progress: Optional[Callable[[int, int], None]] = None,
|
||||
legacy_off: bool = False,
|
||||
) -> List[int]:
|
||||
"""Music phase (ComfyUI defaults). Returns codec ids without the vocab offset."""
|
||||
ids = self.generate_tokens(
|
||||
prefix_embeds, "semantic", max_tokens, seed, temperature, top_p, top_k, repetition_penalty, penalty_window,
|
||||
min_tokens, progress, legacy_off=legacy_off,
|
||||
)
|
||||
return [t - CODEC_OFFSET for t in ids]
|
||||
|
||||
def generate_abc_tokens(
|
||||
self,
|
||||
prefix_embeds: torch.Tensor,
|
||||
seed: int,
|
||||
max_tokens: int = 8192,
|
||||
temperature: float = 0.7,
|
||||
top_p: float = 0.9,
|
||||
top_k: int = 30,
|
||||
repetition_penalty: float = 1.005,
|
||||
penalty_window: int = 100,
|
||||
progress: Optional[Callable[[int, int], None]] = None,
|
||||
) -> List[int]:
|
||||
"""ABC phase (ComfyUI "Generate ABC" defaults): text ids of the sheet, without ABC_END."""
|
||||
return self.generate_tokens(
|
||||
prefix_embeds, "abc", max_tokens, seed, temperature, top_p, top_k, repetition_penalty, penalty_window,
|
||||
min(32, max_tokens), progress,
|
||||
)
|
||||
|
||||
@torch.no_grad()
|
||||
def synthesize(
|
||||
|
||||
@@ -0,0 +1,110 @@
|
||||
"""SheetSage2's score decoder (BART-style, post-LN, learned positions with offset 2) in plain
|
||||
torch, so the m-a-p/SheetSage2 remote code runs on current transformers without BartDecoder.
|
||||
Parameter names follow the released checkpoint (decoder.embed_positions, decoder.layers.N.self_attn.*,
|
||||
encoder_attn.*, fc1/fc2, *_layer_norm, layernorm_embedding). KV cache = per-layer tuples with the
|
||||
batch dimension first, which the upstream generation loop can index."""
|
||||
|
||||
from types import SimpleNamespace
|
||||
from typing import Optional, Tuple
|
||||
|
||||
import torch
|
||||
import torch.nn.functional as F
|
||||
from torch import nn
|
||||
|
||||
POSITION_OFFSET = 2 # BART learned positional embedding convention
|
||||
|
||||
|
||||
class _Attention(nn.Module):
|
||||
def __init__(self, dim: int, heads: int):
|
||||
super().__init__()
|
||||
self.heads = heads
|
||||
self.head_dim = dim // heads
|
||||
self.q_proj = nn.Linear(dim, dim)
|
||||
self.k_proj = nn.Linear(dim, dim)
|
||||
self.v_proj = nn.Linear(dim, dim)
|
||||
self.out_proj = nn.Linear(dim, dim)
|
||||
|
||||
def _split(self, x):
|
||||
b, l, _ = x.shape
|
||||
return x.view(b, l, self.heads, self.head_dim).transpose(1, 2)
|
||||
|
||||
def forward(self, x, kv_source=None, mask=None, past: Optional[Tuple[torch.Tensor, torch.Tensor]] = None, cache_static=False):
|
||||
q = self._split(self.q_proj(x))
|
||||
if cache_static and past is not None:
|
||||
k, v = past # cross attention: memory keys never change
|
||||
else:
|
||||
src = x if kv_source is None else kv_source
|
||||
k, v = self._split(self.k_proj(src)), self._split(self.v_proj(src))
|
||||
if past is not None and not cache_static:
|
||||
k, v = torch.cat([past[0], k], 2), torch.cat([past[1], v], 2)
|
||||
out = F.scaled_dot_product_attention(q, k, v, attn_mask=mask)
|
||||
b, _, l, _ = out.shape
|
||||
return self.out_proj(out.transpose(1, 2).reshape(b, l, -1)), (k, v)
|
||||
|
||||
|
||||
class _Layer(nn.Module):
|
||||
def __init__(self, dim: int, heads: int, ffn: int):
|
||||
super().__init__()
|
||||
self.self_attn = _Attention(dim, heads)
|
||||
self.self_attn_layer_norm = nn.LayerNorm(dim)
|
||||
self.encoder_attn = _Attention(dim, heads)
|
||||
self.encoder_attn_layer_norm = nn.LayerNorm(dim)
|
||||
self.fc1 = nn.Linear(dim, ffn)
|
||||
self.fc2 = nn.Linear(ffn, dim)
|
||||
self.final_layer_norm = nn.LayerNorm(dim)
|
||||
|
||||
def forward(self, x, memory, self_mask, past):
|
||||
self_past = None if past is None else past[:2]
|
||||
cross_past = None if past is None else past[2:]
|
||||
h, self_kv = self.self_attn(x, mask=self_mask, past=self_past)
|
||||
x = self.self_attn_layer_norm(x + h)
|
||||
h, cross_kv = self.encoder_attn(x, kv_source=memory, past=cross_past, cache_static=True)
|
||||
x = self.encoder_attn_layer_norm(x + h)
|
||||
x = self.final_layer_norm(x + self.fc2(F.gelu(self.fc1(x))))
|
||||
return x, self_kv + cross_kv
|
||||
|
||||
|
||||
class ScoreDecoder(nn.Module):
|
||||
"""Drop-in for the BartDecoder the upstream model builds: ``ScoreDecoder(bart_config, embed_tokens=...)``."""
|
||||
|
||||
def __init__(self, config, embed_tokens: nn.Embedding):
|
||||
super().__init__()
|
||||
dim, heads = config.d_model, config.decoder_attention_heads
|
||||
self.embed_tokens = embed_tokens
|
||||
self.embed_positions = nn.Embedding(config.max_position_embeddings + POSITION_OFFSET, dim)
|
||||
self.layernorm_embedding = nn.LayerNorm(dim)
|
||||
self.layers = nn.ModuleList([_Layer(dim, heads, config.decoder_ffn_dim) for _ in range(config.decoder_layers)])
|
||||
self.pad_token_id = config.pad_token_id
|
||||
|
||||
def gradient_checkpointing_disable(self):
|
||||
pass
|
||||
|
||||
def forward(self, input_ids, attention_mask=None, encoder_hidden_states=None, encoder_attention_mask=None,
|
||||
past_key_values=None, use_cache=False, output_hidden_states=False, return_dict=True, **_):
|
||||
b, l = input_ids.shape
|
||||
past_len = 0 if past_key_values is None else past_key_values[0][0].shape[2]
|
||||
positions = torch.arange(past_len, past_len + l, device=input_ids.device) + POSITION_OFFSET
|
||||
x = self.embed_tokens(input_ids) + self.embed_positions(positions)[None]
|
||||
x = self.layernorm_embedding(x)
|
||||
mask = None
|
||||
if l > 1:
|
||||
causal = torch.ones(l, past_len + l, dtype=torch.bool, device=x.device).tril(past_len)
|
||||
mask = causal[None, None]
|
||||
if attention_mask is not None:
|
||||
mask = mask & attention_mask.bool()[:, None, None, :]
|
||||
elif attention_mask is not None and past_len == 0:
|
||||
mask = attention_mask.bool()[:, None, None, :]
|
||||
if mask is not None:
|
||||
mask = torch.zeros(mask.shape, dtype=x.dtype, device=x.device).masked_fill(~mask, float("-inf"))
|
||||
hidden_states = [x] if output_hidden_states else None
|
||||
new_cache = []
|
||||
for i, layer in enumerate(self.layers):
|
||||
x, kv = layer(x, encoder_hidden_states, mask, None if past_key_values is None else past_key_values[i])
|
||||
new_cache.append(kv)
|
||||
if hidden_states is not None:
|
||||
hidden_states.append(x)
|
||||
return SimpleNamespace(
|
||||
last_hidden_state=x,
|
||||
past_key_values=tuple(new_cache) if use_cache else None,
|
||||
hidden_states=tuple(hidden_states) if hidden_states is not None else None,
|
||||
)
|
||||
@@ -0,0 +1,27 @@
|
||||
"""Stand-ins for mir_eval.chord and pretty_midi so m-a-p/SheetSage2's remote code imports without
|
||||
those packages. Installed into sys.modules only when the real packages are absent."""
|
||||
|
||||
import importlib
|
||||
import sys
|
||||
import types
|
||||
|
||||
|
||||
def install_shims():
|
||||
for name in ("mir_eval", "pretty_midi"):
|
||||
try:
|
||||
importlib.import_module(name)
|
||||
continue # real package available; leave it alone
|
||||
except ImportError:
|
||||
pass
|
||||
if name == "mir_eval":
|
||||
from . import mir_eval_chord
|
||||
|
||||
pkg = types.ModuleType("mir_eval")
|
||||
pkg.chord = mir_eval_chord
|
||||
pkg.__path__ = []
|
||||
sys.modules["mir_eval"] = pkg
|
||||
sys.modules["mir_eval.chord"] = mir_eval_chord
|
||||
else:
|
||||
from . import pretty_midi
|
||||
|
||||
sys.modules["pretty_midi"] = pretty_midi
|
||||
@@ -0,0 +1,518 @@
|
||||
"""Chord-label helpers vendored from mir_eval (mir_eval/chord.py), so the SheetSage2 remote code
|
||||
runs without the mir_eval package. Unmodified except for the removed module preamble.
|
||||
|
||||
The MIT License (MIT)
|
||||
|
||||
Copyright (c) 2014 Colin Raffel
|
||||
|
||||
Permission is hereby granted, free of charge, to any person obtaining a copy
|
||||
of this software and associated documentation files (the "Software"), to deal
|
||||
in the Software without restriction, including without limitation the rights
|
||||
to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
|
||||
copies of the Software, and to permit persons to whom the Software is
|
||||
furnished to do so, subject to the following conditions:
|
||||
|
||||
The above copyright notice and this permission notice shall be included in
|
||||
all copies or substantial portions of the Software.
|
||||
|
||||
THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
|
||||
IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
|
||||
FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
|
||||
AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
|
||||
LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
|
||||
OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
|
||||
THE SOFTWARE.
|
||||
"""
|
||||
|
||||
import re
|
||||
import warnings
|
||||
|
||||
import numpy as np
|
||||
|
||||
|
||||
BITMAP_LENGTH = 12
|
||||
NO_CHORD = "N"
|
||||
NO_CHORD_ENCODED = -1, np.array([0] * BITMAP_LENGTH), -1
|
||||
X_CHORD = "X"
|
||||
X_CHORD_ENCODED = -1, np.array([-1] * BITMAP_LENGTH), -1
|
||||
|
||||
|
||||
class InvalidChordException(Exception):
|
||||
r"""Exception class for suspect / invalid chord labels"""
|
||||
|
||||
def __init__(self, message="", chord_label=None):
|
||||
self.message = message
|
||||
self.chord_label = chord_label
|
||||
self.name = self.__class__.__name__
|
||||
super().__init__(message)
|
||||
|
||||
|
||||
# --- Chord Primitives ---
|
||||
def _pitch_classes():
|
||||
r"""Map from pitch class (str) to semitone (int)."""
|
||||
pitch_classes = ["C", "D", "E", "F", "G", "A", "B"]
|
||||
semitones = [0, 2, 4, 5, 7, 9, 11]
|
||||
return {c: s for c, s in zip(pitch_classes, semitones)}
|
||||
|
||||
|
||||
def _scale_degrees():
|
||||
r"""Map scale degrees (str) to semitones (int)."""
|
||||
degrees = ["1", "2", "3", "4", "5", "6", "7", "8", "9", "10", "11", "12", "13"]
|
||||
semitones = [0, 2, 4, 5, 7, 9, 11, 12, 14, 16, 17, 19, 21]
|
||||
return {d: s for d, s in zip(degrees, semitones)}
|
||||
|
||||
|
||||
# Maps pitch classes (strings) to semitone indexes (ints).
|
||||
PITCH_CLASSES = _pitch_classes()
|
||||
|
||||
|
||||
def pitch_class_to_semitone(pitch_class):
|
||||
r"""Convert a pitch class to semitone.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
pitch_class : str
|
||||
Spelling of a given pitch class, e.g. 'C#', 'Gbb'
|
||||
|
||||
Returns
|
||||
-------
|
||||
semitone : int
|
||||
Semitone value of the pitch class.
|
||||
|
||||
"""
|
||||
semitone = 0
|
||||
for idx, char in enumerate(pitch_class):
|
||||
if char == "#" and idx > 0:
|
||||
semitone += 1
|
||||
elif char == "b" and idx > 0:
|
||||
semitone -= 1
|
||||
elif idx == 0:
|
||||
semitone = PITCH_CLASSES.get(char)
|
||||
else:
|
||||
raise InvalidChordException(
|
||||
"Pitch class improperly formed: %s" % pitch_class
|
||||
)
|
||||
return semitone % 12
|
||||
|
||||
|
||||
# Maps scale degrees (strings) to semitone indexes (ints).
|
||||
SCALE_DEGREES = _scale_degrees()
|
||||
|
||||
|
||||
def scale_degree_to_semitone(scale_degree):
|
||||
r"""Convert a scale degree to semitone.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
scale_degree : str
|
||||
Spelling of a relative scale degree, e.g. 'b3', '7', '#5'
|
||||
|
||||
Returns
|
||||
-------
|
||||
semitone : int
|
||||
Relative semitone of the scale degree, wrapped to a single octave
|
||||
|
||||
Raises
|
||||
------
|
||||
InvalidChordException if `scale_degree` is invalid.
|
||||
"""
|
||||
semitone = 0
|
||||
offset = 0
|
||||
if scale_degree.startswith("#"):
|
||||
offset = scale_degree.count("#")
|
||||
scale_degree = scale_degree.strip("#")
|
||||
elif scale_degree.startswith("b"):
|
||||
offset = -1 * scale_degree.count("b")
|
||||
scale_degree = scale_degree.strip("b")
|
||||
|
||||
semitone = SCALE_DEGREES.get(scale_degree, None)
|
||||
if semitone is None:
|
||||
raise InvalidChordException(
|
||||
"Scale degree improperly formed: {}, expected one of {}.".format(
|
||||
scale_degree, list(SCALE_DEGREES.keys())
|
||||
)
|
||||
)
|
||||
return semitone + offset
|
||||
|
||||
|
||||
def scale_degree_to_bitmap(scale_degree, modulo=False, length=BITMAP_LENGTH):
|
||||
"""Create a bitmap representation of a scale degree.
|
||||
|
||||
Note that values in the bitmap may be negative, indicating that the
|
||||
semitone is to be removed.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
scale_degree : str
|
||||
Spelling of a relative scale degree, e.g. 'b3', '7', '#5'
|
||||
modulo : bool, default=True
|
||||
If a scale degree exceeds the length of the bit-vector, modulo the
|
||||
scale degree back into the bit-vector; otherwise it is discarded.
|
||||
length : int, default=12
|
||||
Length of the bit-vector to produce
|
||||
|
||||
Returns
|
||||
-------
|
||||
bitmap : np.ndarray, in [-1, 0, 1], len=`length`
|
||||
Bitmap representation of this scale degree.
|
||||
"""
|
||||
sign = 1
|
||||
if scale_degree.startswith("*"):
|
||||
sign = -1
|
||||
scale_degree = scale_degree.strip("*")
|
||||
edit_map = [0] * length
|
||||
sd_idx = scale_degree_to_semitone(scale_degree)
|
||||
if sd_idx < length or modulo:
|
||||
edit_map[sd_idx % length] = sign
|
||||
return np.array(edit_map)
|
||||
|
||||
|
||||
# Maps quality strings to bitmaps, corresponding to relative pitch class
|
||||
# semitones, i.e. vector[0] is the tonic.
|
||||
QUALITIES = {
|
||||
# 1 2 3 4 5 6 7
|
||||
"maj": [1, 0, 0, 0, 1, 0, 0, 1, 0, 0, 0, 0],
|
||||
"min": [1, 0, 0, 1, 0, 0, 0, 1, 0, 0, 0, 0],
|
||||
"aug": [1, 0, 0, 0, 1, 0, 0, 0, 1, 0, 0, 0],
|
||||
"dim": [1, 0, 0, 1, 0, 0, 1, 0, 0, 0, 0, 0],
|
||||
"sus4": [1, 0, 0, 0, 0, 1, 0, 1, 0, 0, 0, 0],
|
||||
"sus2": [1, 0, 1, 0, 0, 0, 0, 1, 0, 0, 0, 0],
|
||||
"7": [1, 0, 0, 0, 1, 0, 0, 1, 0, 0, 1, 0],
|
||||
"maj7": [1, 0, 0, 0, 1, 0, 0, 1, 0, 0, 0, 1],
|
||||
"min7": [1, 0, 0, 1, 0, 0, 0, 1, 0, 0, 1, 0],
|
||||
"minmaj7": [1, 0, 0, 1, 0, 0, 0, 1, 0, 0, 0, 1],
|
||||
"maj6": [1, 0, 0, 0, 1, 0, 0, 1, 0, 1, 0, 0],
|
||||
"min6": [1, 0, 0, 1, 0, 0, 0, 1, 0, 1, 0, 0],
|
||||
"dim7": [1, 0, 0, 1, 0, 0, 1, 0, 0, 1, 0, 0],
|
||||
"hdim7": [1, 0, 0, 1, 0, 0, 1, 0, 0, 0, 1, 0],
|
||||
"aug7": [1, 0, 0, 0, 1, 0, 0, 0, 1, 0, 1, 0],
|
||||
"maj9": [1, 0, 0, 0, 1, 0, 0, 1, 0, 0, 0, 1],
|
||||
"min9": [1, 0, 0, 1, 0, 0, 0, 1, 0, 0, 1, 0],
|
||||
"9": [1, 0, 0, 0, 1, 0, 0, 1, 0, 0, 1, 0],
|
||||
"min11": [1, 0, 0, 1, 0, 0, 0, 1, 0, 0, 1, 0],
|
||||
"11": [1, 0, 0, 0, 1, 0, 0, 1, 0, 0, 1, 0],
|
||||
"maj13": [1, 0, 0, 0, 1, 0, 0, 1, 0, 0, 0, 1],
|
||||
"min13": [1, 0, 0, 1, 0, 0, 0, 1, 0, 0, 1, 0],
|
||||
"13": [1, 0, 0, 0, 1, 0, 0, 1, 0, 0, 1, 0],
|
||||
"1": [1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
|
||||
"5": [1, 0, 0, 0, 0, 0, 0, 1, 0, 0, 0, 0],
|
||||
"": [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
|
||||
}
|
||||
|
||||
|
||||
def quality_to_bitmap(quality):
|
||||
"""Return the bitmap for a given quality.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
quality : str
|
||||
Chord quality name.
|
||||
|
||||
Returns
|
||||
-------
|
||||
bitmap : np.ndarray
|
||||
Bitmap representation of this quality (12-dim).
|
||||
|
||||
"""
|
||||
if quality not in QUALITIES:
|
||||
raise InvalidChordException(
|
||||
"Unsupported chord quality shorthand: '%s' "
|
||||
"Did you mean to reduce extended chords?" % quality
|
||||
)
|
||||
return np.array(QUALITIES[quality])
|
||||
|
||||
|
||||
# Maps extended chord qualities to the subset above, translating additional
|
||||
# voicings to extensions as a set of scale degrees (strings).
|
||||
# TODO(ejhumphrey): Revisit how minmaj7's are mapped. This is how TMC did it,
|
||||
# but MMV handles it like a separate quality (rather than an add7).
|
||||
EXTENDED_QUALITY_REDUX = {
|
||||
"minmaj7": ("min", {"7"}),
|
||||
"maj9": ("maj7", {"9"}),
|
||||
"min9": ("min7", {"9"}),
|
||||
"9": ("7", {"9"}),
|
||||
"b9": ("7", {"b9"}),
|
||||
"#9": ("7", {"#9"}),
|
||||
"11": ("7", {"9", "11"}),
|
||||
"#11": ("7", {"9", "#11"}),
|
||||
"13": ("7", {"9", "11", "13"}),
|
||||
"b13": ("7", {"9", "11", "b13"}),
|
||||
"min11": ("min7", {"9", "11"}),
|
||||
"maj13": ("maj7", {"9", "11", "13"}),
|
||||
"min13": ("min7", {"9", "11", "13"}),
|
||||
}
|
||||
|
||||
|
||||
def reduce_extended_quality(quality):
|
||||
"""Map an extended chord quality to a simpler one, moving upper voices to
|
||||
a set of scale degree extensions.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
quality : str
|
||||
Extended chord quality to reduce.
|
||||
|
||||
Returns
|
||||
-------
|
||||
base_quality : str
|
||||
New chord quality.
|
||||
extensions : set
|
||||
Scale degrees extensions for the quality.
|
||||
|
||||
"""
|
||||
return EXTENDED_QUALITY_REDUX.get(quality, (quality, set()))
|
||||
|
||||
|
||||
# --- Chord Label Parsing ---
|
||||
# This monster regexp is pulled from the JAMS chord namespace,
|
||||
# which is in turn derived from the context-free grammar of
|
||||
# Harte et al., 2005.
|
||||
CHORD_RE = re.compile(
|
||||
r"""^((N|X)|(([A-G](b*|#*))((:(maj|min|dim|aug|1|5|sus2|sus4|maj6|min6|7|maj7|min7|dim7|hdim7|minmaj7|aug7|9|maj9|min9|11|maj11|min11|13|maj13|min13)(\((\*?((b*|#*)([1-9]|1[0-3]?))(,\*?((b*|#*)([1-9]|1[0-3]?)))*)\))?)|(:\((\*?((b*|#*)([1-9]|1[0-3]?))(,\*?((b*|#*)([1-9]|1[0-3]?)))*)\)))?((/((b*|#*)([1-9]|1[0-3]?)))?)?))$"""
|
||||
) # nopep8
|
||||
|
||||
|
||||
def validate_chord_label(chord_label):
|
||||
"""Test for well-formedness of a chord label.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
chord_label : str
|
||||
Chord label to validate.
|
||||
"""
|
||||
if not CHORD_RE.match(chord_label):
|
||||
raise InvalidChordException("Invalid chord label: " "{}".format(chord_label))
|
||||
pass
|
||||
|
||||
|
||||
def split(chord_label, reduce_extended_chords=False):
|
||||
"""Parse a chord label into its four constituent parts:
|
||||
- root
|
||||
- quality shorthand
|
||||
- scale degrees
|
||||
- bass
|
||||
|
||||
Note: Chords lacking quality AND interval information are major.
|
||||
- If a quality is specified, it is returned.
|
||||
- If an interval is specified WITHOUT a quality, the quality field is
|
||||
empty.
|
||||
|
||||
Some examples::
|
||||
|
||||
'C' -> ['C', 'maj', {}, '1']
|
||||
'G#:min(*b3,*5)/5' -> ['G#', 'min', {'*b3', '*5'}, '5']
|
||||
'A:(3)/6' -> ['A', '', {'3'}, '6']
|
||||
|
||||
Parameters
|
||||
----------
|
||||
chord_label : str
|
||||
A chord label.
|
||||
reduce_extended_chords : bool
|
||||
Whether to map the upper voicings of extended chords (9's, 11's, 13's)
|
||||
to semitone extensions. (Default value = False)
|
||||
|
||||
Returns
|
||||
-------
|
||||
chord_parts : list
|
||||
Split version of the chord label.
|
||||
|
||||
"""
|
||||
chord_label = str(chord_label)
|
||||
validate_chord_label(chord_label)
|
||||
if chord_label == NO_CHORD:
|
||||
return [chord_label, "", set(), ""]
|
||||
|
||||
bass = "1"
|
||||
if "/" in chord_label:
|
||||
chord_label, bass = chord_label.split("/")
|
||||
|
||||
scale_degrees = set()
|
||||
omission = False
|
||||
if "(" in chord_label:
|
||||
chord_label, scale_degrees = chord_label.split("(")
|
||||
omission = "*" in scale_degrees
|
||||
scale_degrees = scale_degrees.strip(")")
|
||||
scale_degrees = {i.strip() for i in scale_degrees.split(",")}
|
||||
|
||||
# Note: Chords lacking quality AND added interval information are major.
|
||||
# If a quality shorthand is specified, it is returned.
|
||||
# If an interval is specified WITHOUT a quality, the quality field is
|
||||
# empty.
|
||||
# Intervals specifying omissions MUST have a quality.
|
||||
if omission and ":" not in chord_label:
|
||||
raise InvalidChordException(
|
||||
"Intervals specifying omissions MUST have a quality."
|
||||
)
|
||||
quality = "" if scale_degrees else "maj"
|
||||
if ":" in chord_label:
|
||||
chord_root, quality_name = chord_label.split(":")
|
||||
# Extended chords (with ":"s) may not explicitly have Major qualities,
|
||||
# so only overwrite the default if the string is not empty.
|
||||
if quality_name:
|
||||
quality = quality_name.lower()
|
||||
else:
|
||||
chord_root = chord_label
|
||||
|
||||
if reduce_extended_chords:
|
||||
quality, addl_scale_degrees = reduce_extended_quality(quality)
|
||||
scale_degrees.update(addl_scale_degrees)
|
||||
|
||||
return [chord_root, quality, scale_degrees, bass]
|
||||
|
||||
|
||||
def join(chord_root, quality="", extensions=None, bass=""):
|
||||
r"""Join the parts of a chord into a complete chord label.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
chord_root : str
|
||||
Root pitch class of the chord, e.g. 'C', 'Eb'
|
||||
quality : str
|
||||
Quality of the chord, e.g. 'maj', 'hdim7'
|
||||
(Default value = '')
|
||||
extensions : list
|
||||
Any added or absent scaled degrees for this chord, e.g. ['4', '\*3']
|
||||
(Default value = None)
|
||||
bass : str
|
||||
Scale degree of the bass note, e.g. '5'.
|
||||
(Default value = '')
|
||||
|
||||
Returns
|
||||
-------
|
||||
chord_label : str
|
||||
A complete chord label.
|
||||
|
||||
"""
|
||||
chord_label = chord_root
|
||||
if quality or extensions:
|
||||
chord_label += ":%s" % quality
|
||||
if extensions:
|
||||
chord_label += "(%s)" % ",".join(extensions)
|
||||
if bass and bass != "1":
|
||||
chord_label += "/%s" % bass
|
||||
validate_chord_label(chord_label)
|
||||
return chord_label
|
||||
|
||||
|
||||
# --- Chords to Numerical Representations ---
|
||||
def encode(chord_label, reduce_extended_chords=False, strict_bass_intervals=False):
|
||||
"""Translate a chord label to numerical representations for evaluation.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
chord_label : str
|
||||
Chord label to encode.
|
||||
reduce_extended_chords : bool
|
||||
Whether to map the upper voicings of extended chords (9's, 11's, 13's)
|
||||
to semitone extensions.
|
||||
(Default value = False)
|
||||
strict_bass_intervals : bool
|
||||
Whether to require that the bass scale degree is present in the chord.
|
||||
(Default value = False)
|
||||
|
||||
Returns
|
||||
-------
|
||||
root_number : int
|
||||
Absolute semitone of the chord's root.
|
||||
semitone_bitmap : np.ndarray, dtype=int
|
||||
12-dim vector of relative semitones in the chord spelling.
|
||||
bass_number : int
|
||||
Relative semitone of the chord's bass note, e.g. 0=root, 7=fifth, etc.
|
||||
"""
|
||||
if chord_label == NO_CHORD:
|
||||
return NO_CHORD_ENCODED
|
||||
if chord_label == X_CHORD:
|
||||
return X_CHORD_ENCODED
|
||||
chord_root, quality, scale_degrees, bass = split(
|
||||
chord_label, reduce_extended_chords=reduce_extended_chords
|
||||
)
|
||||
|
||||
root_number = pitch_class_to_semitone(chord_root)
|
||||
bass_number = scale_degree_to_semitone(bass) % 12
|
||||
|
||||
semitone_bitmap = quality_to_bitmap(quality)
|
||||
semitone_bitmap[0] = 1
|
||||
|
||||
for scale_degree in scale_degrees:
|
||||
semitone_bitmap += scale_degree_to_bitmap(scale_degree, reduce_extended_chords)
|
||||
|
||||
semitone_bitmap = (semitone_bitmap > 0).astype(np.int64)
|
||||
if not semitone_bitmap[bass_number] and strict_bass_intervals:
|
||||
raise InvalidChordException(
|
||||
"Given bass scale degree is absent from this chord: " "%s" % chord_label,
|
||||
chord_label,
|
||||
)
|
||||
else:
|
||||
semitone_bitmap[bass_number] = 1
|
||||
return root_number, semitone_bitmap, bass_number
|
||||
|
||||
|
||||
def encode_many(chord_labels, reduce_extended_chords=False):
|
||||
"""Translate a set of chord labels to numerical representations for sane
|
||||
evaluation.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
chord_labels : list
|
||||
Set of chord labels to encode.
|
||||
reduce_extended_chords : bool
|
||||
Whether to map the upper voicings of extended chords (9's, 11's, 13's)
|
||||
to semitone extensions.
|
||||
(Default value = False)
|
||||
|
||||
Returns
|
||||
-------
|
||||
root_number : np.ndarray, dtype=int
|
||||
Absolute semitone of the chord's root.
|
||||
interval_bitmap : np.ndarray, dtype=int
|
||||
12-dim vector of relative semitones in the given chord quality.
|
||||
bass_number : np.ndarray, dtype=int
|
||||
Relative semitones of the chord's bass notes.
|
||||
|
||||
"""
|
||||
num_items = len(chord_labels)
|
||||
roots, basses = np.zeros([2, num_items], dtype=np.int64)
|
||||
semitones = np.zeros([num_items, 12], dtype=np.int64)
|
||||
local_cache = dict()
|
||||
for i, label in enumerate(chord_labels):
|
||||
result = local_cache.get(label, None)
|
||||
if result is None:
|
||||
result = encode(label, reduce_extended_chords)
|
||||
local_cache[label] = result
|
||||
roots[i], semitones[i], basses[i] = result
|
||||
return roots, semitones, basses
|
||||
|
||||
|
||||
def rotate_bitmap_to_root(bitmap, chord_root):
|
||||
"""Circularly shift a relative bitmap to its absolute pitch classes.
|
||||
|
||||
For clarity, the best explanation is an example. Given 'G:Maj', the root
|
||||
and quality map are as follows::
|
||||
|
||||
root=5
|
||||
quality=[1, 0, 0, 0, 1, 0, 0, 1, 0, 0, 0, 0] # Relative chord shape
|
||||
|
||||
After rotating to the root, the resulting bitmap becomes::
|
||||
|
||||
abs_quality = [0, 0, 1, 0, 0, 0, 0, 1, 0, 0, 0, 1] # G, B, and D
|
||||
|
||||
Parameters
|
||||
----------
|
||||
bitmap : np.ndarray, shape=(12,)
|
||||
Bitmap of active notes, relative to the given root.
|
||||
chord_root : int
|
||||
Absolute pitch class number.
|
||||
|
||||
Returns
|
||||
-------
|
||||
bitmap : np.ndarray, shape=(12,)
|
||||
Absolute bitmap of active pitch classes.
|
||||
|
||||
"""
|
||||
bitmap = np.asarray(bitmap)
|
||||
assert bitmap.ndim == 1, "Currently only 1D bitmaps are supported."
|
||||
idxs = list(np.nonzero(bitmap))
|
||||
idxs[-1] = (idxs[-1] + chord_root) % 12
|
||||
abs_bitmap = np.zeros_like(bitmap)
|
||||
abs_bitmap[tuple(idxs)] = 1
|
||||
return abs_bitmap
|
||||
@@ -0,0 +1,89 @@
|
||||
"""Minimal in-memory stand-in for the parts of pretty_midi that SheetSage2's ABC path uses:
|
||||
PrettyMIDI(resolution=...) / PrettyMIDI(bytes-like), .instruments, .resolution, .write(stream);
|
||||
Instrument(program, name=...) with .notes/.name/.program/.is_drum; Note(velocity, pitch, start, end).
|
||||
|
||||
"Bytes" are our own serialization (the only consumer is this shim), but note times are rounded to
|
||||
MIDI ticks on write exactly as a real file round trip would (120 bpm default tempo: ticks per second
|
||||
= 2 * resolution), because the ABC builder is documented to follow the serialized MIDI timing."""
|
||||
|
||||
import json
|
||||
from io import BytesIO
|
||||
|
||||
_MAGIC = b"AITK-MIDI-JSON\n"
|
||||
|
||||
|
||||
class Note:
|
||||
def __init__(self, velocity, pitch, start, end):
|
||||
self.velocity = int(velocity)
|
||||
self.pitch = int(pitch)
|
||||
self.start = float(start)
|
||||
self.end = float(end)
|
||||
|
||||
def __repr__(self):
|
||||
return f"Note(start={self.start:.6f}, end={self.end:.6f}, pitch={self.pitch}, velocity={self.velocity})"
|
||||
|
||||
|
||||
class Instrument:
|
||||
def __init__(self, program, is_drum=False, name=""):
|
||||
self.program = int(program)
|
||||
self.is_drum = bool(is_drum)
|
||||
self.name = name
|
||||
self.notes = []
|
||||
|
||||
|
||||
class PrettyMIDI:
|
||||
def __init__(self, midi_file=None, resolution=220, initial_tempo=120.0):
|
||||
self.resolution = int(resolution)
|
||||
self.initial_tempo = float(initial_tempo)
|
||||
self.instruments = []
|
||||
if midi_file is not None:
|
||||
if isinstance(midi_file, (bytes, bytearray)):
|
||||
data = bytes(midi_file)
|
||||
elif hasattr(midi_file, "read"):
|
||||
data = midi_file.read()
|
||||
else:
|
||||
with open(str(midi_file), "rb") as f:
|
||||
data = f.read()
|
||||
if not data.startswith(_MAGIC):
|
||||
raise ValueError("pretty_midi shim can only read MIDI written by this shim")
|
||||
payload = json.loads(data[len(_MAGIC) :].decode("utf-8"))
|
||||
self.resolution = int(payload["resolution"])
|
||||
self.initial_tempo = float(payload["initial_tempo"])
|
||||
for inst in payload["instruments"]:
|
||||
instrument = Instrument(inst["program"], inst["is_drum"], inst["name"])
|
||||
instrument.notes = [Note(n["velocity"], n["pitch"], n["start"], n["end"]) for n in inst["notes"]]
|
||||
self.instruments.append(instrument)
|
||||
|
||||
def _ticks_per_second(self):
|
||||
return self.resolution * self.initial_tempo / 60.0
|
||||
|
||||
def _quantize(self, seconds):
|
||||
tps = self._ticks_per_second()
|
||||
return int(round(seconds * tps)) / tps
|
||||
|
||||
def write(self, filename):
|
||||
payload = {
|
||||
"resolution": self.resolution,
|
||||
"initial_tempo": self.initial_tempo,
|
||||
"instruments": [
|
||||
{
|
||||
"program": i.program,
|
||||
"is_drum": i.is_drum,
|
||||
"name": i.name,
|
||||
"notes": [
|
||||
{"velocity": n.velocity, "pitch": n.pitch, "start": self._quantize(n.start), "end": self._quantize(n.end)}
|
||||
for n in i.notes
|
||||
],
|
||||
}
|
||||
for i in self.instruments
|
||||
],
|
||||
}
|
||||
data = _MAGIC + json.dumps(payload).encode("utf-8")
|
||||
if hasattr(filename, "write"):
|
||||
filename.write(data)
|
||||
else:
|
||||
with open(str(filename), "wb") as f:
|
||||
f.write(data)
|
||||
|
||||
def get_end_time(self):
|
||||
return max((n.end for i in self.instruments for n in i.notes), default=0.0)
|
||||
@@ -38,11 +38,26 @@ class YuE2TextTokenizer:
|
||||
def decode(self, ids: List[int]) -> str:
|
||||
return self.tokenizer.decode(ids, skip_special_tokens=True)
|
||||
|
||||
def prefix_head_ids(self, style: str, lyrics: str, cot: str = "off") -> List[int]:
|
||||
"""Prompt up to and including ABC_START (ComfyUI's ``prefix``); the ABC block and
|
||||
``[ABC_END, MUSIC_START]`` follow (see ``abc_tail_ids``)."""
|
||||
prompt = f"{INSTRUCTIONS[cot]}\n[Tags]\n{style}\n[Lyrics]\n{lyrics}\n"
|
||||
return [EOD] + self.encode(prompt) + [ABC_START]
|
||||
|
||||
@staticmethod
|
||||
def abc_tail_ids(abc_ids: List[int]) -> List[int]:
|
||||
return list(abc_ids) + [ABC_END, MUSIC_START]
|
||||
|
||||
def encode_abc(self, abc: str) -> List[int]:
|
||||
ids = self.encode(abc)
|
||||
if any(not 0 <= t < EOD for t in ids):
|
||||
raise ValueError("ABC ids must stay inside the ordinary text vocabulary")
|
||||
return ids
|
||||
|
||||
def prefix_ids(self, style: str, lyrics: str, cot: str = "off", abc: str = "") -> List[int]:
|
||||
"""Everything the AR expert sees before the first codec token."""
|
||||
prompt = f"{INSTRUCTIONS[cot]}\n[Tags]\n{style}\n[Lyrics]\n{lyrics}\n"
|
||||
abc_ids = [] if cot == "off" else self.encode(abc)
|
||||
return [EOD] + self.encode(prompt) + [ABC_START] + abc_ids + [ABC_END, MUSIC_START]
|
||||
abc_ids = [] if cot == "off" else self.encode_abc(abc)
|
||||
return self.prefix_head_ids(style, lyrics, cot) + self.abc_tail_ids(abc_ids)
|
||||
|
||||
|
||||
class TokenizerHead(nn.Module):
|
||||
@@ -75,6 +90,66 @@ def _rebuild_rotary(model: nn.Module) -> int:
|
||||
return n
|
||||
|
||||
|
||||
SHEETSAGE_REPO = "m-a-p/SheetSage2"
|
||||
SHEETSAGE_FILE = "Comfy-Org/YuE2/audio_encoders/sheetsage2_bf16.safetensors"
|
||||
# tied to token_embedding.weight (loaded); everything else in the Comfy file maps 1:1 onto upstream
|
||||
NONPERSISTENT_OK = {"decoder.embed_tokens.weight"}
|
||||
SHEETSAGE_SAMPLE_RATE = 24000
|
||||
|
||||
|
||||
class SheetSage2Transcriber:
|
||||
"""Audio -> ABC lead sheet with the released SheetSage2 (transformers remote code), the same
|
||||
transcriber ComfyUI ships for cover mode. ``full`` keeps chords, ``melody`` drops them."""
|
||||
|
||||
def __init__(self, weights_path: str, repo: str = SHEETSAGE_REPO):
|
||||
"""``weights_path``: the Comfy repack ``audio_encoders/sheetsage2_bf16.safetensors`` (merged
|
||||
weights); ``repo`` supplies config + code (transformers remote code) only."""
|
||||
from .shims import install_shims
|
||||
|
||||
install_shims() # mir_eval.chord / pretty_midi stand-ins unless the real packages exist
|
||||
import sys
|
||||
|
||||
from safetensors.torch import load_file
|
||||
from transformers import AutoConfig
|
||||
from transformers.dynamic_module_utils import get_class_from_dynamic_module
|
||||
|
||||
from .sheetsage_decoder import ScoreDecoder
|
||||
|
||||
cfg = AutoConfig.from_pretrained(repo, trust_remote_code=True)
|
||||
cfg.weights_format = "merged" # the Comfy file carries the LoRA already folded into the encoder
|
||||
cls = get_class_from_dynamic_module(cfg.auto_map["AutoModel"], repo)
|
||||
sys.modules[cls.__module__].BartDecoder = ScoreDecoder # upstream targets transformers 4.45's BartDecoder
|
||||
cls.post_init = lambda self: None # transformers-5 tied-weight bookkeeping rejects upstream's list; weights load explicitly
|
||||
model = cls(cfg)
|
||||
sd = {k: v.to(torch.float32) for k, v in load_file(weights_path).items()}
|
||||
missing, unexpected = model.load_state_dict(sd, strict=False)
|
||||
missing = [k for k in missing if k not in NONPERSISTENT_OK]
|
||||
if missing or unexpected:
|
||||
raise ValueError(f"SheetSage2 weights do not match: missing {missing[:5]} unexpected {unexpected[:5]}")
|
||||
model.output_projection.weight = model.token_embedding.weight # keep the tie after loading
|
||||
self.model = model.eval()
|
||||
self.model.requires_grad_(False)
|
||||
_rebuild_rotary(self.model)
|
||||
|
||||
def to(self, device):
|
||||
self.model.to(device)
|
||||
return self
|
||||
|
||||
@property
|
||||
def device(self):
|
||||
return next(self.model.parameters()).device
|
||||
|
||||
@torch.no_grad()
|
||||
def transcribe(self, waveform: torch.Tensor, sample_rate: int, cot: str = "full") -> str:
|
||||
"""waveform [C, samples] -> ABC text (``cot`` "full" or "melody")."""
|
||||
mono = waveform.float().mean(0) if waveform.dim() == 2 else waveform.float()
|
||||
result = self.model.transcribe(mono.cpu(), sampling_rate=sample_rate, melody_only=cot == "melody")
|
||||
abc = result.get("abc") or ""
|
||||
if not abc.strip():
|
||||
raise RuntimeError(f"SheetSage2 produced no ABC: {result.get('abc_error') or 'unknown reason'}")
|
||||
return abc
|
||||
|
||||
|
||||
class SemanticTokenizer(nn.Module):
|
||||
"""Waveform -> per-frame codec ids (0..32767) at 25 Hz."""
|
||||
|
||||
|
||||
@@ -1,6 +1,7 @@
|
||||
"""YuE2 (m-a-p/YuE2-3B) for ai-toolkit.
|
||||
|
||||
Weights: the Comfy-Org all-in-one checkpoint (``checkpoints/yue2_3b_bf16.safetensors``)
|
||||
Weights: the Comfy-Org all-in-one checkpoint (default ``checkpoints/yue2_3b_int8_convrot.safetensors``,
|
||||
shipped int8 convrot layers attached as-is; ``yue2_3b_bf16.safetensors`` also loads)
|
||||
resolved under MODELS_PATH in ComfyUI's folder layout and downloaded there when missing.
|
||||
The official audio -> semantic-token tokenizer is unreleased; training conditioning comes
|
||||
from the community head (MERT-v2-FullSong + classifier) in
|
||||
@@ -15,8 +16,10 @@ flow loss on the real latents. With ``model_kwargs.ar_loss_weight`` > 0 the same
|
||||
also trains the AR expert with next-token cross entropy over the codec tokens, so one LoRA
|
||||
network on both experts learns composition and rendering together.
|
||||
|
||||
Captions use the ace-step tag format: <CAPTION> (style), <LYRICS>, <DURATION> (max seconds
|
||||
for samples).
|
||||
Prompt modes follow ComfyUI (``model_kwargs.cot``): "full" (default, Generate ABC path: the AR
|
||||
writes a chord-annotated ABC sheet, then the codec tokens), "melody" (sheet without chords) or
|
||||
"off" (no sheet; what Comfy runs when the ABC input is empty). For full/melody the training sheet
|
||||
comes from SheetSage2 at cache time and rides the latent cache as ``abc_ids``.
|
||||
"""
|
||||
|
||||
import os
|
||||
@@ -25,6 +28,7 @@ import re
|
||||
from typing import List, Optional
|
||||
|
||||
import torch
|
||||
import torch.utils.checkpoint
|
||||
from safetensors.torch import load_file
|
||||
from tqdm import tqdm
|
||||
|
||||
@@ -38,12 +42,15 @@ from toolkit.print import print_acc
|
||||
from toolkit.prompt_utils import PromptEmbeds
|
||||
from toolkit.samplers.custom_flowmatch_sampler import CustomFlowMatchEulerDiscreteScheduler
|
||||
|
||||
from .src.model import CODEC_OFFSET, FRAMES_PER_SECOND, MUSIC_END, YuE2Model, merge_nar_lora
|
||||
from .src.model import ABC_END, CODEC_OFFSET, CONTEXT, FRAMES_PER_SECOND, INSTRUCTIONS, MUSIC_END, MUSIC_START, YuE2Model, merge_nar_lora
|
||||
from .src.pipeline import YuE2Pipeline
|
||||
from .src.tokenizer import HEAD_FILE, HEAD_REPO, NAR_LORA_FILE, SemanticTokenizer, YuE2TextTokenizer
|
||||
from .src.tokenizer import HEAD_FILE, HEAD_REPO, NAR_LORA_FILE, SHEETSAGE_FILE, SemanticTokenizer, SheetSage2Transcriber, YuE2TextTokenizer
|
||||
from .src.vae import SAMPLE_RATE, YuE2VAE
|
||||
|
||||
DEFAULT_CHECKPOINT = "Comfy-Org/YuE2/checkpoints/yue2_3b_bf16.safetensors"
|
||||
COT_CODES = {"off": 0, "melody": 1, "full": 2}
|
||||
COT_NAMES = {v: k for k, v in COT_CODES.items()}
|
||||
|
||||
DEFAULT_CHECKPOINT = "Comfy-Org/YuE2/checkpoints/yue2_3b_int8_convrot.safetensors"
|
||||
# community tokenizer files have no ComfyUI folder; they live under <MODELS_PATH>/ai_toolkit/
|
||||
AI_TOOLKIT_DIR = "ai_toolkit"
|
||||
|
||||
@@ -130,10 +137,10 @@ def parse_caption(text: str) -> dict:
|
||||
|
||||
|
||||
class YuE2TextEncoder(torch.nn.Module):
|
||||
"""Prompt prefix -> AR token embeddings. Holds no weights of its own; the
|
||||
embedding table lives in the AR expert (looked up on whatever device it is on)."""
|
||||
"""Prompt -> AR token embeddings of the prefix head (instruction, tags, lyrics, ABC_START).
|
||||
Holds no weights of its own; the embedding table lives in the AR expert."""
|
||||
|
||||
def __init__(self, model: YuE2Model, tokenizer: YuE2TextTokenizer, cot: str = "off"):
|
||||
def __init__(self, model: YuE2Model, tokenizer: YuE2TextTokenizer, cot: str = "full"):
|
||||
super().__init__()
|
||||
self._model = [model]
|
||||
self.tokenizer = tokenizer
|
||||
@@ -146,14 +153,11 @@ class YuE2TextEncoder(torch.nn.Module):
|
||||
|
||||
@property
|
||||
def dtype(self):
|
||||
return self._model[0].ar.model.embed_tokens.weight.dtype
|
||||
|
||||
def prefix_ids(self, style: str, lyrics: str) -> List[int]:
|
||||
return self.tokenizer.prefix_ids(style, lyrics, cot=self.cot)
|
||||
return self._model[0].ar.dtype
|
||||
|
||||
@torch.no_grad()
|
||||
def forward(self, style: str, lyrics: str) -> torch.Tensor:
|
||||
ids = torch.tensor([self.prefix_ids(style, lyrics)], dtype=torch.long)
|
||||
ids = torch.tensor([self.tokenizer.prefix_head_ids(style, lyrics, cot=self.cot)], dtype=torch.long)
|
||||
return self._model[0].ar.embed(ids)
|
||||
|
||||
|
||||
@@ -182,6 +186,15 @@ class YuE2AudioModel(BaseAudioModel):
|
||||
self.merge_nar_lora = bool(kw.get("merge_nar_lora", False))
|
||||
self.nar_lora_path = kw.get("nar_lora_path", f"{HEAD_REPO}/{NAR_LORA_FILE}")
|
||||
self.semantic_head_path = kw.get("semantic_head_path", f"{HEAD_REPO}/{HEAD_FILE}")
|
||||
# ComfyUI's Generate ABC path: "full" (chords) or "melody"; "off" skips the sheet (empty ABC in Comfy)
|
||||
self.cot = str(kw.get("cot", "full"))
|
||||
if self.cot not in INSTRUCTIONS:
|
||||
raise ValueError(f"model_kwargs.cot must be one of {list(INSTRUCTIONS)}")
|
||||
self.abc_max_tokens = int(kw.get("abc_max_tokens", 8192))
|
||||
# fraction of training items fed without the sheet (off-mode prefix), so one LoRA serves both Comfy paths
|
||||
self.abc_dropout = float(kw.get("abc_dropout", 0.5))
|
||||
self.sheetsage_path = kw.get("sheetsage_path", SHEETSAGE_FILE)
|
||||
self.transcriber: Optional[SheetSage2Transcriber] = None
|
||||
self.sample_max_seconds = float(kw.get("sample_max_seconds", 120.0))
|
||||
self.sample_ar_temperature = float(kw.get("sample_ar_temperature", 1.0))
|
||||
# reference sampler value; 1.0 lets a memorized song replay (the penalty knocks it off path)
|
||||
@@ -230,7 +243,7 @@ class YuE2AudioModel(BaseAudioModel):
|
||||
flush()
|
||||
self.model.to(device)
|
||||
self.vae.to(self.vae_device_torch)
|
||||
self.text_encoder = YuE2TextEncoder(self.model, self.tokenizer)
|
||||
self.text_encoder = YuE2TextEncoder(self.model, self.tokenizer, cot=self.cot)
|
||||
self.pipeline = YuE2Pipeline(self.model, self.vae)
|
||||
self.pipeline.do_tiled_decoding = True
|
||||
self.pipeline.use_cuda_graphs = bool(self.model_config.model_kwargs.get("ar_cuda_graphs", True))
|
||||
@@ -243,6 +256,12 @@ class YuE2AudioModel(BaseAudioModel):
|
||||
self.semantic_tokenizer = SemanticTokenizer(head_path, debug=self.debug)
|
||||
return self.semantic_tokenizer
|
||||
|
||||
def _get_transcriber(self) -> SheetSage2Transcriber:
|
||||
if self.transcriber is None:
|
||||
print_acc("Loading SheetSage2 (audio -> ABC)")
|
||||
self.transcriber = SheetSage2Transcriber(resolve_named_file(self.sheetsage_path, component="sheetsage2"))
|
||||
return self.transcriber
|
||||
|
||||
def get_transformer_block_names(self) -> Optional[List[str]]:
|
||||
return ["model.layers"]
|
||||
|
||||
@@ -269,7 +288,8 @@ class YuE2AudioModel(BaseAudioModel):
|
||||
return PromptEmbeds(torch.cat(embeds, 0), attention_mask=torch.cat(masks, 0))
|
||||
|
||||
def encode_audio(self, audio_tensor: torch.Tensor, device=None, dtype=None):
|
||||
"""[B, 2, samples] at 48 kHz -> DTO [B, T, 64] with codec ``tokens`` [B, T]."""
|
||||
"""[B, 2, samples] at 48 kHz -> DTO [B, T, 64] with codec ``tokens`` [B, T] and, unless cot is
|
||||
"off", the SheetSage2 sheet as ``abc_ids`` [B, N] (text ids, -1 padded)."""
|
||||
if device is None:
|
||||
device = self.vae_device_torch
|
||||
if self.vae.device == torch.device("cpu"):
|
||||
@@ -282,7 +302,20 @@ class YuE2AudioModel(BaseAudioModel):
|
||||
tokens = torch.stack([tokenizer.tokenize(wav, SAMPLE_RATE) for wav in audio_tensor]) # [B, T']
|
||||
n = min(latents.shape[1], tokens.shape[1])
|
||||
latents = latents[:, :n].to(dtype=self.torch_dtype if dtype is None else dtype)
|
||||
return DTO(latents, tokens=tokens[:, :n].to(latents.device, torch.int32))
|
||||
extras = {"tokens": tokens[:, :n].to(latents.device, torch.int32)}
|
||||
if self.cot != "off":
|
||||
transcriber = self._get_transcriber()
|
||||
if transcriber.device != self.device_torch:
|
||||
transcriber.to(self.device_torch)
|
||||
sheets = [self.tokenizer.encode_abc(transcriber.transcribe(wav, SAMPLE_RATE, cot=self.cot)) for wav in audio_tensor]
|
||||
width = max(len(x) for x in sheets)
|
||||
abc = torch.full((len(sheets), width), -1, dtype=torch.int32)
|
||||
for i, ids in enumerate(sheets):
|
||||
abc[i, : len(ids)] = torch.tensor(ids, dtype=torch.int32)
|
||||
extras["abc_ids"] = abc.to(latents.device)
|
||||
# which sheet flavor the cache holds; training refuses a mismatch instead of silently using it
|
||||
extras["abc_mode"] = torch.full((len(sheets),), COT_CODES[self.cot], dtype=torch.int32, device=latents.device)
|
||||
return DTO(latents, **extras)
|
||||
|
||||
# ------------------------------------------------------------------
|
||||
# training
|
||||
@@ -330,42 +363,67 @@ class YuE2AudioModel(BaseAudioModel):
|
||||
length = int((prefix_mask > 0.5).sum().item())
|
||||
return prefix_embeds[:length]
|
||||
|
||||
def _ar_info(self, logits: torch.Tensor, ids: torch.Tensor, prefix: torch.Tensor, total: int, whole_song: bool) -> str:
|
||||
"""Readout diagnostics: what the AR loss was computed over."""
|
||||
if not self.debug or self._loss_log_every <= 0 or (self._loss_log_step + 1) % self._loss_log_every != 0:
|
||||
return ""
|
||||
with torch.no_grad():
|
||||
p = logits.softmax(-1)
|
||||
p_target = p.gather(1, ids[:, None]).squeeze(1)
|
||||
return (f"[ar targets {ids.shape[0]} unique {ids.unique().numel()} (song {total}, whole {whole_song}) "
|
||||
f"prefix {prefix.shape[0]} p(END|prefix) {p[0, MUSIC_END].item():.3f} p(target) median {p_target.median().item():.3f}]")
|
||||
|
||||
def _ar_kl(self, ar_embeds: torch.Tensor, logits: torch.Tensor):
|
||||
"""KL(base || lora) per position, mean. Base pass = same prefill with the LoRA inactive, no grad."""
|
||||
net = getattr(self, "_network", None)
|
||||
n = logits.shape[0]
|
||||
was_active = getattr(net, "is_active", None)
|
||||
if net is not None:
|
||||
net.is_active = False
|
||||
try:
|
||||
with torch.no_grad():
|
||||
_, base_hidden = self.model.ar.prefill(ar_embeds, return_hidden=True)
|
||||
base_hidden = base_hidden[0, -n - 1 : -1]
|
||||
finally:
|
||||
def _ar_losses(self, ar_embeds: torch.Tensor, ids: torch.Tensor, prefix: torch.Tensor, total: int, whole_song: bool):
|
||||
"""Next-token CE over ``ids`` (and KL(base || lora) when ``ar_kl_weight`` > 0), computed in
|
||||
position chunks with recomputed logits so the fp32 [n, vocab] tensors never exist at once.
|
||||
Returns (ce, kl_or_None, info_str)."""
|
||||
model = self.model
|
||||
n = ids.shape[0]
|
||||
_, hidden = model.ar.prefill(ar_embeds, return_hidden=True)
|
||||
hidden = hidden[0, -n - 1 : -1]
|
||||
base_hidden = None
|
||||
if self.ar_kl_weight > 0:
|
||||
net = getattr(self, "_network", None)
|
||||
was_active = getattr(net, "is_active", None)
|
||||
if net is not None:
|
||||
net.is_active = was_active
|
||||
kl = logits.new_zeros(())
|
||||
for s0 in range(0, n, 1024): # chunked: two fp32 [n, vocab] log-softmaxes at once would be ~10 GB
|
||||
with torch.no_grad():
|
||||
base_logp = torch.log_softmax(self.model.ar.model.lm_head(base_hidden[s0 : s0 + 1024]).float(), -1)
|
||||
kl = kl + torch.nn.functional.kl_div(torch.log_softmax(logits[s0 : s0 + 1024], -1), base_logp, log_target=True, reduction="sum")
|
||||
return kl / n
|
||||
net.is_active = False
|
||||
try:
|
||||
with torch.no_grad():
|
||||
_, base_hidden = model.ar.prefill(ar_embeds, return_hidden=True)
|
||||
base_hidden = base_hidden[0, -n - 1 : -1]
|
||||
finally:
|
||||
if net is not None:
|
||||
net.is_active = was_active
|
||||
want_info = self.debug and self._loss_log_every > 0 and (self._loss_log_step + 1) % self._loss_log_every == 0
|
||||
lm_head = model.ar.model.lm_head
|
||||
ce_sum = hidden.new_zeros((), dtype=torch.float32)
|
||||
kl_sum = hidden.new_zeros((), dtype=torch.float32) if base_hidden is not None else None
|
||||
p_target, p_end_first = [], None
|
||||
chunk = 2048
|
||||
for s0 in range(0, n, chunk):
|
||||
h = hidden[s0 : s0 + chunk]
|
||||
tgt = ids[s0 : s0 + chunk]
|
||||
if torch.is_grad_enabled():
|
||||
logits = torch.utils.checkpoint.checkpoint(lambda x: lm_head(x).float(), h, use_reentrant=False)
|
||||
else:
|
||||
logits = lm_head(h).float()
|
||||
ce_sum = ce_sum + torch.nn.functional.cross_entropy(logits, tgt, reduction="sum")
|
||||
if base_hidden is not None:
|
||||
with torch.no_grad():
|
||||
base_logp = torch.log_softmax(lm_head(base_hidden[s0 : s0 + chunk]).float(), -1)
|
||||
kl_sum = kl_sum + torch.nn.functional.kl_div(torch.log_softmax(logits, -1), base_logp, log_target=True, reduction="sum")
|
||||
if want_info:
|
||||
with torch.no_grad():
|
||||
p = logits.softmax(-1)
|
||||
if s0 == 0:
|
||||
p_end_first = p[0, MUSIC_END].item()
|
||||
p_target.append(p.gather(1, tgt[:, None]).squeeze(1))
|
||||
del logits
|
||||
info = ""
|
||||
if want_info:
|
||||
p_target = torch.cat(p_target)
|
||||
info = (f"[ar targets {n} unique {ids.unique().numel()} (song {total}, whole {whole_song}) "
|
||||
f"prefix {prefix.shape[0]} p(END|prefix) {p_end_first:.3f} p(target) median {p_target.median().item():.3f}]")
|
||||
return ce_sum / n, (kl_sum / n if kl_sum is not None else None), info
|
||||
|
||||
def _ar_inputs(self, prefix: torch.Tensor, tokens: torch.Tensor, end_token: bool = True):
|
||||
"""One item: prefix (unpadded [L, H]) + codec tokens [+ MUSIC_END] -> embeds [1, L', H], ids [n]."""
|
||||
ids = tokens.long() + CODEC_OFFSET
|
||||
def _ar_inputs(self, prefix: torch.Tensor, abc_ids: torch.Tensor, tokens: torch.Tensor, end_token: bool = True):
|
||||
"""One item: prefix head (unpadded [L, H]) + abc + [ABC_END, MUSIC_START] + codec tokens [+ MUSIC_END]
|
||||
-> embeds [1, L', H] and the ids after the head [n] (the AR targets)."""
|
||||
device = tokens.device
|
||||
parts = [abc_ids.long().to(device), torch.tensor([ABC_END, MUSIC_START], device=device), tokens.long() + CODEC_OFFSET]
|
||||
if end_token:
|
||||
ids = torch.cat([ids, torch.tensor([MUSIC_END], device=tokens.device)])
|
||||
parts.append(torch.tensor([MUSIC_END], device=device))
|
||||
ids = torch.cat(parts)
|
||||
embeds = torch.cat([prefix.to(self.model.ar.dtype), self.model.ar.embed(ids)], dim=0)[None]
|
||||
return embeds, ids
|
||||
|
||||
@@ -416,6 +474,16 @@ class YuE2AudioModel(BaseAudioModel):
|
||||
# `batch` must be a named parameter: predict_noise only forwards it when it sees it in the signature
|
||||
if batch is None or not isinstance(batch.latents, DTO) or batch.latents.get("tokens") is None:
|
||||
raise ValueError("YuE2 training needs codec tokens in the latent cache; enable latent caching")
|
||||
abc_all = batch.latents.get("abc_ids")
|
||||
if self.cot != "off":
|
||||
mode = batch.latents.get("abc_mode")
|
||||
cached = None if mode is None else int(mode.reshape(-1)[0].item())
|
||||
if abc_all is None or cached != COT_CODES[self.cot]:
|
||||
raise ValueError(
|
||||
f"YuE2 cot={self.cot!r} needs a matching ABC sheet in the latent cache "
|
||||
f"(cache has {'none' if cached is None else COT_NAMES.get(cached, cached)!r}): "
|
||||
"delete the dataset's _latent_cache folder so it is rebuilt"
|
||||
)
|
||||
start, end = getattr(batch, "yue2_window", (0, latent_model_input.shape[1]))
|
||||
tokens_all = batch.latents.tokens
|
||||
model = self.model
|
||||
@@ -426,35 +494,48 @@ class YuE2AudioModel(BaseAudioModel):
|
||||
prefix_embeds = text_embeddings.text_embeds.to(device, self.torch_dtype)
|
||||
prefix_mask = text_embeddings.attention_mask.to(device)
|
||||
train_ar = self.ar_loss_weight > 0 and torch.is_grad_enabled()
|
||||
captions = batch.get_caption_list() if hasattr(batch, "get_caption_list") else None
|
||||
preds, lm_losses, kl_losses = [], [], []
|
||||
for i in range(latent_model_input.shape[0]):
|
||||
song = tokens_all[i].to(device)
|
||||
total = song.shape[0]
|
||||
prefix = self._train_prefix(prefix_embeds[i], prefix_mask[i])
|
||||
# NAR conditioning follows the released chunk protocol: prefix + window + MUSIC_END
|
||||
embeds, ids = self._ar_inputs(prefix, song[start:end])
|
||||
if self.cot == "off" or abc_all is None:
|
||||
abc = torch.zeros(0, dtype=torch.long, device=device)
|
||||
elif torch.is_grad_enabled() and captions is not None and random.random() < self.abc_dropout:
|
||||
# sheet dropped: the item trains exactly as an off-mode prompt (off instruction, empty ABC block)
|
||||
parsed = parse_caption(captions[i])
|
||||
ids = torch.tensor([self.tokenizer.prefix_head_ids(parsed["style"], parsed["lyrics"], cot="off")], device=device)
|
||||
prefix = model.ar.embed(ids)[0].to(self.torch_dtype)
|
||||
abc = torch.zeros(0, dtype=torch.long, device=device)
|
||||
else:
|
||||
row = abc_all[i].to(device)
|
||||
abc = row[row >= 0].long()
|
||||
# NAR conditioning follows the released chunk protocol: prefix + abc + window + MUSIC_END
|
||||
embeds, ids = self._ar_inputs(prefix, abc, song[start:end])
|
||||
whole_song = start == 0 and end == total
|
||||
if train_ar and not whole_song:
|
||||
# the language-model loss must see the song from its beginning, or the AR
|
||||
# learns that lyrics and tokens can start anywhere relative to each other
|
||||
limit = total if self.ar_max_tokens <= 0 else min(total, self.ar_max_tokens)
|
||||
ar_embeds, ar_ids = self._ar_inputs(prefix, song[:limit], end_token=limit == total)
|
||||
_, hidden = model.ar.prefill(ar_embeds, return_hidden=True)
|
||||
logits = model.ar.model.lm_head(hidden[0, -ar_ids.shape[0] - 1 : -1]).float()
|
||||
lm_losses.append(torch.nn.functional.cross_entropy(logits, ar_ids))
|
||||
if self.ar_kl_weight > 0:
|
||||
kl_losses.append(self._ar_kl(ar_embeds, logits))
|
||||
self._pending_aux_info = self._ar_info(logits, ar_ids, prefix, total, whole_song)
|
||||
ar_embeds, ar_ids = self._ar_inputs(prefix, abc, song[:limit], end_token=limit == total)
|
||||
ce, kl, info = self._ar_losses(ar_embeds, ar_ids, prefix, total, whole_song)
|
||||
lm_losses.append(ce)
|
||||
if kl is not None:
|
||||
kl_losses.append(kl)
|
||||
self._pending_aux_info = info
|
||||
with torch.no_grad():
|
||||
cache, _ = model.ar.prefill(embeds)
|
||||
elif train_ar:
|
||||
ce, kl, info = self._ar_losses(embeds, ids, prefix, total, whole_song)
|
||||
lm_losses.append(ce)
|
||||
if kl is not None:
|
||||
kl_losses.append(kl)
|
||||
self._pending_aux_info = info
|
||||
with torch.no_grad():
|
||||
cache, _ = model.ar.prefill(embeds)
|
||||
else:
|
||||
cache, hidden = model.ar.prefill(embeds, return_hidden=train_ar)
|
||||
if train_ar:
|
||||
logits = model.ar.model.lm_head(hidden[0, -ids.shape[0] - 1 : -1]).float()
|
||||
lm_losses.append(torch.nn.functional.cross_entropy(logits, ids))
|
||||
if self.ar_kl_weight > 0:
|
||||
kl_losses.append(self._ar_kl(embeds, logits))
|
||||
self._pending_aux_info = self._ar_info(logits, ids, prefix, total, whole_song)
|
||||
cache, _ = model.ar.prefill(embeds)
|
||||
# the flow loss must not train the AR through its KV cache: only next-token CE shapes the AR
|
||||
cache = [(k.detach(), v.detach()) for k, v in cache]
|
||||
pred = model.nar(latent_model_input[i : i + 1].to(device, self.torch_dtype), t[i : i + 1], cache, embeds.shape[1])
|
||||
@@ -484,10 +565,28 @@ class YuE2AudioModel(BaseAudioModel):
|
||||
max_seconds = getattr(gen_config, "duration", None) or parsed["duration"] or self.sample_max_seconds
|
||||
max_tokens = max(1, int(round(max_seconds * FRAMES_PER_SECOND)))
|
||||
length = int((conditional_embeds.attention_mask[0] > 0.5).sum().item())
|
||||
prefix = conditional_embeds.text_embeds[:1, :length].to(self.device_torch, self.torch_dtype)
|
||||
head = conditional_embeds.text_embeds[:1, :length].to(self.device_torch, self.torch_dtype)
|
||||
|
||||
# stage 1 (cot full/melody): the AR writes the ABC sheet after ABC_START, as Comfy's Generate ABC does
|
||||
abc_ids: List[int] = []
|
||||
if self.cot != "off":
|
||||
self._status_update("YuE2: generating ABC sheet")
|
||||
abc_bar = tqdm(total=self.abc_max_tokens, desc=" abc tokens", unit="tok", position=1, leave=False, dynamic_ncols=True)
|
||||
try:
|
||||
abc_ids = pipeline.generate_abc_tokens(
|
||||
head, seed=gen_config.seed, max_tokens=self.abc_max_tokens, progress=lambda done, total: abc_bar.update(done - abc_bar.n)
|
||||
)
|
||||
finally:
|
||||
abc_bar.close()
|
||||
tail = torch.tensor([self.tokenizer.abc_tail_ids(abc_ids)], device=self.device_torch)
|
||||
prefix = torch.cat([head, self.model.ar.embed(tail).to(head.dtype)], dim=1)
|
||||
# one acoustic frame needs two positions plus the boundary tokens (Comfy's budget rule)
|
||||
max_tokens = min(max_tokens, CONTEXT - prefix.shape[1] - 5)
|
||||
if max_tokens < 1:
|
||||
raise ValueError("YuE2 prompt leaves no room for music; shorten the style, lyrics, or ABC")
|
||||
|
||||
# nested bars under generate_images' "Generating Samples" bar (position 0)
|
||||
self._status_update(f"YuE2: generating codec tokens (up to {max_seconds:.0f}s)")
|
||||
self._status_update(f"YuE2: generating codec tokens (up to {max_tokens / FRAMES_PER_SECOND:.0f}s)")
|
||||
token_bar = tqdm(total=max_tokens, desc=" codec tokens", unit="tok", position=1, leave=False, dynamic_ncols=True)
|
||||
|
||||
def token_progress(done, total):
|
||||
@@ -497,7 +596,7 @@ class YuE2AudioModel(BaseAudioModel):
|
||||
try:
|
||||
codec = pipeline.generate_codec_tokens(
|
||||
prefix, max_tokens=max_tokens, seed=gen_config.seed, temperature=self.sample_ar_temperature,
|
||||
repetition_penalty=self.sample_ar_repetition_penalty, progress=token_progress,
|
||||
repetition_penalty=self.sample_ar_repetition_penalty, progress=token_progress, legacy_off=self.cot == "off",
|
||||
)
|
||||
finally:
|
||||
token_bar.close()
|
||||
|
||||
@@ -14,6 +14,7 @@ from toolkit.samplers.custom_flowmatch_sampler import (
|
||||
CustomFlowMatchEulerDiscreteScheduler,
|
||||
)
|
||||
from safetensors.torch import load_file, save_file
|
||||
from toolkit.util.streamed_safetensors import save_file_streamed
|
||||
from toolkit.accelerator import unwrap_model
|
||||
from optimum.quanto import freeze
|
||||
|
||||
@@ -502,7 +503,8 @@ class HidreamO1Model(BaseModel):
|
||||
if not output_path.endswith(".safetensors"):
|
||||
output_path += ".safetensors"
|
||||
meta = get_meta_for_safetensors(meta, name=self.arch)
|
||||
save_file(save_dict, output_path, metadata=meta)
|
||||
# sequential writer: mmap save_file fails with EINVAL on ntfs3 for large files
|
||||
save_file_streamed(save_dict, output_path, metadata=meta)
|
||||
else:
|
||||
transformer.save_pretrained(
|
||||
save_directory=output_path,
|
||||
|
||||
Reference in new issue
Block a user