Upload app.py

app.py

@@ -1,613 +1,426 @@
-#==================================================================================
-# https://huggingface.co/spaces/asigalov61/Godzilla-Piano-Transformer
-#==================================================================================
-
-print('=' * 70)
-print('Godzilla Piano Transformer Gradio App')
-
-print('=' * 70)
-print('Loading core Godzilla Piano Transformer modules...')
+#!/usr/bin/env python
+"""
+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
-
 import time as reqtime
 import datetime
 from pytz import timezone
 
-print('=' * 70)
-print('Loading main Godzilla Piano Transformer modules...')
-
-os.environ['USE_FLASH_ATTENTION'] = '1'
-
 import torch
-
-torch.set_float32_matmul_precision('high')
-torch.backends.cuda.matmul.allow_tf32 = True # allow tf32 on matmul
-torch.backends.cudnn.allow_tf32 = True # allow tf32 on cudnn
-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)
-
-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 *
-
-import random
-
-print('=' * 70)
-print('Loading aux Godzilla Piano Transformer modules...')
-
 import matplotlib.pyplot as plt
-
 import gradio as gr
 import spaces
 
-print('=' * 70)
-print('PyTorch version:', torch.__version__)
-print('=' * 70)
-print('Done!')
-print('Enjoy! :)')
-print('=' * 70)
-
-#==================================================================================
+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
 
-MODEL_CHECKPOINTS = {
-                    'without velocity - 3 epochs': 'Godzilla_Piano_Transformer_No_Velocity_Trained_Model_14075_steps_0.4534_loss_0.8687_acc.pth'
-                    }
+# -----------------------------
+# CONFIGURATION & GLOBALS
+# -----------------------------
+SEP = '=' * 70
+PDT = timezone('US/Pacific')
 
+MODEL_CHECKPOINT = 'Godzilla_Piano_Transformer_No_Velocity_Trained_Model_14075_steps_0.4534_loss_0.8687_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'
 
-PREVIEW_LENGTH = 120 # in tokens
+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()
 
-print('=' * 70)
-print('Instantiating model...')
+# -----------------------------
+# 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
-                              )
+    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('=' * 70)
-print('Loading model checkpoint...')      
-
-model_checkpoint = hf_hub_download(repo_id='asigalov61/Godzilla-Piano-Transformer', 
-                                   filename='Godzilla_Piano_Transformer_No_Velocity_Trained_Model_14075_steps_0.4534_loss_0.8687_acc.pth')
-
-model.load_state_dict(torch.load(model_checkpoint, map_location='cuda', weights_only=True))
-
+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('=' * 70)
-print('Done!')
-print('=' * 70)
-print('Model will use', dtype, 'precision...')
-print('=' * 70)
+print_sep()
+print("Done!")
+print("Model will use", dtype, "precision...")
+print_sep()
 
 model.cuda()
 model.eval()
 
-#==================================================================================
-
-def load_model(model_selector):
-
-    return [[], []]
-
-#==================================================================================
-
-def load_midi(input_midi, model_selector=''):
-
+# -----------------------------
+# HELPER FUNCTIONS
+# -----------------------------
+def render_midi_output(final_composition):
+    """Generate MIDI score, plot, and audio from final composition."""
+    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
+    )
+    fname = save_midi(final_composition)  # The file name is embedded in the saved MIDI.
+    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]
-    
+    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 = []
-    
-    pc = cscore[0]
-    
-    for c in cscore:
-        score.append(max(0, min(127, c[0][1]-pc[0][1])))
-    
-        for n in c:
-            if model_selector == 'with velocity - 3 epochs':
-                score.extend([max(1, min(127, n[2]))+128, max(1, min(127, n[4]))+256, max(1, min(127, n[5]))+384])
-
-            else:
-                score.extend([max(1, min(127, n[2]))+128, max(1, min(127, n[4]))+256])
-    
-        pc = c
-
+    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, model_selector=''):
-
-    song = tokens
-    song_f = []
-    
-    time = 0
-    dur = 0
-    vel = 90
+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
-    channel = 0
-    patch = 0
-
     patches = [0] * 16
 
-    for m in song:
-
-        if 0 <= m < 128:
-            time += m * 32
-
-        elif 128 < m < 256:
-            dur = (m-128) * 32
-
-        elif 256 < m < 384:
-            pitch = (m-256)
-
-            if model_selector == 'without velocity - 3 epochs' or model_selector == 'without velocity - 7 epochs':
-                song_f.append(['note', time, dur, 0, pitch, max(40, pitch), 0])
-
-        elif 384 < m < 512:
-            vel = (m-384)
-
-            if model_selector == 'with velocity - 3 epochs':
-                song_f.append(['note', time, dur, 0, pitch, vel, 0])
-
-    if batch_number == None:
-        fname = 'Godzilla-Piano-Transformer-Music-Composition'
-        
+    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 = 'Godzilla-Piano-Transformer-Music-Composition_'+str(batch_number)
-    
-    data = TMIDIX.Tegridy_ms_SONG_to_MIDI_Converter(song_f,
-                                                  output_signature = 'Godzilla Piano Transformer',
-                                                  output_file_name = fname,
-                                                  track_name='Project Los Angeles',
-                                                  list_of_MIDI_patches=patches,
-                                                  verbose=False
-                                                  )
-
-    return song_f
-
-#==================================================================================
-
+        fname = f"Godzilla-Piano-Transformer-Music-Composition_{timestamp}_Batch_{batch_number}"
+    
+    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
+
+# -----------------------------
+# MUSIC GENERATION FUNCTION (Combined)
+# -----------------------------
 @spaces.GPU
-def generate_music(prime, 
-                   num_gen_tokens,
-                   num_mem_tokens,
-                   num_gen_batches,
-                   model_temperature,
-                   # model_sampling_top_p,
-                   model_state
-                  ):
-
-    if not prime:
-        inputs = [0]
-
-    else:
-        inputs = prime[-num_mem_tokens:]
-
-    print('Generating...')
-    
-    inp = [inputs] * num_gen_batches
-    
-    inp = torch.LongTensor(inp).cuda()
-    
+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,
-                              #filter_logits_fn=top_p,
-                              #filter_kwargs={'thres': 0.96},
-                              temperature=model_temperature,
-                              return_prime=False,
-                              verbose=False)
-    
-    output = out.tolist()
-
-    print('Done!')
-    print('=' * 70)
-            
-    return output
-    
-#==================================================================================
-
-def generate_callback(input_midi, 
-                      num_prime_tokens, 
-                      num_gen_tokens,
-                      num_mem_tokens,
-                      model_temperature,
-                      # model_sampling_top_p,
-                      final_composition, 
-                      generated_batches, 
-                      block_lines,
-                      model_state
-                     ):
-
-    generated_batches = []
-
+        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"))
+    
+    # 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, model_selector=model_state[2])[:num_prime_tokens]
-        midi_score = save_midi(final_composition, model_selector=model_state[2])
-        block_lines.append(midi_score[-1][1] / 1000)
-        
-    batched_gen_tokens = generate_music(final_composition, 
-                                        num_gen_tokens,
-                                        num_mem_tokens,
-                                        NUM_OUT_BATCHES,
-                                        model_temperature,
-                                        # model_sampling_top_p,
-                                        model_state
-                                       )
-    
-    outputs = []
-    
-    for i in range(len(batched_gen_tokens)):
-
-        tokens = batched_gen_tokens[i]
-        
-        # Preview
-        tokens_preview = final_composition[-PREVIEW_LENGTH:]
-        
-        # Save MIDI to a temporary file
-        midi_score = save_midi(tokens_preview + tokens, i, model_selector=model_state[2])
-
-        # MIDI plot
-
+        final_composition = load_midi(input_midi)[:num_prime_tokens]
+        midi_fname = save_midi(final_composition)
+        # Use the last note's time as a marker.
+        midi_score = TMIDIX.Tegridy_ms_SONG_to_MIDI_Converter(
+            final_composition,
+            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(final_composition[-1] 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 = save_midi(preview_tokens + tokens, batch_number=i)
+        plot_kwargs = {'plot_title': f'Batch # {i}', 'return_plt': True}
         if len(final_composition) > PREVIEW_LENGTH:
-            midi_plot = TMIDIX.plot_ms_SONG(midi_score, 
-                                            plot_title='Batch # ' + str(i),
-                                            preview_length_in_notes=int(PREVIEW_LENGTH / 3),
-                                            return_plt=True
-                                           )
-
-        else:
-            midi_plot = TMIDIX.plot_ms_SONG(midi_score, 
-                                            plot_title='Batch # ' + str(i), 
-                                            return_plt=True
-                                           )
-
-        # File name
-        fname = 'Godzilla-Piano-Transformer-Music-Composition_'+str(i)
-        
-        # Save audio to a temporary file
-        midi_audio = midi_to_colab_audio(fname + '.mid', 
-                                        soundfont_path=SOUDFONT_PATH,
-                                        sample_rate=16000,
-                                        output_for_gradio=True
-                                        )
-
-        outputs.append([(16000, midi_audio), midi_plot, tokens])
-        
-    return outputs, final_composition, generated_batches, block_lines
-
-#==================================================================================
-
-def generate_callback_wrapper(input_midi, 
-                              num_prime_tokens, 
-                              num_gen_tokens,
-                              num_mem_tokens,
-                              model_temperature,
-                              # model_sampling_top_p,
-                              final_composition, 
-                              generated_batches, 
-                              block_lines,
-                              model_selector,
-                              model_state
-                             ):
-
-    print('=' * 70)
-    print('Req start time: {:%Y-%m-%d %H:%M:%S}'.format(datetime.datetime.now(PDT)))
-    start_time = reqtime.time()
-    
-    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('Selected model type:', model_selector)
-
-    if not model_state:
-        model_state = load_model(model_selector)
-        model_state.append(model_selector)
-
-    else:
-        if model_selector != model_state[2]:
-            print('=' * 70)
-            print('Switching model...')
-            model_state = load_model(model_selector)
-            model_state.append(model_selector)
-            print('=' * 70)
-    
-    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('Model top_p:', model_sampling_top_p)
-    print('=' * 70)
-    
-    result = generate_callback(input_midi, 
-                                num_prime_tokens, 
-                                num_gen_tokens,
-                                num_mem_tokens,
-                                model_temperature,
-                                # model_sampling_top_p,
-                                final_composition,
-                                generated_batches,
-                                block_lines,
-                                model_state
-                             )
-    
-    generated_batches = [sublist[-1] for sublist in result[0]]
-
-    print('=' * 70)
-    print('Req end time: {:%Y-%m-%d %H:%M:%S}'.format(datetime.datetime.now(PDT)))
-    print('=' * 70)
-    print('Req execution time:', (reqtime.time() - start_time), 'sec')
-    print('*' * 70)    
-    
-    return tuple([result[1], generated_batches, result[3]] + [item for sublist in result[0] for item in sublist[:-1]] + [model_state])
-
-#==================================================================================
-
-def add_batch(batch_number, final_composition, generated_batches, block_lines, model_state=[]):
-
+            plot_kwargs['preview_length_in_notes'] = int(PREVIEW_LENGTH / 3)
+        midi_score = TMIDIX.Tegridy_ms_SONG_to_MIDI_Converter(
+            preview_tokens + tokens,
+            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])
-
-        # Save MIDI to a temporary file
-        midi_score = save_midi(final_composition, model_selector=model_state[2])
-    
-        block_lines.append(midi_score[-1][1] / 1000)
-
-        # MIDI plot
-        midi_plot = TMIDIX.plot_ms_SONG(midi_score, 
-                                        plot_title='Godzilla Piano Transformer Composition',
-                                        block_lines_times_list=block_lines[:-1],
-                                        return_plt=True)
-        
-        # File name
-        fname = 'Godzilla-Piano-Transformer-Music-Composition'
-        
-        # Save audio to a temporary file
-        midi_audio = midi_to_colab_audio(fname + '.mid', 
-                                        soundfont_path=SOUDFONT_PATH,
-                                        sample_rate=16000,
-                                        output_for_gradio=True
-                                        )
-    
-        print('Added batch #', batch_number)
-        print('=' * 70)
-
-        return (16000, midi_audio), midi_plot, fname+'.mid', final_composition, generated_batches, block_lines
-
+        midi_fname = save_midi(final_composition)
+        block_lines.append(final_composition[-1] if final_composition else 0)
+        midi_score = TMIDIX.Tegridy_ms_SONG_to_MIDI_Converter(
+            final_composition,
+            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, model_state=[]):
-
-    if final_composition:
-
-        if len(final_composition) > num_tokens:
-            final_composition = final_composition[:-num_tokens]
+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()
-    
-        # Save MIDI to a temporary file
-        midi_score = save_midi(final_composition, model_selector=model_state[2])
-    
-        # MIDI plot
-        midi_plot = TMIDIX.plot_ms_SONG(midi_score, 
-                                        plot_title='Godzilla Piano Transformer Composition',
-                                        block_lines_times_list=block_lines[:-1],
-                                        return_plt=True)
-    
-        # File name
-        fname = 'Godzilla-Piano-Transformer-Music-Composition'
-        
-        # Save audio to a temporary file
-        midi_audio = midi_to_colab_audio(fname + '.mid', 
-                                        soundfont_path=SOUDFONT_PATH,
-                                        sample_rate=16000,
-                                        output_for_gradio=True
-                                        )
-        
-        print('Removed batch #', batch_number)
-        print('=' * 70)
-        
-        return (16000, midi_audio), midi_plot, fname+'.mid', final_composition, generated_batches, block_lines
-
+        midi_fname = save_midi(final_composition)
+        midi_score = TMIDIX.Tegridy_ms_SONG_to_MIDI_Converter(
+            final_composition,
+            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():
-    return None, None, None, [], []    
-
-#==================================================================================
 
-def reset(final_composition=[], generated_batches=[], block_lines=[], model_state=[]):
-    
-    final_composition = []
-    generated_batches = []
-    block_lines = []
-    model_state = []
-
-    return final_composition, generated_batches, block_lines
-    
-#==================================================================================
-
-def reset_demo(final_composition=[], generated_batches=[], block_lines=[], model_state=[]):
-    
-    final_composition = []
-    generated_batches = []
-    block_lines = []
-    model_state = []
-
-#==================================================================================
-
-PDT = timezone('US/Pacific')
+def clear():
+    """Clear outputs and reset state."""
+    return None, None, None, [], []
 
-print('=' * 70)
-print('App start time: {:%Y-%m-%d %H:%M:%S}'.format(datetime.datetime.now(PDT)))
-print('=' * 70)
+def reset(final_composition=[], generated_batches=[], block_lines=[]):
+    """Reset composition state."""
+    return [], [], []
 
-#==================================================================================
+def reset_demo(final_composition=[], generated_batches=[], block_lines=[]):
+    """Reset state for demo unload."""
+    return [], [], []
 
+# -----------------------------
+# GRADIO INTERFACE SETUP
+# -----------------------------
 with gr.Blocks() as demo:
-
-    #==================================================================================
-
     demo.load(reset_demo)
 
-    #==================================================================================
-
     gr.Markdown("<h1 style='text-align: center; margin-bottom: 1rem'>Godzilla Piano Transformer</h1>")
     gr.Markdown("<h1 style='text-align: center; 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!
-            """)
-    
-    #==================================================================================
-
+        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([])
-    model_state = gr.State([])
-    
-    #==================================================================================
-    
-    gr.Markdown("## Upload seed MIDI or click 'Generate' button for random output")
-    
+
+    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")
+    input_midi.upload(reset, [final_composition, generated_batches, block_lines],
+                      [final_composition, generated_batches, block_lines])
 
-    model_selector = gr.Dropdown(["without velocity - 3 epochs"], 
-                                label="Select model", 
-                               )
-    
+    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")
-    # model_sampling_top_p = gr.Slider(0.1, 1, value=0.96, step=0.01, label="Model sampling top p value")
-    
     generate_btn = gr.Button("Generate", variant="primary")
 
-    gr.Markdown("## Select batch")
-    
+    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}") as tab:
-            
-            audio_output = gr.Audio(label=f"Batch # {i} MIDI Audio", format="mp3", elem_id="midi_audio")
+        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])
-
-    outputs.extend([model_state])
-
-    generate_btn.click(generate_callback_wrapper, 
-                       [input_midi, 
-                        num_prime_tokens, 
-                        num_gen_tokens,
-                        num_mem_tokens,
-                        model_temperature,
-                        # model_sampling_top_p,
-                        final_composition,
-                        generated_batches,
-                        block_lines,
-                        model_selector,
-                        model_state
-                       ], 
-                       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")
-    
+    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", elem_id="midi_audio")
+
+    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, model_state],
-                  [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, model_state], 
-                     [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]
-                   )
-    
-    #==================================================================================
+    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.unload(reset_demo)
 
-#==================================================================================
-
-demo.launch()
-
-#==================================================================================
+demo.launch()