Hybrid Demucs model implementation (#2506)

Summary:
Draft PR with initial model implementation with minor changes from previous implementation

Pull Request resolved: https://github.com/pytorch/audio/pull/2506

Reviewed By: nateanl

Differential Revision: D37762671

Pulled By: skim0514

fbshipit-source-id: b7dc0a6ef725d6ae6d76c23c882623f7d339977c

docs/source/prototype.models.rst

@@ -28,6 +28,13 @@ ConvEmformer
 
   .. automethod:: infer
 
+HDemucs
+~~~~~~~
+
+.. autoclass:: HDemucs
+
+  .. automethod:: forward
+
 References
 ~~~~~~~~~~
 

docs/source/refs.bib

@@ -382,3 +382,9 @@
   journal={arXiv preprint arXiv:1706.08612},
   year={2017}
 }
+@inproceedings{defossez2021hybrid,
+  title={Hybrid Spectrogram and Waveform Source Separation},
+  author={D{\'e}fossez, Alexandre},
+  booktitle={Proceedings of the ISMIR 2021 Workshop on Music Source Separation},
+  year={2021}
+}

test/torchaudio_unittest/prototype/hdemucs_cpu_test.py

+import torch
+from torchaudio_unittest.common_utils import PytorchTestCase
+from torchaudio_unittest.prototype.hdemucs_test_impl import HDemucsTests
+
+
+class HDemucsFloat32CPUTest(HDemucsTests, PytorchTestCase):
+    dtype = torch.float32
+    device = torch.device("cpu")

test/torchaudio_unittest/prototype/hdemucs_gpu_test.py

+import torch
+from torchaudio_unittest.common_utils import PytorchTestCase, skipIfNoCuda
+from torchaudio_unittest.prototype.hdemucs_test_impl import HDemucsTests
+
+
+@skipIfNoCuda
+class HDemucsFloat32GPUTest(HDemucsTests, PytorchTestCase):
+    dtype = torch.float32
+    device = torch.device("cuda")

test/torchaudio_unittest/prototype/hdemucs_test_impl.py

+import torch
+from parameterized import parameterized
+from torchaudio.prototype.models.hdemucs import _HDecLayer, _HEncLayer, HDemucs
+from torchaudio_unittest.common_utils import TestBaseMixin
+
+
+def _get_hdemucs_model(sources):
+    return HDemucs(sources)
+
+
+class HDemucsTests(TestBaseMixin):
+    def _get_inputs(self, duration: int, channels: int, batch_size: int, sample_rate: int):
+        sample = torch.rand(batch_size, channels, duration * sample_rate, dtype=torch.float32, device=self.device)
+        return sample
+
+    @parameterized.expand(
+        [
+            (["bass", "drums", "other", "vocals"],),
+            (["bass", "drums", "other"],),
+            (["bass", "vocals"],),
+            (["vocals"],),
+        ]
+    )
+    def test_hdemucs_output_shape(self, sources):
+        r"""Feed tensors with specific shape to HDemucs and validate
+        that it outputs with a tensor with expected shape.
+        """
+        batch_size = 1
+        duration = 10
+        channels = 2
+        sample_rate = 44100
+
+        model = _get_hdemucs_model(sources).to(self.device).eval()
+        inputs = self._get_inputs(duration, channels, batch_size, sample_rate)
+
+        split_sample = model(inputs)
+
+        assert split_sample.shape == (batch_size, len(sources), channels, duration * sample_rate)
+
+    def test_encoder_output_shape_frequency(self):
+        r"""Feed tensors with specific shape to HDemucs Decoder and validate
+        that it outputs with a tensor with expected shape for frequency domain.
+        """
+        batch_size = 1
+        chin, chout = 4, 48
+        f_bins = 2048
+        t = 800
+        stride = 4
+
+        model = _HEncLayer(chin, chout).to(self.device).eval()
+
+        x = torch.rand(batch_size, chin, f_bins, t, device=self.device, dtype=self.dtype)
+        out = model(x)
+
+        assert out.size() == (batch_size, chout, f_bins / stride, t)
+
+    def test_decoder_output_shape_frequency(self):
+        r"""Feed tensors with specific shape to HDemucs Decoder and validate
+        that it outputs with a tensor with expected shape for frequency domain.
+        """
+        batch_size = 1
+        chin, chout = 96, 48
+        f_bins = 128
+        t = 800
+        stride = 4
+
+        model = _HDecLayer(chin, chout).to(self.device).eval()
+
+        x = torch.rand(batch_size, chin, f_bins, t, device=self.device, dtype=self.dtype)
+        skip = torch.rand(batch_size, chin, f_bins, t, device=self.device, dtype=self.dtype)
+        z, y = model(x, skip, t)
+
+        assert z.size() == (batch_size, chout, f_bins * stride, t)
+        assert y.size() == (batch_size, chin, f_bins, t)
+
+    def test_encoder_output_shape_time(self):
+        r"""Feed tensors with specific shape to HDemucs Decoder and validate
+        that it outputs with a tensor with expected shape for time domain.
+        """
+        batch_size = 1
+        chin, chout = 4, 48
+        t = 800
+        stride = 4
+
+        model = _HEncLayer(chin, chout, freq=False).to(self.device).eval()
+
+        x = torch.rand(batch_size, chin, t, device=self.device, dtype=self.dtype)
+        out = model(x)
+
+        assert out.size() == (batch_size, chout, t / stride)
+
+    def test_decoder_output_shape_time(self):
+        r"""Feed tensors with specific shape to HDemucs Decoder and validate
+        that it outputs with a tensor with expected shape for time domain.
+        """
+        batch_size = 1
+        chin, chout = 96, 48
+        t = 800
+        stride = 4
+
+        model = _HDecLayer(chin, chout, freq=False).to(self.device).eval()
+
+        x = torch.rand(batch_size, chin, t, device=self.device, dtype=self.dtype)
+        skip = torch.rand(batch_size, chin, t, device=self.device, dtype=self.dtype)
+        z, y = model(x, skip, t * stride)
+
+        assert z.size() == (batch_size, chout, t * stride)
+        assert y.size() == (batch_size, chin, t)

torchaudio/prototype/models/__init__.py

 from .conv_emformer import ConvEmformer
 from .conv_tasnet import conv_tasnet_base
+from .hdemucs import HDemucs
 from .rnnt import conformer_rnnt_base, conformer_rnnt_model
 
 __all__ = [
@@ -7,4 +8,5 @@ __all__ = [
     "conformer_rnnt_model",
     "conv_tasnet_base",
     "ConvEmformer",
+    "HDemucs",
 ]

torchaudio/prototype/models/hdemucs.py

+# *****************************************************************************
+# MIT License
+#
+# Copyright (c) Facebook, Inc. and its affiliates.
+#
+# Permission is hereby granted, free of charge, to any person obtaining a copy
+# of this software and associated documentation files (the "Software"), to deal
+# in the Software without restriction, including without limitation the rights
+# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+# copies of the Software, and to permit persons to whom the Software is
+# furnished to do so, subject to the following conditions:
+#
+# The above copyright notice and this permission notice shall be included in all
+# copies or substantial portions of the Software.
+#
+# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+# SOFTWARE.
+# *****************************************************************************
+
+
+import math
+import typing as tp
+from typing import Any, Dict, List, Optional
+
+import torch
+from torch import nn
+from torch.nn import functional as F
+
+
+class _ScaledEmbedding(torch.nn.Module):
+    r"""Make continuous embeddings and boost learning rate
+
+    Args:
+        num_embeddings (int): number of embeddings
+        embedding_dim (int): embedding dimensions
+        scale (float, optional): amount to scale learning rate (Default: 10.0)
+        smooth (bool, optional): choose to apply smoothing (Default: ``False``)
+    """
+
+    def __init__(self, num_embeddings: int, embedding_dim: int, scale: float = 10.0, smooth: bool = False):
+        super().__init__()
+        self.embedding = nn.Embedding(num_embeddings, embedding_dim)
+        if smooth:
+            weight = torch.cumsum(self.embedding.weight.data, dim=0)
+            # when summing gaussian, scale raises as sqrt(n), so we normalize by that.
+            weight = weight / torch.arange(1, num_embeddings + 1).sqrt()[:, None]
+            self.embedding.weight.data[:] = weight
+        self.embedding.weight.data /= scale
+        self.scale = scale
+
+    @property
+    def weight(self) -> torch.Tensor:
+        return self.embedding.weight * self.scale
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        r"""Forward pass for embedding with scale.
+        Args:
+            x (torch.Tensor): input tensor of shape `(num_embeddings)`
+
+        Returns:
+            (Tensor):
+                Embedding output of shape `(num_embeddings, embedding_dim)`
+        """
+        out = self.embedding(x) * self.scale
+        return out
+
+
+class _HEncLayer(torch.nn.Module):
+
+    r"""Encoder layer. This used both by the time and the frequency branch.
+    Args:
+        chin (int): number of input channels.
+        chout (int): number of output channels.
+        kernel_size (int, optional): Kernel size for encoder (Default: 8)
+        stride (int, optional): Stride for encoder layer (Default: 4)
+        norm_groups (int, optional): number of groups for group norm. (Default: 4)
+        empty (bool, optional): used to make a layer with just the first conv. this is used
+            before merging the time and freq. branches. (Default: ``False``)
+        freq (bool, optional): boolean for whether conv layer is for frequency domain (Default: ``True``)
+        norm_type (string, optional): Norm type, either ``group_norm `` or ``none`` (Default: ``group_norm``)
+        context (int, optional): context size for the 1x1 conv. (Default: 0)
+        dconv_kw (Dict[str, Any] , optional): dictionary of kwargs for the DConv class. (Default: ``None``)
+        pad (bool, optional): true to pad the input. Padding is done so that the output size is
+            always the input size / stride. (Default: ``True``)
+    """
+
+    def __init__(
+        self,
+        chin: int,
+        chout: int,
+        kernel_size: int = 8,
+        stride: int = 4,
+        norm_groups: int = 4,
+        empty: bool = False,
+        freq: bool = True,
+        norm_type: str = "group_norm",
+        context: int = 0,
+        dconv_kw: Dict[str, Any] = None,
+        pad: bool = True,
+    ):
+        super().__init__()
+        if dconv_kw is None:
+            dconv_kw = {}
+        norm_fn = lambda d: nn.Identity()  # noqa
+        if norm_type == "group_norm":
+            norm_fn = lambda d: nn.GroupNorm(norm_groups, d)  # noqa
+        pad_val = kernel_size // 4 if pad else 0
+        klass = nn.Conv1d
+        self.freq = freq
+        self.kernel_size = kernel_size
+        self.stride = stride
+        self.empty = empty
+        self.pad = pad_val
+        if freq:
+            kernel_size = [kernel_size, 1]
+            stride = [stride, 1]
+            pad_val = [pad_val, 0]
+            klass = nn.Conv2d
+        self.conv = klass(chin, chout, kernel_size, stride, pad_val)
+        self.norm1 = norm_fn(chout)
+
+        if self.empty:
+            self.rewrite = nn.Identity()
+            self.norm2 = nn.Identity()
+            self.dconv = nn.Identity()
+        else:
+            self.rewrite = klass(chout, 2 * chout, 1 + 2 * context, 1, context)
+            self.norm2 = norm_fn(2 * chout)
+            self.dconv = _DConv(chout, **dconv_kw)
+
+    def forward(self, x: torch.Tensor, inject: Optional[torch.Tensor] = None) -> torch.Tensor:
+        r"""Forward pass for encoding layer.
+
+        Size depends on whether frequency or time
+
+        Args:
+            x (torch.Tensor): tensor input of shape `(B, C, F, T)` for frequency and shape
+                `(B, C, T)` for time
+            inject (torch.Tensor, optional): on last layer, combine frequency and time branches through inject param,
+                same shape as x (default: ``None``)
+
+        Returns:
+            Tensor
+                output tensor after encoder layer of shape `(B, C, F / stride, T)` for frequency
+                    and shape `(B, C, ceil(T / stride))` for time
+        """
+
+        if not self.freq and x.dim() == 4:
+            B, C, Fr, T = x.shape
+            x = x.view(B, -1, T)
+
+        if not self.freq:
+            le = x.shape[-1]
+            if not le % self.stride == 0:
+                x = F.pad(x, (0, self.stride - (le % self.stride)))
+        y = self.conv(x)
+        if self.empty:
+            return y
+        if inject is not None:
+            assert inject.shape[-1] == y.shape[-1], "injection shapes do not align"
+            if inject.dim() == 3 and y.dim() == 4:
+                inject = inject[:, :, None]
+            y = y + inject
+        y = F.gelu(self.norm1(y))
+        if self.freq:
+            B, C, Fr, T = y.shape
+            y = y.permute(0, 2, 1, 3).reshape(-1, C, T)
+            y = self.dconv(y)
+            y = y.view(B, Fr, C, T).permute(0, 2, 1, 3)
+        else:
+            y = self.dconv(y)
+        z = self.norm2(self.rewrite(y))
+        z = F.glu(z, dim=1)
+        return z
+
+
+class _HDecLayer(torch.nn.Module):
+    r"""Decoder layer. This used both by the time and the frequency branches.
+    Args:
+        chin (int): number of input channels.
+        chout (int): number of output channels.
+        last (bool, optional): whether current layer is final layer (Default: ``False``)
+        kernel_size (int, optional): Kernel size for encoder (Default: 8)
+        stride (int): Stride for encoder layer (Default: 4)
+        norm_groups (int, optional): number of groups for group norm. (Default: 1)
+        empty (bool, optional): used to make a layer with just the first conv. this is used
+            before merging the time and freq. branches. (Default: ``False``)
+        freq (bool, optional): boolean for whether conv layer is for frequency (Default: ``True``)
+        norm_type (str, optional): Norm type, either ``group_norm `` or ``none`` (Default: ``group_norm``)
+        context (int, optional): context size for the 1x1 conv. (Default: 1)
+        dconv_kw (Dict[str, Any] , optional): dictionary of kwargs for the DConv class. (Default: ``None``)
+        pad (bool, optional): true to pad the input. Padding is done so that the output size is
+            always the input size / stride. (Default: ``True``)
+    """
+
+    def __init__(
+        self,
+        chin: int,
+        chout: int,
+        last: bool = False,
+        kernel_size: int = 8,
+        stride: int = 4,
+        norm_groups: int = 1,
+        empty: bool = False,
+        freq: bool = True,
+        norm_type: str = "group_norm",
+        context: int = 1,
+        dconv_kw: Dict[str, Any] = None,
+        pad: bool = True,
+    ):
+        super().__init__()
+        if dconv_kw is None:
+            dconv_kw = {}
+        norm_fn = lambda d: nn.Identity()  # noqa
+        if norm_type == "group_norm":
+            norm_fn = lambda d: nn.GroupNorm(norm_groups, d)  # noqa
+        if pad:
+            assert (kernel_size - stride) % 2 == 0, "kernel size and stride do not align"
+            pad = (kernel_size - stride) // 2
+        else:
+            pad = 0
+        self.pad = pad
+        self.last = last
+        self.freq = freq
+        self.chin = chin
+        self.empty = empty
+        self.stride = stride
+        self.kernel_size = kernel_size
+        klass = nn.Conv1d
+        klass_tr = nn.ConvTranspose1d
+        if freq:
+            kernel_size = [kernel_size, 1]
+            stride = [stride, 1]
+            klass = nn.Conv2d
+            klass_tr = nn.ConvTranspose2d
+        self.conv_tr = klass_tr(chin, chout, kernel_size, stride)
+        self.norm2 = norm_fn(chout)
+        if self.empty:
+            self.rewrite = nn.Identity()
+            self.norm1 = nn.Identity()
+            return
+        self.rewrite = klass(chin, 2 * chin, 1 + 2 * context, 1, context)
+        self.norm1 = norm_fn(2 * chin)
+
+    def forward(self, x: torch.Tensor, skip: Optional[torch.Tensor], length):
+        r"""Forward pass for decoding layer.
+
+        Size depends on whether frequency or time
+
+        Args:
+            x (torch.Tensor): tensor input of shape `(B, C, F, T)` for frequency and shape
+                `(B, C, T)` for time
+            skip (torch.Tensor, optional): on first layer, separate frequency and time branches using param
+                (default: ``None``)
+            length (int): Size of tensor for output
+
+        Returns:
+            (Tensor, Tensor):
+                Tensor
+                    output tensor after decoder layer of shape `(B, C, F * stride, T)` for frequency domain except last
+                        frequency layer shape is `(B, C, kernel_size, T)`. Shape is `(B, C, stride * T)`
+                        for time domain.
+                Tensor
+                    contains the output just before final transposed convolution, which is used when the
+                        freq. and time branch separate. Otherwise, does not matter. Shape is
+                        `(B, C, F, T)` for frequency and `(B, C, T)` for time.
+        """
+        if self.freq and x.dim() == 3:
+            B, C, T = x.shape
+            x = x.view(B, self.chin, -1, T)
+
+        if not self.empty:
+            x = x + skip
+            y = F.glu(self.norm1(self.rewrite(x)), dim=1)
+        else:
+            y = x
+            assert skip is None, "skip must be none when empty is true."
+        z = self.norm2(self.conv_tr(y))
+        if self.freq:
+            if self.pad:
+                z = z[..., self.pad : -self.pad, :]
+        else:
+            z = z[..., self.pad : self.pad + length]
+            assert z.shape[-1] == length, "Last index of z must be equal to length"
+        if not self.last:
+            z = F.gelu(z)
+
+        return z, y
+
+
+class HDemucs(torch.nn.Module):
+    r"""
+    Hybrid Demucs model from *Hybrid Spectrogram and Waveform Source Separation* [:footcite:`defossez2021hybrid`].
+
+    Args:
+        sources (List[str]): list of source names.
+        audio_channels (int, optional): input/output audio channels. (Default: 2)
+        channels (int, optional): initial number of hidden channels. (Default: 48)
+        growth (int, optional): increase the number of hidden channels by this factor at each layer. (Default: 2)
+        nfft (int, optional): number of fft bins. Note that changing this requires careful computation of
+            various shape parameters and will not work out of the box for hybrid models. (Default: 4096)
+        depth (int, optional): number of layers in encoder and decoder (Default: 6)
+        freq_emb (float, optional): add frequency embedding after the first frequency layer if > 0,
+            the actual value controls the weight of the embedding. (Default: 0.2)
+        emb_scale (int, optional): equivalent to scaling the embedding learning rate (Default: 10)
+        emb_smooth (bool, optional): initialize the embedding with a smooth one (with respect to frequencies).
+            (Default: ``True``)
+        kernel_size (int, optional): kernel_size for encoder and decoder layers. (Default: 8)
+        time_stride (int, optional): stride for the final time layer, after the merge. (Default: 2)
+        stride (int, optional): stride for encoder and decoder layers. (Default: 4)
+        context (int, optional): context for 1x1 conv in the decoder. (Default: 4)
+        context_enc (int, optional): context for 1x1 conv in the encoder. (Default: 0)
+        norm_starts (int, optional): layer at which group norm starts being used.
+            decoder layers are numbered in reverse order. (Default: 4)
+        norm_groups (int, optional): number of groups for group norm. (Default: 4)
+        dconv_depth (int, optional): depth of residual DConv branch. (Default: 2)
+        dconv_comp (int, optional): compression of DConv branch. (Default: 4)
+        dconv_attn (int, optional): adds attention layers in DConv branch starting at this layer. (Default: 4)
+        dconv_lstm (int, optional): adds a LSTM layer in DConv branch starting at this layer. (Default: 4)
+        dconv_init (float, optional): initial scale for the DConv branch LayerScale. (Default: 1e-4)
+        sample_rate (int, optional): sample rate, serving as metadata not actually used (Default: 44100)
+        segment (int, optional): segment size (Default: 40)
+    """
+
+    def __init__(
+        self,
+        sources: List[str],
+        audio_channels: int = 2,
+        channels: int = 48,
+        growth: int = 2,
+        nfft: int = 4096,
+        depth: int = 6,
+        freq_emb: float = 0.2,
+        emb_scale: int = 10,
+        emb_smooth: bool = True,
+        kernel_size: int = 8,
+        time_stride: int = 2,
+        stride: int = 4,
+        context: int = 1,
+        context_enc: int = 0,
+        norm_starts: int = 4,
+        norm_groups: int = 4,
+        dconv_depth: int = 2,
+        dconv_comp: int = 4,
+        dconv_attn: int = 4,
+        dconv_lstm: int = 4,
+        dconv_init: float = 1e-4,
+        sample_rate: int = 44100,
+        segment: int = 4 * 10,
+    ):
+        super().__init__()
+        self.depth = depth
+        self.nfft = nfft
+        self.audio_channels = audio_channels
+        self.sources = sources
+        self.kernel_size = kernel_size
+        self.context = context
+        self.stride = stride
+        self.channels = channels
+        self.sample_rate = sample_rate
+        self.segment = segment
+
+        self.hop_length = self.nfft // 4
+        self.freq_emb = None
+
+        self.freq_encoder = nn.ModuleList()
+        self.freq_decoder = nn.ModuleList()
+
+        self.time_encoder = nn.ModuleList()
+        self.time_decoder = nn.ModuleList()
+
+        chin = audio_channels
+        chin_z = chin * 2  # number of channels for the freq branch
+        chout = channels
+        chout_z = channels
+        freqs = self.nfft // 2
+
+        for index in range(self.depth):
+            lstm = index >= dconv_lstm
+            attn = index >= dconv_attn
+            norm_type = "group_norm" if index >= norm_starts else "none"
+            freq = freqs > 1
+            stri = stride
+            ker = kernel_size
+            if not freq:
+                assert freqs == 1, "when freq is false, freqs must be 1"
+                ker = time_stride * 2
+                stri = time_stride
+
+            pad = True
+            last_freq = False
+            if freq and freqs <= kernel_size:
+                ker = freqs
+                pad = False
+                last_freq = True
+
+            kw = {
+                "kernel_size": ker,
+                "stride": stri,
+                "freq": freq,
+                "pad": pad,
+                "norm_type": norm_type,
+                "norm_groups": norm_groups,
+                "dconv_kw": {
+                    "lstm": lstm,
+                    "attn": attn,
+                    "depth": dconv_depth,
+                    "compress": dconv_comp,
+                    "init": dconv_init,
+                },
+            }
+            kwt = dict(kw)
+            kwt["freq"] = 0
+            kwt["kernel_size"] = kernel_size
+            kwt["stride"] = stride
+            kwt["pad"] = True
+            kw_dec = dict(kw)
+
+            if last_freq:
+                chout_z = max(chout, chout_z)
+                chout = chout_z
+
+            enc = _HEncLayer(chin_z, chout_z, context=context_enc, **kw)
+            if freq:
+                if last_freq is True and nfft == 2048:
+                    kwt["stride"] = 2
+                    kwt["kernel_size"] = 4
+                tenc = _HEncLayer(chin, chout, context=context_enc, empty=last_freq, **kwt)
+                self.time_encoder.append(tenc)
+
+            self.freq_encoder.append(enc)
+            if index == 0:
+                chin = self.audio_channels * len(self.sources)
+                chin_z = chin * 2
+            dec = _HDecLayer(chout_z, chin_z, last=index == 0, context=context, **kw_dec)
+            if freq:
+                tdec = _HDecLayer(chout, chin, empty=last_freq, last=index == 0, context=context, **kwt)
+                self.time_decoder.insert(0, tdec)
+            self.freq_decoder.insert(0, dec)
+
+            chin = chout
+            chin_z = chout_z
+            chout = int(growth * chout)
+            chout_z = int(growth * chout_z)
+            if freq:
+                if freqs <= kernel_size:
+                    freqs = 1
+                else:
+                    freqs //= stride
+            if index == 0 and freq_emb:
+                self.freq_emb = _ScaledEmbedding(freqs, chin_z, smooth=emb_smooth, scale=emb_scale)
+                self.freq_emb_scale = freq_emb
+
+        _rescale_module(self)
+
+    def _spec(self, x):
+        hl = self.hop_length
+        nfft = self.nfft
+        x0 = x  # noqa
+
+        # We re-pad the signal in order to keep the property
+        # that the size of the output is exactly the size of the input
+        # divided by the stride (here hop_length), when divisible.
+        # This is achieved by padding by 1/4th of the kernel size (here nfft).
+        # which is not supported by torch.stft.
+        # Having all convolution operations follow this convention allow to easily
+        # align the time and frequency branches later on.
+        assert hl == nfft // 4, "hop length must be nfft // 4"
+        le = int(math.ceil(x.shape[-1] / hl))
+        pad = hl // 2 * 3
+        x = F.pad(x, [pad, pad + le * hl - x.shape[-1]], mode="reflect")
+
+        z = _spectro(x, nfft, hl)[..., :-1, :]
+        assert z.shape[-1] == le + 4, "spectrogram's last dimension must be 4 + input size divided by stride"
+        z = z[..., 2 : 2 + le]
+        return z
+
+    def _ispec(self, z, length=None):
+        hl = self.hop_length
+        z = F.pad(z, [0, 0, 0, 1])
+        z = F.pad(z, [2, 2])
+        pad = hl // 2 * 3
+        le = hl * int(math.ceil(length / hl)) + 2 * pad
+        x = _ispectro(z, hl, length=le)
+        x = x[..., pad : pad + length]
+        return x
+
+    def _magnitude(self, z):
+        # move the complex dimension to the channel one.
+        B, C, Fr, T = z.shape
+        m = torch.view_as_real(z).permute(0, 1, 4, 2, 3)
+        m = m.reshape(B, C * 2, Fr, T)
+        return m
+
+    def _mask(self, m):
+        # `m` is a full spectrogram and `z` is ignored.
+        B, S, C, Fr, T = m.shape
+        out = m.view(B, S, -1, 2, Fr, T).permute(0, 1, 2, 4, 5, 3)
+        out = torch.view_as_complex(out.contiguous())
+        return out
+
+    def forward(self, input: torch.Tensor):
+
+        r"""HDemucs forward call
+
+        Args:
+            input (torch.Tensor): input mixed tensor of shape `(batch_size, channel, num_frames)`
+
+        Returns:
+            Tensor
+                output tensor split into sources of shape `(batch_size, num_sources, channel, num_frames)`
+        """
+
+        if len(input.shape) == 2:
+            input = input.view(1, self.audio_channels, -1)
+
+        x = input
+        length = x.shape[-1]
+
+        z = self._spec(input)
+        mag = self._magnitude(z)
+        x = mag
+
+        B, C, Fq, T = x.shape
+
+        # unlike previous Demucs, we always normalize because it is easier.
+        mean = x.mean(dim=(1, 2, 3), keepdim=True)
+        std = x.std(dim=(1, 2, 3), keepdim=True)
+        x = (x - mean) / (1e-5 + std)
+        # x will be the freq. branch input.
+
+        # Prepare the time branch input.
+        xt = input
+        meant = xt.mean(dim=(1, 2), keepdim=True)
+        stdt = xt.std(dim=(1, 2), keepdim=True)
+        xt = (xt - meant) / (1e-5 + stdt)
+
+        saved = []  # skip connections, freq.
+        saved_t = []  # skip connections, time.
+        lengths: List[int] = []  # saved lengths to properly remove padding, freq branch.
+        lengths_t: List[int] = []  # saved lengths for time branch.
+
+        for idx, encode in enumerate(self.freq_encoder):
+            lengths.append(x.shape[-1])
+            inject = None
+            if idx < len(self.time_encoder):
+                # we have not yet merged branches.
+                lengths_t.append(xt.shape[-1])
+                tenc = self.time_encoder[idx]
+                xt = tenc(xt)
+                if not tenc.empty:
+                    # save for skip connection
+                    saved_t.append(xt)
+                else:
+                    # tenc contains just the first conv., so that now time and freq.
+                    # branches have the same shape and can be merged.
+                    inject = xt
+            x = encode(x, inject)
+            if idx == 0 and self.freq_emb is not None:
+                # add frequency embedding to allow for non equivariant convolutions
+                # over the frequency axis.
+                frs = torch.arange(x.shape[-2], device=x.device)
+                emb = self.freq_emb(frs).t()[None, :, :, None].expand_as(x)
+                x = x + self.freq_emb_scale * emb
+
+            saved.append(x)
+
+        x = torch.zeros_like(x)
+        xt = torch.zeros_like(x)
+        # initialize everything to zero (signal will go through u-net skips).
+
+        for idx, decode in enumerate(self.freq_decoder):
+            skip = saved.pop(-1)
+            x, pre = decode(x, skip, lengths.pop(-1))
+            # `pre` contains the output just before final transposed convolution,
+            # which is used when the freq. and time branch separate.
+            offset = self.depth - len(self.time_decoder)
+            if idx >= offset:
+                tdec = self.time_decoder[idx - offset]
+                length_t = lengths_t.pop(-1)
+                if tdec.empty:
+                    assert pre.shape[2] == 1, "tdec empty is true, pre shape does not match " + pre.shape
+                    pre = pre[:, :, 0]
+                    xt, _ = tdec(pre, None, length_t)
+                else:
+                    skip = saved_t.pop(-1)
+                    xt, _ = tdec(xt, skip, length_t)
+
+        assert len(saved) == 0, "saved is not empty"
+        assert len(lengths_t) == 0, "lengths_t is not empty"
+        assert len(saved_t) == 0, "saved_t is not empty"
+
+        S = len(self.sources)
+        x = x.view(B, S, -1, Fq, T)
+        x = x * std[:, None] + mean[:, None]
+
+        zout = self._mask(x)
+        x = self._ispec(zout, length)
+
+        xt = xt.view(B, S, -1, length)
+        xt = xt * stdt[:, None] + meant[:, None]
+        x = xt + x
+        return x
+
+
+class _DConv(torch.nn.Module):
+    r"""
+    New residual branches in each encoder layer.
+    This alternates dilated convolutions, potentially with LSTMs and attention.
+    Also before entering each residual branch, dimension is projected on a smaller subspace,
+    e.g. of dim `channels // compress`.
+
+    Args:
+        channels (int): input/output channels for residual branch.
+        compress (float, optional): amount of channel compression inside the branch. (default: 4)
+        depth (int, optional): number of layers in the residual branch. Each layer has its own
+            projection, and potentially LSTM and attention.(default: 2)
+        init (float, optional): initial scale for LayerNorm. (default: 1e-4)
+        norm_type (bool, optional): Norm type, either ``group_norm `` or ``none`` (Default: ``group_norm``)
+        attn (bool, optional): use LocalAttention. (Default: ``False``)
+        heads (int, optional): number of heads for the LocalAttention.  (default: 4)
+        ndecay (int, optional): number of decay controls in the LocalAttention. (default: 4)
+        lstm (bool, optional): use LSTM. (Default: ``False``)
+        kernel_size (int, optional): kernel size for the (dilated) convolutions. (default: 3)
+    """
+
+    def __init__(
+        self,
+        channels: int,
+        compress: float = 4,
+        depth: int = 2,
+        init: float = 1e-4,
+        norm_type: str = "group_norm",
+        attn: bool = False,
+        heads: int = 4,
+        ndecay: int = 4,
+        lstm: bool = False,
+        kernel_size: int = 3,
+    ):
+
+        super().__init__()
+        assert kernel_size % 2 == 1, "kernel size should not be divisible by 2"
+        self.channels = channels
+        self.compress = compress
+        self.depth = abs(depth)
+        dilate = depth > 0
+
+        norm_fn: tp.Callable[[int], nn.Module]
+        norm_fn = lambda d: nn.Identity()  # noqa
+        if norm_type == "group_norm":
+            norm_fn = lambda d: nn.GroupNorm(1, d)  # noqa
+
+        hidden = int(channels / compress)
+
+        act = nn.GELU
+
+        self.layers = nn.ModuleList([])
+        for d in range(self.depth):
+            dilation = pow(2, d) if dilate else 1
+            padding = dilation * (kernel_size // 2)
+            mods = [
+                nn.Conv1d(channels, hidden, kernel_size, dilation=dilation, padding=padding),
+                norm_fn(hidden),
+                act(),
+                nn.Conv1d(hidden, 2 * channels, 1),
+                norm_fn(2 * channels),
+                nn.GLU(1),
+                _LayerScale(channels, init),
+            ]
+            if attn:
+                mods.insert(3, _LocalState(hidden, heads=heads, ndecay=ndecay))
+            if lstm:
+                mods.insert(3, _BLSTM(hidden, layers=2, skip=True))
+            layer = nn.Sequential(*mods)
+            self.layers.append(layer)
+
+    def forward(self, x):
+        r"""DConv forward call
+
+        Args:
+            x (torch.Tensor): input tensor for convolution
+
+        Returns:
+            Tensor
+                Output after being run through layers.
+        """
+        for layer in self.layers:
+            x = x + layer(x)
+        return x
+
+
+class _BLSTM(torch.nn.Module):
+    r"""
+    BiLSTM with same hidden units as input dim.
+    If `max_steps` is not None, input will be splitting in overlapping
+    chunks and the LSTM applied separately on each chunk.
+    Args:
+        dim (int): dimensions at LSTM layer.
+        layers (int, optional): number of LSTM layers. (default: 1)
+        skip (bool, optional): (default: ``False``)
+    """
+
+    def __init__(self, dim, layers: int = 1, skip: bool = False):
+        super().__init__()
+        self.max_steps = 200
+        self.lstm = nn.LSTM(bidirectional=True, num_layers=layers, hidden_size=dim, input_size=dim)
+        self.linear = nn.Linear(2 * dim, dim)
+        self.skip = skip
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        r"""BLSTM forward call
+
+        Args:
+            x (torch.Tensor): input tensor for BLSTM shape is `(batch_size, dim, time_steps)`
+
+        Returns:
+            Tensor
+                Output after being run through bidirectional LSTM. Shape is `(batch_size, dim, time_steps)`
+        """
+        B, C, T = x.shape
+        y = x
+        framed = False
+        width = 0
+        stride = 0
+        nframes = 0
+        if self.max_steps is not None and T > self.max_steps:
+            width = self.max_steps
+            stride = width // 2
+            frames = _unfold(x, width, stride)
+            nframes = frames.shape[2]
+            framed = True
+            x = frames.permute(0, 2, 1, 3).reshape(-1, C, width)
+
+        x = x.permute(2, 0, 1)
+
+        x = self.lstm(x)[0]
+        x = self.linear(x)
+        x = x.permute(1, 2, 0)
+        if framed:
+            out = []
+            frames = x.reshape(B, -1, C, width)
+            limit = stride // 2
+            for k in range(nframes):
+                if k == 0:
+                    out.append(frames[:, k, :, :-limit])
+                elif k == nframes - 1:
+                    out.append(frames[:, k, :, limit:])
+                else:
+                    out.append(frames[:, k, :, limit:-limit])
+            out = torch.cat(out, -1)
+            out = out[..., :T]
+            x = out
+        if self.skip:
+            x = x + y
+
+        return x
+
+
+class _LocalState(nn.Module):
+    """Local state allows to have attention based only on data (no positional embedding),
+    but while setting a constraint on the time window (e.g. decaying penalty term).
+    Also a failed experiments with trying to provide some frequency based attention.
+    """
+
+    def __init__(self, channels: int, heads: int = 4, ndecay: int = 4):
+        r"""
+        Args:
+            channels (int): Size of Conv1d layers.
+            heads (int, optional):  (default: 4)
+            ndecay (int, optional): (default: 4)
+        """
+        super(_LocalState, self).__init__()
+        assert channels % heads == 0, "Channels must be divisible by heads."
+        self.heads = heads
+        self.ndecay = ndecay
+        self.content = nn.Conv1d(channels, channels, 1)
+        self.query = nn.Conv1d(channels, channels, 1)
+        self.key = nn.Conv1d(channels, channels, 1)
+
+        self.query_decay = nn.Conv1d(channels, heads * ndecay, 1)
+        if ndecay:
+            # Initialize decay close to zero (there is a sigmoid), for maximum initial window.
+            self.query_decay.weight.data *= 0.01
+            assert self.query_decay.bias is not None, "bias must not be None"
+            self.query_decay.bias.data[:] = -2
+        self.proj = nn.Conv1d(channels + heads * 0, channels, 1)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        r"""LocalState forward call
+
+        Args:
+            x (torch.Tensor): input tensor for LocalState
+
+        Returns:
+            Tensor
+                Output after being run through LocalState layer.
+        """
+        B, C, T = x.shape
+        heads = self.heads
+        indexes = torch.arange(T, device=x.device, dtype=x.dtype)
+        # left index are keys, right index are queries
+        delta = indexes[:, None] - indexes[None, :]
+
+        queries = self.query(x).view(B, heads, -1, T)
+        keys = self.key(x).view(B, heads, -1, T)
+        # t are keys, s are queries
+        dots = torch.einsum("bhct,bhcs->bhts", keys, queries)
+        dots /= math.sqrt(keys.shape[2])
+        if self.ndecay:
+            decays = torch.arange(1, self.ndecay + 1, device=x.device, dtype=x.dtype)
+            decay_q = self.query_decay(x).view(B, heads, -1, T)
+            decay_q = torch.sigmoid(decay_q) / 2
+            decay_kernel = -decays.view(-1, 1, 1) * delta.abs() / math.sqrt(self.ndecay)
+            dots += torch.einsum("fts,bhfs->bhts", decay_kernel, decay_q)
+
+        # Kill self reference.
+        dots.masked_fill_(torch.eye(T, device=dots.device, dtype=torch.bool), -100)
+        weights = torch.softmax(dots, dim=2)
+
+        content = self.content(x).view(B, heads, -1, T)
+        result = torch.einsum("bhts,bhct->bhcs", weights, content)
+        result = result.reshape(B, -1, T)
+        return x + self.proj(result)
+
+
+class _LayerScale(nn.Module):
+    """Layer scale from [Touvron et al 2021] (https://arxiv.org/pdf/2103.17239.pdf).
+    This rescales diagonally residual outputs close to 0 initially, then learnt.
+    """
+
+    def __init__(self, channels: int, init: float = 0):
+        r"""
+        Args:
+            channels (int): Size of  rescaling
+            init (float, optional): Scale to default to (default: 0)
+        """
+        super().__init__()
+        self.scale = nn.Parameter(torch.zeros(channels, requires_grad=True))
+        self.scale.data[:] = init
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        r"""LayerScale forward call
+
+        Args:
+            x (torch.Tensor): input tensor for LayerScale
+
+        Returns:
+            Tensor
+                Output after rescaling tensor.
+        """
+        return self.scale[:, None] * x
+
+
+def _unfold(a: torch.Tensor, kernel_size: int, stride: int) -> torch.Tensor:
+    """Given input of size [*OT, T], output Tensor of size [*OT, F, K]
+    with K the kernel size, by extracting frames with the given stride.
+    This will pad the input so that `F = ceil(T / K)`.
+    see https://github.com/pytorch/pytorch/issues/60466
+    """
+    shape = list(a.shape[:-1])
+    length = int(a.shape[-1])
+    n_frames = math.ceil(length / stride)
+    tgt_length = (n_frames - 1) * stride + kernel_size
+    a = F.pad(input=a, pad=[0, tgt_length - length])
+    strides = [a.stride(dim) for dim in range(a.dim())]
+    assert strides[-1] == 1, "data should be contiguous"
+    strides = strides[:-1] + [stride, 1]
+    shape.append(n_frames)
+    shape.append(kernel_size)
+    return a.as_strided(shape, strides)
+
+
+def _rescale_module(module):
+    r"""
+    Rescales initial weight scale for all models within the module.
+    """
+    for sub in module.modules():
+        if isinstance(sub, (nn.Conv1d, nn.ConvTranspose1d, nn.Conv2d, nn.ConvTranspose2d)):
+            std = sub.weight.std().detach()
+            scale = (std / 0.1) ** 0.5
+            sub.weight.data /= scale
+            if sub.bias is not None:
+                sub.bias.data /= scale
+
+
+def _spectro(x: torch.Tensor, n_fft: int = 512, hop_length: int = 0, pad: int = 0) -> torch.Tensor:
+    other = list(x.shape[:-1])
+    length = int(x.shape[-1])
+    x = x.reshape(-1, length)
+    z = torch.stft(
+        x,
+        n_fft * (1 + pad),
+        hop_length,
+        window=torch.hann_window(n_fft).to(x),
+        win_length=n_fft,
+        normalized=True,
+        center=True,
+        return_complex=True,
+        pad_mode="reflect",
+    )
+    _, freqs, frame = z.shape
+    other.extend([freqs, frame])
+    return z.view(other)
+
+
+def _ispectro(z: torch.Tensor, hop_length: int = 0, length: int = 0, pad: int = 0) -> torch.Tensor:
+    other = list(z.shape[:-2])
+    freqs = int(z.shape[-2])
+    frames = int(z.shape[-1])
+
+    n_fft = 2 * freqs - 2
+    z = z.view(-1, freqs, frames)
+    win_length = n_fft // (1 + pad)
+    x = torch.istft(
+        z,
+        n_fft,
+        hop_length,
+        window=torch.hann_window(win_length).to(z.real),
+        win_length=win_length,
+        normalized=True,
+        length=length,
+        center=True,
+    )
+    _, length = x.shape
+    other.append(length)
+    return x.view(other)