Upload app.py

app.py

@@ -10,6 +10,7 @@ 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
 
 in_space = os.getenv("SYSTEM") == "spaces"
 
@@ -22,7 +23,7 @@ ckpt_load = 'model.pt'
 
 start = "000000000000\n"
 num_samples = 1
-max_new_tokens = 564
+max_new_tokens = 384
 
 seed = random.randint(1, 100000)
 torch.manual_seed(seed)
@@ -58,9 +59,9 @@ model.to(device)
 if compile:
     model = torch.compile(model)
 
-tokenizer = re.compile(r'000000000000|\d{1}|\n')
+tokenizer = re.compile(r'000000000000|\d{2}|\n')
 
-meta_path = os.path.join('', 'meta.pkl')
+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)
@@ -131,7 +132,6 @@ def generate_midi(temperature, top_k):
     return midi_events
 
 
-
 def write_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))
@@ -152,19 +152,21 @@ def write_midi(midi_events, bpm):
     print(f"Generated: {midi_path}")
 
 
-def render_wav(midi_file):
-    
+def render_wav(midi_file, uploaded_sf2=None):
     sf2_dir = 'sf2_kits'
     audio_format = 's16'
     sample_rate = '44100'
     gain = '2.0'
 
-    sf2_files = [f for f in os.listdir(sf2_dir) if f.endswith('.sf2')]
-    if not sf2_files:
-        raise ValueError("No SoundFont (.sf2) file found in directory.")
+    if uploaded_sf2:
+        sf2_file = uploaded_sf2
+    else:
+        sf2_files = [f for f in os.listdir(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))
 
-    sf2_file = os.path.join(sf2_dir, random.choice(sf2_files))
-    print(sf2_file)
+    print(f"Using SoundFont: {sf2_file}")
     output_wav = os.path.join(temp_dir, 'output.wav')
 
     with open(os.devnull, 'w') as devnull:
@@ -177,23 +179,7 @@ def render_wav(midi_file):
     return output_wav
 
 
-def render_sfx(wav_raw, settings):
-    wav_fx = wav_raw
-
-    for setting in 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 wav_fx
-
-
-def generate_and_return_files(bpm, temperature, top_k):
+def generate_and_return_files(bpm, temperature, top_k, uploaded_sf2=None):
     midi_events = generate_midi(temperature, top_k)  
     if not midi_events:
         return "Error generating MIDI.", None, None
@@ -201,7 +187,7 @@ def generate_and_return_files(bpm, temperature, top_k):
     write_midi(midi_events, bpm)
     
     midi_file = os.path.join(temp_dir, 'output.mid')
-    wav_raw = render_wav(midi_file)
+    wav_raw = render_wav(midi_file, uploaded_sf2)
     wav_fx = os.path.join(temp_dir, 'output_fx.wav')
 
     sfx_settings = [
@@ -226,22 +212,45 @@ def generate_and_return_files(bpm, temperature, top_k):
     return midi_file, wav_fx
 
 
+custom_css = """
+#generate-btn {
+    background-color: #6366f1 !important;
+    color: white !important;
+    border: none !important;
+    font-size: 16px;
+    padding: 10px 20px;
+    border-radius: 5px;
+    cursor: pointer;
+}
+#generate-btn:hover {
+    background-color: #4f51c5 !important;
+}
+"""
+
+with gr.Blocks(css=custom_css, theme="soft") as iface:
+    gr.Markdown("<h1 style='font-weight: bold; text-align: center;'>nanoMPC - AI Midi Drum Sequencer</h1>")
+    gr.Markdown("<p style='text-align:center;'>nanoMPC is a tiny transformer model that generates MIDI drum beats inspired by Lo-Fi, Boom Bap and other styles of Hip Hop.</p>")
+    
+    with gr.Row():
+        with gr.Column(scale=1):
+            bpm = gr.Slider(minimum=50, maximum=200, step=1, value=90, 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=256, step=1, value=128, label="Top-k")
+            soundfont = gr.File(label="Optional: Upload SoundFont (preset=0, bank=0)")
+        
+        with gr.Column(scale=1):
+            midi_file = gr.File(label="MIDI File Output")
+            audio_file = gr.Audio(label="Generated Audio Output", type="filepath")
+            generate_button = gr.Button("Generate", elem_id="generate-btn")
+    
+    generate_button.click(
+        fn=generate_and_return_files,
+        inputs=[bpm, temperature, top_k, soundfont],
+        outputs=[midi_file, audio_file]
+    )
+
+iface.launch(share=True)
+
+
+
 
-iface = gr.Interface(
-    fn=generate_and_return_files,
-    inputs=[
-        gr.Slider(minimum=50, maximum=200, step=1, value=87, label="bpm"),
-        gr.Slider(minimum=0.1, maximum=2.0, step=0.1, value=1.0, label="temperature"),
-        gr.Slider(minimum=4, maximum=128, step=1, value=16, label="top_k")
-    ],
-    outputs=[
-        gr.File(label="MIDI File"),
-        gr.Audio(label="Generated Audio", type="filepath")
-    ],
-    title="<h1 style='font-weight: bold; text-align: center;'>nanoMPC - AI Midi Drum Sequencer</h1>",
-    description="<p style='text-align:center;'>nanoMPC is a tiny transformer model that generates MIDI drum beats inspired by Lo-Fi, Boom Bap and other styles of Hip Hop.</p>",
-    theme="soft",
-    allow_flagging="never",
-)
-
-iface.launch()