Merge pull request #52 from lucidrains/pope

pope

Author: lucidrains

.github/workflows/test.yml

@@ -0,0 +1,27 @@
+name: test
+on: [push, pull_request]
+
+jobs:
+  test:
+    runs-on: ubuntu-latest
+    strategy:
+      matrix:
+        python-version: ["3.10"]
+
+    steps:
+    - uses: actions/checkout@v4
+
+    - name: Set up Python ${{ matrix.python-version }}
+      uses: actions/setup-python@v5
+      with:
+        python-version: ${{ matrix.python-version }}
+
+    - name: Install dependencies
+      run: |
+        python -m pip install --upgrade pip
+        pip install .
+        pip install pytest
+
+    - name: Test with pytest
+      run: |
+        pytest tests/test_roformer.py

bs_roformer/bs_roformer.py

@@ -12,6 +12,7 @@
 from beartype import beartype
 
 from rotary_embedding_torch import RotaryEmbedding
+from PoPE_pytorch import PoPE, flash_attn_with_pope
 
 from einops import rearrange, pack, unpack
 
@@ -73,6 +74,7 @@ def __init__(
         dim_head = 64,
         dropout = 0.,
         rotary_embed = None,
+        pope_embed = None,
         flash = True,
         learned_value_residual_mix = False
     ):
@@ -82,6 +84,9 @@ def __init__(
         dim_inner = heads * dim_head
 
         self.rotary_embed = rotary_embed
+        self.pope_embed = pope_embed
+
+        assert not (exists(rotary_embed) and exists(pope_embed)), 'cannot have both rotary and pope embeddings'
 
         self.attend = Attend(flash = flash, dropout = dropout)
 
@@ -111,11 +116,14 @@ def forward(self, x, value_residual = None):
             assert exists(value_residual)
             v = v.lerp(value_residual, mix)
 
-        if exists(self.rotary_embed):
+        if exists(self.pope_embed):
+            out = flash_attn_with_pope(q, k, v, pos_emb = self.pope_embed(q.shape[-2]), softmax_scale = self.scale)
+        elif 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)
+            out = self.attend(q, k, v)
+        else:
+            out = self.attend(q, k, v)
 
         gates = self.to_gates(x)
         out = out * rearrange(gates, 'b n h -> b h n 1').sigmoid()
@@ -137,6 +145,7 @@ def __init__(
         ff_mult = 4,
         norm_output = True,
         rotary_embed = None,
+        pope_embed = None,
         flash_attn = True,
         add_value_residual = False,
         num_residual_streams = 1,
@@ -150,7 +159,7 @@ def __init__(
 
         for _ in range(depth):
             self.layers.append(ModuleList([
-                init_hyper_conn(dim = dim, branch = Attention(dim = dim, dim_head = dim_head, heads = heads, dropout = attn_dropout, rotary_embed = rotary_embed, flash = flash_attn, learned_value_residual_mix = add_value_residual)),
+                init_hyper_conn(dim = dim, branch = Attention(dim = dim, dim_head = dim_head, heads = heads, dropout = attn_dropout, rotary_embed = rotary_embed, pope_embed = pope_embed, flash = flash_attn, learned_value_residual_mix = add_value_residual)),
                 init_hyper_conn(dim = dim, branch = FeedForward(dim = dim, mult = ff_mult, dropout = ff_dropout))
             ]))
 
@@ -306,7 +315,8 @@ def __init__(
         multi_stft_resolutions_window_sizes: tuple[int, ...] = (4096, 2048, 1024, 512, 256),
         multi_stft_hop_size = 147,
         multi_stft_normalized = False,
-        multi_stft_window_fn: Callable = torch.hann_window
+        multi_stft_window_fn: Callable = torch.hann_window,
+        use_pope = False
     ):
         super().__init__()
 
@@ -328,18 +338,24 @@ def __init__(
             num_residual_streams = num_residual_streams,
             num_residual_fracs = num_residual_fracs,
             mc_hyper_conn_sinkhorn_iters = mc_hyper_conn_sinkhorn_iters,
-            norm_output = False,
+            norm_output = False
         )
 
-        time_rotary_embed = RotaryEmbedding(dim = dim_head)
-        freq_rotary_embed = RotaryEmbedding(dim = dim_head)
+        if use_pope:
+            time_pope_embed = PoPE(dim = dim_head, heads = heads)
+            freq_pope_embed = PoPE(dim = dim_head, heads = heads)
+            time_rotary_embed = freq_rotary_embed = None
+        else:
+            time_rotary_embed = RotaryEmbedding(dim = dim_head)
+            freq_rotary_embed = RotaryEmbedding(dim = dim_head)
+            time_pope_embed = freq_pope_embed = None
 
         for layer_index in range(depth):
             is_first = layer_index == 0
 
             self.layers.append(nn.ModuleList([
-                Transformer(depth = time_transformer_depth, rotary_embed = time_rotary_embed, add_value_residual = not is_first, **transformer_kwargs),
-                Transformer(depth = freq_transformer_depth, rotary_embed = freq_rotary_embed, add_value_residual = not is_first, **transformer_kwargs)
+                Transformer(depth = time_transformer_depth, rotary_embed = time_rotary_embed, pope_embed = time_pope_embed, add_value_residual = not is_first, **transformer_kwargs),
+                Transformer(depth = freq_transformer_depth, rotary_embed = freq_rotary_embed, pope_embed = freq_pope_embed, add_value_residual = not is_first, **transformer_kwargs)
             ]))
 
         self.final_norm = RMSNorm(dim)

bs_roformer/mel_band_roformer.py

@@ -12,6 +12,7 @@
 from beartype import beartype
 
 from rotary_embedding_torch import RotaryEmbedding
+from PoPE_pytorch import PoPE, flash_attn_with_pope
 
 from einops import rearrange, pack, unpack, reduce, repeat
 from einops.layers.torch import Rearrange
@@ -84,6 +85,7 @@ def __init__(
         dim_head = 64,
         dropout = 0.,
         rotary_embed = None,
+        pope_embed = None,
         flash = True,
         add_value_residual = False
     ):
@@ -93,6 +95,9 @@ def __init__(
         dim_inner = heads * dim_head
 
         self.rotary_embed = rotary_embed
+        self.pope_embed = pope_embed
+
+        assert not (exists(rotary_embed) and exists(pope_embed)), 'cannot have both rotary and pope embeddings'
 
         self.attend = Attend(flash = flash, dropout = dropout)
 
@@ -124,11 +129,14 @@ def forward(self, x, value_residual = None):
             assert exists(value_residual)
             v = v.lerp(mix, value_residual)
 
-        if exists(self.rotary_embed):
+        if exists(self.pope_embed):
+            out = flash_attn_with_pope(q, k, v, pos_emb = self.pope_embed(q.shape[-2]), softmax_scale = self.scale)
+        elif 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)
+            out = self.attend(q, k, v)
+        else:
+            out = self.attend(q, k, v)
 
         gates = self.to_gates(x)
         out = out * rearrange(gates, 'b n h -> b h n 1').sigmoid()
@@ -217,6 +225,7 @@ def __init__(
         ff_mult = 4,
         norm_output = True,
         rotary_embed = None,
+        pope_embed = None,
         flash_attn = True,
         linear_attn = False,
         add_value_residual = False,
@@ -234,7 +243,7 @@ def __init__(
             if linear_attn:
                 attn = LinearAttention(dim = dim, dim_head = dim_head, heads = heads, dropout = attn_dropout, flash = flash_attn, add_value_residual = add_value_residual)
             else:
-                attn = Attention(dim = dim, dim_head = dim_head, heads = heads, dropout = attn_dropout, rotary_embed = rotary_embed, flash = flash_attn, add_value_residual = add_value_residual)
+                attn = Attention(dim = dim, dim_head = dim_head, heads = heads, dropout = attn_dropout, rotary_embed = rotary_embed, pope_embed = pope_embed, flash = flash_attn, add_value_residual = add_value_residual)
 
             ff = FeedForward(dim = dim, mult = ff_mult, dropout = ff_dropout)
 
@@ -387,7 +396,8 @@ def __init__(
         match_input_audio_length = False, # if True, pad output tensor to match length of input tensor
         add_value_residual = True,
         num_residual_streams = 4,
-        num_residual_fracs = 1
+        num_residual_fracs = 1,
+        use_pope = False
     ):
         super().__init__()
 
@@ -407,8 +417,14 @@ def __init__(
             num_residual_fracs = num_residual_fracs
         )
 
-        time_rotary_embed = RotaryEmbedding(dim = dim_head)
-        freq_rotary_embed = RotaryEmbedding(dim = dim_head)
+        if use_pope:
+            time_pope_embed = PoPE(dim = dim_head, heads = heads)
+            freq_pope_embed = PoPE(dim = dim_head, heads = heads)
+            time_rotary_embed = freq_rotary_embed = None
+        else:
+            time_rotary_embed = RotaryEmbedding(dim = dim_head)
+            freq_rotary_embed = RotaryEmbedding(dim = dim_head)
+            time_pope_embed = freq_pope_embed = None
 
         linear_flash_attn = default(linear_flash_attn, flash_attn)
 
@@ -421,8 +437,8 @@ def __init__(
 
             self.layers.append(nn.ModuleList([
                 Transformer(depth = linear_transformer_depth, linear_attn = True, flash_attn = linear_flash_attn, add_value_residual = add_value_residual and not is_first, **transformer_kwargs) if linear_transformer_depth > 0 else None,
-                Transformer(depth = time_transformer_depth, rotary_embed = time_rotary_embed, flash_attn = flash_attn, add_value_residual = add_value_residual and not is_first, **transformer_kwargs),
-                Transformer(depth = freq_transformer_depth, rotary_embed = freq_rotary_embed, flash_attn = flash_attn, add_value_residual = add_value_residual and not is_first, **transformer_kwargs)
+                Transformer(depth = time_transformer_depth, rotary_embed = time_rotary_embed, pope_embed = time_pope_embed, flash_attn = flash_attn, add_value_residual = add_value_residual and not is_first, **transformer_kwargs),
+                Transformer(depth = freq_transformer_depth, rotary_embed = freq_rotary_embed, pope_embed = freq_pope_embed, flash_attn = flash_attn, add_value_residual = add_value_residual and not is_first, **transformer_kwargs)
             ]))
 
         self.stft_window_fn = partial(default(stft_window_fn, torch.hann_window), stft_win_length)

setup.py

@@ -3,7 +3,7 @@
 setup(
   name = 'BS-RoFormer',
   packages = find_packages(exclude=[]),
-  version = '1.0.6',
+  version = '1.1.0',
   license='MIT',
   description = 'BS-RoFormer - Band-Split Rotary Transformer for SOTA Music Source Separation',
   author = 'Phil Wang',
@@ -22,6 +22,7 @@
     'einops>=0.8.0',
     'hyper-connections>=0.4.4',
     'librosa',
+    'PoPE-pytorch>=0.0.15',
     'rotary-embedding-torch>=0.3.6',
     'torch>=2.0',
   ],

tests/test_roformer.py

@@ -0,0 +1,44 @@
+import torch
+import pytest
+from bs_roformer import BSRoformer, MelBandRoformer
+from PoPE_pytorch import PoPE
+
+@pytest.mark.parametrize('use_pope', [True, False])
+def test_bs_roformer(use_pope):
+    model = BSRoformer(
+        dim = 512,
+        depth = 1,
+        time_transformer_depth = 1,
+        freq_transformer_depth = 1,
+        use_pope = use_pope
+    )
+
+    dummy_audio = torch.randn(1, 1, 44100)
+    out = model(dummy_audio)
+    
+    assert out.shape[0] == dummy_audio.shape[0]
+    assert abs(out.shape[-1] - dummy_audio.shape[-1]) < 1024
+
+    # verify pope presence
+    has_pope = any(isinstance(m, PoPE) for m in model.modules())
+    assert has_pope == use_pope
+
+@pytest.mark.parametrize('use_pope', [True, False])
+def test_mel_band_roformer(use_pope):
+    model = MelBandRoformer(
+        dim = 512,
+        depth = 1,
+        time_transformer_depth = 1,
+        freq_transformer_depth = 1,
+        use_pope = use_pope
+    )
+
+    dummy_audio = torch.randn(1, 1, 44100)
+    out = model(dummy_audio)
+    
+    assert out.shape[0] == dummy_audio.shape[0]
+    assert abs(out.shape[-1] - dummy_audio.shape[-1]) < 1024
+
+    # verify pope presence
+    has_pope = any(isinstance(m, PoPE) for m in model.modules())
+    assert has_pope == use_pope