File size: 17,049 Bytes
a8c245b
 
 
 
12a4852
 
 
 
 
a3ef88a
 
5f04547
 
 
7dd2cd1
 
 
 
a3ef88a
 
 
4154701
a3ef88a
12a4852
 
 
 
a3ef88a
12a4852
 
 
 
 
465603c
9236ae7
053cdcc
83355ca
12a4852
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
a3ef88a
12a4852
 
 
 
 
 
 
7b3dcbc
12a4852
 
 
 
a3ef88a
12a4852
 
 
 
 
018d705
 
 
 
 
 
 
40e4ca3
018d705
 
12a4852
 
 
 
 
 
 
 
 
018d705
 
 
12a4852
 
 
 
 
 
 
018d705
12a4852
 
 
 
018d705
 
 
 
12a4852
 
 
 
 
5244019
12a4852
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
a3ef88a
12a4852
 
 
e38b485
39917f5
7253d65
e38b485
 
a3ef88a
39917f5
12a4852
 
 
 
 
 
 
 
 
 
39917f5
a3ef88a
12a4852
 
 
 
 
 
 
a3ef88a
39917f5
 
12a4852
 
 
 
 
 
 
 
 
 
 
 
b5a3cae
 
12a4852
a3ef88a
b5a3cae
 
 
12a4852
 
 
 
 
 
 
 
 
e0b1ad3
12a4852
 
 
 
4154701
12a4852
 
 
 
 
a3ef88a
12a4852
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
3fd4c9e
 
 
 
 
 
 
 
 
 
 
 
 
12a4852
 
a3ef88a
12a4852
bd45f96
12a4852
3fd4c9e
bd45f96
12a4852
 
 
 
 
 
3b8a3fc
12a4852
 
 
 
 
 
 
5244019
12a4852
7b3dcbc
2abc7eb
e0b1ad3
5244019
12a4852
 
 
 
 
 
5dfe7a2
12a4852
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
8586ed2
 
5244019
3b8a3fc
3524d42
de20229
12a4852
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
8586ed2
 
a3ef88a
12a4852
 
 
 
 
8586ed2
5244019
 
6eeabd7
12a4852
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
8586ed2
 
a3ef88a
12a4852
 
 
8edbfbe
12a4852
 
 
8edbfbe
12a4852
 
 
a3ef88a
1446958
da2ccaf
 
14b9312
12a4852
 
 
 
 
 
 
 
 
 
8586ed2
 
 
12a4852
 
a3ef88a
12a4852
 
dd9d99a
12a4852
5d5971c
ed1250f
5d5971c
a3ef88a
 
 
12a4852
8586ed2
12a4852
a3ef88a
12a4852
 
a3ef88a
 
12a4852
 
 
 
 
 
 
 
 
a3ef88a
 
d6f53f9
12a4852
 
a3ef88a
 
 
12a4852
 
 
 
 
 
 
 
 
 
 
 
5c91d9e
12a4852
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
391
392
393
394
395
396
397
398
399
400
401
402
403
404
405
406
407
408
409
410
411
412
413
414
415
416
417
418
419
420
421
422
423
424
425
426
427
428
429
430
431
432
433
434
435
436
437
438
439
440
#====================================================================
# https://huggingface.co/spaces/asigalov61/Godzilla-Piano-Transformer
#====================================================================

"""
Godzilla Piano Transformer Gradio App - Single Model, Simplified Version
Fast 807M 4k solo Piano music transformer trained on 1.14M+ MIDIs (2.7M+ samples)
Using only one model: "without velocity - 3 epochs"
"""

import os

os.environ["HF_HUB_ENABLE_HF_TRANSFER"] = "1"

import time as reqtime
import datetime
from pytz import timezone

import torch
import matplotlib.pyplot as plt
import gradio as gr
import spaces

from huggingface_hub import hf_hub_download
import TMIDIX
from midi_to_colab_audio import midi_to_colab_audio
from x_transformer_2_3_1 import TransformerWrapper, AutoregressiveWrapper, Decoder

# -----------------------------
# CONFIGURATION & GLOBALS
# -----------------------------
SEP = '=' * 70
PDT = timezone('US/Pacific')

MODEL_CHECKPOINT = 'Godzilla_Piano_Transformer_No_Velocity_Trained_Model_21113_steps_0.3454_loss_0.895_acc.pth'
SOUDFONT_PATH = 'SGM-v2.01-YamahaGrand-Guit-Bass-v2.7.sf2'
NUM_OUT_BATCHES = 12
PREVIEW_LENGTH = 120  # in tokens

# -----------------------------
# PRINT START-UP INFO
# -----------------------------
def print_sep():
    print(SEP)

print_sep()
print("Godzilla Piano Transformer Gradio App")
print_sep()
print("Loading modules...")

# -----------------------------
# ENVIRONMENT & PyTorch Settings
# -----------------------------
os.environ['USE_FLASH_ATTENTION'] = '1'

torch.set_float32_matmul_precision('high')
torch.backends.cuda.matmul.allow_tf32 = True
torch.backends.cudnn.allow_tf32 = True
torch.backends.cuda.enable_mem_efficient_sdp(True)
torch.backends.cuda.enable_math_sdp(True)
torch.backends.cuda.enable_flash_sdp(True)
torch.backends.cuda.enable_cudnn_sdp(True)

print_sep()
print("PyTorch version:", torch.__version__)
print("Done loading modules!")
print_sep()

# -----------------------------
# MODEL INITIALIZATION
# -----------------------------
print_sep()
print("Instantiating model...")

device_type = 'cuda'
dtype = 'bfloat16'
ptdtype = {'bfloat16': torch.bfloat16, 'float16': torch.float16}[dtype]
ctx = torch.amp.autocast(device_type=device_type, dtype=ptdtype)

SEQ_LEN = 4096
PAD_IDX = 384

model = TransformerWrapper(
    num_tokens=PAD_IDX + 1,
    max_seq_len=SEQ_LEN,
    attn_layers=Decoder(
        dim=2048,
        depth=16,
        heads=32,
        rotary_pos_emb=True,
        attn_flash=True
    )
)
model = AutoregressiveWrapper(model, ignore_index=PAD_IDX, pad_value=PAD_IDX)

print_sep()
print("Loading model checkpoint...")
checkpoint = hf_hub_download(
    repo_id='asigalov61/Godzilla-Piano-Transformer',
    filename=MODEL_CHECKPOINT
)
model.load_state_dict(torch.load(checkpoint, map_location='cuda', weights_only=True))
model = torch.compile(model, mode='max-autotune')
print_sep()
print("Done!")
print("Model will use", dtype, "precision...")
print_sep()

model.cuda()
model.eval()

# -----------------------------
# HELPER FUNCTIONS
# -----------------------------
def render_midi_output(final_composition):
    """Generate MIDI score, plot, and audio from final composition."""
    fname, midi_score = save_midi(final_composition)
    time_val = midi_score[-1][1] / 1000  # seconds marker from last note
    midi_plot = TMIDIX.plot_ms_SONG(
        midi_score,
        plot_title='Godzilla Piano Transformer Composition',
        block_lines_times_list=[],
        return_plt=True
    )
    midi_audio = midi_to_colab_audio(
        fname + '.mid',
        soundfont_path=SOUDFONT_PATH,
        sample_rate=16000,
        output_for_gradio=True
    )
    return (16000, midi_audio), midi_plot, fname + '.mid', time_val

# -----------------------------
# MIDI PROCESSING FUNCTIONS
# -----------------------------
def load_midi(input_midi):
    """Process the input MIDI file and create a token sequence using without velocity logic."""
    raw_score = TMIDIX.midi2single_track_ms_score(input_midi.name)
    escore_notes = TMIDIX.advanced_score_processor(
        raw_score, return_enhanced_score_notes=True, apply_sustain=True
    )[0]
    sp_escore_notes = TMIDIX.solo_piano_escore_notes(escore_notes)
    zscore = TMIDIX.recalculate_score_timings(sp_escore_notes)
    zscore = TMIDIX.augment_enhanced_score_notes(zscore, timings_divider=32)
    fscore = TMIDIX.fix_escore_notes_durations(zscore)
    cscore = TMIDIX.chordify_score([1000, fscore])
    
    score = []
    prev_chord = cscore[0]
    for chord in cscore:
        # Time difference token.
        score.append(max(0, min(127, chord[0][1] - prev_chord[0][1])))
        for note in chord:
            score.extend([
                max(1, min(127, note[2])) + 128,
                max(1, min(127, note[4])) + 256
            ])
        prev_chord = chord
    return score

def save_midi(tokens, batch_number=None):
    """Convert token sequence back to a MIDI score and write it using TMIDIX (without velocity).
       The output MIDI file name incorporates a date-time stamp.
    """
    song_events = []
    time_marker = 0
    duration = 0
    pitch = 0
    patches = [0] * 16

    for token in tokens:
        if 0 <= token < 128:
            time_marker += token * 32
        elif 128 <= token < 256:
            duration = (token - 128) * 32
        elif 256 <= token < 384:
            pitch = token - 256
            song_events.append(['note', time_marker, duration, 0, pitch, max(40, pitch), 0])
        # No velocity tokens are used.
    
    # Generate a time stamp using the PDT timezone.
    timestamp = datetime.datetime.now(PDT).strftime("%Y%m%d_%H%M%S")
    
    '''if batch_number is None:
        fname = f"Godzilla-Piano-Transformer-Music-Composition_{timestamp}"
    else:
        fname = f"Godzilla-Piano-Transformer-Music-Composition_{timestamp}_Batch_{batch_number}"'''

    fname = f"Godzilla-Piano-Transformer-Music-Composition"
    
    TMIDIX.Tegridy_ms_SONG_to_MIDI_Converter(
        song_events,
        output_signature='Godzilla Piano Transformer',
        output_file_name=fname,
        track_name='Project Los Angeles',
        list_of_MIDI_patches=patches,
        verbose=False
    )
    return fname, song_events

# -----------------------------
# MUSIC GENERATION FUNCTION (Combined)
# -----------------------------
@spaces.GPU
def generate_music(prime, num_gen_tokens, num_mem_tokens, num_gen_batches, model_temperature):
    """Generate music tokens given prime tokens and parameters."""
    inputs = prime[-num_mem_tokens:] if prime else [0]
    print("Generating...")
    inp = torch.LongTensor([inputs] * num_gen_batches).cuda()
    with ctx:
        out = model.generate(
            inp,
            num_gen_tokens,
            temperature=model_temperature,
            return_prime=False,
            verbose=False
        )
    print("Done!")
    print_sep()
    return out.tolist()

def generate_music_and_state(input_midi, num_prime_tokens, num_gen_tokens, num_mem_tokens,
                             model_temperature, final_composition, generated_batches, block_lines):
    """
    Generate tokens using the model, update the composition state, and prepare outputs.
    This function combines seed loading, token generation, and UI output packaging.
    """
    print_sep()
    print("Request start time:", datetime.datetime.now(PDT).strftime("%Y-%m-%d %H:%M:%S"))

    print('=' * 70)
    if input_midi is not None:
        fn = os.path.basename(input_midi.name)
        fn1 = fn.split('.')[0]
        print('Input file name:', fn)

    print('Num prime tokens:', num_prime_tokens)
    print('Num gen tokens:', num_gen_tokens)
    print('Num mem tokens:', num_mem_tokens)

    print('Model temp:', model_temperature)
    print('=' * 70)
    
    # Load seed from MIDI if there is no existing composition.
    if not final_composition and input_midi is not None:
        final_composition = load_midi(input_midi)[:num_prime_tokens]
        midi_fname, midi_score = save_midi(final_composition)
        # Use the last note's time as a marker.
        TMIDIX.Tegridy_ms_SONG_to_MIDI_Converter(
            midi_score,
            output_signature='Godzilla Piano Transformer',
            output_file_name=midi_fname,
            track_name='Project Los Angeles',
            list_of_MIDI_patches=[0]*16,
            verbose=False
        )
        block_lines.append(midi_score[-1][1] / 1000 if final_composition else 0)
    
    batched_gen_tokens = generate_music(final_composition, num_gen_tokens, num_mem_tokens,
                                        NUM_OUT_BATCHES, model_temperature)
    
    output_batches = []
    for i, tokens in enumerate(batched_gen_tokens):
        preview_tokens = final_composition[-PREVIEW_LENGTH:]
        midi_fname, midi_score = save_midi(preview_tokens + tokens, batch_number=i)
        plot_kwargs = {'plot_title': f'Batch # {i}', 'return_plt': True}
        if len(final_composition) > PREVIEW_LENGTH:
            plot_kwargs['preview_length_in_notes'] = len([t for t in preview_tokens if t > 256])
        TMIDIX.Tegridy_ms_SONG_to_MIDI_Converter(
            midi_score,
            output_signature='Godzilla Piano Transformer',
            output_file_name=midi_fname,
            track_name='Project Los Angeles',
            list_of_MIDI_patches=[0]*16,
            verbose=False
        )
        midi_plot = TMIDIX.plot_ms_SONG(midi_score, **plot_kwargs)
        midi_audio = midi_to_colab_audio(midi_fname + '.mid',
                                         soundfont_path=SOUDFONT_PATH,
                                         sample_rate=16000,
                                         output_for_gradio=True)
        output_batches.append([(16000, midi_audio), midi_plot, tokens])
    
    # Update generated_batches (for use by add/remove functions)
    generated_batches = batched_gen_tokens
    
    print("Request end time:", datetime.datetime.now(PDT).strftime("%Y-%m-%d %H:%M:%S"))
    print_sep()
    
    # Flatten outputs: states then audio and plots for each batch.
    outputs_flat = []
    for batch in output_batches:
        outputs_flat.extend([batch[0], batch[1]])
    return [final_composition, generated_batches, block_lines] + outputs_flat

# -----------------------------
# BATCH HANDLING FUNCTIONS
# -----------------------------
def add_batch(batch_number, final_composition, generated_batches, block_lines):
    """Add tokens from the specified batch to the final composition and update outputs."""
    if generated_batches:
        final_composition.extend(generated_batches[batch_number])
        midi_fname, midi_score = save_midi(final_composition)
        block_lines.append(midi_score[-1][1] / 1000 if final_composition else 0)
        TMIDIX.Tegridy_ms_SONG_to_MIDI_Converter(
            midi_score,
            output_signature='Godzilla Piano Transformer',
            output_file_name=midi_fname,
            track_name='Project Los Angeles',
            list_of_MIDI_patches=[0]*16,
            verbose=False
        )
        midi_plot = TMIDIX.plot_ms_SONG(
            midi_score,
            plot_title='Godzilla Piano Transformer Composition',
            block_lines_times_list=block_lines[:-1],
            return_plt=True
        )
        midi_audio = midi_to_colab_audio(midi_fname + '.mid',
                                         soundfont_path=SOUDFONT_PATH,
                                         sample_rate=16000,
                                         output_for_gradio=True)
        print("Added batch #", batch_number)
        print_sep()
        return (16000, midi_audio), midi_plot, midi_fname + '.mid', final_composition, generated_batches, block_lines
    else:
        return None, None, None, [], [], []

def remove_batch(batch_number, num_tokens, final_composition, generated_batches, block_lines):
    """Remove tokens from the final composition and update outputs."""
    if final_composition and len(final_composition) > num_tokens:
        final_composition = final_composition[:-num_tokens]
        if block_lines:
            block_lines.pop()
        midi_fname, midi_score = save_midi(final_composition)
        TMIDIX.Tegridy_ms_SONG_to_MIDI_Converter(
            midi_score,
            output_signature='Godzilla Piano Transformer',
            output_file_name=midi_fname,
            track_name='Project Los Angeles',
            list_of_MIDI_patches=[0]*16,
            verbose=False
        )
        midi_plot = TMIDIX.plot_ms_SONG(
            midi_score,
            plot_title='Godzilla Piano Transformer Composition',
            block_lines_times_list=block_lines[:-1],
            return_plt=True
        )
        midi_audio = midi_to_colab_audio(midi_fname + '.mid',
                                         soundfont_path=SOUDFONT_PATH,
                                         sample_rate=16000,
                                         output_for_gradio=True)
        print("Removed batch #", batch_number)
        print_sep()
        return (16000, midi_audio), midi_plot, midi_fname + '.mid', final_composition, generated_batches, block_lines
    else:
        return None, None, None, [], [], []

def clear():
    """Clear outputs and reset state."""
    return None, None, None, [], []

def reset(final_composition=[], generated_batches=[], block_lines=[]):
    """Reset composition state."""
    return [], [], []

# -----------------------------
# GRADIO INTERFACE SETUP
# -----------------------------
with gr.Blocks() as demo:

    gr.Markdown("<h1 style='text-align: left; margin-bottom: 1rem'>Godzilla Piano Transformer</h1>")
    gr.Markdown("<h1 style='text-align: left; margin-bottom: 1rem'>Fast 807M 4k solo Piano music transformer trained on 1.14M+ MIDIs (2.7M+ samples)</h1>")
    gr.HTML("""
        Check out <a href="https://huggingface.co/datasets/asigalov61/Godzilla-Piano">Godzilla Piano dataset</a> on Hugging Face
        <p>
            <a href="https://huggingface.co/spaces/asigalov61/Godzilla-Piano-Transformer?duplicate=true">
                <img src="https://huggingface.co/datasets/huggingface/badges/resolve/main/duplicate-this-space-md.svg" alt="Duplicate in Hugging Face">
            </a>
        </p>
        for faster execution and endless generation!
    """)

    # Global state variables for composition
    final_composition = gr.State([])
    generated_batches = gr.State([])
    block_lines = gr.State([])

    gr.Markdown("## Upload seed MIDI or click 'Generate' for a random output")
    input_midi = gr.File(label="Input MIDI", file_types=[".midi", ".mid", ".kar"])
    input_midi.upload(reset, [final_composition, generated_batches, block_lines],
                      [final_composition, generated_batches, block_lines])

    gr.Markdown("## Generate")
    num_prime_tokens = gr.Slider(15, 3072, value=3072, step=1, label="Number of prime tokens")
    num_gen_tokens = gr.Slider(15, 1024, value=512, step=1, label="Number of tokens to generate")
    num_mem_tokens = gr.Slider(15, 4096, value=4096, step=1, label="Number of memory tokens")
    model_temperature = gr.Slider(0.1, 1, value=0.9, step=0.01, label="Model temperature")
    generate_btn = gr.Button("Generate", variant="primary")

    gr.Markdown("## Batch Previews")
    outputs = [final_composition, generated_batches, block_lines]
    # Two outputs (audio and plot) for each batch
    for i in range(NUM_OUT_BATCHES):
        with gr.Tab(f"Batch # {i}"):
            audio_output = gr.Audio(label=f"Batch # {i} MIDI Audio", format="mp3")
            plot_output = gr.Plot(label=f"Batch # {i} MIDI Plot")
            outputs.extend([audio_output, plot_output])
    generate_btn.click(
        generate_music_and_state,
        [input_midi, num_prime_tokens, num_gen_tokens, num_mem_tokens, model_temperature,
         final_composition, generated_batches, block_lines],
        outputs
    )

    gr.Markdown("## Add/Remove Batch")
    batch_number = gr.Slider(0, NUM_OUT_BATCHES - 1, value=0, step=1, label="Batch number to add/remove")
    add_btn = gr.Button("Add batch", variant="primary")
    remove_btn = gr.Button("Remove batch", variant="stop")
    clear_btn = gr.ClearButton()

    final_audio_output = gr.Audio(label="Final MIDI audio", format="mp3")
    final_plot_output = gr.Plot(label="Final MIDI plot")
    final_file_output = gr.File(label="Final MIDI file")

    add_btn.click(
        add_batch,
        [batch_number, final_composition, generated_batches, block_lines],
        [final_audio_output, final_plot_output, final_file_output, final_composition, generated_batches, block_lines]
    )
    remove_btn.click(
        remove_batch,
        [batch_number, num_gen_tokens, final_composition, generated_batches, block_lines],
        [final_audio_output, final_plot_output, final_file_output, final_composition, generated_batches, block_lines]
    )
    clear_btn.click(clear, inputs=None,
                    outputs=[final_audio_output, final_plot_output, final_file_output, final_composition, block_lines])

demo.launch()