mirror of
https://github.com/Comfy-Org/ComfyUI.git
synced 2026-09-21 13:38:08 -05:00
Support Yue2 music model. (#16250)
This commit is contained in:
comfyanonymous
@@ -1,10 +1,13 @@
|
||||
from .wav2vec2 import Wav2Vec2Model
|
||||
from .whisper import WhisperLargeV3
|
||||
from .sheetsage2 import SheetSage2
|
||||
from .sheetsage2_abc import events_to_abc
|
||||
import comfy.model_management
|
||||
import comfy.ops
|
||||
import comfy.utils
|
||||
import logging
|
||||
import torchaudio
|
||||
import torch
|
||||
|
||||
|
||||
class AudioEncoderModel():
|
||||
@@ -13,6 +16,9 @@ class AudioEncoderModel():
|
||||
offload_device = comfy.model_management.text_encoder_offload_device()
|
||||
self.dtype = comfy.model_management.text_encoder_dtype(self.load_device)
|
||||
model_type = config.pop("model_type")
|
||||
self.model_sample_rate = config.pop("model_sample_rate", 16000)
|
||||
if model_type == "sheetsage2":
|
||||
self.dtype = torch.bfloat16 if comfy.model_management.should_use_bf16(self.load_device) else torch.float32
|
||||
model_config = dict(config)
|
||||
model_config.update({
|
||||
"dtype": self.dtype,
|
||||
@@ -24,9 +30,10 @@ class AudioEncoderModel():
|
||||
self.model = Wav2Vec2Model(**model_config)
|
||||
elif model_type == "whisper3":
|
||||
self.model = WhisperLargeV3(**model_config)
|
||||
elif model_type == "sheetsage2":
|
||||
self.model = SheetSage2(**model_config)
|
||||
self.model.eval()
|
||||
self.patcher = comfy.model_patcher.CoreModelPatcher(self.model, load_device=self.load_device, offload_device=offload_device)
|
||||
self.model_sample_rate = 16000
|
||||
comfy.model_management.archive_model_dtypes(self.model)
|
||||
|
||||
def load_sd(self, sd):
|
||||
@@ -46,6 +53,17 @@ class AudioEncoderModel():
|
||||
return outputs
|
||||
|
||||
|
||||
class SheetSage2AudioEncoder(AudioEncoderModel):
|
||||
def generate_abc(self, audio, sample_rate, melody_only=True):
|
||||
audio = torchaudio.functional.resample(audio.float().mean(dim=1), sample_rate, self.model_sample_rate)
|
||||
comfy.model_management.load_model_gpu(self.patcher)
|
||||
scores = []
|
||||
for waveform in audio:
|
||||
events = self.model.transcribe(waveform[None].to(self.load_device))
|
||||
scores.append(events_to_abc(events, waveform.shape[-1] / self.model_sample_rate, melody_only=melody_only))
|
||||
return scores
|
||||
|
||||
|
||||
def load_audio_encoder_from_sd(sd, prefix=""):
|
||||
sd = comfy.utils.state_dict_prefix_replace(sd, {"wav2vec2.": ""})
|
||||
if "encoder.layer_norm.bias" in sd: #wav2vec2
|
||||
@@ -79,10 +97,12 @@ def load_audio_encoder_from_sd(sd, prefix=""):
|
||||
config = {
|
||||
"model_type": "whisper3",
|
||||
}
|
||||
elif "encoder.feature_extractor.mel_mean" in sd and "decoder.layernorm_embedding.weight" in sd and "layer_weight" in sd:
|
||||
config = {"model_type": "sheetsage2", "model_sample_rate": 24000}
|
||||
else:
|
||||
raise RuntimeError("ERROR: audio encoder not supported.")
|
||||
|
||||
audio_encoder = AudioEncoderModel(config)
|
||||
audio_encoder = SheetSage2AudioEncoder(config) if config["model_type"] == "sheetsage2" else AudioEncoderModel(config)
|
||||
m, u = audio_encoder.load_sd(sd)
|
||||
if len(m) > 0:
|
||||
logging.warning("missing audio encoder: {}".format(m))
|
||||
|
||||
@@ -0,0 +1,182 @@
|
||||
"""MERT-v2 ConvNeXt/Conformer audio encoder."""
|
||||
|
||||
import torch
|
||||
from torch import nn
|
||||
from torch.nn import functional as F
|
||||
|
||||
import comfy.ops
|
||||
import comfy.quant_ops
|
||||
from comfy.ldm.modules.attention import optimized_attention_for_device
|
||||
|
||||
|
||||
class MelFrontend(nn.Module):
|
||||
def __init__(self, device=None):
|
||||
super().__init__()
|
||||
self.register_buffer("mel_mean", torch.empty(128, device=device, dtype=torch.float32))
|
||||
self.register_buffer("mel_std", torch.empty(128, device=device, dtype=torch.float32))
|
||||
self.spectrogram = nn.Module()
|
||||
self.spectrogram.register_buffer("window", torch.empty(2048, device=device, dtype=torch.float32))
|
||||
self.mel_scale = nn.Module()
|
||||
self.mel_scale.register_buffer("fb", torch.empty(1025, 128, device=device, dtype=torch.float32))
|
||||
|
||||
def forward(self, waveform):
|
||||
window = comfy.ops.cast_to_input(self.spectrogram.window, waveform)
|
||||
spectrum = torch.stft(waveform, n_fft=2048, hop_length=240, win_length=2048,
|
||||
window=window, return_complex=True).abs().square()
|
||||
mel = spectrum.transpose(-1, -2) @ comfy.ops.cast_to_input(self.mel_scale.fb, waveform)
|
||||
mel = 10.0 * mel.clamp_min(1e-10).log10()
|
||||
mean = comfy.ops.cast_to_input(self.mel_mean, waveform)
|
||||
std = comfy.ops.cast_to_input(self.mel_std, waveform)
|
||||
return (mel[:, :-1] - mean) / std.clamp_min(1e-5)
|
||||
|
||||
|
||||
class Transpose(nn.Module):
|
||||
def forward(self, x):
|
||||
return x.transpose(1, 2)
|
||||
|
||||
|
||||
class GlobalResponseNorm(nn.Module):
|
||||
def __init__(self, dim, device=None, dtype=None):
|
||||
super().__init__()
|
||||
self.weight = nn.Parameter(torch.empty(1, 1, dim, device=device, dtype=dtype))
|
||||
self.bias = nn.Parameter(torch.empty(1, 1, dim, device=device, dtype=dtype))
|
||||
|
||||
def forward(self, x):
|
||||
magnitude = torch.linalg.vector_norm(x, dim=1, keepdim=True)
|
||||
normalized = magnitude / (magnitude.mean(dim=-1, keepdim=True) + 1e-6)
|
||||
weight = comfy.ops.cast_to_input(self.weight, x)
|
||||
bias = comfy.ops.cast_to_input(self.bias, x)
|
||||
return weight * (x * normalized) + bias + x
|
||||
|
||||
|
||||
class ConvNextLayer(nn.Module):
|
||||
def __init__(self, dim, device=None, dtype=None, operations=None):
|
||||
super().__init__()
|
||||
self.depthwise_block = nn.Sequential(
|
||||
Transpose(), operations.Conv1d(dim, dim, 7, padding=3, groups=dim, device=device, dtype=dtype), Transpose(),
|
||||
)
|
||||
self.pointwise_block = nn.Sequential(
|
||||
operations.LayerNorm(dim, eps=1e-6, device=device, dtype=dtype),
|
||||
operations.Linear(dim, 4 * dim, device=device, dtype=dtype), nn.GELU(),
|
||||
GlobalResponseNorm(4 * dim, device=device, dtype=dtype),
|
||||
operations.Linear(4 * dim, dim, device=device, dtype=dtype),
|
||||
)
|
||||
|
||||
def forward(self, x):
|
||||
return x + self.pointwise_block(self.depthwise_block(x))
|
||||
|
||||
|
||||
class ConvNextBlock(nn.Module):
|
||||
def __init__(self, in_channels, out_channels, stride, depth, device=None, dtype=None, operations=None):
|
||||
super().__init__()
|
||||
self.resampling_layer = nn.Identity()
|
||||
if in_channels != out_channels or stride > 1:
|
||||
self.resampling_layer = nn.Sequential(
|
||||
operations.LayerNorm(in_channels, eps=1e-6, device=device, dtype=dtype), Transpose(),
|
||||
operations.Conv1d(in_channels, out_channels, 2, stride=stride, device=device, dtype=dtype), Transpose(),
|
||||
)
|
||||
self.convnext_layers = nn.Sequential(*[
|
||||
ConvNextLayer(out_channels, device=device, dtype=dtype, operations=operations) for _ in range(depth)
|
||||
])
|
||||
|
||||
def forward(self, x):
|
||||
return self.convnext_layers(self.resampling_layer(x))
|
||||
|
||||
|
||||
class Attention(nn.Module):
|
||||
def __init__(self, dim, heads, device=None, dtype=None, operations=None):
|
||||
super().__init__()
|
||||
self.heads = heads
|
||||
self.query_proj = operations.Linear(dim, dim, device=device, dtype=dtype)
|
||||
self.key_proj = operations.Linear(dim, dim, device=device, dtype=dtype)
|
||||
self.value_proj = operations.Linear(dim, dim, device=device, dtype=dtype)
|
||||
self.out_proj = operations.Linear(dim, dim, device=device, dtype=dtype)
|
||||
|
||||
def forward(self, x, positions, attention):
|
||||
shape = (x.shape[0], x.shape[1], self.heads, -1)
|
||||
q = self.query_proj(x).reshape(shape).transpose(1, 2)
|
||||
k = self.key_proj(x).reshape(shape).transpose(1, 2)
|
||||
v = self.value_proj(x).reshape(shape).transpose(1, 2)
|
||||
q, k = comfy.quant_ops.ck.apply_rope_split_half(q, k, positions)
|
||||
return self.out_proj(attention(q, k, v, self.heads, skip_reshape=True))
|
||||
|
||||
|
||||
class FeedForward(nn.Module):
|
||||
def __init__(self, dim, intermediate, device=None, dtype=None, operations=None):
|
||||
super().__init__()
|
||||
self.w_1 = operations.Linear(dim, intermediate, device=device, dtype=dtype)
|
||||
self.w_2 = operations.Linear(intermediate, dim, device=device, dtype=dtype)
|
||||
|
||||
def forward(self, x):
|
||||
return self.w_2(F.gelu(self.w_1(x)))
|
||||
|
||||
|
||||
class ConvolutionModule(nn.Module):
|
||||
def __init__(self, dim, device=None, dtype=None, operations=None):
|
||||
super().__init__()
|
||||
self.layer_norm = operations.LayerNorm(dim, eps=1e-5, device=device, dtype=dtype)
|
||||
self.conv_block = nn.Sequential(
|
||||
Transpose(), operations.Conv1d(dim, dim * 2, 1, bias=False, device=device, dtype=dtype), nn.GLU(dim=1),
|
||||
operations.Conv1d(dim, dim, 31, padding=15, groups=dim, bias=False, device=device, dtype=dtype),
|
||||
nn.Sequential(Transpose(), operations.LayerNorm(dim, eps=1e-5, device=device, dtype=dtype), Transpose()),
|
||||
nn.GELU(), operations.Conv1d(dim, dim, 1, bias=False, device=device, dtype=dtype), Transpose(),
|
||||
)
|
||||
|
||||
def forward(self, x):
|
||||
return self.conv_block(self.layer_norm(x))
|
||||
|
||||
|
||||
class ConformerBlock(nn.Module):
|
||||
def __init__(self, dim, intermediate, heads, device=None, dtype=None, operations=None):
|
||||
super().__init__()
|
||||
self.ffn1_layer_norm = operations.LayerNorm(dim, eps=1e-5, device=device, dtype=dtype)
|
||||
self.ffn1 = FeedForward(dim, intermediate, device=device, dtype=dtype, operations=operations)
|
||||
self.attn_layer_norm = operations.LayerNorm(dim, eps=1e-5, device=device, dtype=dtype)
|
||||
self.attn = Attention(dim, heads, device=device, dtype=dtype, operations=operations)
|
||||
self.conv_module = ConvolutionModule(dim, device=device, dtype=dtype, operations=operations)
|
||||
self.ffn2_layer_norm = operations.LayerNorm(dim, eps=1e-5, device=device, dtype=dtype)
|
||||
self.ffn2 = FeedForward(dim, intermediate, device=device, dtype=dtype, operations=operations)
|
||||
self.final_layer_norm = operations.LayerNorm(dim, eps=1e-5, device=device, dtype=dtype)
|
||||
|
||||
def forward(self, x, positions, attention):
|
||||
x = x + 0.5 * self.ffn1(self.ffn1_layer_norm(x))
|
||||
x = x + self.attn(self.attn_layer_norm(x), positions, attention)
|
||||
x = x + self.conv_module(x)
|
||||
x = x + 0.5 * self.ffn2(self.ffn2_layer_norm(x))
|
||||
return self.final_layer_norm(x)
|
||||
|
||||
|
||||
class MERT2(nn.Module):
|
||||
def __init__(self, dim=1024, intermediate=4096, heads=16, layers=24, channels=(128, 512, 1024),
|
||||
depths=(3, 4, 5), device=None, dtype=None, operations=None):
|
||||
super().__init__()
|
||||
self.head_dim = dim // heads
|
||||
self.feature_extractor = MelFrontend(device=device)
|
||||
channels = (128, *channels)
|
||||
self.subsampling_module = nn.Sequential(*[
|
||||
ConvNextBlock(channels[i], channels[i + 1], (1, 2, 2)[i], depths[i], device=device, dtype=dtype, operations=operations)
|
||||
for i in range(3)
|
||||
])
|
||||
self.layers = nn.ModuleList([
|
||||
ConformerBlock(dim, intermediate, heads, device=device, dtype=dtype, operations=operations) for _ in range(layers)
|
||||
])
|
||||
|
||||
def position_embeddings(self, x):
|
||||
inverse = 1.0 / (10000 ** (torch.arange(0, self.head_dim, 2, device=x.device, dtype=torch.float32) / self.head_dim))
|
||||
angles = torch.arange(x.shape[1], device=x.device, dtype=torch.float32)[:, None] * inverse
|
||||
cos, sin = angles.cos().to(x.dtype), angles.sin().to(x.dtype)
|
||||
return torch.stack((cos, -sin, sin, cos), dim=-1).reshape(1, 1, x.shape[1], -1, 2, 2).float()
|
||||
|
||||
def forward(self, mel, layer_weight, output_hidden_states=False):
|
||||
x = self.subsampling_module(mel)
|
||||
weights = comfy.ops.cast_to_input(layer_weight, x).softmax(dim=0)
|
||||
mixed = x * weights[0]
|
||||
states = [x] if output_hidden_states else None
|
||||
positions = self.position_embeddings(x)
|
||||
attention = optimized_attention_for_device(x.device)
|
||||
for weight, layer in zip(weights[1:], self.layers):
|
||||
x = layer(x, positions, attention)
|
||||
mixed = mixed + x * weight
|
||||
if output_hidden_states:
|
||||
states.append(x)
|
||||
return mixed, states
|
||||
@@ -0,0 +1,537 @@
|
||||
"""SheetSage2 audio-to-score generation with the released event vocabulary."""
|
||||
|
||||
import copy
|
||||
import logging
|
||||
|
||||
import numpy as np
|
||||
import torch
|
||||
from torch import nn
|
||||
from torch.nn import functional as F
|
||||
|
||||
import comfy.model_management
|
||||
import comfy.model_prefetch
|
||||
import comfy.ops
|
||||
import comfy.utils
|
||||
from comfy.audio_encoders.mert2 import MERT2
|
||||
from comfy.ldm.modules.attention import optimized_attention_for_device
|
||||
from comfy.text_encoders.llama import FixedKV
|
||||
|
||||
|
||||
class DecoderAttention(nn.Module):
|
||||
def __init__(self, dim, heads, device=None, dtype=None, operations=None):
|
||||
super().__init__()
|
||||
self.heads = heads
|
||||
self.q_proj = operations.Linear(dim, dim, device=device, dtype=dtype)
|
||||
self.k_proj = operations.Linear(dim, dim, device=device, dtype=dtype)
|
||||
self.v_proj = operations.Linear(dim, dim, device=device, dtype=dtype)
|
||||
self.out_proj = operations.Linear(dim, dim, device=device, dtype=dtype)
|
||||
|
||||
def project(self, x, projection):
|
||||
return projection(x).reshape(x.shape[0], x.shape[1], self.heads, -1).transpose(1, 2)
|
||||
|
||||
def forward(self, x, attention, mask=None, cache=None, memory=None):
|
||||
q = self.project(x, self.q_proj)
|
||||
if memory is not None:
|
||||
k, v = memory
|
||||
else:
|
||||
k, v = self.project(x, self.k_proj), self.project(x, self.v_proj)
|
||||
length = x.shape[1]
|
||||
if isinstance(cache, FixedKV):
|
||||
key, value = k.transpose(1, 2), v.transpose(1, 2)
|
||||
if length == 1 and cache.index > 0:
|
||||
position = cache.position.view(-1, 1, 1, 1).expand_as(key)
|
||||
cache.key.scatter_(1, position, key)
|
||||
cache.value.scatter_(1, position, value)
|
||||
valid = torch.arange(cache.key.shape[1], device=x.device)[None] < cache.seqlen[:, None]
|
||||
mask = torch.zeros(valid.shape, device=x.device, dtype=x.dtype).masked_fill_(~valid, -torch.inf)[:, None, None]
|
||||
out = attention(q, cache.key.transpose(1, 2), cache.value.transpose(1, 2), self.heads, mask=mask, skip_reshape=True)
|
||||
return self.out_proj(out), cache
|
||||
cache.key[:, :length].copy_(key)
|
||||
cache.value[:, :length].copy_(value)
|
||||
elif cache is not None:
|
||||
key, value, index = cache
|
||||
key[:, :, index:index + length].copy_(k)
|
||||
value[:, :, index:index + length].copy_(v)
|
||||
k, v = key[:, :, :index + length], value[:, :, :index + length]
|
||||
cache = key, value, index + length
|
||||
out = attention(q, k, v, self.heads, mask=mask, skip_reshape=True)
|
||||
return self.out_proj(out), cache
|
||||
|
||||
|
||||
class DecoderLayer(nn.Module):
|
||||
def __init__(self, dim, intermediate, heads, device=None, dtype=None, operations=None):
|
||||
super().__init__()
|
||||
self.self_attn = DecoderAttention(dim, heads, device=device, dtype=dtype, operations=operations)
|
||||
self.self_attn_layer_norm = operations.LayerNorm(dim, eps=1e-5, device=device, dtype=dtype)
|
||||
self.encoder_attn = DecoderAttention(dim, heads, device=device, dtype=dtype, operations=operations)
|
||||
self.encoder_attn_layer_norm = operations.LayerNorm(dim, eps=1e-5, device=device, dtype=dtype)
|
||||
self.fc1 = operations.Linear(dim, intermediate, device=device, dtype=dtype)
|
||||
self.fc2 = operations.Linear(intermediate, dim, device=device, dtype=dtype)
|
||||
self.final_layer_norm = operations.LayerNorm(dim, eps=1e-5, device=device, dtype=dtype)
|
||||
|
||||
def forward(self, x, attention, mask, cache):
|
||||
self_cache, memory = cache
|
||||
out, self_cache = self.self_attn(x, attention, mask=mask, cache=self_cache)
|
||||
x = self.self_attn_layer_norm(x + out)
|
||||
out, _ = self.encoder_attn(x, attention, memory=memory)
|
||||
x = self.encoder_attn_layer_norm(x + out)
|
||||
return self.final_layer_norm(x + self.fc2(F.gelu(self.fc1(x)))), (self_cache, memory)
|
||||
|
||||
|
||||
class Decoder(nn.Module):
|
||||
def __init__(self, dim, intermediate, heads, layers, max_tokens, device=None, dtype=None, operations=None):
|
||||
super().__init__()
|
||||
self.embed_positions = operations.Embedding(max_tokens + 2, dim, device=device, dtype=dtype)
|
||||
self.layernorm_embedding = operations.LayerNorm(dim, eps=1e-5, device=device, dtype=dtype)
|
||||
self.layers = nn.ModuleList([
|
||||
DecoderLayer(dim, intermediate, heads, device=device, dtype=dtype, operations=operations) for _ in range(layers)
|
||||
])
|
||||
|
||||
def forward(self, x, positions, cache):
|
||||
length = x.shape[1]
|
||||
fixed = isinstance(cache[0][0], FixedKV)
|
||||
index = cache[0][0].index if fixed else cache[0][0][2]
|
||||
mask = None
|
||||
if length > 1:
|
||||
mask = torch.full((length, index + length), -torch.inf, device=x.device, dtype=x.dtype).triu_(index + 1)
|
||||
x = self.layernorm_embedding(x + self.embed_positions(positions, out_dtype=x.dtype))
|
||||
graph = fixed and length == 1 and index > 0
|
||||
if graph:
|
||||
x = x.clone()
|
||||
attention = optimized_attention_for_device(x.device, mask=mask is not None or graph, small_input=True)
|
||||
queue = comfy.model_prefetch.make_prefetch_queue(list(self.layers), x.device, {"prefetch_dynamic_vbars": True})
|
||||
for i, layer in enumerate(self.layers):
|
||||
if fixed:
|
||||
cache[i][0].prepare(length)
|
||||
|
||||
def core():
|
||||
nonlocal x
|
||||
out, cache[i] = layer(x, attention, mask, cache[i])
|
||||
if graph:
|
||||
x.copy_(out)
|
||||
else:
|
||||
x = out
|
||||
|
||||
comfy.model_prefetch.prefetch_queue_pop(queue, x.device, layer, x.dtype, core=core,
|
||||
enable_graph=graph, malloc_scope="block")
|
||||
if fixed:
|
||||
cache[i][0].advance(length)
|
||||
comfy.model_prefetch.prefetch_queue_pop(queue, x.device, None, malloc_scope="block")
|
||||
return x
|
||||
|
||||
|
||||
class SheetSage2(nn.Module):
|
||||
def __init__(self, dim=512, intermediate=2048, heads=8, layers=6, max_tokens=5120,
|
||||
mert_config=None, device=None, dtype=None, operations=None):
|
||||
super().__init__()
|
||||
mert_config = {} if mert_config is None else mert_config
|
||||
self.dtype = dtype
|
||||
self.max_tokens = max_tokens
|
||||
self.encoder = MERT2(**mert_config, device=device, dtype=dtype, operations=operations)
|
||||
self.layer_weight = nn.Parameter(torch.empty(len(self.encoder.layers) + 1, device=device, dtype=dtype))
|
||||
self.encoder_projection = operations.Linear(mert_config.get("dim", 1024), dim, device=device, dtype=dtype)
|
||||
self.tokenizer = ScoreTokenizer()
|
||||
self.token_embedding = operations.Embedding(self.tokenizer.n_tokens, dim, device=device, dtype=dtype)
|
||||
self.decoder = Decoder(dim, intermediate, heads, layers, max_tokens, device=device, dtype=dtype, operations=operations)
|
||||
self.output_projection = operations.Linear(dim, self.tokenizer.n_tokens, bias=False, device=device, dtype=dtype)
|
||||
|
||||
def get_dynamic_vram__units(self):
|
||||
return list(self.decoder.layers), []
|
||||
|
||||
def encode(self, waveform, output_hidden_states=False):
|
||||
# The released encoder attends to the entire 300-second window, including its padding.
|
||||
waveform = F.pad(waveform, (0, max(0, 300 * 24000 - waveform.shape[-1])))
|
||||
mel = self.encoder.feature_extractor(waveform.float()).to(self.dtype)
|
||||
mixed, states = self.encoder(mel, self.layer_weight, output_hidden_states=output_hidden_states)
|
||||
return self.encoder_projection(mixed), states
|
||||
|
||||
def forward(self, audio):
|
||||
return self.encode(audio.mean(dim=1), output_hidden_states=True)
|
||||
|
||||
def init_cache(self, memory):
|
||||
batch, _, dim = memory.shape
|
||||
fixed = comfy.model_prefetch.malloc_graph_enabled(memory.device)
|
||||
cache = []
|
||||
for layer in self.decoder.layers:
|
||||
heads = layer.self_attn.heads
|
||||
shape = (batch, self.max_tokens, heads, dim // heads) if fixed else (batch, heads, self.max_tokens, dim // heads)
|
||||
# Fixed attention includes masked future slots, whose values must remain finite.
|
||||
key = torch.zeros(shape, device=memory.device, dtype=memory.dtype) if fixed else torch.empty(shape, device=memory.device, dtype=memory.dtype)
|
||||
value = torch.zeros_like(key) if fixed else torch.empty_like(key)
|
||||
if fixed:
|
||||
self_cache = FixedKV(key, value, 0, torch.empty(batch, device=memory.device, dtype=torch.long),
|
||||
torch.zeros(batch, device=memory.device, dtype=torch.int32))
|
||||
else:
|
||||
self_cache = key, value, 0
|
||||
cross = layer.encoder_attn
|
||||
cache.append((self_cache, (cross.project(memory, cross.k_proj), cross.project(memory, cross.v_proj))))
|
||||
return cache
|
||||
|
||||
def decode(self, ids, positions, cache):
|
||||
x = self.token_embedding(ids, out_dtype=self.dtype)
|
||||
return self.output_projection(self.decoder(x, positions, cache)[:, -1:])
|
||||
|
||||
def generate_tokens(self, memory, stop_seconds, prefix=None):
|
||||
tokenizer = self.tokenizer
|
||||
tokens = tokenizer.prompt_prefix() if prefix is None else list(prefix)
|
||||
state = PromptGrammarState(tokenizer)
|
||||
for token in tokens[tokens.index(tokenizer.out_token) + 1:]:
|
||||
state.update(token)
|
||||
cache = self.init_cache(memory)
|
||||
device = memory.device
|
||||
ids = torch.tensor([tokens], device=device, dtype=torch.long)
|
||||
positions = torch.arange(2, len(tokens) + 2, device=device)[None]
|
||||
logits = self.decode(ids, positions, cache)
|
||||
ids = torch.empty((1, 1), device=device, dtype=torch.long)
|
||||
positions = torch.full((1, 1), len(tokens) + 2, device=device, dtype=torch.long)
|
||||
fixed = isinstance(cache[0][0], FixedKV)
|
||||
progress = comfy.utils.ProgressBar(self.max_tokens - len(tokens))
|
||||
try:
|
||||
for step in comfy.utils.model_trange(self.max_tokens - len(tokens), desc="SheetSage2 transcription", unit="token"):
|
||||
comfy.model_management.throw_exception_if_processing_interrupted()
|
||||
scores = logits[0, -1].float().masked_fill(~state.allowed(device), -torch.inf)
|
||||
next_id = scores.argmax()
|
||||
token = next_id.item()
|
||||
tokens.append(token)
|
||||
progress.update_absolute(step + 1)
|
||||
if state.update(token):
|
||||
break
|
||||
if tokenizer.time_token_start <= token < tokenizer.time_token_end and tokenizer.token_to_time_id(token) / tokenizer.time_hz >= stop_seconds:
|
||||
tokens.append(tokenizer.eos_token)
|
||||
break
|
||||
if len(tokens) == self.max_tokens:
|
||||
logging.warning("SheetSage2 reached its token limit; the transcription may be incomplete.")
|
||||
tokens.append(tokenizer.eos_token)
|
||||
break
|
||||
ids.copy_(next_id)
|
||||
if fixed:
|
||||
comfy.model_prefetch.malloc_graph_begin(device)
|
||||
logits.copy_(self.decode(ids, positions, cache))
|
||||
if fixed:
|
||||
comfy.model_prefetch.malloc_graph_end()
|
||||
positions.add_(1)
|
||||
finally:
|
||||
comfy.model_prefetch.cleanup_prefetch_queues()
|
||||
return tokens
|
||||
|
||||
def transcribe(self, waveform):
|
||||
duration = waveform.shape[-1] / 24000
|
||||
stitched = []
|
||||
for window in sliding_window_plan(duration):
|
||||
comfy.model_management.throw_exception_if_processing_interrupted()
|
||||
start = window["start"]
|
||||
prefix, base = overlap_prefix(stitched, self.tokenizer, start, window["prefix_end"])
|
||||
if prefix is not None and len(prefix) >= self.max_tokens - 128:
|
||||
raise ValueError("SheetSage2 overlap fills the token context; transcribe shorter audio sections.")
|
||||
segment = waveform[:, round(start * 24000):round(window["end"] * 24000)]
|
||||
memory, _ = self.encode(segment)
|
||||
stop = window["generation_stop"] if window["generation_stop"] is not None else min(duration - start, 300.0)
|
||||
tokens = self.generate_tokens(memory, stop, prefix)
|
||||
decoded = self.tokenizer.decode_sequence(tokens)
|
||||
lookup = event_time_map(decoded, 300.0)
|
||||
stitched.extend(stitched_window_events(decoded, lookup, start, window["accept_start"],
|
||||
window["accept_end"], duration, global_subbeat_base=base))
|
||||
stitched.sort(key=lambda event: (event["time"], event["global_subbeat"]))
|
||||
return stitched
|
||||
|
||||
|
||||
CHROMATIC_SHARPS = ("C", "C#", "D", "D#", "E", "F", "F#", "G", "G#", "A", "A#", "B")
|
||||
STRUCTURE_LABELS = (
|
||||
"silence", "intro", "outro", "verse", "chorus", "bridge", "pre-chorus", "post-chorus", "interlude",
|
||||
"fade-out", "loop", "rap", "preshot", "irregular", "instrumental", "intro and verse", "pre-chorus and chorus",
|
||||
"verse and pre-chorus", "solo", "theme", "development", "variation", "pre-outro",
|
||||
)
|
||||
DURATION_TEMPLATES = (1, 2, 3, 4, 6, 8, 12, 16, 24, 32, 48, 64, 96, 128, 192, 256, 384, 512, 768, 1024, 1536, 2048, 3072, 4096)
|
||||
EVENT_FIELDS = ("timestamp", "rhythm", "structure", "key", "chord", "melody")
|
||||
FIELD_TO_INDEX = {name: index for index, name in enumerate(EVENT_FIELDS)}
|
||||
|
||||
|
||||
class ScoreTokenizer:
|
||||
pad_token, sos_token, eos_token, out_token = 0, 1, 2, 3
|
||||
time_hz = 100
|
||||
|
||||
def __init__(self):
|
||||
self.meter_pairs = tuple((numerator, denominator) for numerator in range(1, 33) for denominator in (1, 2, 4, 8, 16, 32))
|
||||
self.full_chord_labels = ["N"]
|
||||
inversions = {"maj": ("/2", "/3", "/5"), "min": ("/2", "/b3", "/5"),
|
||||
"maj7": ("/3", "/5", "/7"), "min7": ("/b3", "/5", "/b7"), "7": ("/3", "/5", "/b7")}
|
||||
for quality in ("maj", "min", "dim", "aug", "maj7", "min7", "7", "hdim7", "dim7", "minmaj7", "sus2", "sus4", "sus4(b7)", "maj6", "min6"):
|
||||
for root in CHROMATIC_SHARPS:
|
||||
self.full_chord_labels.extend(f"{root}:{quality}{inversion}" for inversion in (*inversions.get(quality, ()), ""))
|
||||
offset = 260
|
||||
self.ranges = []
|
||||
for name, count in (("subbeat_shift", 257), ("time", 30000), ("meter", 192), ("eighth_position", 256),
|
||||
("structure", len(STRUCTURE_LABELS)), ("key", 24), ("majmin_chord", 25),
|
||||
("full_chord", len(self.full_chord_labels)), ("pitch", 256), ("duration", len(DURATION_TEMPLATES))):
|
||||
setattr(self, f"{name}_token_start", offset)
|
||||
setattr(self, f"{name}_token_end", offset + count)
|
||||
self.ranges.append((name.replace("full_chord", "chord_full").replace("majmin_chord", "chord_majmin"), offset, offset + count))
|
||||
offset += count
|
||||
self.n_tokens = offset
|
||||
|
||||
def prompt_prefix(self):
|
||||
return [self.sos_token, 4, 5, 6, 7, 9, 11, self.out_token]
|
||||
|
||||
def token_type(self, token):
|
||||
for name, start, end in self.ranges:
|
||||
if start <= token < end:
|
||||
return name
|
||||
return {0: "pad", 1: "sos", 2: "eos", 3: "out"}.get(token, "prompt")
|
||||
|
||||
def token_to_time_id(self, token):
|
||||
return token - self.time_token_start
|
||||
|
||||
def decode_field(self, field, tokens):
|
||||
if field == "timestamp":
|
||||
return self.token_to_time_id(tokens[0]) / self.time_hz
|
||||
if field == "rhythm":
|
||||
rhythm = {}
|
||||
for token in tokens:
|
||||
if self.token_type(token) == "meter":
|
||||
rhythm["meter"] = self.meter_pairs[token - self.meter_token_start]
|
||||
else:
|
||||
rhythm["eighth_position"] = token - self.eighth_position_token_start
|
||||
return rhythm
|
||||
if field == "structure":
|
||||
return STRUCTURE_LABELS[tokens[0] - self.structure_token_start]
|
||||
if field == "key":
|
||||
index = tokens[0] - self.key_token_start
|
||||
return f"{CHROMATIC_SHARPS[index % 12]}:{'minor' if index >= 12 else 'major'}"
|
||||
if field == "chord":
|
||||
return self.full_chord_labels[tokens[0] - self.full_chord_token_start]
|
||||
notes, index = [], 0
|
||||
while index < len(tokens):
|
||||
pitch = tokens[index] - self.pitch_token_start
|
||||
index += 1
|
||||
duration = 0
|
||||
if index < len(tokens) and self.token_type(tokens[index]) == "duration":
|
||||
duration = tokens[index] - self.duration_token_start
|
||||
index += 1
|
||||
notes.append({"pitch": pitch % 128, "track": int(pitch >= 128),
|
||||
"duration_bin": duration, "duration_steps": DURATION_TEMPLATES[duration]})
|
||||
return notes
|
||||
|
||||
def decode_sequence(self, tokens):
|
||||
fields = {"time": "timestamp", "meter": "rhythm", "eighth_position": "rhythm", "structure": "structure",
|
||||
"key": "key", "chord_full": "chord", "pitch": "melody", "duration": "melody"}
|
||||
position, subbeat, events = tokens.index(self.out_token) + 1, 0, []
|
||||
while position < len(tokens) and tokens[position] != self.eos_token:
|
||||
if self.token_type(tokens[position]) != "subbeat_shift":
|
||||
raise ValueError("SheetSage2 produced an event without a beat position.")
|
||||
while position < len(tokens) and self.token_type(tokens[position]) == "subbeat_shift":
|
||||
subbeat += tokens[position] - self.subbeat_shift_token_start
|
||||
position += 1
|
||||
payload = {}
|
||||
while position < len(tokens) and self.token_type(tokens[position]) not in ("subbeat_shift", "eos"):
|
||||
token = tokens[position]
|
||||
payload.setdefault(fields[self.token_type(token)], []).append(token)
|
||||
position += 1
|
||||
if payload:
|
||||
events.append({"subbeat": subbeat, "tokens_by_field": payload,
|
||||
"values": {field: self.decode_field(field, values) for field, values in payload.items()}})
|
||||
return {"events": events}
|
||||
|
||||
def encode_events(self, events):
|
||||
tokens, previous = self.prompt_prefix(), 0
|
||||
for event in events:
|
||||
shift = event["subbeat"] - previous
|
||||
while shift > 256:
|
||||
tokens.append(self.subbeat_shift_token_end - 1)
|
||||
shift -= 256
|
||||
tokens.append(self.subbeat_shift_token_start + shift)
|
||||
previous = event["subbeat"]
|
||||
for field in EVENT_FIELDS:
|
||||
tokens.extend(event["tokens_by_field"].get(field, ()))
|
||||
return tokens
|
||||
|
||||
|
||||
def sliding_window_plan(duration, window_seconds=300.0, overlap_seconds=200.0, lookahead_seconds=100.0):
|
||||
hop = window_seconds - overlap_seconds
|
||||
start, accepted = 0.0, 0.0
|
||||
result = []
|
||||
while True:
|
||||
last = start + window_seconds >= duration - 1e-6
|
||||
accept_end = duration if last else start + window_seconds - lookahead_seconds
|
||||
result.append(dict(start=start, end=min(duration, start + window_seconds),
|
||||
accept_start=accepted, accept_end=accept_end, prefix_end=accepted,
|
||||
generation_stop=None if last else window_seconds - lookahead_seconds))
|
||||
if last:
|
||||
return result
|
||||
accepted = accept_end
|
||||
start = min(start + hop, duration - window_seconds)
|
||||
|
||||
|
||||
def overlap_prefix(stitched, tokenizer, start, prefix_end):
|
||||
events = [event for event in stitched if start - 1e-4 <= event["time"] < prefix_end - 1e-4]
|
||||
events.sort(key=lambda event: (event["global_subbeat"], event["time"]))
|
||||
first = next((i for i, event in enumerate(events) if "timestamp" in event["values"] or "rhythm" in event["values"]), None)
|
||||
if first is None:
|
||||
return None, 0
|
||||
events = copy.deepcopy(events[first:])
|
||||
base = events[0]["global_subbeat"]
|
||||
context = {}
|
||||
for event in stitched:
|
||||
if event["time"] > events[0]["time"] + 1e-6:
|
||||
continue
|
||||
for field in ("structure", "key", "chord"):
|
||||
if event["tokens_by_field"].get(field):
|
||||
context[field] = event["tokens_by_field"][field]
|
||||
for token in event["tokens_by_field"].get("rhythm", ()):
|
||||
if tokenizer.token_type(token) == "meter":
|
||||
context["meter"] = token
|
||||
for event in events:
|
||||
event["subbeat"] = max(0, event["global_subbeat"] - base)
|
||||
if "timestamp" in event["tokens_by_field"]:
|
||||
time_id = min(29999, max(0, round((event["time"] - start) * tokenizer.time_hz)))
|
||||
event["tokens_by_field"]["timestamp"] = [tokenizer.time_token_start + time_id]
|
||||
first_fields = events[0]["tokens_by_field"]
|
||||
for field in ("structure", "key", "chord"):
|
||||
if field not in first_fields and field in context:
|
||||
first_fields[field] = list(context[field])
|
||||
rhythm = first_fields.get("rhythm", [])
|
||||
if any(tokenizer.token_type(token) == "eighth_position" for token in rhythm) and not any(tokenizer.token_type(token) == "meter" for token in rhythm) and "meter" in context:
|
||||
first_fields["rhythm"] = [context["meter"], *rhythm]
|
||||
return tokenizer.encode_events(events), base
|
||||
|
||||
|
||||
def event_time_map(decoded, target_seconds):
|
||||
anchors = sorted({event["subbeat"]: event["values"]["timestamp"] for event in decoded["events"] if "timestamp" in event["values"]}.items())
|
||||
if not anchors:
|
||||
return lambda step: min(target_seconds, max(0.0, step * 0.125))
|
||||
steps, times = np.asarray(anchors, dtype=np.float64).T
|
||||
period = float(np.median(np.diff(times) / np.maximum(np.diff(steps), 1))) if len(anchors) > 1 else 0.125
|
||||
if not np.isfinite(period) or period <= 0:
|
||||
period = 0.125
|
||||
|
||||
def lookup(step):
|
||||
if step <= steps[0]:
|
||||
return float(np.clip(times[0] + (step - steps[0]) * period, 0, target_seconds))
|
||||
if step >= steps[-1]:
|
||||
return float(np.clip(times[-1] + (step - steps[-1]) * period, 0, target_seconds))
|
||||
return float(np.interp(step, steps, times))
|
||||
|
||||
return lookup
|
||||
|
||||
|
||||
def stitched_window_events(decoded, lookup, start, accept_start, accept_end, duration, global_subbeat_base=0):
|
||||
accepted = []
|
||||
for source in decoded["events"]:
|
||||
time = start + lookup(source["subbeat"])
|
||||
if time < accept_start - 1e-4 or time >= accept_end - 1e-4 or time >= duration - 1e-4:
|
||||
continue
|
||||
event = copy.deepcopy(source)
|
||||
event["time"] = float(np.clip(time, 0, duration))
|
||||
event["global_subbeat"] = global_subbeat_base + event["subbeat"]
|
||||
if "timestamp" in event["values"]:
|
||||
event["values"]["timestamp"] = event["time"]
|
||||
for note in event["values"].get("melody", ()):
|
||||
note["end_time"] = min(duration, max(event["time"] + 0.04, start + lookup(event["subbeat"] + note["duration_steps"])))
|
||||
accepted.append(event)
|
||||
return accepted
|
||||
|
||||
|
||||
class PromptGrammarState:
|
||||
def __init__(self, tokenizer):
|
||||
self.tokenizer = tokenizer
|
||||
self.generated_events = 0
|
||||
self.in_shift = True
|
||||
self.shift_run = 0
|
||||
self.payload_count = 0
|
||||
self.last_field_index = -1
|
||||
self.incomplete = None
|
||||
|
||||
def _allow_field_starts(self, allowed):
|
||||
tokenizer = self.tokenizer
|
||||
if self.last_field_index < FIELD_TO_INDEX["timestamp"]:
|
||||
allowed[tokenizer.time_token_start : tokenizer.time_token_end] = True
|
||||
if self.last_field_index < FIELD_TO_INDEX["rhythm"]:
|
||||
allowed[tokenizer.meter_token_start : tokenizer.meter_token_end] = True
|
||||
allowed[
|
||||
tokenizer.eighth_position_token_start : tokenizer.eighth_position_token_end
|
||||
] = True
|
||||
if self.last_field_index < FIELD_TO_INDEX["structure"]:
|
||||
allowed[
|
||||
tokenizer.structure_token_start : tokenizer.structure_token_end
|
||||
] = True
|
||||
if self.last_field_index < FIELD_TO_INDEX["key"]:
|
||||
allowed[tokenizer.key_token_start : tokenizer.key_token_end] = True
|
||||
if self.last_field_index < FIELD_TO_INDEX["chord"]:
|
||||
allowed[
|
||||
tokenizer.full_chord_token_start : tokenizer.full_chord_token_end
|
||||
] = True
|
||||
if self.last_field_index <= FIELD_TO_INDEX["melody"]:
|
||||
allowed[tokenizer.pitch_token_start : tokenizer.pitch_token_end] = True
|
||||
|
||||
def allowed(self, device):
|
||||
tokenizer = self.tokenizer
|
||||
allowed = torch.zeros(tokenizer.n_tokens, dtype=torch.bool, device=device)
|
||||
can_end = self.payload_count > 0
|
||||
|
||||
if can_end:
|
||||
allowed[tokenizer.eos_token] = True
|
||||
if self.payload_count > 0 or self.in_shift:
|
||||
if self.shift_run < 4:
|
||||
allowed[
|
||||
tokenizer.subbeat_shift_token_start : tokenizer.subbeat_shift_token_end
|
||||
] = True
|
||||
|
||||
if self.incomplete == "rhythm_after_meter":
|
||||
allowed[
|
||||
tokenizer.eighth_position_token_start : tokenizer.eighth_position_token_end
|
||||
] = True
|
||||
return allowed
|
||||
|
||||
if self.incomplete == "melody_after_pitch":
|
||||
allowed[tokenizer.duration_token_start : tokenizer.duration_token_end] = True
|
||||
allowed[tokenizer.pitch_token_start : tokenizer.pitch_token_end] = True
|
||||
return allowed
|
||||
|
||||
self._allow_field_starts(allowed)
|
||||
return allowed
|
||||
|
||||
def update(self, token):
|
||||
tokenizer = self.tokenizer
|
||||
token = int(token)
|
||||
token_type = tokenizer.token_type(token)
|
||||
if token == tokenizer.eos_token:
|
||||
return True
|
||||
if token_type == "subbeat_shift":
|
||||
if not self.in_shift and self.payload_count > 0:
|
||||
self.generated_events += 1
|
||||
self.payload_count = 0
|
||||
self.last_field_index = -1
|
||||
self.incomplete = None
|
||||
self.in_shift = True
|
||||
self.shift_run += 1
|
||||
return False
|
||||
|
||||
self.in_shift = False
|
||||
self.shift_run = 0
|
||||
self.payload_count += 1
|
||||
if token_type == "time":
|
||||
self.last_field_index = FIELD_TO_INDEX["timestamp"]
|
||||
self.incomplete = None
|
||||
elif token_type == "meter":
|
||||
self.last_field_index = FIELD_TO_INDEX["rhythm"]
|
||||
self.incomplete = "rhythm_after_meter"
|
||||
elif token_type == "eighth_position":
|
||||
self.last_field_index = FIELD_TO_INDEX["rhythm"]
|
||||
self.incomplete = None
|
||||
elif token_type == "structure":
|
||||
self.last_field_index = FIELD_TO_INDEX["structure"]
|
||||
self.incomplete = None
|
||||
elif token_type == "key":
|
||||
self.last_field_index = FIELD_TO_INDEX["key"]
|
||||
self.incomplete = None
|
||||
elif token_type == "chord_full":
|
||||
self.last_field_index = FIELD_TO_INDEX["chord"]
|
||||
self.incomplete = None
|
||||
elif token_type == "pitch":
|
||||
self.last_field_index = FIELD_TO_INDEX["melody"]
|
||||
self.incomplete = "melody_after_pitch"
|
||||
elif token_type == "duration":
|
||||
self.last_field_index = FIELD_TO_INDEX["melody"]
|
||||
self.incomplete = None
|
||||
else:
|
||||
raise RuntimeError(f"Unexpected prompt token type {token_type!r}")
|
||||
return False
|
||||
@@ -0,0 +1,921 @@
|
||||
"""SheetSage2 timed events to two-voice ABC, adapted from m-a-p/SheetSage2."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import math
|
||||
import re
|
||||
from collections import Counter
|
||||
from dataclasses import dataclass, replace
|
||||
from fractions import Fraction
|
||||
from typing import Sequence
|
||||
|
||||
import numpy as np
|
||||
|
||||
|
||||
_SUPPORTED_DURATION_UNITS = frozenset({1, 2, 3, 4, 6, 8, 12, 16, 24, 32, 48})
|
||||
_MUSIC_ELEMENT_RE = re.compile(
|
||||
r'"(?P<quoted>[^"]*)"'
|
||||
r"|\[K:(?P<key>[^\]]+)\]"
|
||||
r"|(?P<note>[_=^]*[A-Ga-gz][,']*)(?P<duration>\d*)(?P<tie>-?)"
|
||||
)
|
||||
|
||||
|
||||
def interval_rows(events, field, duration):
|
||||
rows = [[event["time"], duration, event["values"][field]] for event in events if field in event["values"]]
|
||||
for previous, current in zip(rows, rows[1:]):
|
||||
previous[1] = current[0]
|
||||
return [row for row in rows if row[1] > row[0]]
|
||||
|
||||
|
||||
def events_to_abc(events, duration, melody_only=True):
|
||||
beats, meter = [], None
|
||||
notes = {"Vocal": [], "Ins": []}
|
||||
for event in events:
|
||||
rhythm = event["values"].get("rhythm", {})
|
||||
meter = rhythm.get("meter", meter)
|
||||
eighth = rhythm.get("eighth_position")
|
||||
if eighth is not None and meter is not None:
|
||||
position = Fraction(eighth * meter[1], 8)
|
||||
if position.denominator != 1 or not 0 <= position < meter[0]:
|
||||
raise BeatGridError(f"Beat position {eighth} is outside the decoded {meter[0]}/{meter[1]} grid.")
|
||||
beats.append(BeatEvent(event["time"], int(position) + 1, meter[0], meter[1]))
|
||||
for note in event["values"].get("melody", ()):
|
||||
end = min(duration, note["end_time"])
|
||||
if end > event["time"]:
|
||||
notes[VOICE_IDS[note["track"]]].append([event["time"], end, note["pitch"]])
|
||||
if len(beats) < 2:
|
||||
raise BeatGridError("SheetSage2 needs at least two decoded beats to produce ABC.")
|
||||
if any(current.time <= previous.time for previous, current in zip(beats, beats[1:])):
|
||||
raise BeatGridError("SheetSage2 decoded beat times are not increasing.")
|
||||
period = float(np.median(np.diff([beat.time for beat in beats[-9:]])))
|
||||
while beats[-1].time < duration - 1e-6:
|
||||
previous = beats[-1]
|
||||
beats.append(BeatEvent(previous.time + period, previous.beat_id % previous.declared_numerator + 1,
|
||||
previous.declared_numerator, previous.denominator))
|
||||
intervals = {}
|
||||
for field in ("key", "structure", "chord"):
|
||||
intervals[field] = [[max(beats[0].time, start), min(beats[-1].time, end), value]
|
||||
for start, end, value in interval_rows(events, field, duration)
|
||||
if end > beats[0].time and start < beats[-1].time]
|
||||
if not intervals["key"]:
|
||||
raise AbcRebuildError("SheetSage2 did not decode a key for the ABC score.")
|
||||
keys = [(start, end, key_symbol_to_abc(key)) for start, end, key in intervals["key"]]
|
||||
measures, diagnostics = infer_measures(beats)
|
||||
times, quarters, denominators = _build_grid(beats, measures)
|
||||
voices = {}
|
||||
for voice, track in notes.items():
|
||||
track.sort(key=lambda note: (note[0], note[2], note[1]))
|
||||
for previous, current in zip(track, track[1:]):
|
||||
if previous[1] > current[0] + 1e-6:
|
||||
previous[1] = current[0]
|
||||
voices[voice] = _notes_to_arr([note for note in track if note[1] > note[0] + 1e-6], times, voice)
|
||||
score = RebuiltAbcScore(
|
||||
beats=beats, measures=measures, subbeat_times=times, subbeat_quarters=quarters,
|
||||
subbeat_denominators=denominators,
|
||||
key_arr=_fill_intervals(keys, times, default=keys[0][2], dtype="<U16"),
|
||||
chord_arr=np.full(len(times), "N", dtype="<U64") if melody_only else _fill_intervals(intervals["chord"], times, default="N", dtype="<U64"),
|
||||
structure_events=_structure_events(intervals["structure"], times), voice_arrs=voices, diagnostics=diagnostics,
|
||||
)
|
||||
return score_to_abc(score)
|
||||
|
||||
|
||||
SUBBEAT_DIVISION = 4
|
||||
|
||||
VOICE_IDS = ("Vocal", "Ins")
|
||||
|
||||
NO_CHORDS = frozenset({"N", "X", "?"})
|
||||
|
||||
class AbcRebuildError(ValueError):
|
||||
"""Base class for deterministic reconstruction failures."""
|
||||
|
||||
class BeatGridError(AbcRebuildError):
|
||||
pass
|
||||
|
||||
class ChordSymbolError(AbcRebuildError):
|
||||
pass
|
||||
|
||||
class MelodyVoiceError(AbcRebuildError):
|
||||
pass
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class BeatEvent:
|
||||
time: float
|
||||
beat_id: int
|
||||
declared_numerator: int
|
||||
denominator: int
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class Measure:
|
||||
index: int
|
||||
start_beat: int
|
||||
end_beat: int
|
||||
numerator: int
|
||||
denominator: int
|
||||
pickup: bool = False
|
||||
partial: bool = False
|
||||
inferred: bool = False
|
||||
notated_numerator: int | None = None
|
||||
notated_denominator: int | None = None
|
||||
pad_before: bool = False
|
||||
|
||||
@property
|
||||
def beat_count(self) -> int:
|
||||
return self.end_beat - self.start_beat
|
||||
|
||||
@property
|
||||
def start_t(self) -> int:
|
||||
return self.start_beat * SUBBEAT_DIVISION
|
||||
|
||||
@property
|
||||
def end_t(self) -> int:
|
||||
return self.end_beat * SUBBEAT_DIVISION
|
||||
|
||||
@property
|
||||
def abc_numerator(self) -> int:
|
||||
return self.notated_numerator or self.numerator
|
||||
|
||||
@property
|
||||
def abc_denominator(self) -> int:
|
||||
return self.notated_denominator or self.denominator
|
||||
|
||||
@dataclass
|
||||
class RebuiltAbcScore:
|
||||
beats: list[BeatEvent]
|
||||
measures: list[Measure]
|
||||
subbeat_times: np.ndarray
|
||||
subbeat_quarters: np.ndarray
|
||||
subbeat_denominators: np.ndarray
|
||||
key_arr: np.ndarray
|
||||
chord_arr: np.ndarray
|
||||
structure_events: list[tuple[int, str]]
|
||||
voice_arrs: dict[str, np.ndarray]
|
||||
diagnostics: list[str]
|
||||
subbeat_div: int = SUBBEAT_DIVISION
|
||||
|
||||
@dataclass
|
||||
class MeasureGroup:
|
||||
measures: list[Measure]
|
||||
structure_labels: list[str]
|
||||
meter_changed: bool
|
||||
key_changed: bool
|
||||
|
||||
_QUALITY_TO_ABC = {
|
||||
"maj": "",
|
||||
"min": "m",
|
||||
"dim": "dim",
|
||||
"aug": "aug",
|
||||
"7": "7",
|
||||
"maj7": "maj7",
|
||||
"min7": "m7",
|
||||
"dim7": "dim7",
|
||||
"hdim7": "m7b5",
|
||||
"sus4": "sus4",
|
||||
"sus2": "sus2",
|
||||
"maj6": "6",
|
||||
"min6": "m6",
|
||||
"sus4(b7)": "7sus4",
|
||||
# abc2midi and SymMusic both accept the parenthesized major seventh.
|
||||
# Common aliases such as mmaj7/mM7 trigger abc2midi diagnostics.
|
||||
"minmaj7": "m(maj7)",
|
||||
}
|
||||
|
||||
_NATURAL_PITCH_CLASS = {
|
||||
"C": 0,
|
||||
"D": 2,
|
||||
"E": 4,
|
||||
"F": 5,
|
||||
"G": 7,
|
||||
"A": 9,
|
||||
"B": 11,
|
||||
}
|
||||
|
||||
_LETTERS = "CDEFGAB"
|
||||
|
||||
_SHARP_PITCH_NAMES = ("C", "C#", "D", "D#", "E", "F", "F#", "G", "G#", "A", "A#", "B")
|
||||
|
||||
_FLAT_PITCH_NAMES = ("C", "Db", "D", "Eb", "E", "F", "Gb", "G", "Ab", "A", "Bb", "B")
|
||||
|
||||
_ROOT_RE = re.compile(r"^(?P<letter>[A-G])(?P<accidental>#{0,2}|b{0,2})$")
|
||||
|
||||
_BASS_DEGREE_RE = re.compile(r"^(?P<accidental>#{0,2}|b{0,2})(?P<degree>[1-9]|1[0-3])$")
|
||||
|
||||
_KEY_SIGNATURE_ACCIDENTALS = {
|
||||
"C": 0,
|
||||
"G": 1,
|
||||
"D": 2,
|
||||
"A": 3,
|
||||
"E": 4,
|
||||
"B": 5,
|
||||
"F#": 6,
|
||||
"C#": 7,
|
||||
"F": -1,
|
||||
"Bb": -2,
|
||||
"Eb": -3,
|
||||
"Ab": -4,
|
||||
"Db": -5,
|
||||
"Gb": -6,
|
||||
"Cb": -7,
|
||||
"Am": 0,
|
||||
"Em": 1,
|
||||
"Bm": 2,
|
||||
"F#m": 3,
|
||||
"C#m": 4,
|
||||
"G#m": 5,
|
||||
"D#m": 6,
|
||||
"A#m": 7,
|
||||
"Dm": -1,
|
||||
"Gm": -2,
|
||||
"Cm": -3,
|
||||
"Fm": -4,
|
||||
"Bbm": -5,
|
||||
"Ebm": -6,
|
||||
"Abm": -7,
|
||||
}
|
||||
|
||||
_KEY_RELATIVE_PITCH_NAMES = {
|
||||
7: ("B#", "C#", "C##", "D#", "D##", "E#", "F#", "F##", "G#", "G##", "A#", "B"),
|
||||
6: ("B#", "C#", "C##", "D#", "E", "E#", "F#", "F##", "G#", "G##", "A#", "B"),
|
||||
5: ("B#", "C#", "C##", "D#", "E", "E#", "F#", "F##", "G#", "A", "A#", "B"),
|
||||
4: ("B#", "C#", "D", "D#", "E", "E#", "F#", "F##", "G#", "A", "A#", "B"),
|
||||
3: ("B#", "C#", "D", "D#", "E", "E#", "F#", "G", "G#", "A", "A#", "B"),
|
||||
2: ("C", "C#", "D", "D#", "E", "E#", "F#", "G", "G#", "A", "A#", "B"),
|
||||
1: ("C", "C#", "D", "D#", "E", "F", "F#", "G", "G#", "A", "A#", "B"),
|
||||
0: ("C", "C#", "D", "D#", "E", "F", "F#", "G", "G#", "A", "Bb", "B"),
|
||||
-1: ("C", "C#", "D", "Eb", "E", "F", "F#", "G", "G#", "A", "Bb", "B"),
|
||||
-2: ("C", "C#", "D", "Eb", "E", "F", "F#", "G", "Ab", "A", "Bb", "B"),
|
||||
-3: ("C", "Db", "D", "Eb", "E", "F", "F#", "G", "Ab", "A", "Bb", "B"),
|
||||
-4: ("C", "Db", "D", "Eb", "E", "F", "Gb", "G", "Ab", "A", "Bb", "B"),
|
||||
-5: ("C", "Db", "D", "Eb", "E", "F", "Gb", "G", "Ab", "A", "Bb", "Cb"),
|
||||
-6: ("C", "Db", "D", "Eb", "Fb", "F", "Gb", "G", "Ab", "A", "Bb", "Cb"),
|
||||
-7: ("C", "Db", "D", "Eb", "Fb", "F", "Gb", "G", "Ab", "Bbb", "Bb", "Cb"),
|
||||
}
|
||||
|
||||
def _mode_with_first_tiebreak(values: Sequence[int]) -> int:
|
||||
counts = Counter(values)
|
||||
maximum = max(counts.values())
|
||||
return next(value for value in values if counts[value] == maximum)
|
||||
|
||||
def infer_measures(beats: Sequence[BeatEvent]) -> tuple[list[Measure], list[str]]:
|
||||
"""Infer self-consistent measures from actual downbeat boundaries."""
|
||||
|
||||
downbeat_indices = [index for index, beat in enumerate(beats) if beat.beat_id == 1]
|
||||
if not downbeat_indices:
|
||||
raise BeatGridError("No downbeat (beat ID 1) exists in the beat lab")
|
||||
spans: list[tuple[int, int, bool, bool]] = []
|
||||
if downbeat_indices[0] > 0:
|
||||
spans.append((0, downbeat_indices[0], True, False))
|
||||
spans.extend(
|
||||
(start, end, False, False)
|
||||
for start, end in zip(downbeat_indices, downbeat_indices[1:])
|
||||
)
|
||||
if downbeat_indices[-1] < len(beats) - 1:
|
||||
# Exported beat labs use their last row as the score end boundary. If
|
||||
# that row is not a downbeat, the final bar is intentionally truncated.
|
||||
spans.append((downbeat_indices[-1], len(beats) - 1, False, True))
|
||||
if not spans:
|
||||
raise BeatGridError("No positive-length measure exists between downbeats")
|
||||
|
||||
diagnostics = []
|
||||
measures = []
|
||||
for measure_index, (start, end, pickup, partial) in enumerate(spans):
|
||||
events = list(beats[start:end])
|
||||
# Downbeat spans still define measures when the model skips a beat ID.
|
||||
beat_count = len(events)
|
||||
|
||||
denominators = [event.denominator for event in events]
|
||||
denominator = _mode_with_first_tiebreak(denominators)
|
||||
declared_numerators = [event.declared_numerator for event in events]
|
||||
declared_numerator = _mode_with_first_tiebreak(declared_numerators)
|
||||
numerator_conflict = any(value != beat_count for value in declared_numerators)
|
||||
denominator_conflict = any(value != denominator for value in denominators)
|
||||
pad_final_partial = (
|
||||
partial
|
||||
and len(set(declared_numerators)) == 1
|
||||
and not denominator_conflict
|
||||
and declared_numerator >= beat_count
|
||||
)
|
||||
inferred = pickup or partial or numerator_conflict or denominator_conflict
|
||||
if pad_final_partial and declared_numerator > beat_count:
|
||||
diagnostics.append(
|
||||
f"measure {measure_index}: padded final {beat_count}/{denominator} span "
|
||||
f"to declared {declared_numerator}/{denominator} with trailing rest"
|
||||
)
|
||||
elif numerator_conflict:
|
||||
diagnostics.append(
|
||||
f"measure {measure_index}: inferred {beat_count}/{denominator} from downbeat span; "
|
||||
f"declared numerators were {declared_numerators}"
|
||||
)
|
||||
if denominator_conflict:
|
||||
diagnostics.append(
|
||||
f"measure {measure_index}: placed denominator {denominator} at the measure boundary; "
|
||||
f"row declarations were {denominators}"
|
||||
)
|
||||
measures.append(
|
||||
Measure(
|
||||
index=measure_index,
|
||||
start_beat=start,
|
||||
end_beat=end,
|
||||
numerator=beat_count,
|
||||
denominator=denominator,
|
||||
pickup=pickup,
|
||||
partial=partial,
|
||||
inferred=inferred,
|
||||
notated_numerator=(
|
||||
declared_numerator if pad_final_partial else beat_count
|
||||
),
|
||||
)
|
||||
)
|
||||
if len(measures) >= 2:
|
||||
first = measures[0]
|
||||
following = measures[1]
|
||||
first_duration = first.numerator / first.denominator
|
||||
following_duration = (
|
||||
following.abc_numerator / following.abc_denominator
|
||||
)
|
||||
if first_duration < following_duration:
|
||||
measures[0] = replace(
|
||||
first,
|
||||
inferred=True,
|
||||
notated_numerator=following.abc_numerator,
|
||||
notated_denominator=following.abc_denominator,
|
||||
pad_before=True,
|
||||
)
|
||||
diagnostics.append(
|
||||
f"measure 0: padded leading {first.numerator}/{first.denominator} span "
|
||||
f"to {following.abc_numerator}/{following.abc_denominator} "
|
||||
f"with preceding rest"
|
||||
)
|
||||
return measures, diagnostics
|
||||
|
||||
def _build_grid(beats: Sequence[BeatEvent], measures: Sequence[Measure]):
|
||||
interval_denominators = np.zeros(len(beats) - 1, dtype=np.int32)
|
||||
for measure in measures:
|
||||
interval_denominators[measure.start_beat:measure.end_beat] = measure.denominator
|
||||
if np.any(interval_denominators == 0):
|
||||
raise BeatGridError("Downbeat spans do not cover every beat interval")
|
||||
|
||||
subbeat_times = []
|
||||
subbeat_denominators = []
|
||||
quarter_positions = [0.0]
|
||||
current_quarter = 0.0
|
||||
for index in range(len(beats) - 1):
|
||||
start = beats[index].time
|
||||
end = beats[index + 1].time
|
||||
denominator = int(interval_denominators[index])
|
||||
times = np.linspace(start, end, SUBBEAT_DIVISION + 1)[:-1]
|
||||
subbeat_times.extend(float(value) for value in times)
|
||||
subbeat_denominators.extend([denominator] * SUBBEAT_DIVISION)
|
||||
quarter_step = 4.0 / denominator / SUBBEAT_DIVISION
|
||||
for _ in range(SUBBEAT_DIVISION):
|
||||
current_quarter += quarter_step
|
||||
quarter_positions.append(current_quarter)
|
||||
subbeat_times.append(beats[-1].time)
|
||||
subbeat_denominators.append(int(interval_denominators[-1]))
|
||||
return (
|
||||
np.asarray(subbeat_times, dtype=np.float64),
|
||||
np.asarray(quarter_positions, dtype=np.float64),
|
||||
np.asarray(subbeat_denominators, dtype=np.int32),
|
||||
)
|
||||
|
||||
def _subbeat_boundaries(subbeat_times: np.ndarray) -> np.ndarray:
|
||||
return (subbeat_times[:-1] + subbeat_times[1:]) / 2
|
||||
|
||||
def _quantize_time(time: float, subbeat_times: np.ndarray) -> int:
|
||||
return int(np.searchsorted(_subbeat_boundaries(subbeat_times), float(time)))
|
||||
|
||||
def _fill_intervals(rows, subbeat_times, *, default, dtype):
|
||||
result = np.full(len(subbeat_times), default, dtype=dtype)
|
||||
for start, end, value in rows:
|
||||
start_t = _quantize_time(start, subbeat_times)
|
||||
end_t = _quantize_time(end, subbeat_times)
|
||||
start_t = max(0, min(start_t, len(result) - 1))
|
||||
end_t = max(0, min(end_t, len(result) - 1))
|
||||
if start_t == end_t == len(result) - 1:
|
||||
continue
|
||||
if end_t <= start_t:
|
||||
raise AbcRebuildError(
|
||||
f"Interval {start:.6f}-{end:.6f} ({value}) is shorter than the ABC subbeat grid"
|
||||
)
|
||||
result[start_t:end_t] = value
|
||||
if len(result) > 1:
|
||||
result[-1] = result[-2]
|
||||
return result
|
||||
|
||||
def _structure_events(rows, subbeat_times):
|
||||
events = []
|
||||
for start, _, label in rows:
|
||||
t = _quantize_time(start, subbeat_times)
|
||||
t = max(0, min(t, len(subbeat_times) - 1))
|
||||
events.append((t, label))
|
||||
return events
|
||||
|
||||
def _notes_to_arr(notes, subbeat_times, voice_id):
|
||||
result = np.zeros(len(subbeat_times), dtype=np.int32)
|
||||
boundaries = _subbeat_boundaries(subbeat_times)
|
||||
for note in sorted(notes, key=lambda item: (item[0], item[1], item[2])):
|
||||
start_t = int(np.searchsorted(boundaries, note[0]))
|
||||
end_t = int(np.searchsorted(boundaries, note[1]))
|
||||
start_t = max(0, min(start_t, len(result) - 1))
|
||||
end_t = max(0, min(end_t, len(result) - 1))
|
||||
if start_t == end_t == len(result) - 1:
|
||||
continue
|
||||
if end_t <= start_t:
|
||||
raise MelodyVoiceError(
|
||||
f"{voice_id}: note pitch={note[2]} at {note[0]:.6f}-{note[1]:.6f} "
|
||||
"cannot be represented on the decoded subbeat grid"
|
||||
)
|
||||
if np.any(result[start_t:end_t] != 0):
|
||||
raise MelodyVoiceError(
|
||||
f"{voice_id}: overlapping quantized melody notes at subbeats {start_t}:{end_t}"
|
||||
)
|
||||
sustain = note[2] * 2 + 2
|
||||
result[start_t:end_t] = sustain
|
||||
result[start_t] = sustain + 1
|
||||
return result
|
||||
|
||||
def _pitch_class(root: str) -> tuple[int, str, str]:
|
||||
match = _ROOT_RE.fullmatch(root)
|
||||
if match is None:
|
||||
raise ChordSymbolError(f"Invalid pitch spelling {root!r}")
|
||||
letter = match.group("letter")
|
||||
accidental = match.group("accidental")
|
||||
offset = accidental.count("#") - accidental.count("b")
|
||||
return (_NATURAL_PITCH_CLASS[letter] + offset) % 12, letter, accidental
|
||||
|
||||
def portable_pitch_name(root: str, *, preserve_double: bool = False) -> str:
|
||||
pitch_class, _, accidental = _pitch_class(root)
|
||||
if preserve_double or len(accidental) <= 1:
|
||||
return root
|
||||
names = _SHARP_PITCH_NAMES if accidental.startswith("#") else _FLAT_PITCH_NAMES
|
||||
return names[pitch_class]
|
||||
|
||||
def _bass_degree_to_pitch(root: str, degree_text: str) -> str:
|
||||
if _ROOT_RE.fullmatch(degree_text):
|
||||
return portable_pitch_name(degree_text, preserve_double=True)
|
||||
match = _BASS_DEGREE_RE.fullmatch(degree_text)
|
||||
if match is None:
|
||||
raise ChordSymbolError(f"Invalid chord bass degree {degree_text!r}")
|
||||
root_pc, root_letter, root_accidental = _pitch_class(root)
|
||||
degree = int(match.group("degree"))
|
||||
degree_accidental = match.group("accidental")
|
||||
scale_semitones = (0, 2, 4, 5, 7, 9, 11)
|
||||
interval = scale_semitones[(degree - 1) % 7] + 12 * ((degree - 1) // 7)
|
||||
interval += degree_accidental.count("#") - degree_accidental.count("b")
|
||||
target_pc = (root_pc + interval) % 12
|
||||
|
||||
target_letter_index = (_LETTERS.index(root_letter) + degree - 1) % 7
|
||||
target_letter = _LETTERS[target_letter_index]
|
||||
natural_pc = _NATURAL_PITCH_CLASS[target_letter]
|
||||
difference = (target_pc - natural_pc + 6) % 12 - 6
|
||||
if difference in {-2, -1, 0, 1, 2}:
|
||||
accidental = {-2: "bb", -1: "b", 0: "", 1: "#", 2: "##"}[difference]
|
||||
return target_letter + accidental
|
||||
names = _SHARP_PITCH_NAMES if "#" in (root_accidental + degree_accidental) else _FLAT_PITCH_NAMES
|
||||
return names[target_pc]
|
||||
|
||||
def chord_symbol_to_abc(chord: str) -> str | None:
|
||||
chord = chord.strip()
|
||||
if chord in NO_CHORDS:
|
||||
return None
|
||||
if ":" not in chord:
|
||||
raise ChordSymbolError(f"Chord {chord!r} is missing the ':' quality separator")
|
||||
root, descriptor = chord.split(":", 1)
|
||||
if "/" in descriptor:
|
||||
quality, bass_degree = descriptor.split("/", 1)
|
||||
else:
|
||||
quality, bass_degree = descriptor, None
|
||||
if quality not in _QUALITY_TO_ABC:
|
||||
raise ChordSymbolError(
|
||||
f"Unsupported chord quality {quality!r} in {chord!r}; refusing to rewrite it as major"
|
||||
)
|
||||
chord_root = portable_pitch_name(root, preserve_double=True)
|
||||
text = chord_root + _QUALITY_TO_ABC[quality]
|
||||
if bass_degree:
|
||||
text += "/" + _bass_degree_to_pitch(root, bass_degree)
|
||||
return text
|
||||
|
||||
def key_symbol_to_abc(key: str) -> str:
|
||||
key = key.strip()
|
||||
if ":" in key:
|
||||
root, mode = key.split(":", 1)
|
||||
if mode not in {"major", "minor"}:
|
||||
raise AbcRebuildError(f"Unsupported key mode {mode!r} in {key!r}")
|
||||
elif key.endswith("m"):
|
||||
root, mode = key[:-1], "minor"
|
||||
else:
|
||||
root, mode = key, "major"
|
||||
root_pc, _, accidental = _pitch_class(root)
|
||||
candidate = portable_pitch_name(root) + ("m" if mode == "minor" else "")
|
||||
if candidate in _KEY_SIGNATURE_ACCIDENTALS:
|
||||
return candidate
|
||||
names = _FLAT_PITCH_NAMES if "b" in accidental else _SHARP_PITCH_NAMES
|
||||
candidate = names[root_pc] + ("m" if mode == "minor" else "")
|
||||
if candidate not in _KEY_SIGNATURE_ACCIDENTALS:
|
||||
fallback_names = _SHARP_PITCH_NAMES if names is _FLAT_PITCH_NAMES else _FLAT_PITCH_NAMES
|
||||
candidate = fallback_names[root_pc] + ("m" if mode == "minor" else "")
|
||||
if candidate not in _KEY_SIGNATURE_ACCIDENTALS:
|
||||
raise AbcRebuildError(f"Cannot encode portable ABC key for {key!r}")
|
||||
return candidate
|
||||
|
||||
def get_key_accidentals(key: str) -> list[int]:
|
||||
try:
|
||||
count = _KEY_SIGNATURE_ACCIDENTALS[key]
|
||||
except KeyError as exc:
|
||||
raise AbcRebuildError(f"Unsupported ABC key signature {key!r}") from exc
|
||||
accidentals = [0] * 7
|
||||
order = "FCGDAEB" if count > 0 else "BEADGCF"
|
||||
for letter in order[:abs(count)]:
|
||||
accidentals[_LETTERS.index(letter)] = 1 if count > 0 else -1
|
||||
return accidentals
|
||||
|
||||
def note_to_abc(note: int, key_accidentals: Sequence[int], measure_accidentals: dict) -> str:
|
||||
"""Use key-relative spelling and write only bar-state changes.
|
||||
|
||||
The two target parsers propagate an accidental to the same note letter in
|
||||
every octave until the next barline. ``measure_accidentals`` is therefore
|
||||
keyed by letter and reset by the caller for every bar (and after an inline
|
||||
key change). This preserves pitches across parsers while still omitting
|
||||
repeated accidental marks. The key-relative spelling can use double
|
||||
accidentals in remote keys; MIDI G is F## in G# minor, for example.
|
||||
"""
|
||||
|
||||
accidental_count = sum(key_accidentals)
|
||||
try:
|
||||
pitch_name = _KEY_RELATIVE_PITCH_NAMES[accidental_count][note % 12]
|
||||
except KeyError as exc:
|
||||
raise AbcRebuildError(
|
||||
f"Unsupported key signature accidental count {accidental_count}"
|
||||
) from exc
|
||||
letter = pitch_name[0]
|
||||
accidental = pitch_name[1:]
|
||||
accidental_number = {"": 0, "#": 1, "##": 2, "b": -1, "bb": -2}[accidental]
|
||||
octave = (note - 60) // 12
|
||||
# Cb and B# cross the MIDI octave boundary even though their written note
|
||||
# letter does not.
|
||||
if note % 12 == 11 and accidental_number == -1:
|
||||
octave += 1
|
||||
elif note % 12 == 0 and accidental_number == 1:
|
||||
octave -= 1
|
||||
scale_index = _LETTERS.index(letter)
|
||||
current_accidental = measure_accidentals.get(
|
||||
scale_index,
|
||||
key_accidentals[scale_index],
|
||||
)
|
||||
accidental_text = ""
|
||||
if current_accidental != accidental_number:
|
||||
measure_accidentals[scale_index] = accidental_number
|
||||
accidental_text = {-2: "__", -1: "_", 0: "=", 1: "^", 2: "^^"}[
|
||||
accidental_number
|
||||
]
|
||||
|
||||
if octave > 0:
|
||||
letter = letter.lower()
|
||||
if octave > 1:
|
||||
letter += "'" * (octave - 1)
|
||||
elif octave < 0:
|
||||
letter += "," * abs(octave)
|
||||
return accidental_text + letter
|
||||
|
||||
def abc_unit_denominator(score: RebuiltAbcScore) -> int:
|
||||
values = [
|
||||
denominator * score.subbeat_div
|
||||
for measure in score.measures
|
||||
for denominator in (measure.denominator, measure.abc_denominator)
|
||||
]
|
||||
denominator = math.lcm(*values)
|
||||
if denominator > 1024:
|
||||
raise AbcRebuildError(f"Required ABC unit length 1/{denominator} is unreasonably small")
|
||||
return denominator
|
||||
|
||||
def _measure_actual_units(measure: Measure, unit_denominator: int) -> int:
|
||||
return measure.numerator * unit_denominator // measure.denominator
|
||||
|
||||
def _measure_abc_units(measure: Measure, unit_denominator: int) -> int:
|
||||
return measure.abc_numerator * unit_denominator // measure.abc_denominator
|
||||
|
||||
def _measure_padding_units(measure: Measure, unit_denominator: int) -> int:
|
||||
return (
|
||||
_measure_abc_units(measure, unit_denominator)
|
||||
- _measure_actual_units(measure, unit_denominator)
|
||||
)
|
||||
|
||||
def _duration_units(score: RebuiltAbcScore, start_t: int, end_t: int, unit_denominator: int) -> int:
|
||||
units = 0
|
||||
for denominator in score.subbeat_denominators[start_t:end_t]:
|
||||
divisor = int(denominator) * score.subbeat_div
|
||||
if unit_denominator % divisor:
|
||||
raise AbcRebuildError(
|
||||
f"ABC L:1/{unit_denominator} cannot express a 1/{divisor} subbeat exactly"
|
||||
)
|
||||
units += unit_denominator // divisor
|
||||
return units
|
||||
|
||||
def estimate_tempo(score: RebuiltAbcScore) -> float:
|
||||
seconds = score.subbeat_times[-1] - score.subbeat_times[0]
|
||||
quarter_notes = score.subbeat_quarters[-1] - score.subbeat_quarters[0]
|
||||
if seconds <= 0 or quarter_notes <= 0:
|
||||
raise AbcRebuildError("Cannot estimate tempo from a zero-duration score")
|
||||
return float(quarter_notes / seconds * 60.0)
|
||||
|
||||
def _continues_pitch(value: int, next_value: int) -> bool:
|
||||
if value <= 0:
|
||||
return False
|
||||
pitch = value // 2 - 1
|
||||
return next_value == pitch * 2 + 2
|
||||
|
||||
def _same_note_segment(value: int, next_value: int) -> bool:
|
||||
if value == 0:
|
||||
return next_value == 0
|
||||
pitch = value // 2 - 1
|
||||
return next_value == pitch * 2 + 2
|
||||
|
||||
def _split_duration_units(duration: int) -> list[int]:
|
||||
"""Split a duration into values accepted by strict music parsers."""
|
||||
|
||||
if duration <= 0:
|
||||
raise AbcRebuildError(f"Cannot serialize non-positive duration {duration}")
|
||||
result = []
|
||||
remaining = int(duration)
|
||||
while remaining:
|
||||
if remaining in _SUPPORTED_DURATION_UNITS:
|
||||
result.append(remaining)
|
||||
break
|
||||
candidates = [
|
||||
value
|
||||
for value in _SUPPORTED_DURATION_UNITS
|
||||
if value < remaining
|
||||
]
|
||||
if not candidates:
|
||||
raise AbcRebuildError(
|
||||
f"Duration {duration} cannot be split into representable ABC values"
|
||||
)
|
||||
chunk = max(candidates)
|
||||
result.append(chunk)
|
||||
remaining -= chunk
|
||||
return result
|
||||
|
||||
def _duration_text(duration: int) -> str:
|
||||
return "" if duration == 1 else str(duration)
|
||||
|
||||
def _render_duration_tokens(
|
||||
prefix: str,
|
||||
note_text: str,
|
||||
duration: int,
|
||||
*,
|
||||
tie_out: bool,
|
||||
) -> list[str]:
|
||||
chunks = _split_duration_units(duration)
|
||||
tokens = []
|
||||
for index, chunk in enumerate(chunks):
|
||||
continues = note_text != "z" and (
|
||||
index + 1 < len(chunks) or tie_out
|
||||
)
|
||||
tokens.append(
|
||||
(prefix if index == 0 else "")
|
||||
+ note_text
|
||||
+ _duration_text(chunk)
|
||||
+ ("-" if continues else "")
|
||||
)
|
||||
return tokens
|
||||
|
||||
def _render_voice_measure(
|
||||
score: RebuiltAbcScore,
|
||||
voice_id: str,
|
||||
measure: Measure,
|
||||
unit_denominator: int,
|
||||
) -> str:
|
||||
voice = score.voice_arrs[voice_id]
|
||||
show_chords = voice_id == "Vocal"
|
||||
measure_accidentals = {}
|
||||
current_key = str(score.key_arr[measure.start_t])
|
||||
key_accidentals = get_key_accidentals(current_key)
|
||||
parts = []
|
||||
padding = _measure_padding_units(measure, unit_denominator)
|
||||
if padding < 0:
|
||||
raise AbcRebuildError(
|
||||
f"Measure {measure.index}: notated meter is shorter than its decoded span"
|
||||
)
|
||||
leading_padding = padding if measure.pad_before else 0
|
||||
trailing_padding = 0 if measure.pad_before else padding
|
||||
t = measure.start_t
|
||||
while t < measure.end_t:
|
||||
change_points = [measure.end_t]
|
||||
for probe in range(t + 1, measure.end_t):
|
||||
if not _same_note_segment(int(voice[t]), int(voice[probe])):
|
||||
change_points.append(probe)
|
||||
break
|
||||
for probe in range(t + 1, measure.end_t):
|
||||
if score.key_arr[probe] != score.key_arr[probe - 1]:
|
||||
change_points.append(probe)
|
||||
break
|
||||
if show_chords:
|
||||
for probe in range(t + 1, measure.end_t):
|
||||
if score.chord_arr[probe] != score.chord_arr[probe - 1]:
|
||||
change_points.append(probe)
|
||||
break
|
||||
next_t = min(change_points)
|
||||
|
||||
prefix = ""
|
||||
key = str(score.key_arr[t])
|
||||
if t > measure.start_t and key != current_key:
|
||||
current_key = key
|
||||
key_accidentals = get_key_accidentals(current_key)
|
||||
measure_accidentals = {}
|
||||
prefix += f"[K:{current_key}]"
|
||||
|
||||
if show_chords and (t == measure.start_t or score.chord_arr[t] != score.chord_arr[t - 1]):
|
||||
chord = str(score.chord_arr[t])
|
||||
chord_text = chord_symbol_to_abc(chord)
|
||||
if chord_text is not None:
|
||||
prefix += f'"{chord_text}"'
|
||||
|
||||
value = int(voice[t])
|
||||
if value == 0:
|
||||
note_text = "z"
|
||||
else:
|
||||
note_text = note_to_abc(value // 2 - 1, key_accidentals, measure_accidentals)
|
||||
duration = _duration_units(score, t, next_t, unit_denominator)
|
||||
if t == measure.start_t and leading_padding:
|
||||
if value == 0 and not prefix:
|
||||
duration += leading_padding
|
||||
else:
|
||||
parts.extend(
|
||||
_render_duration_tokens(
|
||||
"",
|
||||
"z",
|
||||
leading_padding,
|
||||
tie_out=False,
|
||||
)
|
||||
)
|
||||
leading_padding = 0
|
||||
if value == 0 and next_t == measure.end_t and trailing_padding:
|
||||
duration += trailing_padding
|
||||
trailing_padding = 0
|
||||
if duration <= 0:
|
||||
raise AbcRebuildError(f"Non-positive ABC duration at subbeats {t}:{next_t}")
|
||||
tie_out = (
|
||||
value > 0
|
||||
and next_t < len(voice)
|
||||
and _continues_pitch(value, int(voice[next_t]))
|
||||
)
|
||||
parts.extend(
|
||||
_render_duration_tokens(
|
||||
prefix,
|
||||
note_text,
|
||||
duration,
|
||||
tie_out=tie_out,
|
||||
)
|
||||
)
|
||||
t = next_t
|
||||
if leading_padding:
|
||||
raise AbcRebuildError(
|
||||
f"Measure {measure.index}: leading rest padding was not serialized"
|
||||
)
|
||||
if trailing_padding:
|
||||
parts.extend(
|
||||
_render_duration_tokens(
|
||||
"",
|
||||
"z",
|
||||
trailing_padding,
|
||||
tie_out=False,
|
||||
)
|
||||
)
|
||||
return "".join(parts)
|
||||
|
||||
def _is_compressible_full_rest(rendered_measure: str) -> bool:
|
||||
"""Whether a rendered measure can be losslessly replaced by ABC ``Z``."""
|
||||
|
||||
cursor = 0
|
||||
saw_note = False
|
||||
for match in _MUSIC_ELEMENT_RE.finditer(rendered_measure):
|
||||
if rendered_measure[cursor:match.start()]:
|
||||
return False
|
||||
cursor = match.end()
|
||||
if match.group("quoted") is not None or match.group("key") is not None:
|
||||
return False
|
||||
saw_note = True
|
||||
if match.group("note") != "z" or match.group("tie"):
|
||||
return False
|
||||
return saw_note and cursor == len(rendered_measure)
|
||||
|
||||
def _render_voice_group(
|
||||
score: RebuiltAbcScore,
|
||||
voice_id: str,
|
||||
measures: list[Measure],
|
||||
unit_denominator: int,
|
||||
) -> str:
|
||||
rendered = [
|
||||
_render_voice_measure(
|
||||
score,
|
||||
voice_id,
|
||||
measure,
|
||||
unit_denominator,
|
||||
)
|
||||
for measure in measures
|
||||
]
|
||||
parts = []
|
||||
index = 0
|
||||
while index < len(rendered):
|
||||
if not _is_compressible_full_rest(rendered[index]):
|
||||
parts.append(rendered[index] + "|")
|
||||
index += 1
|
||||
continue
|
||||
end = index + 1
|
||||
while (
|
||||
end < len(rendered)
|
||||
and _is_compressible_full_rest(rendered[end])
|
||||
):
|
||||
end += 1
|
||||
count = end - index
|
||||
parts.append("Z" + (str(count) if count > 1 else "") + "|")
|
||||
index = end
|
||||
return "".join(parts)
|
||||
|
||||
def _sanitize_structure_label(value: str) -> str:
|
||||
return " ".join(str(value).split())
|
||||
|
||||
def _measure_groups(score: RebuiltAbcScore) -> list[MeasureGroup]:
|
||||
first_measure = score.measures[0]
|
||||
active_meter = (
|
||||
first_measure.abc_numerator,
|
||||
first_measure.abc_denominator,
|
||||
)
|
||||
active_key = str(score.key_arr[first_measure.start_t])
|
||||
active_structure = ""
|
||||
groups: list[MeasureGroup] = []
|
||||
|
||||
for measure in score.measures:
|
||||
meter = (measure.abc_numerator, measure.abc_denominator)
|
||||
key = str(score.key_arr[measure.start_t])
|
||||
meter_changed = meter != active_meter
|
||||
key_changed = key != active_key
|
||||
new_structure_labels = []
|
||||
for t, label in score.structure_events:
|
||||
if not measure.start_t <= t < measure.end_t:
|
||||
continue
|
||||
clean_label = _sanitize_structure_label(label)
|
||||
if clean_label and clean_label != active_structure:
|
||||
new_structure_labels.append(clean_label)
|
||||
active_structure = clean_label
|
||||
|
||||
start_group = (
|
||||
not groups
|
||||
or len(groups[-1].measures) >= 4
|
||||
or meter_changed
|
||||
or key_changed
|
||||
or bool(new_structure_labels)
|
||||
)
|
||||
if start_group:
|
||||
groups.append(
|
||||
MeasureGroup(
|
||||
measures=[measure],
|
||||
structure_labels=new_structure_labels,
|
||||
meter_changed=meter_changed,
|
||||
key_changed=key_changed,
|
||||
)
|
||||
)
|
||||
else:
|
||||
groups[-1].measures.append(measure)
|
||||
|
||||
active_meter = meter
|
||||
active_key = str(score.key_arr[measure.end_t - 1])
|
||||
return groups
|
||||
|
||||
def score_to_abc(score: RebuiltAbcScore) -> str:
|
||||
unit_denominator = abc_unit_denominator(score)
|
||||
first_measure = score.measures[0]
|
||||
first_key = str(score.key_arr[first_measure.start_t])
|
||||
lines = [
|
||||
"X:1",
|
||||
"T:",
|
||||
f"M:{first_measure.abc_numerator}/{first_measure.abc_denominator}",
|
||||
f"L:1/{unit_denominator}",
|
||||
f"Q:1/4={int(round(estimate_tempo(score)))}",
|
||||
'V: Vocal clef=treble name="Vocal Melody" snm="Vocal"',
|
||||
'V: Ins clef=treble name="Ins Melody" snm="Inst."',
|
||||
f"K:{first_key}",
|
||||
]
|
||||
for group in _measure_groups(score):
|
||||
lines.extend(f"% {label}" for label in group.structure_labels)
|
||||
first_group_measure = group.measures[0]
|
||||
for voice_id in VOICE_IDS:
|
||||
lines.append(f"V: {voice_id}")
|
||||
if group.meter_changed:
|
||||
lines.append(
|
||||
f"M:{first_group_measure.abc_numerator}/"
|
||||
f"{first_group_measure.abc_denominator}"
|
||||
)
|
||||
if group.key_changed:
|
||||
lines.append(
|
||||
f"K:{score.key_arr[first_group_measure.start_t]}"
|
||||
)
|
||||
lines.append(
|
||||
_render_voice_group(
|
||||
score,
|
||||
voice_id,
|
||||
group.measures,
|
||||
unit_denominator,
|
||||
)
|
||||
)
|
||||
text = "\n".join(lines) + "\n"
|
||||
return text
|
||||
@@ -1024,6 +1024,12 @@ class ACEAudio15(LatentFormat):
|
||||
latent_dimensions = 1
|
||||
temporal_downscale_ratio = 1764
|
||||
|
||||
class YuE2(LatentFormat):
|
||||
latent_channels = 64
|
||||
latent_dimensions = 1
|
||||
temporal_downscale_ratio = 1920
|
||||
|
||||
|
||||
class MiniMaxMusic3(LatentFormat):
|
||||
latent_channels = 128
|
||||
latent_dimensions = 1
|
||||
|
||||
@@ -157,7 +157,7 @@ class DecoderBlock(nn.Module):
|
||||
else:
|
||||
upsample_layer = WNConvTranspose1d(in_channels=in_channels,
|
||||
out_channels=out_channels,
|
||||
kernel_size=2*stride, stride=stride, padding=math.ceil(stride/2))
|
||||
kernel_size=2*stride, stride=stride, padding=math.ceil(stride/2), output_padding=stride % 2)
|
||||
|
||||
self.layers = nn.Sequential(
|
||||
get_activation("snake" if use_snake else "elu", antialias=antialias_activation, channels=in_channels),
|
||||
@@ -261,16 +261,20 @@ class AudioOobleckVAE(nn.Module):
|
||||
use_snake=True,
|
||||
antialias_activation=False,
|
||||
use_nearest_upsample=False,
|
||||
final_tanh=False):
|
||||
final_tanh=False,
|
||||
sample_latent=True):
|
||||
super().__init__()
|
||||
self.encoder = OobleckEncoder(in_channels, channels, latent_dim * 2, c_mults, strides, use_snake, antialias_activation)
|
||||
self.decoder = OobleckDecoder(in_channels, channels, latent_dim, c_mults, strides, use_snake, antialias_activation,
|
||||
use_nearest_upsample=use_nearest_upsample, final_tanh=final_tanh)
|
||||
self.bottleneck = VAEBottleneck()
|
||||
self.sample_latent = sample_latent
|
||||
|
||||
def encode(self, x):
|
||||
return self.bottleneck.encode(self.encoder(x))
|
||||
encoded = self.encoder(x)
|
||||
if not self.sample_latent:
|
||||
return encoded.chunk(2, dim=1)[0]
|
||||
return self.bottleneck.encode(encoded)
|
||||
|
||||
def decode(self, x):
|
||||
return self.decoder(self.bottleneck.decode(x))
|
||||
|
||||
|
||||
@@ -0,0 +1,87 @@
|
||||
"""YuE2 acoustic transformer. Adapted from M·A·P YuE2 (Apache-2.0)."""
|
||||
|
||||
import torch
|
||||
from torch import nn
|
||||
|
||||
import comfy.model_management
|
||||
import comfy.model_prefetch
|
||||
import comfy.ops
|
||||
from comfy.ldm.modules.attention import optimized_attention_for_device
|
||||
from comfy.ldm.modules.diffusionmodules.util import timestep_embedding
|
||||
from comfy.text_encoders.llama import Qwen3_8BConfig, RMSNorm, TransformerBlock, precompute_freqs_cis
|
||||
|
||||
|
||||
def model_config(**overrides):
|
||||
return Qwen3_8BConfig(**{
|
||||
"vocab_size": 184704, "hidden_size": 2048, "intermediate_size": 6144,
|
||||
"num_hidden_layers": 28, "num_attention_heads": 16, "num_key_value_heads": 8,
|
||||
"max_position_embeddings": 24576, "merged_qkv": True, "merged_mlp": True,
|
||||
**overrides,
|
||||
})
|
||||
|
||||
|
||||
class TimestepEmbedder(nn.Module):
|
||||
def __init__(self, hidden_size, dtype, device, operations):
|
||||
super().__init__()
|
||||
self.mlp = nn.Sequential(
|
||||
operations.Linear(256, hidden_size, dtype=dtype, device=device),
|
||||
nn.SiLU(),
|
||||
operations.Linear(hidden_size, hidden_size, dtype=dtype, device=device),
|
||||
)
|
||||
|
||||
def forward(self, t, dtype):
|
||||
return self.mlp(timestep_embedding(t, 256).to(dtype))
|
||||
|
||||
|
||||
class AudioPositionEmbedding(nn.Module):
|
||||
def __init__(self, frames, hidden_size, dtype, device):
|
||||
super().__init__()
|
||||
self.register_buffer("pe", torch.empty(frames, hidden_size, dtype=dtype, device=device))
|
||||
|
||||
def forward(self, length, x):
|
||||
return comfy.ops.cast_to_input(self.pe[:length], x)
|
||||
|
||||
|
||||
class YuE2(nn.Module):
|
||||
def __init__(self, dtype=None, device=None, operations=None, **kwargs):
|
||||
super().__init__()
|
||||
self.dtype = dtype
|
||||
self.config = model_config(**kwargs.get("config", {}))
|
||||
config = self.config
|
||||
self.model = nn.Module()
|
||||
self.model.layers = nn.ModuleList([
|
||||
TransformerBlock(config, i, device=device, dtype=dtype, ops=operations)
|
||||
for i in range(config.num_hidden_layers)
|
||||
])
|
||||
self.model.norm = RMSNorm(config.hidden_size, config.rms_norm_eps, device=device, dtype=dtype)
|
||||
self.vae2llm = operations.Linear(64, config.hidden_size, dtype=dtype, device=device)
|
||||
self.llm2vae = operations.Linear(config.hidden_size, 64, dtype=dtype, device=device)
|
||||
self.time_embedder = TimestepEmbedder(config.hidden_size, dtype, device, operations)
|
||||
self.latent_pos_embed = AudioPositionEmbedding(config.max_position_embeddings, config.hidden_size, dtype, device)
|
||||
|
||||
def forward(self, x, timestep, context, yue2_chunks, transformer_options={}, **kwargs):
|
||||
batch, channels, frames = x.shape
|
||||
if frames != yue2_chunks[-1][1]:
|
||||
raise ValueError("YuE2 latent duration must match the seconds output of YuE2 Text Encode.")
|
||||
config = self.config
|
||||
time = self.time_embedder(timestep.to(x.dtype), x.dtype)[:, None]
|
||||
output = torch.empty_like(x)
|
||||
attention = optimized_attention_for_device(x.device)
|
||||
for start, end, kv_start, kv_end in yue2_chunks:
|
||||
comfy.model_management.throw_exception_if_processing_interrupted()
|
||||
ar_length = kv_end - kv_start
|
||||
length = end - start + 2
|
||||
state = torch.nn.functional.pad(x[..., start:end].transpose(1, 2), (0, 0, 1, 1))
|
||||
state = self.vae2llm(state) + time + self.latent_pos_embed(length, x)[None]
|
||||
positions = torch.arange(ar_length, ar_length + length, device=x.device)[None]
|
||||
rope = precompute_freqs_cis(config.head_dim, positions, config.rope_theta, device=x.device)
|
||||
prefix = context[:, kv_start:kv_end].reshape(batch, ar_length, config.num_hidden_layers, 2, config.num_key_value_heads, config.head_dim)
|
||||
prefix = prefix.permute(2, 3, 0, 4, 1, 5)
|
||||
prefetch = comfy.model_prefetch.make_prefetch_queue(list(self.model.layers), x.device, transformer_options)
|
||||
for index, layer in enumerate(self.model.layers):
|
||||
comfy.model_prefetch.prefetch_queue_pop(prefetch, x.device, layer, state.dtype)
|
||||
state, _ = layer(state, freqs_cis=rope, optimized_attention=attention,
|
||||
past_key_value=(prefix[index, 0], prefix[index, 1], ar_length))
|
||||
comfy.model_prefetch.prefetch_queue_pop(prefetch, x.device, None)
|
||||
output[..., start:end] = self.llm2vae(self.model.norm(state))[:, 1:-1].transpose(1, 2)
|
||||
return output
|
||||
@@ -23,6 +23,7 @@ import logging
|
||||
import comfy.ldm.lightricks.av_model
|
||||
import comfy.ldm.minimax.model
|
||||
import comfy.ldm.minimax_music.dit
|
||||
import comfy.ldm.yue2.model
|
||||
import comfy.nested_tensor
|
||||
import comfy.ldm.lightricks.symmetric_patchifier
|
||||
import comfy.context_windows
|
||||
@@ -2551,6 +2552,25 @@ class ACEStep15(BaseModel):
|
||||
out['refer_audio'] = comfy.conds.CONDRegular(refer_audio)
|
||||
return out
|
||||
|
||||
class YuE2(BaseModel):
|
||||
def __init__(self, model_config, model_type=ModelType.FLOW, device=None):
|
||||
super().__init__(model_config, model_type, device=device, unet_model=comfy.ldm.yue2.model.YuE2)
|
||||
|
||||
def extra_conds(self, **kwargs):
|
||||
out = super().extra_conds(**kwargs)
|
||||
context = kwargs["cross_attn"].to(device=kwargs["device"], dtype=self.get_dtype_inference())
|
||||
out["c_crossattn"] = comfy.conds.CONDRegular(context)
|
||||
out["yue2_chunks"] = comfy.conds.CONDConstant(kwargs["yue2_chunks"])
|
||||
return out
|
||||
|
||||
def extra_conds_shapes(self, **kwargs):
|
||||
return {"c_crossattn": kwargs["cross_attn"].shape}
|
||||
|
||||
def memory_required(self, input_shape, cond_shapes={}):
|
||||
context_size = sum(math.prod(shape) for shape in cond_shapes.get("c_crossattn", []))
|
||||
return super().memory_required(input_shape, cond_shapes) + context_size * comfy.model_management.dtype_size(self.get_dtype_inference())
|
||||
|
||||
|
||||
class MiniMaxMusic3(BaseModel):
|
||||
def __init__(self, model_config, model_type=ModelType.FLOW, device=None):
|
||||
super().__init__(model_config, model_type, device=device, unet_model=comfy.ldm.minimax_music.dit.MiniMaxMusic3DiT)
|
||||
|
||||
@@ -1152,6 +1152,12 @@ def detect_unet_config(state_dict, key_prefix, metadata=None):
|
||||
"text_dim": 4096,
|
||||
}
|
||||
|
||||
if all(key_prefix + key in state_dict for key in (
|
||||
"vae2llm.weight", "llm2vae.weight", "latent_pos_embed.pe",
|
||||
"model.layers.0.self_attn.qkv_proj.weight", "time_embedder.mlp.0.weight",
|
||||
)):
|
||||
return {"audio_model": "yue2"}
|
||||
|
||||
if '{}input_blocks.0.0.weight'.format(key_prefix) not in state_dict_keys:
|
||||
return None
|
||||
|
||||
|
||||
+16
-1
@@ -79,6 +79,7 @@ import comfy.text_encoders.qwen35
|
||||
import comfy.text_encoders.qwen3vl
|
||||
import comfy.text_encoders.minimax
|
||||
import comfy.text_encoders.minimax_music
|
||||
import comfy.text_encoders.yue2
|
||||
import comfy.ldm.minimax.vae
|
||||
import comfy.ldm.minimax.audio_vae
|
||||
import comfy.text_encoders.boogu
|
||||
@@ -695,6 +696,7 @@ class VAE:
|
||||
decoder_config={'target': "comfy.ldm.modules.diffusionmodules.model.Decoder", 'params': decoder_ddconfig if decoder_ddconfig is not None else ddconfig})
|
||||
elif "decoder.layers.1.layers.0.beta" in sd:
|
||||
config = {}
|
||||
yue2_vae = "decoder.layers.6.layers.1.weight_v" in sd or "decoder.layers.6.layers.1.parametrizations.weight.original1" in sd
|
||||
param_key = None
|
||||
self.upscale_ratio = 2048
|
||||
self.downscale_ratio = 2048
|
||||
@@ -709,6 +711,9 @@ class VAE:
|
||||
self.upscale_ratio = 1920
|
||||
self.downscale_ratio = 1920
|
||||
|
||||
if yue2_vae:
|
||||
config.update(channels=64, c_mults=[1, 2, 4, 8, 16, 32], strides=[2, 2, 4, 4, 5, 6],
|
||||
sample_latent=False)
|
||||
self.first_stage_model = AudioOobleckVAE(**config)
|
||||
self.memory_used_encode = lambda shape, dtype: (1000 * shape[2]) * model_management.dtype_size(dtype)
|
||||
self.memory_used_decode = lambda shape, dtype: (1000 * shape[2] * 2048) * model_management.dtype_size(dtype)
|
||||
@@ -720,6 +725,10 @@ class VAE:
|
||||
self.process_input = lambda audio: audio
|
||||
self.working_dtypes = [torch.float16, torch.bfloat16, torch.float32]
|
||||
self.disable_offload = True
|
||||
if yue2_vae:
|
||||
self.audio_sample_rate = 48000
|
||||
self.upscale_ratio = self.downscale_ratio = 1920
|
||||
self.memory_used_decode = lambda shape, dtype: (1500 * shape[-1] * 1920) * model_management.dtype_size(dtype)
|
||||
elif "blocks.2.blocks.3.stack.5.weight" in sd or "decoder.blocks.2.blocks.3.stack.5.weight" in sd or "layers.4.layers.1.attn_block.attn.qkv.weight" in sd or "encoder.layers.4.layers.1.attn_block.attn.qkv.weight" in sd: #genmo mochi vae
|
||||
if "blocks.2.blocks.3.stack.5.weight" in sd:
|
||||
sd = comfy.utils.state_dict_prefix_replace(sd, {"": "decoder."})
|
||||
@@ -1552,6 +1561,7 @@ class CLIPType(Enum):
|
||||
JOYIMAGE = 33
|
||||
MAGE = 34
|
||||
MINIMAX = 35
|
||||
YUE2 = 36
|
||||
|
||||
|
||||
|
||||
@@ -1741,7 +1751,12 @@ def load_text_encoder_state_dicts(state_dicts=[], embedding_directory=None, clip
|
||||
clip_target.params = {}
|
||||
if len(clip_data) == 1:
|
||||
te_model = detect_te_model(clip_data[0])
|
||||
if clip_type == CLIPType.MINIMAX and "model.audio_decoder.projection.weight" in clip_data[0]:
|
||||
if clip_type == CLIPType.YUE2 and "yue2_tokenizer_json" in clip_data[0]:
|
||||
tokenizer_data["yue2_tokenizer_json"] = clip_data[0].pop("yue2_tokenizer_json")
|
||||
detect = comfy.text_encoders.hunyuan_video.llama_detect(clip_data[0])
|
||||
clip_target.clip = comfy.text_encoders.yue2.te(**detect)
|
||||
clip_target.tokenizer = comfy.text_encoders.yue2.YuE2Tokenizer
|
||||
elif clip_type == CLIPType.MINIMAX and "model.audio_decoder.projection.weight" in clip_data[0]:
|
||||
tokenizer_data["tokenizer_json"] = clip_data[0].pop("tokenizer_json", None)
|
||||
quant = comfy.utils.detect_layer_quantization(clip_data[0], "")
|
||||
if quant is not None:
|
||||
|
||||
@@ -33,6 +33,7 @@ import comfy.text_encoders.mage_flow
|
||||
import comfy.text_encoders.joyimage
|
||||
import comfy.text_encoders.anima
|
||||
import comfy.text_encoders.ace15
|
||||
import comfy.text_encoders.yue2
|
||||
import comfy.text_encoders.longcat_image
|
||||
import comfy.text_encoders.ernie
|
||||
import comfy.text_encoders.cogvideo
|
||||
@@ -2265,6 +2266,27 @@ class ACEStep15(supported_models_base.BASE):
|
||||
|
||||
return supported_models_base.ClipTarget(comfy.text_encoders.ace15.ACE15Tokenizer, comfy.text_encoders.ace15.te(**detect))
|
||||
|
||||
class YuE2(supported_models_base.BASE):
|
||||
unet_config = {"audio_model": "yue2"}
|
||||
unet_extra_config = {}
|
||||
latent_format = latent_formats.YuE2
|
||||
supported_inference_dtypes = [torch.bfloat16, torch.float32]
|
||||
sampling_settings = {"multiplier": 1.0}
|
||||
memory_usage_factor = 4.0
|
||||
vae_key_prefix = ["vae."]
|
||||
text_encoder_key_prefix = ["text_encoders."]
|
||||
|
||||
def get_model(self, state_dict, prefix="", device=None):
|
||||
return model_base.YuE2(self, device=device)
|
||||
|
||||
def model_type(self, state_dict, prefix=""):
|
||||
return model_base.ModelType.FLOW
|
||||
|
||||
def clip_target(self, state_dict={}):
|
||||
detect = comfy.text_encoders.hunyuan_video.llama_detect(state_dict, self.text_encoder_key_prefix[0])
|
||||
return supported_models_base.ClipTarget(comfy.text_encoders.yue2.YuE2Tokenizer, comfy.text_encoders.yue2.te(**detect))
|
||||
|
||||
|
||||
class MiniMaxMusic3(supported_models_base.BASE):
|
||||
unet_config = {
|
||||
"audio_model": "minimax_music3",
|
||||
@@ -2583,6 +2605,7 @@ models = [
|
||||
ACEStep,
|
||||
ACEStep15,
|
||||
MiniMaxMusic3,
|
||||
YuE2,
|
||||
Omnigen2,
|
||||
Boogu,
|
||||
MageFlow,
|
||||
|
||||
@@ -0,0 +1,266 @@
|
||||
"""YuE2 score/semantic generation and acoustic prefix conditioning."""
|
||||
|
||||
import logging
|
||||
|
||||
import torch
|
||||
from tokenizers import Tokenizer
|
||||
|
||||
import comfy.model_management
|
||||
import comfy.model_prefetch
|
||||
import comfy.ops
|
||||
import comfy.utils
|
||||
from comfy.ldm.yue2.model import model_config
|
||||
from comfy.text_encoders.llama import FixedKV, Llama2_
|
||||
|
||||
|
||||
EOD = 151643
|
||||
ABC_START, ABC_END = 151847, 151848
|
||||
MUSIC_START, MUSIC_END = 151851, 151852
|
||||
CODEC_OFFSET, CODEC_SIZE = 151853, 32768
|
||||
CONTEXT = 24576
|
||||
FRAMES_PER_SECOND = 25
|
||||
INSTRUCTIONS = {
|
||||
"off": "Generate music with codec tokens from the given conditions.",
|
||||
"melody": "Generate a melody-only ABC transcription without chord symbols, then generate music with codec tokens from the given conditions.",
|
||||
"full": "Generate a chord-annotated ABC transcription, then generate music with codec tokens from the given conditions.",
|
||||
}
|
||||
|
||||
|
||||
def distribution(logits, history, step, phase, temperature, top_p, top_k, repetition_penalty, penalty_window, min_tokens, legacy_off=False):
|
||||
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)
|
||||
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[..., 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)
|
||||
counts.scatter_add_(-1, recent, torch.ones_like(recent, dtype=scores.dtype))
|
||||
penalty = repetition_penalty ** counts
|
||||
scores = torch.where(scores < 0, scores * penalty, scores / penalty)
|
||||
if temperature == 0:
|
||||
return scores
|
||||
scores /= temperature
|
||||
threshold = scores.topk(min(top_k, scores.shape[-1])).values[..., -1, None]
|
||||
scores.masked_fill_(scores < threshold, -torch.inf)
|
||||
if top_p < 1:
|
||||
values, indices = scores.sort(descending=True)
|
||||
probabilities = values.softmax(-1)
|
||||
removed = probabilities.cumsum(-1) - probabilities > top_p
|
||||
removed[..., :3 if legacy_off else 1] = False
|
||||
values.masked_fill_(removed, -torch.inf)
|
||||
scores = values.scatter(-1, indices, values)
|
||||
return scores
|
||||
|
||||
|
||||
def chunk_ranges(frames, prefix_tokens, context=CONTEXT):
|
||||
size = (context - prefix_tokens - 3) // 2
|
||||
if frames < 1 or size < 1:
|
||||
raise ValueError("YuE2 needs music tokens and enough context for at least one acoustic frame.")
|
||||
return [(start, min(start + size, frames)) for start in range(0, frames, size)]
|
||||
|
||||
|
||||
class YuE2Tokenizer:
|
||||
def __init__(self, embedding_directory=None, tokenizer_data={}):
|
||||
data = tokenizer_data["yue2_tokenizer_json"]
|
||||
if torch.is_tensor(data):
|
||||
data = data.numpy().tobytes()
|
||||
self.tokenizer_json = data
|
||||
self.tokenizer = Tokenizer.from_str(data.decode("utf-8"))
|
||||
|
||||
def tokenize_with_weights(self, text, return_word_ids=False, **kwargs):
|
||||
cot = kwargs.get("cot", "full")
|
||||
prompt = f"{INSTRUCTIONS[cot]}\n[Tags]\n{text}\n[Lyrics]\n{kwargs.get('lyrics', '')}\n"
|
||||
return {
|
||||
"prefix": [EOD] + self.tokenizer.encode(prompt).ids + [ABC_START],
|
||||
"negative": [EOD] + self.tokenizer.encode(INSTRUCTIONS[cot]).ids,
|
||||
"abc_ids": self.tokenizer.encode(kwargs.get("abc", "")).ids,
|
||||
"cot": cot,
|
||||
"seed": kwargs.get("seed", 0),
|
||||
"max_tokens": kwargs.get("max_tokens", 9000),
|
||||
"temperature": kwargs.get("temperature", 1.0),
|
||||
"top_p": kwargs.get("top_p", 0.95),
|
||||
"top_k": kwargs.get("top_k", 100),
|
||||
"repetition_penalty": kwargs.get("repetition_penalty", 1.2),
|
||||
"cfg_scale": kwargs.get("cfg_scale", 1.01 if cot == "off" else 1.0),
|
||||
}
|
||||
|
||||
def state_dict(self):
|
||||
return {"yue2_tokenizer_json": torch.frombuffer(bytearray(self.tokenizer_json), dtype=torch.uint8)}
|
||||
|
||||
def decode(self, ids, skip_special_tokens=True):
|
||||
return self.tokenizer.decode(ids, skip_special_tokens=skip_special_tokens)
|
||||
|
||||
|
||||
class YuE2TEModel(torch.nn.Module):
|
||||
def __init__(self, device="cpu", dtype=None, model_options={}, config=None):
|
||||
super().__init__()
|
||||
self.config = model_config(**{"fixed_kv": True, **(config or {})})
|
||||
operations = model_options.get("custom_operations", comfy.ops.manual_cast)
|
||||
quant = model_options.get("quantization_metadata")
|
||||
if quant is not None and "custom_operations" not in model_options:
|
||||
operations = comfy.ops.mixed_precision_ops(quant, dtype)
|
||||
self.model = Llama2_(self.config, device=device, dtype=dtype, ops=operations)
|
||||
self.model.prefetch_dynamic_vbars = True
|
||||
self.model.graph_dynamic_vbar_blocks = True
|
||||
self.dtypes = {dtype}
|
||||
self.execution_device = device
|
||||
|
||||
def get_dynamic_vram__units(self):
|
||||
return self.model.get_dynamic_vram__units()
|
||||
|
||||
def set_clip_options(self, options):
|
||||
self.execution_device = options.get("execution_device", self.execution_device)
|
||||
|
||||
def reset_clip_options(self):
|
||||
pass
|
||||
|
||||
def load_sd(self, state_dict):
|
||||
return self.load_state_dict(state_dict, strict=False, assign=getattr(self, "can_assign_sd", False))
|
||||
|
||||
def memory_estimation_function(self, tokens, device=None):
|
||||
config = self.config
|
||||
abc_length = 0 if tokens["cot"] == "off" else len(tokens["abc_ids"])
|
||||
length = min(config.max_position_embeddings, len(tokens["prefix"]) + abc_length + tokens["max_tokens"] + 2)
|
||||
branches = 1 if tokens["cfg_scale"] == 1.0 else 2
|
||||
dtype = torch.bfloat16 if comfy.model_management.should_use_bf16(device) else torch.float32
|
||||
cache = branches * 2 * config.num_hidden_layers * config.num_key_value_heads * config.head_dim * length
|
||||
prefill = branches * (length * length + length * (config.intermediate_size * 3 + config.hidden_size * 8))
|
||||
return (cache + prefill) * comfy.model_management.dtype_size(dtype)
|
||||
|
||||
def _prefill(self, prefixes, capacity, dtype):
|
||||
length = max(map(len, prefixes))
|
||||
ids = torch.tensor([[0] * (length - len(prefix)) + prefix for prefix in prefixes], device=self.execution_device, dtype=torch.long)
|
||||
mask = positions = None
|
||||
if any(len(prefix) != length for prefix in prefixes):
|
||||
mask = torch.ones((len(prefixes), capacity), device=self.execution_device, dtype=torch.long)
|
||||
for index, prefix in enumerate(prefixes):
|
||||
mask[index, :length - len(prefix)] = 0
|
||||
positions = mask[:, :length].cumsum(-1).sub_(1).clamp_min_(0)
|
||||
cache = self.model.init_kv_cache(len(prefixes), capacity, self.execution_device, dtype)
|
||||
output = self.model(ids, attention_mask=mask[:, :length] if mask is not None else None,
|
||||
position_ids=positions, past_key_values=cache, dtype=dtype)
|
||||
return self.model.lm_head(output[0][:, -1]), output[2], mask
|
||||
|
||||
def _generate(self, prefix, seed, max_tokens, phase, dtype, negative=None, cfg_scale=1.0, legacy_off=False, **sampling):
|
||||
if max(len(prefix), len(negative or [])) + max_tokens > self.config.max_position_embeddings:
|
||||
raise ValueError("YuE2 prompt plus generation budget exceeds the model context; reduce the token budget or prompt length.")
|
||||
device = self.execution_device
|
||||
rng_device = device if torch.device(device).type != "mps" else "cpu"
|
||||
generator = torch.Generator(device=rng_device).manual_seed(seed)
|
||||
prefixes = [prefix] if cfg_scale == 1.0 else [prefix, negative]
|
||||
prefix_length = max(map(len, prefixes))
|
||||
logits, cache, mask = self._prefill(prefixes, prefix_length + max_tokens, dtype)
|
||||
fixed_kv = isinstance(cache[0], FixedKV)
|
||||
decode_tokens = torch.empty((len(prefixes), 1), device=device, dtype=torch.long)
|
||||
positions = torch.tensor([[len(p)] for p in prefixes], device=device, dtype=torch.long)
|
||||
history = []
|
||||
end = ABC_END if phase == "abc" else MUSIC_END
|
||||
progress = comfy.utils.ProgressBar(max_tokens)
|
||||
try:
|
||||
for step in comfy.utils.model_trange(max_tokens, desc="YuE2 ABC sampling" if phase == "abc" else "YuE2 music sampling", unit="token"):
|
||||
comfy.model_management.throw_exception_if_processing_interrupted()
|
||||
guided = logits if cfg_scale == 1.0 else logits[1:] + cfg_scale * (logits[:1] - logits[1:])
|
||||
scores = distribution(guided, history, step, phase, legacy_off=legacy_off, **sampling)
|
||||
if sampling["temperature"] == 0:
|
||||
next_id = scores.argmax(-1, keepdim=True)
|
||||
else:
|
||||
probabilities = scores.softmax(-1).to(rng_device)
|
||||
next_id = torch.multinomial(probabilities, 1, generator=generator).to(device)
|
||||
decode_tokens.copy_(next_id)
|
||||
token = next_id.item()
|
||||
progress.update_absolute(step + 1)
|
||||
if token == end:
|
||||
return history, False
|
||||
history.append(token)
|
||||
if step + 1 < max_tokens:
|
||||
# Keep decode allocations stable; sampling has a changing history window.
|
||||
if fixed_kv:
|
||||
comfy.model_prefetch.malloc_graph_begin(device)
|
||||
output = self.model(decode_tokens, past_key_values=cache, dtype=dtype, position_ids=positions,
|
||||
attention_mask=mask[:, :prefix_length + step + 1] if mask is not None and not fixed_kv else None)
|
||||
logits.copy_(self.model.lm_head(output[0][:, -1]))
|
||||
cache = output[2]
|
||||
del output
|
||||
if fixed_kv:
|
||||
comfy.model_prefetch.malloc_graph_end()
|
||||
positions.add_(1)
|
||||
finally:
|
||||
# Each phase has different KV buffers and may change the CFG batch size.
|
||||
comfy.model_prefetch.cleanup_prefetch_queues()
|
||||
logging.warning("YuE2 %s reached its token budget; increase the limit for a complete result.", phase)
|
||||
return history, True
|
||||
|
||||
def _acoustic_conditioning(self, prefix, tokens, dtype):
|
||||
config = self.config
|
||||
ranges = chunk_ranges(len(tokens), len(prefix), config.max_position_embeddings)
|
||||
total = sum(len(prefix) + end - start + 1 for start, end in ranges)
|
||||
# A normal [batch, tokens, features] conditioning tensor, with each layer's KV in features.
|
||||
output = torch.empty((1, total, config.num_hidden_layers, 2, config.num_key_value_heads, config.head_dim),
|
||||
device=comfy.model_management.intermediate_device(), dtype=dtype)
|
||||
chunks = []
|
||||
offset = 0
|
||||
for start, end in ranges:
|
||||
comfy.model_management.throw_exception_if_processing_interrupted()
|
||||
ids = prefix + tokens[start:end] + [MUSIC_END]
|
||||
_, cache, _ = self._prefill([ids], len(ids), dtype)
|
||||
for index, kv in enumerate(cache):
|
||||
if isinstance(kv, FixedKV):
|
||||
key, value = kv.key, kv.value
|
||||
else:
|
||||
key, value, _ = kv
|
||||
key, value = key.transpose(1, 2), value.transpose(1, 2)
|
||||
output[:, offset:offset + len(ids), index, 0].copy_(key)
|
||||
output[:, offset:offset + len(ids), index, 1].copy_(value)
|
||||
chunks.append((start, end, offset, offset + len(ids)))
|
||||
offset += len(ids)
|
||||
del cache
|
||||
return output.flatten(2), tuple(chunks)
|
||||
|
||||
def generate(self, tokens, do_sample=True, max_length=256, temperature=1.0, top_k=50, top_p=0.95, repetition_penalty=1.0, seed=None, **kwargs):
|
||||
dtype = torch.bfloat16 if comfy.model_management.should_use_bf16(self.execution_device) else torch.float32
|
||||
ids, _ = self._generate(
|
||||
tokens["prefix"], tokens["seed"] if seed is None else seed, max_length, "abc", dtype,
|
||||
temperature=temperature if do_sample else 0, top_p=top_p, top_k=top_k,
|
||||
repetition_penalty=repetition_penalty, penalty_window=100, min_tokens=min(32, max_length),
|
||||
)
|
||||
return ids
|
||||
|
||||
def encode_token_weights(self, tokens):
|
||||
device = self.execution_device
|
||||
dtype = torch.bfloat16 if comfy.model_management.should_use_bf16(device) else torch.float32
|
||||
prefix = tokens["prefix"]
|
||||
abc_ids = tokens["abc_ids"]
|
||||
cot = tokens["cot"]
|
||||
if cot == "off":
|
||||
abc_ids = []
|
||||
prefix = prefix + abc_ids + [ABC_END, MUSIC_START]
|
||||
negative = tokens["negative"] + ([MUSIC_START] if cot == "off" else [ABC_START] + abc_ids + [ABC_END, MUSIC_START])
|
||||
semantic, semantic_truncated = self._generate(
|
||||
prefix, tokens["seed"], tokens["max_tokens"], "semantic", dtype,
|
||||
negative=negative, cfg_scale=tokens["cfg_scale"], legacy_off=cot == "off",
|
||||
temperature=tokens["temperature"], top_p=tokens["top_p"], top_k=tokens["top_k"],
|
||||
repetition_penalty=tokens["repetition_penalty"], penalty_window=50,
|
||||
min_tokens=min(200, tokens["max_tokens"]),
|
||||
)
|
||||
conditioning, chunks = self._acoustic_conditioning(prefix, semantic, dtype)
|
||||
return conditioning, None, {
|
||||
"yue2_chunks": chunks, "yue2_abc_ids": abc_ids, "yue2_frames": len(semantic),
|
||||
"yue2_truncated": semantic_truncated,
|
||||
}
|
||||
|
||||
|
||||
def te(dtype_llama=None, llama_quantization_metadata=None):
|
||||
class YuE2TEModel_(YuE2TEModel):
|
||||
def __init__(self, device="cpu", dtype=None, model_options={}):
|
||||
dtype = comfy.model_management.pick_weight_dtype(dtype_llama, dtype, device)
|
||||
if llama_quantization_metadata is not None:
|
||||
model_options = {**model_options, "quantization_metadata": llama_quantization_metadata}
|
||||
super().__init__(device=device, dtype=dtype, model_options=model_options)
|
||||
return YuE2TEModel_
|
||||
@@ -50,12 +50,34 @@ class AudioEncoderEncode(io.ComfyNode):
|
||||
return io.NodeOutput(output)
|
||||
|
||||
|
||||
class SheetSage2AudioToABC(io.ComfyNode):
|
||||
@classmethod
|
||||
def define_schema(cls):
|
||||
return io.Schema(
|
||||
node_id="SheetSage2AudioToABC",
|
||||
display_name="SheetSage2 Audio to ABC",
|
||||
category="model/conditioning/yue2",
|
||||
description="Transcribes vocal and instrumental melodies from music into ABC notation. Connect abc output to YuE2 Generate Music node and use the matching mode.",
|
||||
inputs=[
|
||||
io.AudioEncoder.Input("audio_encoder"),
|
||||
io.Audio.Input("audio"),
|
||||
io.Combo.Input("mode", options=["melody", "full"], tooltip="full: generates melody and chords; melody: generates melody only, recommended for covers."),
|
||||
],
|
||||
outputs=[io.String.Output(display_name="abc", is_output_list=True)],
|
||||
)
|
||||
|
||||
@classmethod
|
||||
def execute(cls, audio_encoder, audio, mode):
|
||||
return io.NodeOutput(audio_encoder.generate_abc(audio["waveform"], audio["sample_rate"], melody_only=mode == "melody"))
|
||||
|
||||
|
||||
class AudioEncoder(ComfyExtension):
|
||||
@override
|
||||
async def get_node_list(self) -> list[type[io.ComfyNode]]:
|
||||
return [
|
||||
AudioEncoderLoader,
|
||||
AudioEncoderEncode,
|
||||
SheetSage2AudioToABC,
|
||||
]
|
||||
|
||||
|
||||
|
||||
@@ -0,0 +1,98 @@
|
||||
import torch
|
||||
from typing_extensions import override
|
||||
|
||||
import comfy.model_management
|
||||
from comfy.text_encoders.yue2 import FRAMES_PER_SECOND
|
||||
from comfy_api.latest import ComfyExtension, io
|
||||
|
||||
|
||||
class YuE2GenerateABC(io.ComfyNode):
|
||||
@classmethod
|
||||
def define_schema(cls):
|
||||
return io.Schema(
|
||||
node_id="YuE2GenerateABC",
|
||||
display_name="YuE2 Generate ABC",
|
||||
category="model/conditioning/yue2",
|
||||
description="Generates the ABC notation of a song from style and lyrics. Connect abc to YuE2 Generate Music node.",
|
||||
inputs=[
|
||||
io.Clip.Input("clip"),
|
||||
io.String.Input("style", multiline=True, dynamic_prompts=True),
|
||||
io.String.Input("lyrics", multiline=True, dynamic_prompts=True),
|
||||
io.Int.Input("seed", default=0, min=0, max=0xffffffffffffffff, control_after_generate=True),
|
||||
io.Combo.Input("mode", options=["full", "melody"], tooltip="full: generates melody and chords; melody: generates melody only, recommended for covers."),
|
||||
io.Int.Input("max_abc_tokens", default=8192, min=1, max=20000, advanced=True),
|
||||
],
|
||||
outputs=[io.String.Output(display_name="abc")],
|
||||
)
|
||||
|
||||
@classmethod
|
||||
def execute(cls, clip, style, lyrics, seed, mode, max_abc_tokens):
|
||||
tokens = clip.tokenize(style, lyrics=lyrics, cot=mode, seed=seed, max_tokens=max_abc_tokens)
|
||||
ids = clip.generate(tokens, max_length=max_abc_tokens, temperature=0.7, top_p=0.9, top_k=30, repetition_penalty=1.005, seed=seed)
|
||||
return io.NodeOutput(clip.decode(ids))
|
||||
|
||||
|
||||
class YuE2GenerateMusic(io.ComfyNode):
|
||||
@classmethod
|
||||
def define_schema(cls):
|
||||
return io.Schema(
|
||||
node_id="YuE2GenerateMusic",
|
||||
display_name="YuE2 Generate Music",
|
||||
category="model/conditioning/yue2",
|
||||
description="Generates music tokens and acoustic conditioning from style, lyrics, and an ABC notation. Provide the generated seconds to the Empty YuE2 Latent Audio node. An empty ABC input ignores the selected mode.",
|
||||
inputs=[
|
||||
io.Clip.Input("clip"),
|
||||
io.String.Input("style", multiline=True, dynamic_prompts=True),
|
||||
io.String.Input("lyrics", multiline=True, dynamic_prompts=True),
|
||||
io.String.Input("abc", default="", multiline=True, tooltip="Connect the ABC generator or supply an edited score. Leave empty to use off mode automatically."),
|
||||
io.Int.Input("seed", default=0, min=0, max=0xffffffffffffffff, control_after_generate=True),
|
||||
io.Combo.Input("mode", options=["full", "melody"], tooltip="full: generates melody and chords; melody: generates melody only, recommended for covers."),
|
||||
io.Float.Input("max_duration", default=360.0, min=0.04, max=360.0, step=0.04, tooltip="Maximum duration; generation can stop earlier. The release uses a 360-second budget."),
|
||||
io.Float.Input("temperature", default=1.0, min=0.0, max=5.0, step=0.05, advanced=True),
|
||||
io.Float.Input("top_p", default=0.95, min=0.01, max=1.0, step=0.01, advanced=True),
|
||||
io.Int.Input("top_k", default=100, min=1, max=32768, advanced=True),
|
||||
io.Float.Input("repetition_penalty", default=1.2, min=0.01, max=10.0, step=0.01, advanced=True),
|
||||
],
|
||||
outputs=[io.Conditioning.Output(), io.Float.Output(display_name="seconds")],
|
||||
)
|
||||
|
||||
@classmethod
|
||||
def execute(cls, clip, style, lyrics, seed, mode, max_duration, temperature, top_p, top_k, repetition_penalty, abc=""):
|
||||
if not abc.strip():
|
||||
mode = "off"
|
||||
tokens = clip.tokenize(style, lyrics=lyrics, cot=mode, seed=seed, abc=abc,
|
||||
max_tokens=max(1, round(max_duration * FRAMES_PER_SECOND)),
|
||||
temperature=temperature, top_p=top_p, top_k=top_k, repetition_penalty=repetition_penalty)
|
||||
conditioning = clip.encode_from_tokens_scheduled(tokens)
|
||||
return io.NodeOutput(conditioning, conditioning[0][1]["yue2_frames"] / FRAMES_PER_SECOND)
|
||||
|
||||
|
||||
class EmptyYuE2LatentAudio(io.ComfyNode):
|
||||
@classmethod
|
||||
def define_schema(cls):
|
||||
return io.Schema(
|
||||
node_id="EmptyYuE2LatentAudio",
|
||||
display_name="Empty YuE2 Latent Audio",
|
||||
category="model/latent/yue2",
|
||||
inputs=[
|
||||
io.Float.Input("seconds", default=120.0, min=0.04, max=1000.0, step=0.04),
|
||||
io.Int.Input("batch_size", default=1, min=1, max=4096),
|
||||
],
|
||||
outputs=[io.Latent.Output()],
|
||||
)
|
||||
|
||||
@classmethod
|
||||
def execute(cls, seconds, batch_size):
|
||||
latent = torch.zeros((batch_size, 64, max(1, round(seconds * FRAMES_PER_SECOND))),
|
||||
device=comfy.model_management.intermediate_device(), dtype=comfy.model_management.intermediate_dtype())
|
||||
return io.NodeOutput({"samples": latent, "type": "audio", "downscale_ratio_temporal": 1920})
|
||||
|
||||
|
||||
class YuE2Extension(ComfyExtension):
|
||||
@override
|
||||
async def get_node_list(self):
|
||||
return [YuE2GenerateABC, YuE2GenerateMusic, EmptyYuE2LatentAudio]
|
||||
|
||||
|
||||
async def comfy_entrypoint():
|
||||
return YuE2Extension()
|
||||
@@ -1009,7 +1009,7 @@ class CLIPLoader:
|
||||
@classmethod
|
||||
def INPUT_TYPES(s):
|
||||
return {"required": { "clip_name": (folder_paths.get_filename_list("text_encoders"), ),
|
||||
"type": (["stable_diffusion", "stable_cascade", "sd3", "stable_audio", "mochi", "ltxv", "pixart", "cosmos", "lumina2", "wan", "hidream", "chroma", "ace", "omnigen2", "qwen_image", "hunyuan_image", "flux2", "ovis", "longcat_image", "cogvideox", "lens", "pixeldit", "ideogram4", "boogu", "krea2", "joyimage", "mage", "minimax"], ),
|
||||
"type": (["stable_diffusion", "stable_cascade", "sd3", "stable_audio", "mochi", "ltxv", "pixart", "cosmos", "lumina2", "wan", "hidream", "chroma", "ace", "omnigen2", "qwen_image", "hunyuan_image", "flux2", "ovis", "longcat_image", "cogvideox", "lens", "pixeldit", "ideogram4", "boogu", "krea2", "joyimage", "mage", "minimax", "yue2"], ),
|
||||
},
|
||||
"optional": {
|
||||
"device": (["default", "cpu"], {"advanced": True}),
|
||||
@@ -2458,6 +2458,7 @@ async def init_builtin_extra_nodes():
|
||||
"nodes_lt_upsampler.py",
|
||||
"nodes_lt_audio.py",
|
||||
"nodes_minimax_music.py",
|
||||
"nodes_yue2.py",
|
||||
"nodes_minimax_h3.py",
|
||||
"nodes_lt.py",
|
||||
"nodes_lt_keyframes.py",
|
||||
|
||||
Reference in New Issue
Block a user