Spaces:
Running
on
Zero
Running
on
Zero
#==================================================================== | |
# 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) | |
# ----------------------------- | |
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() |