init

.gitignore

@@ -0,0 +1,13 @@
+output/
+*__pycache__/
+samples*/
+runs/
+checkpoints/
+master_ip
+logs/
+*.DS_Store
+.idea
+tools/
+.vscode/
+convert_*
+*.pt

__init__.py

@@ -0,0 +1,3 @@
+from .nodes import NODE_CLASS_MAPPINGS, NODE_DISPLAY_NAME_MAPPINGS
+
+__all__ = ["NODE_CLASS_MAPPINGS", "NODE_DISPLAY_NAME_MAPPINGS"]

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
+}

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)

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",
+}

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 = ""

requirements.txt

@@ -0,0 +1,3 @@
+librosa
+rotary_embedding_torch
+einops