Update app.py

app.py

@@ -75,19 +75,18 @@ 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 = 1794
-
-model = TransformerWrapper(
-        num_tokens = PAD_IDX+1,
-        max_seq_len = SEQ_LEN,
-        attn_layers = Decoder(dim = 2048,
-                              depth = 4,
-                              heads = 32,
-                              rotary_pos_emb = True,
-                              attn_flash = True
-                              )
-)
+SEQ_LEN = 1536
+PAD_IDX = 708
+
+model = TransformerWrapper(num_tokens = PAD_IDX+1,
+                           max_seq_len = SEQ_LEN,
+                           attn_layers = Decoder(dim = 2048,
+                                                 depth = 8,
+                                                 heads = 32,
+                                                 rotary_pos_emb = True,
+                                                 attn_flash = True
+                                                 )
+                           )
 
 model = AutoregressiveWrapper(model, ignore_index=PAD_IDX, pad_value=PAD_IDX)
 
@@ -100,6 +99,9 @@ model.load_state_dict(torch.load(model_checkpoint, map_location=device_type, wei
 
 model = torch.compile(model, mode='max-autotune')
 
+model.to(device_type)
+model.eval()
+
 print('=' * 70)
 print('Done!')
 print('=' * 70)
@@ -113,47 +115,46 @@ def load_midi(input_midi, melody_patch=-1, use_nth_note=1):
     raw_score = TMIDIX.midi2single_track_ms_score(input_midi)
 
     escore_notes = TMIDIX.advanced_score_processor(raw_score, return_enhanced_score_notes=True)[0]
-    escore_notes = TMIDIX.augment_enhanced_score_notes(escore_notes, timings_divider=32)
     
-    sp_escore_notes = TMIDIX.solo_piano_escore_notes(escore_notes, keep_drums=False)
-
-    if melody_patch == -1:
-        zscore = TMIDIX.recalculate_score_timings(sp_escore_notes)
-
-    else:
-        mel_score = [e for e in sp_escore_notes if e[6] == melody_patch]
-
-        if mel_score:
-            zscore = TMIDIX.recalculate_score_timings(mel_score)
-
-        else:
-            zscore = TMIDIX.recalculate_score_timings(sp_escore_notes)
+    sp_escore_notes = TMIDIX.solo_piano_escore_notes(escore_notes)
+    zscore = TMIDIX.recalculate_score_timings(sp_escore_notes)
     
-    cscore = TMIDIX.chordify_score([1000, zscore])[:MAX_MELODY_NOTES:use_nth_note]
+    escore = TMIDIX.augment_enhanced_score_notes(zscore, timings_divider=32)
     
-    score = []
+    escore = TMIDIX.fix_escore_notes_durations(escore)
     
-    score_list = []
+    cscore = TMIDIX.chordify_score([1000, escore])
+    
+    score = []
+    chords = []
     
     pc = cscore[0]
     
     for c in cscore:
-        score.append(max(0, min(127, c[0][1]-pc[0][1])))
     
-        scl = [[max(0, min(127, c[0][1]-pc[0][1]))]]
+        tones_chord = sorted(set([e[4] % 12 for e in c]))
     
-        n = c[0]
-        
-        score.extend([max(1, min(127, n[2]))+128, max(1, min(127, n[4]))+256])
-        scl.append([max(1, min(127, n[2]))+128, max(1, min(127, n[4]))+256])
+        if tones_chord not in TMIDIX.ALL_CHORDS_SORTED:
+            tones_chord = TMIDIX.check_and_fix_tones_chord(tones_chord, use_full_chords=False)
+    
+        chord_tok = TMIDIX.ALL_CHORDS_SORTED.index(tones_chord)
+        chords.append(chord_tok+384)
+    
+        score.append(chord_tok+384)
+        score.append(max(0, min(127, c[0][1]-pc[0][1])))
     
-        score_list.append(scl)
+        for n in c:
+            score.extend([max(1, min(127, n[2]))+128, max(1, min(127, n[4]))+256])
     
         pc = c
     
-    score_list.append(scl)
+    print('Done!')
+    print('=' * 70)
+    print('Score has', len(chords), 'chords')
+    print('Score hss', len(score), 'tokens')
+    print('=' * 70)
 
-    return score, score_list
+    return score, chords
 
 #==================================================================================
 
@@ -268,8 +269,7 @@ def Generate_Accompaniment(input_midi,
     print('=' * 70)
     print('Generating...')
 
-    model.to(device_type)
-    model.eval()
+
 
     #==================================================================