Update app.py

app.py

@@ -234,7 +234,7 @@ def first_note_idx(score_tokens, patch):
 
 @spaces.GPU
 def Inpaint_MIDI_Template(midi_template_idx, 
-                          inpaint_acc,
+                          inpainting_mode,
                           model_temperature,
                           model_sampling_top_p 
                          ):
@@ -242,31 +242,35 @@ def Inpaint_MIDI_Template(midi_template_idx,
     #===============================================================================
 
     def inpaint(melody_chords,
-                inpaint_MIDI_patch=0,
+                inpaint_MIDI_patch=[0],
                 number_of_prime_tokens=0,
                 number_of_memory_tokens=4096,
                 temperature=1.0,
                 model_sampling_top_p_value=0.96,
                 verbose=False
                ):
+
+        #=====================================================================
     
         if verbose:
             print('=' * 70)
             print('Giant Music Transformer Inpainting Model Generator')
             print('=' * 70)
+
+        #=====================================================================
         
         out2 = []
         
         for m in melody_chords[:number_of_prime_tokens]:
           out2.append(m)
-        
-        torch.cuda.empty_cache()
+
+        #=====================================================================
         
         for i in tqdm.tqdm(range(number_of_prime_tokens, len(melody_chords))):
         
             cpatch = (melody_chords[i]-2304) // 129
         
-            if 2304 <= melody_chords[i] < 18945 and (cpatch) == inpaint_MIDI_patch:
+            if 2304 <= melody_chords[i] < 18945 and cpatch in inpaint_MIDI_patch:
         
                 inp = torch.LongTensor(out2[-number_of_memory_tokens:]).cuda()
         
@@ -294,7 +298,8 @@ def Inpaint_MIDI_Template(midi_template_idx,
         
             else:
                 out2.append(melody_chords[i])
-        
+
+        #=====================================================================
         
         song_f = []
         
@@ -344,6 +349,8 @@ def Inpaint_MIDI_Template(midi_template_idx,
                 pitch = (ss-2304) % 129
         
                 song_f.append(['note', time, dur, channel, pitch, vel, patch ])
+
+        #=====================================================================
                     
         if verbose:
             print('Done!')
@@ -360,7 +367,7 @@ def Inpaint_MIDI_Template(midi_template_idx,
     print('Requested settings:')
     print('=' * 70)
     print('MIDI template idx:', midi_template_idx)
-    print('Inpaint acc:', inpaint_acc)
+    print('Parts to inpaint:', inpainting_mode)
     print('Model temperature:', model_temperature)
     print('Model sampling top p:', model_sampling_top_p)
     print('=' * 70)
@@ -383,7 +390,7 @@ def Inpaint_MIDI_Template(midi_template_idx,
     mt_idx = midi_templates.index(midi_template)
     mt_md5 = midi_template[0]
     
-    events_matrix = toks_to_score(midi_template[1])[:MAX_NOTES_TO_INPAINT]
+    inp_score = toks_to_score(midi_template[1])[:MAX_NOTES_TO_INPAINT]
 
     print('=' * 70)
     print('Selected MIDI template idx:', mt_idx)
@@ -392,84 +399,43 @@ def Inpaint_MIDI_Template(midi_template_idx,
     #==================================================================
 
     print('=' * 70)
-    print('Creating scores...')
+    print('Prepping patches...')
 
     mel_pat = 40
     acc_pat = 24
     bass_pat = 35
-    
-    mel_score = []
-    acc_score = []
-    bass_score = []
-    drums_score = []
-    
-    fnote = False
-    
-    for e in events_matrix:
-        
-        if e[6] == mel_pat:
-            mel_score.append(e)
-            fnote = True
-    
-        elif e[6] == acc_pat:
-            acc_score.append(e)
-    
-        elif e[6] == bass_pat:
-            bass_score.append(e)
-    
-        elif e[6] == 128:
-            drums_score.append(e)
-    
-        if not fnote:
-            mel_score.append(e)
-    
-    print('=' * 70)
 
-    #==================================================================
-    
-    print('Inpainting melody...')
+    ipatches = []
 
-    mel_drums_tokens = score_to_toks(mel_score+drums_score)
-    mel_drums_score = inpaint(mel_drums_tokens, mel_pat)
+    if 'Melody' in inpainting_mode:
+        ipatches.append(mel_pat)
 
-    print('Done!')
-    print('=' * 70)
+    if 'Accompaniment' in inpainting_mode:
+        ipatches.append(acc_pat)
 
-    #==================================================================
+    if 'Base' in inpainting_mode:
+        ipatches.append(bass_pat)
 
-    if inpaint_acc:
+    if not inpainting_mode:
+        ipatches.append(mel_pat)
 
-        print('Inpainting accompaniment...')
-    
-        mel_drums_acc_tokens = score_to_toks(mel_drums_score+acc_score)
-        pidx = first_note_idx(mel_drums_acc_tokens, mel_pat)
-        mel_drums_acc_score = inpaint(mel_drums_acc_tokens, acc_pat, number_of_prime_tokens=pidx)
-
-        mel_drums_acc_bass_tokens = score_to_toks(mel_drums_acc_score+bass_score)
-    
-        print('Done!')
-        print('=' * 70)
-
-    else:
-        print('Skipping accompaniment inpainting...')
-        mel_drums_acc_bass_tokens = score_to_toks(mel_drums_score+bass_score)
-        print('=' * 70)
+    print('=' * 70)
 
     #==================================================================
     
-    print('Inpainting base...')
+    print('Inpainting...')
 
-    pidx = first_note_idx(mel_drums_acc_bass_tokens, mel_pat)
-    mel_drums_acc_bass_score = inpaint(mel_drums_acc_bass_tokens, bass_pat, number_of_prime_tokens=pidx)
+    inp_score_tokens = score_to_toks(inp_score)
+    out_score = inpaint(inp_score_tokens, ipatches, 300)
 
     print('Done!')
     print('=' * 70)
-    
+
     #==================================================================
 
     print('Patching final score...')
 
-    output_score, patches, overflow_patches = TMIDIX.patch_enhanced_score_notes(mel_drums_acc_bass_score)
+    output_score, patches, overflow_patches = TMIDIX.patch_enhanced_score_notes(out_score)
 
     print('Done!')
     print('=' * 70)
@@ -560,7 +526,7 @@ with gr.Blocks() as demo:
     gr.Markdown("## MIDI Templates options")
 
     midi_template_idx = gr.Slider(-1, 16383, value=-1, label="Desired MIDI template number")
-    inpaint_acc = gr.Checkbox(value=True, label="Inpaint accompaniment")
+    inpainting_mode = gr.CheckboxGroup(['Melody', 'Accompaniment', 'Base'], value=['Melody'], label="Which MIDI template parts to inpaint")
 
     gr.Markdown("## Model options")
 
@@ -577,7 +543,7 @@ with gr.Blocks() as demo:
     
     generate_btn.click(Inpaint_MIDI_Template, 
                        [midi_template_idx, 
-                        inpaint_acc,
+                        inpainting_mode,
                         model_temperature,
                         model_sampling_top_p                        
                        ],