From 004dd7b25c0b06c8b75ca77a59fbe04938dc0fe6 Mon Sep 17 00:00:00 2001 From: kijai <40791699+kijai@users.noreply.github.com> Date: Sat, 23 Aug 2025 20:29:44 +0300 Subject: [PATCH] init --- .gitignore | 13 + __init__.py | 3 + example_workflows/melband_example.json | 237 +++++++++++ model/mel_band_roformer.py | 545 +++++++++++++++++++++++++ nodes.py | 175 ++++++++ pyproject.toml | 15 + readme.md | 0 requirements.txt | 3 + 8 files changed, 991 insertions(+) create mode 100644 .gitignore create mode 100644 __init__.py create mode 100644 example_workflows/melband_example.json create mode 100644 model/mel_band_roformer.py create mode 100644 nodes.py create mode 100644 pyproject.toml create mode 100644 readme.md create mode 100644 requirements.txt diff --git a/.gitignore b/.gitignore new file mode 100644 index 0000000..d32eacc --- /dev/null +++ b/.gitignore @@ -0,0 +1,13 @@ +output/ +*__pycache__/ +samples*/ +runs/ +checkpoints/ +master_ip +logs/ +*.DS_Store +.idea +tools/ +.vscode/ +convert_* +*.pt \ No newline at end of file diff --git a/__init__.py b/__init__.py new file mode 100644 index 0000000..2e96bd6 --- /dev/null +++ b/__init__.py @@ -0,0 +1,3 @@ +from .nodes import NODE_CLASS_MAPPINGS, NODE_DISPLAY_NAME_MAPPINGS + +__all__ = ["NODE_CLASS_MAPPINGS", "NODE_DISPLAY_NAME_MAPPINGS"] \ No newline at end of file diff --git a/example_workflows/melband_example.json b/example_workflows/melband_example.json new file mode 100644 index 0000000..1f32eb3 --- /dev/null +++ b/example_workflows/melband_example.json @@ -0,0 +1,237 @@ +{ + "id": "8b7a9a57-2303-4ef5-9fc2-bf41713bd1fc", + "revision": 0, + "last_node_id": 324, + "last_link_id": 579, + "nodes": [ + { + "id": 318, + "type": "MelRoFormerSampler", + "pos": [ + 2710.03759765625, + -1660.3865966796875 + ], + "size": [ + 240.40696716308594, + 47.67810821533203 + ], + "flags": {}, + "order": 2, + "mode": 0, + "inputs": [ + { + "name": "model", + "type": "MELROFORMERMODEL", + "link": 566 + }, + { + "name": "audio", + "type": "AUDIO", + "link": 579 + } + ], + "outputs": [ + { + "name": "vocals", + "type": "AUDIO", + "links": [ + 569 + ] + }, + { + "name": "instruments", + "type": "AUDIO", + "links": [ + 575 + ] + } + ], + "properties": { + "Node name for S&R": "MelRoFormerSampler" + }, + "widgets_values": [] + }, + { + "id": 319, + "type": "PreviewAudio", + "pos": [ + 3017.6923828125, + -1672.13427734375 + ], + "size": [ + 383.5520935058594, + 88 + ], + "flags": {}, + "order": 3, + "mode": 0, + "inputs": [ + { + "name": "audio", + "type": "AUDIO", + "link": 569 + } + ], + "outputs": [], + "properties": { + "cnr_id": "comfy-core", + "ver": "0.3.51", + "Node name for S&R": "PreviewAudio" + }, + "widgets_values": [] + }, + { + "id": 323, + "type": "PreviewAudio", + "pos": [ + 3018.808349609375, + -1498.72900390625 + ], + "size": [ + 376.83953857421875, + 91.35621643066406 + ], + "flags": {}, + "order": 4, + "mode": 0, + "inputs": [ + { + "name": "audio", + "type": "AUDIO", + "link": 575 + } + ], + "outputs": [], + "properties": { + "cnr_id": "comfy-core", + "ver": "0.3.51", + "Node name for S&R": "PreviewAudio" + }, + "widgets_values": [] + }, + { + "id": 317, + "type": "LoadAudio", + "pos": [ + 2201.01220703125, + -1506.000244140625 + ], + "size": [ + 274.080078125, + 136 + ], + "flags": {}, + "order": 0, + "mode": 0, + "inputs": [], + "outputs": [ + { + "name": "AUDIO", + "type": "AUDIO", + "links": [ + 579 + ] + } + ], + "properties": { + "cnr_id": "comfy-core", + "ver": "0.3.51", + "Node name for S&R": "LoadAudio" + }, + "widgets_values": [ + "0321. Alphaville - Big In Japan.mp3", + null, + null + ] + }, + { + "id": 315, + "type": "MelRoFormerModelLoader", + "pos": [ + 2188.627197265625, + -1667.1629638671875 + ], + "size": [ + 458.6287536621094, + 63.593719482421875 + ], + "flags": {}, + "order": 1, + "mode": 0, + "inputs": [], + "outputs": [ + { + "name": "model", + "type": "MELROFORMERMODEL", + "links": [ + 566 + ] + } + ], + "properties": { + "Node name for S&R": "MelRoFormerModelLoader" + }, + "widgets_values": [ + "MelRoFormer\\MelBandRoformer_fp16.safetensors" + ] + } + ], + "links": [ + [ + 566, + 315, + 0, + 318, + 0, + "MELROFORMERMODEL" + ], + [ + 569, + 318, + 0, + 319, + 0, + "AUDIO" + ], + [ + 575, + 318, + 1, + 323, + 0, + "AUDIO" + ], + [ + 579, + 317, + 0, + 318, + 1, + "AUDIO" + ] + ], + "groups": [], + "config": {}, + "extra": { + "ds": { + "scale": 1.1918176537728358, + "offset": [ + -1914.6136194563167, + 1969.9978180568914 + ] + }, + "frontendVersion": "1.26.3", + "node_versions": { + "ComfyUI-WanVideoWrapper": "0a11c67a0c0062b534178920a0d6dcaa75e7b5fe", + "comfy-core": "0.3.43", + "audio-separation-nodes-comfyui": "31a4567726e035097cc2d1f767767908a6fda2ea", + "ComfyUI-KJNodes": "f7eb33abc80a2aded1b46dff0dd14d07856a7d50", + "comfyui-videohelpersuite": "a7ce59e381934733bfae03b1be029756d6ce936d" + }, + "VHS_latentpreview": true, + "VHS_latentpreviewrate": 0, + "VHS_MetadataImage": true, + "VHS_KeepIntermediate": true + }, + "version": 0.4 +} \ No newline at end of file diff --git a/model/mel_band_roformer.py b/model/mel_band_roformer.py new file mode 100644 index 0000000..57c123c --- /dev/null +++ b/model/mel_band_roformer.py @@ -0,0 +1,545 @@ +# https://github.com/KimberleyJensen/Mel-Band-Roformer-Vocal-Model/blob/main/models/mel_band_roformer/mel_band_roformer.py + +from functools import partial + +import torch +from torch import nn +from torch.nn import Module, ModuleList +import torch.nn.functional as F + +from rotary_embedding_torch import RotaryEmbedding + +from einops import rearrange, pack, unpack, reduce, repeat + +from librosa import filters + +# helper functions + +def exists(val): + return val is not None + + +def default(v, d): + return v if exists(v) else d + + +def pack_one(t, pattern): + return pack([t], pattern) + + +def unpack_one(t, ps, pattern): + return unpack(t, ps, pattern)[0] + + +def pad_at_dim(t, pad, dim=-1, value=0.): + dims_from_right = (- dim - 1) if dim < 0 else (t.ndim - dim - 1) + zeros = ((0, 0) * dims_from_right) + return F.pad(t, (*zeros, *pad), value=value) + + +# norm + +class RMSNorm(Module): + def __init__(self, dim): + super().__init__() + self.scale = dim ** 0.5 + self.gamma = nn.Parameter(torch.ones(dim)) + + def forward(self, x): + return F.normalize(x, dim=-1) * self.scale * self.gamma + + +# attention + +class FeedForward(Module): + def __init__( + self, + dim, + mult=4, + dropout=0. + ): + super().__init__() + dim_inner = int(dim * mult) + self.net = nn.Sequential( + RMSNorm(dim), + nn.Linear(dim, dim_inner), + nn.GELU(), + nn.Dropout(dropout), + nn.Linear(dim_inner, dim), + nn.Dropout(dropout) + ) + + def forward(self, x): + return self.net(x) + + +class Attention(Module): + def __init__( + self, + dim, + heads=8, + dim_head=64, + dropout=0., + rotary_embed=None, + ): + super().__init__() + self.heads = heads + self.scale = dim_head ** -0.5 + dim_inner = heads * dim_head + + self.rotary_embed = rotary_embed + + self.attend = F.scaled_dot_product_attention + + self.norm = RMSNorm(dim) + self.to_qkv = nn.Linear(dim, dim_inner * 3, bias=False) + + self.to_gates = nn.Linear(dim, heads) + + self.to_out = nn.Sequential( + nn.Linear(dim_inner, dim, bias=False), + nn.Dropout(dropout) + ) + + def forward(self, x): + x = self.norm(x) + + q, k, v = rearrange(self.to_qkv(x), 'b n (qkv h d) -> qkv b h n d', qkv=3, h=self.heads) + + if exists(self.rotary_embed): + q = self.rotary_embed.rotate_queries_or_keys(q) + k = self.rotary_embed.rotate_queries_or_keys(k) + + out = self.attend(q, k, v) + + gates = self.to_gates(x) + out = out * rearrange(gates, 'b n h -> b h n 1').sigmoid() + + out = rearrange(out, 'b h n d -> b n (h d)') + return self.to_out(out) + + +class Transformer(Module): + def __init__( + self, + *, + dim, + depth, + dim_head=64, + heads=8, + attn_dropout=0., + ff_dropout=0., + ff_mult=4, + norm_output=True, + rotary_embed=None, + flash_attn=True + ): + super().__init__() + self.layers = ModuleList([]) + + for _ in range(depth): + self.layers.append(ModuleList([ + Attention(dim=dim, dim_head=dim_head, heads=heads, dropout=attn_dropout, rotary_embed=rotary_embed), + FeedForward(dim=dim, mult=ff_mult, dropout=ff_dropout) + ])) + + self.norm = RMSNorm(dim) if norm_output else nn.Identity() + + def forward(self, x): + + for attn, ff in self.layers: + x = attn(x) + x + x = ff(x) + x + + return self.norm(x) + + +# bandsplit module + +class BandSplit(Module): + def __init__( + self, + dim, + dim_inputs + ): + super().__init__() + self.dim_inputs = dim_inputs + self.to_features = ModuleList([]) + + for dim_in in dim_inputs: + net = nn.Sequential( + RMSNorm(dim_in), + nn.Linear(dim_in, dim) + ) + + self.to_features.append(net) + + def forward(self, x): + x = x.split(self.dim_inputs, dim=-1) + + outs = [] + for split_input, to_feature in zip(x, self.to_features): + split_output = to_feature(split_input) + outs.append(split_output) + + return torch.stack(outs, dim=-2) + + +def MLP( + dim_in, + dim_out, + dim_hidden=None, + depth=1, + activation=nn.Tanh +): + dim_hidden = default(dim_hidden, dim_in) + + net = [] + dims = (dim_in, *((dim_hidden,) * depth), dim_out) + + for ind, (layer_dim_in, layer_dim_out) in enumerate(zip(dims[:-1], dims[1:])): + is_last = ind == (len(dims) - 2) + + net.append(nn.Linear(layer_dim_in, layer_dim_out)) + + if is_last: + continue + + net.append(activation()) + + return nn.Sequential(*net) + + +class MaskEstimator(Module): + def __init__( + self, + dim, + dim_inputs, + depth, + mlp_expansion_factor=4 + ): + super().__init__() + self.dim_inputs = dim_inputs + self.to_freqs = ModuleList([]) + dim_hidden = dim * mlp_expansion_factor + + for dim_in in dim_inputs: + net = [] + + mlp = nn.Sequential( + MLP(dim, dim_in * 2, dim_hidden=dim_hidden, depth=depth), + nn.GLU(dim=-1) + ) + + self.to_freqs.append(mlp) + + def forward(self, x): + x = x.unbind(dim=-2) + + outs = [] + + for band_features, mlp in zip(x, self.to_freqs): + freq_out = mlp(band_features) + outs.append(freq_out) + + return torch.cat(outs, dim=-1) + + +# main class + +class MelBandRoformer(Module): + def __init__( + self, + dim, + *, + depth, + stereo=False, + num_stems=1, + time_transformer_depth=2, + freq_transformer_depth=2, + num_bands=60, + dim_head=64, + heads=8, + attn_dropout=0.1, + ff_dropout=0.1, + flash_attn=True, + dim_freqs_in=1025, + sample_rate=44100, # needed for mel filter bank from librosa + stft_n_fft=2048, + stft_hop_length=512, + # 10ms at 44100Hz, from sections 4.1, 4.4 in the paper - @faroit recommends // 2 or // 4 for better reconstruction + stft_win_length=2048, + stft_normalized=False, + stft_window_fn = None, + mask_estimator_depth=1, + multi_stft_resolution_loss_weight=1., + multi_stft_resolutions_window_sizes = (4096, 2048, 1024, 512, 256), + multi_stft_hop_size=147, + multi_stft_normalized=False, + multi_stft_window_fn = torch.hann_window, + match_input_audio_length=False, # if True, pad output tensor to match length of input tensor + ): + super().__init__() + + self.stereo = stereo + self.audio_channels = 2 if stereo else 1 + self.num_stems = num_stems + + self.layers = ModuleList([]) + + transformer_kwargs = dict( + dim=dim, + heads=heads, + dim_head=dim_head, + attn_dropout=attn_dropout, + ff_dropout=ff_dropout, + flash_attn=flash_attn + ) + + time_rotary_embed = RotaryEmbedding(dim=dim_head) + freq_rotary_embed = RotaryEmbedding(dim=dim_head) + + for _ in range(depth): + self.layers.append(nn.ModuleList([ + Transformer(depth=time_transformer_depth, rotary_embed=time_rotary_embed, **transformer_kwargs), + Transformer(depth=freq_transformer_depth, rotary_embed=freq_rotary_embed, **transformer_kwargs) + ])) + + self.stft_window_fn = partial(default(stft_window_fn, torch.hann_window), stft_win_length) + + self.stft_kwargs = dict( + n_fft=stft_n_fft, + hop_length=stft_hop_length, + win_length=stft_win_length, + normalized=stft_normalized + ) + + freqs = torch.stft(torch.randn(1, 4096), **self.stft_kwargs, return_complex=True).shape[1] + + # create mel filter bank + # with librosa.filters.mel as in section 2 of paper + + mel_filter_bank_numpy = filters.mel(sr=sample_rate, n_fft=stft_n_fft, n_mels=num_bands) + + mel_filter_bank = torch.from_numpy(mel_filter_bank_numpy) + + # for some reason, it doesn't include the first freq? just force a value for now + + mel_filter_bank[0][0] = 1. + + # In some systems/envs we get 0.0 instead of ~1.9e-18 in the last position, + # so let's force a positive value + + mel_filter_bank[-1, -1] = 1. + + # binary as in paper (then estimated masks are averaged for overlapping regions) + + freqs_per_band = mel_filter_bank > 0 + assert freqs_per_band.any(dim=0).all(), 'all frequencies need to be covered by all bands for now' + + repeated_freq_indices = repeat(torch.arange(freqs), 'f -> b f', b=num_bands) + freq_indices = repeated_freq_indices[freqs_per_band] + + if stereo: + freq_indices = repeat(freq_indices, 'f -> f s', s=2) + freq_indices = freq_indices * 2 + torch.arange(2) + freq_indices = rearrange(freq_indices, 'f s -> (f s)') + + self.register_buffer('freq_indices', freq_indices, persistent=False) + self.register_buffer('freqs_per_band', freqs_per_band, persistent=False) + + num_freqs_per_band = reduce(freqs_per_band, 'b f -> b', 'sum') + num_bands_per_freq = reduce(freqs_per_band, 'b f -> f', 'sum') + + self.register_buffer('num_freqs_per_band', num_freqs_per_band, persistent=False) + self.register_buffer('num_bands_per_freq', num_bands_per_freq, persistent=False) + + # band split and mask estimator + + freqs_per_bands_with_complex = tuple(2 * f * self.audio_channels for f in num_freqs_per_band.tolist()) + + self.band_split = BandSplit( + dim=dim, + dim_inputs=freqs_per_bands_with_complex + ) + + self.mask_estimators = nn.ModuleList([]) + + for _ in range(num_stems): + mask_estimator = MaskEstimator( + dim=dim, + dim_inputs=freqs_per_bands_with_complex, + depth=mask_estimator_depth + ) + + self.mask_estimators.append(mask_estimator) + + # for the multi-resolution stft loss + + self.multi_stft_resolution_loss_weight = multi_stft_resolution_loss_weight + self.multi_stft_resolutions_window_sizes = multi_stft_resolutions_window_sizes + self.multi_stft_n_fft = stft_n_fft + self.multi_stft_window_fn = multi_stft_window_fn + + self.multi_stft_kwargs = dict( + hop_length=multi_stft_hop_size, + normalized=multi_stft_normalized + ) + + self.match_input_audio_length = match_input_audio_length + + def forward( + self, + raw_audio, + target=None, + return_loss_breakdown=False + ): + """ + einops + + b - batch + f - freq + t - time + s - audio channel (1 for mono, 2 for stereo) + n - number of 'stems' + c - complex (2) + d - feature dimension + """ + + device = raw_audio.device + + if raw_audio.ndim == 2: + raw_audio = rearrange(raw_audio, 'b t -> b 1 t') + + batch, channels, raw_audio_length = raw_audio.shape + + istft_length = raw_audio_length if self.match_input_audio_length else None + + assert (not self.stereo and channels == 1) or ( + self.stereo and channels == 2), 'stereo needs to be set to True if passing in audio signal that is stereo (channel dimension of 2). also need to be False if mono (channel dimension of 1)' + + # to stft + + raw_audio, batch_audio_channel_packed_shape = pack_one(raw_audio, '* t') + + stft_window = self.stft_window_fn(device=device) + + stft_repr = torch.stft(raw_audio, **self.stft_kwargs, window=stft_window, return_complex=True) + stft_repr = torch.view_as_real(stft_repr) + + stft_repr = unpack_one(stft_repr, batch_audio_channel_packed_shape, '* f t c') + stft_repr = rearrange(stft_repr, + 'b s f t c -> b (f s) t c') # merge stereo / mono into the frequency, with frequency leading dimension, for band splitting + + # index out all frequencies for all frequency ranges across bands ascending in one go + + batch_arange = torch.arange(batch, device=device)[..., None] + + # account for stereo + + x = stft_repr[batch_arange, self.freq_indices] + + # fold the complex (real and imag) into the frequencies dimension + + x = rearrange(x, 'b f t c -> b t (f c)') + + x = self.band_split(x) + + # axial / hierarchical attention + + for time_transformer, freq_transformer in self.layers: + x = rearrange(x, 'b t f d -> b f t d') + x, ps = pack([x], '* t d') + + x = time_transformer(x) + + x, = unpack(x, ps, '* t d') + x = rearrange(x, 'b f t d -> b t f d') + x, ps = pack([x], '* f d') + + x = freq_transformer(x) + + x, = unpack(x, ps, '* f d') + + num_stems = len(self.mask_estimators) + + masks = torch.stack([fn(x) for fn in self.mask_estimators], dim=1) + masks = rearrange(masks, 'b n t (f c) -> b n f t c', c=2) + + # modulate frequency representation + + stft_repr = rearrange(stft_repr, 'b f t c -> b 1 f t c') + + # complex number multiplication + + stft_repr = torch.view_as_complex(stft_repr) + masks = torch.view_as_complex(masks) + + masks = masks.type(stft_repr.dtype) + + # need to average the estimated mask for the overlapped frequencies + + scatter_indices = repeat(self.freq_indices, 'f -> b n f t', b=batch, n=num_stems, t=stft_repr.shape[-1]) + + stft_repr_expanded_stems = repeat(stft_repr, 'b 1 ... -> b n ...', n=num_stems) + masks_summed = torch.zeros_like(stft_repr_expanded_stems).scatter_add_(2, scatter_indices, masks) + + denom = repeat(self.num_bands_per_freq, 'f -> (f r) 1', r=channels) + + masks_averaged = masks_summed / denom.clamp(min=1e-8) + + # modulate stft repr with estimated mask + + stft_repr = stft_repr * masks_averaged + + # istft + + stft_repr = rearrange(stft_repr, 'b n (f s) t -> (b n s) f t', s=self.audio_channels) + + recon_audio = torch.istft(stft_repr, **self.stft_kwargs, window=stft_window, return_complex=False, + length=istft_length) + + recon_audio = rearrange(recon_audio, '(b n s) t -> b n s t', b=batch, s=self.audio_channels, n=num_stems) + + if num_stems == 1: + recon_audio = rearrange(recon_audio, 'b 1 s t -> b s t') + + # if a target is passed in, calculate loss for learning + + if not exists(target): + return recon_audio + + if self.num_stems > 1: + assert target.ndim == 4 and target.shape[1] == self.num_stems + + if target.ndim == 2: + target = rearrange(target, '... t -> ... 1 t') + + target = target[..., :recon_audio.shape[-1]] # protect against lost length on istft + + loss = F.l1_loss(recon_audio, target) + + multi_stft_resolution_loss = 0. + + for window_size in self.multi_stft_resolutions_window_sizes: + res_stft_kwargs = dict( + n_fft=max(window_size, self.multi_stft_n_fft), # not sure what n_fft is across multi resolution stft + win_length=window_size, + return_complex=True, + window=self.multi_stft_window_fn(window_size, device=device), + **self.multi_stft_kwargs, + ) + + recon_Y = torch.stft(rearrange(recon_audio, '... s t -> (... s) t'), **res_stft_kwargs) + target_Y = torch.stft(rearrange(target, '... s t -> (... s) t'), **res_stft_kwargs) + + multi_stft_resolution_loss = multi_stft_resolution_loss + F.l1_loss(recon_Y, target_Y) + + weighted_multi_resolution_loss = multi_stft_resolution_loss * self.multi_stft_resolution_loss_weight + + total_loss = loss + weighted_multi_resolution_loss + + if not return_loss_breakdown: + return total_loss + + return total_loss, (loss, multi_stft_resolution_loss) diff --git a/nodes.py b/nodes.py new file mode 100644 index 0000000..3239090 --- /dev/null +++ b/nodes.py @@ -0,0 +1,175 @@ +import os +import torch +import torch.nn.functional as F +from tqdm import tqdm + +import librosa +import folder_paths + +from .model.mel_band_roformer import MelBandRoformer + +script_directory = os.path.dirname(os.path.abspath(__file__)) + +from comfy import model_management as mm +from comfy.utils import load_torch_file, ProgressBar +device = mm.get_torch_device() +offload_device = mm.unet_offload_device() + +def get_windowing_array(window_size, fade_size, device): + fadein = torch.linspace(0, 1, fade_size) + fadeout = torch.linspace(1, 0, fade_size) + window = torch.ones(window_size) + window[-fade_size:] *= fadeout + window[:fade_size] *= fadein + return window.to(device) + +class MelBandRoFormerModelLoader: + @classmethod + def INPUT_TYPES(s): + return { + "required": { + "model_name": (folder_paths.get_filename_list("diffusion_models"), {"tooltip": "These models are loaded from the 'ComfyUI/models/diffusion_models' -folder",}), + }, + } + + RETURN_TYPES = ("MELROFORMERMODEL",) + RETURN_NAMES = ("model", ) + FUNCTION = "loadmodel" + CATEGORY = "Mel-Band RoFormer" + + def loadmodel(self, model_name): + model_config = { + "dim": 384, + "depth": 6, + "stereo": True, + "num_stems": 1, + "time_transformer_depth": 1, + "freq_transformer_depth": 1, + "num_bands": 60, + "dim_head": 64, + "heads": 8, + "attn_dropout": 0, + "ff_dropout": 0, + "flash_attn": True, + "dim_freqs_in": 1025, + "sample_rate": 44100, # needed for mel filter bank from librosa + "stft_n_fft": 2048, + "stft_hop_length": 441, + "stft_win_length": 2048, + "stft_normalized": False, + "mask_estimator_depth": 2, + "multi_stft_resolution_loss_weight": 1.0, + "multi_stft_resolutions_window_sizes": (4096, 2048, 1024, 512, 256), + "multi_stft_hop_size": 147, + "multi_stft_normalized": False, + } + model = MelBandRoformer(**model_config).eval() + model_path = folder_paths.get_full_path_or_raise("diffusion_models", model_name) + model.load_state_dict(load_torch_file(model_path), strict=True) + + return (model,) + +class MelBandRoFormerSampler: + @classmethod + def INPUT_TYPES(s): + return { + "required": { + "model": ("MELROFORMERMODEL",), + "audio": ("AUDIO",), + }, + } + + RETURN_TYPES = ("AUDIO","AUDIO",) + RETURN_NAMES = ("vocals", "instruments") + FUNCTION = "process" + CATEGORY = "Mel-Band RoFormer" + + def process(self, model, audio): + + audio_input = audio["waveform"] + sample_rate = audio["sample_rate"] + + B, audio_channels, audio_length = audio_input.shape + + sr = 44100 + + if audio_channels == 1: + raise NotImplementedError("Mono input is not supported yet.") + + if sample_rate != sr: + print(f"Resampling input {sample_rate} to {sr}") + audio_np = audio_input.cpu().numpy() + resampled = librosa.resample(audio_np, orig_sr=sample_rate, target_sr=sr, axis=-1) + audio_input = torch.from_numpy(resampled) + audio_input = original_audio = audio_input[0] + + C = 352800 + N = 2 + step = C // N + fade_size = C // 10 + border = C - step + + if audio_length > 2 * border and border > 0: + audio_input = F.pad(audio_input, (border, border), mode='reflect') + + windowing_array = get_windowing_array(C, fade_size, device) + + + audio_input = audio_input.to(device) + vocals = torch.zeros_like(audio_input, dtype=torch.float32).to(device) + counter = torch.zeros_like(audio_input, dtype=torch.float32).to(device) + + total_length = audio_input.shape[1] + num_chunks = (total_length + step - 1) // step + + model.to(device) + + comfy_pbar = ProgressBar(num_chunks) + + for i in tqdm(range(0, total_length, step), desc="Processing chunks"): + part = audio_input[:, i:i + C] + length = part.shape[-1] + if length < C: + if length > C // 2 + 1: + part = F.pad(input=part, pad=(0, C - length), mode='reflect') + else: + part = F.pad(input=part, pad=(0, C - length, 0, 0), mode='constant', value=0) + + x = model(part.unsqueeze(0))[0] + + window = windowing_array.clone() + if i == 0: + window[:fade_size] = 1 + elif i + C >= total_length: + window[-fade_size:] = 1 + + vocals[..., i:i+length] += x[..., :length] * window[..., :length] + counter[..., i:i+length] += window[..., :length] + comfy_pbar.update(1) + + model.to(offload_device) + + estimated_sources = vocals / counter + + if audio_length > 2 * border and border > 0: + estimated_sources = estimated_sources[..., border:-border] + + vocals_out = { + "waveform": estimated_sources.unsqueeze(0).cpu(), + "sample_rate": sr, + } + instruments_out = { + "waveform": (original_audio.to(device) - estimated_sources).unsqueeze(0).cpu(), + "sample_rate": sr, + } + + return (vocals_out, instruments_out) + +NODE_CLASS_MAPPINGS = { + "MelBandRoFormerModelLoader": MelBandRoFormerModelLoader, + "MelBandRoFormerSampler": MelBandRoFormerSampler, +} +NODE_DISPLAY_NAME_MAPPINGS = { + "MelBandRoFormerModelLoader": "Mel-Band RoFormer Model Loader", + "MelBandRoFormerSampler": "Mel-Band RoFormer Sampler", +} diff --git a/pyproject.toml b/pyproject.toml new file mode 100644 index 0000000..edb8a90 --- /dev/null +++ b/pyproject.toml @@ -0,0 +1,15 @@ +[project] +name = "ComfyUI-MelRoFormer" +description = "ComfyUI wrapper nodes for WanVideo" +version = "1.0.0" +license = {file = "LICENSE"} +dependencies = ["librosa", "rotary_embedding_torch", "einops"] + +[project.urls] +Repository = "https://github.com/kijai/ComfyUI-MelRoFormer" +# Used by Comfy Registry https://comfyregistry.org + +[tool.comfy] +PublisherId = "kijai" +DisplayName = "ComfyUI-MelRoFormer" +Icon = "" diff --git a/readme.md b/readme.md new file mode 100644 index 0000000..e69de29 diff --git a/requirements.txt b/requirements.txt new file mode 100644 index 0000000..639a040 --- /dev/null +++ b/requirements.txt @@ -0,0 +1,3 @@ +librosa +rotary_embedding_torch +einops \ No newline at end of file