some scaffold, copy in essentials (attention, what else)

Author: lucidrains

Diff for: .github/workflows/python-publish.yml

@@ -0,0 +1,38 @@
+
+  
+# This workflow will upload a Python Package using Twine when a release is created
+# For more information see: https://help.github.com/en/actions/language-and-framework-guides/using-python-with-github-actions#publishing-to-package-registries
+
+# This workflow uses actions that are not certified by GitHub.
+# They are provided by a third-party and are governed by
+# separate terms of service, privacy policy, and support
+# documentation.
+
+name: Upload Python Package
+
+on:
+  release:
+    types: [published]
+
+jobs:
+  deploy:
+
+    runs-on: ubuntu-latest
+
+    steps:
+    - uses: actions/checkout@v2
+    - name: Set up Python
+      uses: actions/setup-python@v2
+      with:
+        python-version: '3.x'
+    - name: Install dependencies
+      run: |
+        python -m pip install --upgrade pip
+        pip install build
+    - name: Build package
+      run: python -m build
+    - name: Publish package
+      uses: pypa/gh-action-pypi-publish@27b31702a0e7fc50959f5ad993c78deac1bdfc29
+      with:
+        user: __token__
+        password: ${{ secrets.PYPI_API_TOKEN }}

Diff for: musiclm_pytorch/__init__.py

@@ -0,0 +1 @@
+from musiclm_pytorch.musiclm_pytorch import MuLaN

Diff for: musiclm_pytorch/musiclm_pytorch.py

@@ -0,0 +1,119 @@
+import torch
+import torch.nn.functional as F
+from torch import nn, einsum
+
+from einops import rearrange, repeat, reduce
+
+# functions
+
+def exists(val):
+    return val is not None
+
+# biasless layernorm
+
+class LayerNorm(nn.Module):
+    def __init__(self, dim):
+        super().__init__()
+        self.gamma = nn.Parameter(torch.ones(dim))
+        self.register_buffer("beta", torch.zeros(dim))
+
+    def forward(self, x):
+        return F.layer_norm(x, x.shape[-1:], self.gamma, self.beta)
+
+# attention
+
+class Attention(nn.Module):
+    def __init__(
+        self,
+        dim,
+        causal = False,
+        dim_head = 64,
+        heads = 8,
+        num_null_kv = 0,
+        dropout = 0.1
+    ):
+        super().__init__()
+        self.heads = heads
+        self.scale = dim_head ** -0.5
+        self.causal = causal
+        inner_dim = dim_head * heads
+
+        self.norm = LayerNorm(dim)
+
+        self.attn_dropout = nn.Dropout(dropout)
+
+        self.to_q = nn.Linear(dim, inner_dim, bias = False)
+        self.to_kv = nn.Linear(dim, inner_dim * 2, bias = False)
+
+        self.to_out = nn.Sequential(
+            nn.Linear(inner_dim, dim, bias = False),
+            nn.Dropout(dropout)
+        )
+
+    def forward(
+        self,
+        x,
+        context = None,
+        mask = None,
+        attn_bias = None,
+        prefix_context = None,
+        prefix_context_mask = None
+    ):
+        b, n, _, device = *x.shape, x.device
+
+        # prenorm
+
+        x = self.norm(x)
+
+        # project for queries, keys, values
+
+        q, k, v = self.to_q(x), *self.to_kv(x).chunk(2, dim = -1)
+
+        # split for multi-headed attention
+
+        q, k, v = map(lambda t: rearrange(t, 'b n (h d) -> b h n d', h = self.heads), (q, k, v))
+
+        q = q * self.scale
+
+        # similarities
+
+        sim = einsum('b h i d, b h j d -> b h i j', q, k)
+
+        if exists(mask):
+            mask = rearrange(mask, 'b j -> b 1 1 j')
+            sim = sim.masked_fill(~mask, -torch.finfo(sim.dtype).max)
+
+        if self.causal:
+            i, j = sim.shape[-2:]
+            causal_mask = torch.ones((i, j), dtype = torch.bool, device = x.device).triu(j - i + 1)
+            sim = sim.masked_fill(causal_mask, -torch.finfo(sim.dtype).max)
+
+        # attention
+
+        attn = sim.softmax(dim = -1)
+        attn = self.attn_dropout(attn)
+
+        # aggregate
+
+        out = einsum('b h i j, b h j d -> b h i d', attn, v)
+
+        # merge heads
+
+        out = rearrange(out, 'b h n d -> b n (h d)')
+        return self.to_out(out)
+
+# main classes
+
+class MuLaN(nn.Module):
+    def __init__(self):
+        super().__init__()
+
+    def forward(self, x):
+        return x
+
+class MusicLM(nn.Module):
+    def __init__(self):
+        super().__init__()
+
+    def forward(self, x):
+        return x

Diff for: setup.py

@@ -0,0 +1,32 @@
+from setuptools import setup, find_packages
+
+setup(
+  name = 'musiclm-pytorch',
+  packages = find_packages(exclude=[]),
+  version = '0.0.1',
+  license='MIT',
+  description = 'MusicLM - AudioLM + Audio CLIP to text to music synthesis',
+  author = 'Phil Wang',
+  author_email = '[email protected]',
+  long_description_content_type = 'text/markdown',
+  url = 'https://github.com/lucidrains/musiclm-pytorch',
+  keywords = [
+    'artificial intelligence',
+    'deep learning',
+    'transformers',
+    'attention mechanism',
+    'text to music',
+    'contrastive learning'
+  ],
+  install_requires=[
+    'einops>=0.4',
+    'torch>=1.6',
+  ],
+  classifiers=[
+    'Development Status :: 4 - Beta',
+    'Intended Audience :: Developers',
+    'Topic :: Scientific/Engineering :: Artificial Intelligence',
+    'License :: OSI Approved :: MIT License',
+    'Programming Language :: Python :: 3.6',
+  ],
+)