adding pad1d

demucs/hdemucs.py

@@ -8,6 +8,7 @@
 """
 from copy import deepcopy
 import math
+import typing as tp
 
 from openunmix.filtering import wiener
 import torch
@@ -19,6 +20,18 @@
 from .spec import spectro, ispectro
 
 
+def pad1d(x: torch.Tensor, paddings: tp.Tuple[int, int], mode: str = 'zero', value: float = 0.):
+    """Tiny wrapper around F.pad, just to allow for reflect padding on small input.
+    If this is the case, we insert extra 0 padding to the right before the reflection happen."""
+    length = x.shape[-1]
+    padding_left, padding_right = paddings
+    if mode == 'reflect':
+        max_pad = max(padding_left, padding_right)
+        if length <= max_pad:
+            x = F.pad(x, (0, max_pad - length + 1))
+    return F.pad(x, paddings, mode, value)
+
+
 class ScaledEmbedding(nn.Module):
     """
     Boost learning rate for embeddings (with `scale`).
@@ -580,9 +593,9 @@ def _spec(self, x):
             le = int(math.ceil(x.shape[-1] / hl))
             pad = hl // 2 * 3
             if not self.hybrid_old:
-                x = F.pad(x, (pad, pad + le * hl - x.shape[-1]), mode='reflect')
+                x = pad1d(x, (pad, pad + le * hl - x.shape[-1]), mode='reflect')
             else:
-                x = F.pad(x, (pad, pad + le * hl - x.shape[-1]))
+                x = pad1d(x, (pad, pad + le * hl - x.shape[-1]))
 
         z = spectro(x, nfft, hl)[..., :-1, :]
         if self.hybrid: