Update app.py

app.py

@@ -23,7 +23,7 @@ pipe = transformers.pipeline(
 )
 print('Done')
 
-default_prefix = '''pitch duration wait
+default_prefix = '''pitch duration wait velocity instrument
 
 71 1310 0
 48 330 350
@@ -47,35 +47,50 @@ default_prefix = '''pitch duration wait
 default_prefix_len = default_prefix.count('\n') - 2
 
 
+
 def postprocess(txt, path):
-    # saves the text representation to a midi file
+    # assert txt.startswith(prompt)
     txt = txt.split('\n\n')[-1]
 
-    notes = []
+    # track = symusic.core.TrackSecond()
+    tracks = {}
+
     now = 0
     # we need to ignore the invalid output by the model
     try:
         for line in txt.split('\n'):
-            pitch, duration, wait = [int(x) for x in line.split()]
+            pitch, duration, wait, velocity, instrument = line.split()
+            pitch, duration, wait, velocity = [int(x) for x in [pitch, duration, wait, velocity]]
+            if instrument not in tracks:
+                tracks[instrument] = symusic.core.TrackSecond()
+                if instrument != 'drum':
+                    tracks[instrument].program = int(instrument)
+                else:
+                    tracks[instrument].is_drum = True
             # Eg. Note(time=7.47, duration=5.25, pitch=43, velocity=64, ttype='Quarter')
-            notes.append(symusic.core.NoteSecond(
+            tracks[instrument].notes.append(symusic.core.NoteSecond(
                 time=now/1000,
                 duration=duration/1000,
                 pitch=int(pitch),
-                velocity=80,
+                velocity=int(velocity * 4),
             ))
             now += wait
     except Exception as e:
-        print('Ignored error:', e)
+        print('Postprocess: Ignored error:', e)
+    
+    print(f'Postprocess: Got {sum(len(track.notes) for track in tracks.values())} notes')
 
     try:
-        track = symusic.core.TrackSecond()
-        track.notes = symusic.core.NoteSecondList(notes)
+        # track = symusic.core.TrackSecond()
+        # track.notes = symusic.core.NoteSecondList(notes)
         score = symusic.Score(ttype='Second')
-        score.tracks.append(track)
+        # score.tracks.append(track)
+        score.tracks.extend(tracks.values())
         score.dump_midi(path)
     except Exception as e:
-        print('Ignored error:', e)
+        print('Postprocess: Ignored postprocessing error:', e)
+
+
 
 with gr.Blocks() as demo:
     chatbot_box = gr.Chatbot(type="messages", render_markdown=False, sanitize_html=False)