Upload 18 files

.gitattributes

@@ -33,3 +33,4 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 *.zip filter=lfs diff=lfs merge=lfs -text
 *.zst filter=lfs diff=lfs merge=lfs -text
 *tfevents* filter=lfs diff=lfs merge=lfs -text
+sf2_kits/drum_breaks.sf2 filter=lfs diff=lfs merge=lfs -text

LICENSE

@@ -0,0 +1,43 @@
+Model License and Usage Terms
+-----------------------------
+
+Copyright © 2024 Patchbanks. All rights reserved.
+
+The outputs generated by this model are granted royalty-free for personal use 
+and commercial entertainment purposes, including but not limited to songs, albums, 
+videos, or other forms of digital media released by an individual. However, these 
+outputs may not be used, resold, or licensed for commercial production, which 
+includes, but is not limited to, profiting from sound libraries, audio and MIDI 
+sample packs, stock media, subscription platforms, or generative AI platforms. 
+Additionally, the collection of generated output data for the purpose of training 
+AI models or machine learning systems is strictly prohibited.
+
+The model weights, along with any derivatives or modifications, are also restricted 
+from commercial use and may not be sold, licensed, deployed on cloud servers, 
+or otherwise distributed for commercial purposes.
+
+
+Model Architecture License
+---------------------------
+
+MIT License
+
+Copyright (c) 2022 Andrej Karpathy
+
+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.

app.py

@@ -0,0 +1,286 @@
+import os
+import pickle
+import torch
+import random
+import subprocess
+import re
+import pretty_midi
+import gradio as gr
+from contextlib import nullcontext
+from model import GPTConfig, GPT
+from pedalboard import Pedalboard, Reverb, Compressor, Gain, Limiter
+from pedalboard.io import AudioFile
+import gradio as gr
+import spaces
+
+in_space = os.getenv("SYSTEM") == "spaces"
+
+temp_dir = 'temp'
+os.makedirs(temp_dir, exist_ok=True)
+
+init_from = 'resume'
+out_dir = 'checkpoints'
+ckpt_load = 'model.pt'
+
+start = "000000000000\n"
+num_samples = 1
+max_new_tokens = 768
+
+seed = random.randint(1, 100000)
+torch.manual_seed(seed)
+device = 'cpu' if torch.cuda.is_available() else 'cpu'
+dtype = 'bfloat16' if torch.cuda.is_available() and torch.cuda.is_bf16_supported() else 'float16'
+compile = False
+exec(open('configurator.py').read())
+
+
+torch.backends.cuda.matmul.allow_tf32 = True
+torch.backends.cudnn.allow_tf32 = True
+device_type = 'cpu' if 'cuda' in device else 'cpu'
+ptdtype = {'float32': torch.float32, 'bfloat16': torch.bfloat16, 'float16': torch.float16}[dtype]
+ctx = nullcontext() if device_type == 'cpu' else torch.amp.autocast(device_type=device_type, dtype=ptdtype)
+
+if init_from == 'resume':
+    ckpt_path = os.path.join(out_dir, ckpt_load)
+    checkpoint = torch.load(ckpt_path, map_location=device, weights_only=True)
+    gptconf = GPTConfig(**checkpoint['model_args'])
+    model = GPT(gptconf)
+    state_dict = checkpoint['model']
+    unwanted_prefix = '_orig_mod.'
+    for k, v in list(state_dict.items()):
+        if k.startswith(unwanted_prefix):
+            state_dict[k[len(unwanted_prefix):]] = state_dict.pop(k)
+    model.load_state_dict(state_dict)
+elif init_from.startswith('gpt2'):
+    model = GPT.from_pretrained(init_from, dict(dropout=0.0))
+
+model.eval()
+model.to(device)
+if compile:
+    model = torch.compile(model)
+
+tokenizer = re.compile(r'000000000000|\d{2}|\n')
+
+meta_path = os.path.join('data', checkpoint['config']['dataset'], 'meta.pkl')
+with open(meta_path, 'rb') as f:
+    meta = pickle.load(f)
+    stoi = meta.get('stoi', None)
+    itos = meta.get('itos', None)
+
+def encode(text):
+    matches = tokenizer.findall(text)
+    return [stoi[c] for c in matches]
+
+def decode(encoded):
+    return ''.join([itos[i] for i in encoded])
+
+def clear_midi(dir):
+    for file in os.listdir(dir):
+        if file.endswith('.mid'):
+            os.remove(os.path.join(dir, file))
+
+clear_midi(temp_dir)
+
+
+def generate_midi(temperature, top_k):
+    start_ids = encode(start)
+    x = (torch.tensor(start_ids, dtype=torch.long, device=device)[None, ...])
+        
+    midi_events = []
+    seq_count = 0
+
+    with torch.no_grad():
+        for _ in range(num_samples):
+            sequence = []
+            y = model.generate(x, max_new_tokens, temperature=temperature, top_k=top_k)
+            tkn_seq = decode(y[0].tolist())
+            lines = tkn_seq.splitlines()
+
+            for event in lines:
+                if event.startswith(start.strip()):
+                    if sequence:
+                        midi_events.append(sequence)
+                        sequence = []
+                    seq_count += 1
+                elif event.strip() == "":
+                    continue
+                else:
+                    try:
+                        p = int(event[0:2])
+                        v = int(event[2:4])
+                        s = int(event[4:8])
+                        e = int(event[8:12])
+                    except ValueError:
+                        p, v, s, e = 0, 0, 0, 0
+                    sequence.append({'file_name': f'nanompc_{seq_count:02d}', 'pitch': p, 'velocity': v, 'start': s, 'end': e})
+
+            if sequence:
+                midi_events.append(sequence)
+
+    round_bars = []
+    
+    for sequence in midi_events:
+        filtered_sequence = []
+        for event in sequence:
+            if event['start'] < 1536 and event['end'] <= 1536:
+                filtered_sequence.append(event)
+        if filtered_sequence:
+            round_bars.append(filtered_sequence)
+
+    midi_events = round_bars
+
+    for track in midi_events:
+        track.sort(key=lambda x: x['start'])
+        unique_notes = []
+        
+        for note in track:
+            if not any(abs(note['start'] - n['start']) < 12 and note['pitch'] == n['pitch'] for n in unique_notes):
+                unique_notes.append(note)
+        
+        track[:] = unique_notes 
+
+    return midi_events
+
+
+def write_single_midi(midi_events, bpm):
+    midi_data = pretty_midi.PrettyMIDI(initial_tempo=bpm, resolution=96)
+    midi_data.time_signature_changes.append(pretty_midi.containers.TimeSignature(4, 4, 0))
+    instrument = pretty_midi.Instrument(0)
+    midi_data.instruments.append(instrument)
+
+    for event in midi_events[0]:
+        pitch = event['pitch']
+        velocity = event['velocity']
+        start = midi_data.tick_to_time(event['start'])
+        end = midi_data.tick_to_time(event['end'])
+        note = pretty_midi.Note(pitch=pitch, velocity=velocity, start=start, end=end)
+        instrument.notes.append(note)
+
+    midi_path = os.path.join(temp_dir, 'output.mid')
+    midi_data.write(midi_path)
+    print(f"Generated: {midi_path}")
+
+
+def render_wav(midi_file, uploaded_sf2=None, output_level='2.0'):
+    sf2_dir = 'sf2_kits'
+    audio_format = 's16'
+    sample_rate = '44100'
+    gain = str(output_level)
+
+    if uploaded_sf2:
+        sf2_file = uploaded_sf2
+    else:
+        sf2_files = [f for f in os.listdir(os.path.join(sf2_dir)) if f.endswith('.sf2')]
+        if not sf2_files:
+            raise ValueError("No SoundFont (.sf2) file found in directory.")
+        sf2_file = os.path.join(sf2_dir, random.choice(sf2_files))
+
+    output_wav = os.path.join(temp_dir, 'output.wav')
+
+    with open(os.devnull, 'w') as devnull:
+        command = [
+            'fluidsynth', '-ni', sf2_file, midi_file, '-F', output_wav, '-r', str(sample_rate), 
+            '-o', f'audio.file.format={audio_format}', '-g', str(gain)
+        ]
+        subprocess.call(command, stdout=devnull, stderr=devnull)
+
+    return output_wav
+
+
+def generate_and_return_files(bpm, temperature, top_k, uploaded_sf2=None, output_level='2.0'):
+    midi_events = generate_midi(temperature, top_k)  
+    if not midi_events:
+        return "Error generating MIDI.", None, None
+    
+    write_single_midi(midi_events, bpm)
+    
+    midi_file = os.path.join(temp_dir, 'output.mid')
+    wav_raw = render_wav(midi_file, uploaded_sf2, output_level)
+    wav_fx = os.path.join(temp_dir, 'output_fx.wav')
+
+    sfx_settings = [
+        {
+            'board': Pedalboard([
+                Reverb(room_size=0.01, wet_level=random.uniform(0.005, 0.01), dry_level=0.75, width=1.0),
+                Compressor(threshold_db=-3.0, ratio=8.0, attack_ms=0.0, release_ms=300.0),
+            ])
+        }
+    ]
+
+    for setting in sfx_settings:
+        board = setting['board']
+
+        with AudioFile(wav_raw) as f:
+            with AudioFile(wav_fx, 'w', f.samplerate, f.num_channels) as o:
+                while f.tell() < f.frames:
+                    chunk = f.read(int(f.samplerate))
+                    effected = board(chunk, f.samplerate, reset=False)
+                    o.write(effected)
+
+    return midi_file, wav_fx
+
+
+custom_css = """
+#container {
+  max-width: 1200px !important;
+  margin: 0 auto !important;
+}
+#generate-btn {
+  font-size: 18px;
+  color: white;
+  padding: 10px 20px;
+  border: none;
+  border-radius: 5px;
+  cursor: pointer;
+  background: linear-gradient(90deg, hsla(268, 90%, 70%, 1) 0%, hsla(260, 72%, 74%, 1) 50%, hsla(247, 73%, 65%, 1) 100%);
+  transition: background 1s ease;
+}
+#generate-btn:hover {
+  color: white;
+  background: linear-gradient(90deg, hsla(268, 90%, 62%, 1) 0%, hsla(260, 70%, 70%, 1) 50%, hsla(247, 73%, 55%, 1) 100%);
+}
+#container .prose {
+  text-align: center !important;
+}
+#container h1 {
+  font-weight: bold;
+  font-size: 40px;
+  margin: 0px;
+}
+#container p {
+  font-size: 18px;
+  text-align: center;
+}
+
+"""
+
+with gr.Blocks(
+    css=custom_css,
+    theme=gr.themes.Default(
+        font=[gr.themes.GoogleFont("Roboto"), "sans-serif"],
+        primary_hue="violet",
+        secondary_hue="violet"
+    )
+) as iface:
+    with gr.Column(elem_id="container"):
+        gr.Markdown("<h1>Neural Breaks</h1>")
+        gr.Markdown("<p>Neural Breaks is a generative MIDI model trained **on** dynamic transcriptions of funk and soul drum breaks.</p>")
+
+        bpm = gr.Slider(minimum=50, maximum=200, step=1, value=100, label="BPM")
+        temperature = gr.Slider(minimum=0.1, maximum=2.0, step=0.1, value=1.0, label="Temperature")
+        top_k = gr.Slider(minimum=4, maximum=16, step=1, value=8, label="Top-k")
+        output_level = gr.Slider(minimum=0, maximum=3, step=0.10, value=2.0, label="Output Gain")
+        generate_button = gr.Button("Generate", elem_id="generate-btn")
+        midi_file = gr.File(label="MIDI Output")
+        audio_file = gr.Audio(label="Audio Output", type="filepath")
+        soundfont = gr.File(label="Optional: Upload SoundFont (preset=0, bank=0)")
+
+        generate_button.click(
+            fn=generate_and_return_files,
+            inputs=[bpm, temperature, top_k, soundfont, output_level],
+            outputs=[midi_file, audio_file]
+        )
+
+        gr.Markdown("<p style='font-size: 16px;'>Developed by <a href='https://www.patchbanks.com/' target='_blank'><strong>Patchbanks</strong></a></p>")
+
+iface.launch(share=True)

checkpoints/model.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:dc9998fce7cc7b59f949df3ec0dacd12fcc0baf15c8cb9e5f0fa62922e45030d
+size 24327014

config/eval_gpt2.py

@@ -0,0 +1,8 @@
+# evaluate the base gpt2
+# n_layer=12, n_head=12, n_embd=768
+# 124M parameters
+batch_size = 8
+eval_iters = 500 # use more iterations to get good estimate
+eval_only = True
+wandb_log = False
+init_from = 'gpt2'

config/eval_gpt2_large.py

@@ -0,0 +1,8 @@
+# evaluate the base gpt2
+# n_layer=36, n_head=20, n_embd=1280
+# 774M parameters
+batch_size = 8
+eval_iters = 500 # use more iterations to get good estimate
+eval_only = True
+wandb_log = False
+init_from = 'gpt2-large'

config/eval_gpt2_medium.py

@@ -0,0 +1,8 @@
+# evaluate the base gpt2
+# n_layer=24, n_head=16, n_embd=1024
+# 350M parameters
+batch_size = 8
+eval_iters = 500 # use more iterations to get good estimate
+eval_only = True
+wandb_log = False
+init_from = 'gpt2-medium'

config/eval_gpt2_xl.py

@@ -0,0 +1,8 @@
+# evaluate the base gpt2
+# n_layer=48, n_head=25, n_embd=1600
+# 1558M parameters
+batch_size = 8
+eval_iters = 500 # use more iterations to get good estimate
+eval_only = True
+wandb_log = False
+init_from = 'gpt2-xl'

config/finetune_shakespeare.py

@@ -0,0 +1,25 @@
+import time
+
+out_dir = 'out-shakespeare'
+eval_interval = 5
+eval_iters = 40
+wandb_log = False # feel free to turn on
+wandb_project = 'shakespeare'
+wandb_run_name = 'ft-' + str(time.time())
+
+dataset = 'shakespeare'
+init_from = 'gpt2-xl' # this is the largest GPT-2 model
+
+# only save checkpoints if the validation loss improves
+always_save_checkpoint = False
+
+# the number of examples per iter:
+# 1 batch_size * 32 grad_accum * 1024 tokens = 32,768 tokens/iter
+# shakespeare has 301,966 tokens, so 1 epoch ~= 9.2 iters
+batch_size = 1
+gradient_accumulation_steps = 32
+max_iters = 20
+
+# finetune at constant LR
+learning_rate = 3e-5
+decay_lr = False

config/train_gpt2.py

@@ -0,0 +1,25 @@
+# config for training GPT-2 (124M) down to very nice loss of ~2.85 on 1 node of 8X A100 40GB
+# launch as the following (e.g. in a screen session) and wait ~5 days:
+# $ torchrun --standalone --nproc_per_node=8 train.py config/train_gpt2.py
+
+wandb_log = True
+wandb_project = 'owt'
+wandb_run_name='gpt2-124M'
+
+# these make the total batch size be ~0.5M
+# 12 batch size * 1024 block size * 5 gradaccum * 8 GPUs = 491,520
+batch_size = 12
+block_size = 1024
+gradient_accumulation_steps = 5 * 8
+
+# this makes total number of tokens be 300B
+max_iters = 600000
+lr_decay_iters = 600000
+
+# eval stuff
+eval_interval = 1000
+eval_iters = 200
+log_interval = 10
+
+# weight decay
+weight_decay = 1e-1

config/train_shakespeare_char.py

@@ -0,0 +1,37 @@
+# train a miniature character-level shakespeare model
+# good for debugging and playing on macbooks and such
+
+out_dir = 'out-shakespeare-char'
+eval_interval = 250 # keep frequent because we'll overfit
+eval_iters = 200
+log_interval = 10 # don't print too too often
+
+# we expect to overfit on this small dataset, so only save when val improves
+always_save_checkpoint = False
+
+wandb_log = False # override via command line if you like
+wandb_project = 'shakespeare-char'
+wandb_run_name = 'mini-gpt'
+
+dataset = 'shakespeare_char'
+gradient_accumulation_steps = 1
+batch_size = 64
+block_size = 256 # context of up to 256 previous characters
+
+# baby GPT model :)
+n_layer = 6
+n_head = 6
+n_embd = 384
+dropout = 0.2
+
+learning_rate = 1e-3 # with baby networks can afford to go a bit higher
+max_iters = 5000
+lr_decay_iters = 5000 # make equal to max_iters usually
+min_lr = 1e-4 # learning_rate / 10 usually
+beta2 = 0.99 # make a bit bigger because number of tokens per iter is small
+
+warmup_iters = 100 # not super necessary potentially
+
+# on macbook also add
+# device = 'cpu'  # run on cpu only
+# compile = False # do not torch compile the model

configurator.py

@@ -0,0 +1,57 @@
+"""
+Poor Man's Configurator. Probably a terrible idea. Example usage:
+$ python train.py config/override_file.py --batch_size=32
+this will first run config/override_file.py, then override batch_size to 32
+
+The code in this file will be run as follows from e.g. train.py:
+>>> exec(open('configurator.py').read())
+
+So it's not a Python module, it's just shuttling this code away from train.py
+The code in this script then overrides the globals()
+
+I know people are not going to love this, I just really dislike configuration
+complexity and having to prepend config. to every single variable. If someone
+comes up with a better simple Python solution I am all ears.
+"""
+
+import sys
+from ast import literal_eval
+import argparse
+
+parser = argparse.ArgumentParser(description="nanoMPC")
+parser.add_argument("--bpm", type=int, default=90, help="Beats per minute")
+parser.add_argument("--num_samples", type=int, default=1, help="Number of samples")
+args, unknown_args = parser.parse_known_args()  # Capture unknown args for configurator
+
+bpm = args.bpm
+num_samples = args.num_samples
+
+for arg in unknown_args:
+    if arg.startswith('--'):
+        print(f"Skipping command-line argument: {arg}")
+        continue
+
+    if '=' not in arg:
+        config_file = arg
+        print(f"Overriding config with {config_file}:")
+        with open(config_file) as f:
+            print(f.read())
+        exec(open(config_file).read())
+    else:
+        key, val = arg.split('=')
+        key = key[2:]
+
+        if key in globals():
+            if key in ['bpm', 'num_samples']:
+                continue
+            try:
+                attempt = literal_eval(val)
+            except (SyntaxError, ValueError):
+                attempt = val
+            assert type(attempt) == type(globals()[key])
+            print(f"Overriding: {key} = {attempt}")
+            globals()[key] = attempt
+        else:
+            raise ValueError(f"Unknown config key: {key}")
+
+

data/neural_breaks/meta.pkl

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:20108664c20f9b9e9393eb222b8d5f33b947c2c448e2c06f33e4f6fc557f0fbb
+size 1184

model.py

@@ -0,0 +1,301 @@
+import inspect
+from dataclasses import dataclass
+from torch.utils.data import DataLoader, Dataset
+from torch.nn import functional as F
+import torch.nn as nn
+import torch
+import math
+import numpy as np
+
+
+class RMSNorm(torch.nn.Module):
+    def __init__(self, dim: int, eps: float):
+        super().__init__()
+        self.eps = eps
+        self.weight = nn.Parameter(torch.ones(dim))
+
+    def _norm(self, x):
+        return x * torch.rsqrt(x.pow(2).mean(-1, keepdim=True) + self.eps)
+
+    def forward(self, x):
+        output = self._norm(x.float()).type_as(x)
+        return output * self.weight
+
+
+class CausalSelfAttention(nn.Module):
+    def __init__(self, config):
+        super().__init__()
+        assert config.n_embd % config.n_head == 0
+        self.config = config  # Store the config object
+
+        self.n_head = config.n_head
+        self.n_embd = config.n_embd
+        self.dropout = config.dropout
+
+        self.c_attn = nn.Linear(config.n_embd, 3 * config.n_embd, bias=config.bias)
+        self.c_proj = nn.Linear(config.n_embd, config.n_embd, bias=config.bias)
+        self.attn_dropout = nn.Dropout(config.dropout)
+        self.resid_dropout = nn.Dropout(config.dropout)
+
+        self.rel_attn_bias = nn.Embedding(config.block_size * 2 - 1, config.n_head)
+
+    def forward(self, x):
+        B, T, C = x.size()
+        q, k, v = self.c_attn(x).split(self.n_embd, dim=2)
+        q = q.view(B, T, self.n_head, C // self.n_head).transpose(1, 2)
+        k = k.view(B, T, self.n_head, C // self.n_head).transpose(1, 2)
+        v = v.view(B, T, self.n_head, C // self.n_head).transpose(1, 2)
+
+        if hasattr(torch.nn.functional, 'scaled_dot_product_attention'):
+            attn_logits = torch.nn.functional.scaled_dot_product_attention(q, k, v, attn_mask=None, dropout_p=self.dropout if self.training else 0, is_causal=True)
+        else:
+            attn_logits = (q @ k.transpose(-2, -1)) / math.sqrt(C // self.n_head)
+            max_rpe = self.config.block_size // 2  # Use config object
+            rpe_matrix = self.generate_rpe(T, max_rpe).to(x.device)
+            rpe_embeddings = self.rel_attn_bias(rpe_matrix).transpose(1, 2).unsqueeze(0)
+            attn_logits = attn_logits + rpe_embeddings
+            attn_logits = attn_logits.masked_fill(self.bias[:,:,:T,:T] == 0, float('-inf'))
+            attn_logits = F.softmax(attn_logits, dim=-1)
+            attn_logits = self.attn_dropout(attn_logits)
+            attn_logits = attn_logits @ v
+
+        y = attn_logits.transpose(1, 2).contiguous().view(B, T, C)
+        y = self.resid_dropout(self.c_proj(y))
+        return y
+
+    def generate_rpe(self, length, max_rpe):
+        range_vec = torch.arange(length)
+        range_mat = range_vec.unsqueeze(0) - range_vec.unsqueeze(1)
+        range_mat_clipped = torch.clamp(range_mat, -max_rpe, max_rpe)
+        final_mat = range_mat_clipped + max_rpe
+        return final_mat
+
+
+class MLP(nn.Module):
+
+    def __init__(self, config):
+        super().__init__()
+        self.c_fc    = nn.Linear(config.n_embd, 4 * config.n_embd, bias=config.bias)
+        self.gelu    = nn.GELU()
+        self.c_proj  = nn.Linear(4 * config.n_embd, config.n_embd, bias=config.bias)
+        self.dropout = nn.Dropout(config.dropout)
+
+    def forward(self, x):
+        x = self.c_fc(x)
+        x = self.gelu(x)
+        x = self.c_proj(x)
+        x = self.dropout(x)
+        return x
+
+class Block(nn.Module):
+    def __init__(self, config):
+        super().__init__()
+        self.ln_1 = RMSNorm(config.n_embd, eps=1e-5)
+        self.attn = CausalSelfAttention(config)
+        self.ln_2 = RMSNorm(config.n_embd, eps=1e-5)
+        self.mlp = MLP(config)
+
+    def forward(self, x):
+        x = x + self.attn(self.ln_1(x))
+        x = x + self.mlp(self.ln_2(x))
+        return x
+
+@dataclass
+class GPTConfig:
+    block_size: int = 1024
+    vocab_size: int = 50304
+    n_layer: int = 12
+    n_head: int = 12
+    n_embd: int = 768
+    dropout: float = 0.0
+    bias: bool = True
+
+class GPT(nn.Module):
+    def __init__(self, config):
+        super().__init__()
+        assert config.vocab_size is not None
+        assert config.block_size is not None
+        self.config = config
+
+        self.transformer = nn.ModuleDict(dict(
+            wte = nn.Embedding(config.vocab_size, config.n_embd),
+            wpe = nn.Embedding(config.block_size, config.n_embd),
+            drop = nn.Dropout(config.dropout),
+            h = nn.ModuleList([Block(config) for _ in range(config.n_layer)]),
+            ln_f = RMSNorm(config.n_embd, eps=1e-5),
+        ))
+        self.lm_head = nn.Linear(config.n_embd, config.vocab_size, bias=False)
+        self.transformer.wte.weight = self.lm_head.weight
+
+        self.apply(self._init_weights)
+        for pn, p in self.named_parameters():
+            if pn.endswith('c_proj.weight'):
+                torch.nn.init.normal_(p, mean=0.0, std=0.02/math.sqrt(2 * config.n_layer))
+
+        #print("number of parameters: %.2fM" % (self.get_num_params()/1e6,))
+
+    def get_num_params(self, non_embedding=True):
+        n_params = sum(p.numel() for p in self.parameters())
+        if non_embedding:
+            n_params -= self.transformer.wpe.weight.numel()
+        return n_params
+
+    def _init_weights(self, module):
+        if isinstance(module, nn.Linear):
+            torch.nn.init.normal_(module.weight, mean=0.0, std=0.02)
+            if module.bias is not None:
+                torch.nn.init.zeros_(module.bias)
+        elif isinstance(module, nn.Embedding):
+            torch.nn.init.normal_(module.weight, mean=0.0, std=0.02)
+
+    def forward(self, idx, targets=None, noise_pct=0.1):
+        device = idx.device
+        b, t = idx.size()
+        assert t <= self.config.block_size, f"Cannot forward sequence of length {t}, block size is only {self.config.block_size}"
+        pos = torch.arange(0, t, dtype=torch.long, device=device)
+
+        tok_emb = self.transformer.wte(idx)
+        pos_emb = self.transformer.wpe(pos)
+
+        # add noise to the input
+        if noise_pct > 0.0:
+            noise_std = torch.std(tok_emb) * noise_pct
+            noise = torch.randn_like(tok_emb) * noise_std
+            tok_emb = tok_emb + noise
+
+        x = self.transformer.drop(tok_emb + pos_emb)
+        for block in self.transformer.h:
+            x = block(x)
+        x = self.transformer.ln_f(x)
+
+        if targets is not None:
+
+            logits = self.lm_head(x)
+            loss = F.cross_entropy(logits.view(-1, logits.size(-1)), targets.view(-1), ignore_index=-1)
+        else:
+
+            logits = self.lm_head(x[:, [-1], :])
+            loss = None
+
+        return logits, loss
+
+
+    def crop_block_size(self, block_size):
+        assert block_size <= self.config.block_size
+        self.config.block_size = block_size
+        self.transformer.wpe.weight = nn.Parameter(self.transformer.wpe.weight[:block_size])
+        for block in self.transformer.h:
+            if hasattr(block.attn, 'bias'):
+                block.attn.bias = block.attn.bias[:,:,:block_size,:block_size]
+
+    @classmethod
+    def from_pretrained(cls, model_type, override_args=None):
+        assert model_type in {'gpt2', 'gpt2-medium', 'gpt2-large', 'gpt2-xl'}
+        override_args = override_args or {}
+
+        assert all(k == 'dropout' for k in override_args)
+        from transformers import GPT2LMHeadModel
+        print("loading weights from pretrained gpt: %s" % model_type)
+
+        config_args = {
+            'gpt2':         dict(n_layer=12, n_head=12, n_embd=768),
+            'gpt2-medium':  dict(n_layer=24, n_head=16, n_embd=1024),
+            'gpt2-large':   dict(n_layer=36, n_head=20, n_embd=1280),
+            'gpt2-xl':      dict(n_layer=48, n_head=25, n_embd=1600),
+        }[model_type]
+        print("forcing vocab_size=50257, block_size=1024, bias=True")
+        config_args['vocab_size'] = 50257
+        config_args['block_size'] = 1024
+        config_args['bias'] = True
+
+        if 'dropout' in override_args:
+            print(f"overriding dropout rate to {override_args['dropout']}")
+            config_args['dropout'] = override_args['dropout']
+
+        config = GPTConfig(**config_args)
+        model = GPT(config)
+        sd = model.state_dict()
+        sd_keys = sd.keys()
+        sd_keys = [k for k in sd_keys if not k.endswith('.attn.bias')]
+
+        model_hf = GPT2LMHeadModel.from_pretrained(model_type)
+        sd_hf = model_hf.state_dict()
+
+        sd_keys_hf = sd_hf.keys()
+        sd_keys_hf = [k for k in sd_keys_hf if not k.endswith('.attn.masked_bias')]
+        sd_keys_hf = [k for k in sd_keys_hf if not k.endswith('.attn.bias')]
+        transposed = ['attn.c_attn.weight', 'attn.c_proj.weight', 'mlp.c_fc.weight', 'mlp.c_proj.weight']
+
+        assert len(sd_keys_hf) == len(sd_keys), f"mismatched keys: {len(sd_keys_hf)} != {len(sd_keys)}"
+        for k in sd_keys_hf:
+            if any(k.endswith(w) for w in transposed):
+
+                assert sd_hf[k].shape[::-1] == sd[k].shape
+                with torch.no_grad():
+                    sd[k].copy_(sd_hf[k].t())
+            else:
+
+                assert sd_hf[k].shape == sd[k].shape
+                with torch.no_grad():
+                    sd[k].copy_(sd_hf[k])
+
+        return model
+
+    def configure_optimizers(self, weight_decay, learning_rate, betas, device_type):
+
+        param_dict = {pn: p for pn, p in self.named_parameters()}
+
+        param_dict = {pn: p for pn, p in param_dict.items() if p.requires_grad}
+
+        decay_params = [p for n, p in param_dict.items() if p.dim() >= 2]
+        nodecay_params = [p for n, p in param_dict.items() if p.dim() < 2]
+        optim_groups = [
+            {'params': decay_params, 'weight_decay': weight_decay},
+            {'params': nodecay_params, 'weight_decay': 0.0}
+        ]
+        num_decay_params = sum(p.numel() for p in decay_params)
+        num_nodecay_params = sum(p.numel() for p in nodecay_params)
+        print(f"num decayed parameter tensors: {len(decay_params)}, with {num_decay_params:,} parameters")
+        print(f"num non-decayed parameter tensors: {len(nodecay_params)}, with {num_nodecay_params:,} parameters")
+
+        fused_available = 'fused' in inspect.signature(torch.optim.AdamW).parameters
+        use_fused = fused_available and device_type == 'cuda'
+        extra_args = dict(fused=True) if use_fused else dict()
+        optimizer = torch.optim.AdamW(optim_groups, lr=learning_rate, betas=betas, **extra_args)
+        print(f"using fused AdamW: {use_fused}")
+
+        return optimizer
+
+    def estimate_mfu(self, fwdbwd_per_iter, dt):
+        """ estimate model flops utilization (MFU) in units of A100 bfloat16 peak FLOPS """
+
+        N = self.get_num_params()
+        cfg = self.config
+        L, H, Q, T = cfg.n_layer, cfg.n_head, cfg.n_embd//cfg.n_head, cfg.block_size
+        flops_per_token = 6*N + 12*L*H*Q*T
+        flops_per_fwdbwd = flops_per_token * T
+        flops_per_iter = flops_per_fwdbwd * fwdbwd_per_iter
+
+        flops_achieved = flops_per_iter * (1.0/dt)
+        flops_promised = 312e12
+        mfu = flops_achieved / flops_promised
+        return mfu
+
+    @torch.no_grad()
+    def generate(self, idx, max_new_tokens, temperature=1.0, top_k=None):
+        for _ in range(max_new_tokens):
+            idx_cond = idx if idx.size(1) <= self.config.block_size else idx[:, -self.config.block_size:]
+            logits, _ = self(idx_cond)
+            logits = logits[:, -1, :] / temperature
+            if top_k is not None:
+                v, _ = torch.topk(logits, min(top_k, logits.size(-1)))
+                logits[logits < v[:, [-1]]] = -float('Inf')
+            probs = F.softmax(logits, dim=-1)
+            idx_next = torch.multinomial(probs, num_samples=1)
+            
+            if idx_next.item() == 0: # stop token
+                break
+            
+            idx = torch.cat((idx, idx_next), dim=1)
+
+        return idx

packages.txt

@@ -0,0 +1 @@
+fluidsynth

requirements.txt

@@ -0,0 +1,4 @@
+pretty_midi==0.2.10
+pedalboard==0.9.3
+torch
+gradio

sf2_kits/drum_breaks.sf2

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:4fc32e717f84b48194933c9b929d3727306dede46589b76af9c2b1077b4221dd
+size 31717174