passing the midi audio properly

app.py

@@ -117,27 +117,27 @@ def generate_midi(seed, use_chords, chord_progression, bpm):
     # Clean up temporary MIDI file
     os.remove(midi_filename)
 
-    return wav_filename
+    # Load the generated audio
+    song, sr = torchaudio.load(wav_filename)
+
+    # Clean up temporary MIDI file
+    os.remove(midi_filename)
+    os.remove(wav_filename)
+
+    return song.numpy(), sr
 
 @spaces.GPU(duration=120)
-def generate_music(midi_audio, prompt_duration, musicgen_model, num_iterations, bpm):
-    if isinstance(midi_audio, tuple):
-        wav_filename, sample_rate = midi_audio
-        song, sr = torchaudio.load(wav_filename)
-    else:
-        # Assuming midi_audio is a numpy array
-        with tempfile.NamedTemporaryFile(suffix=".wav", delete=False) as temp_file:
-            temp_filename = temp_file.name
-            torchaudio.save(temp_filename, torch.from_numpy(midi_audio), sample_rate=44100)
-            song, sr = torchaudio.load(temp_filename)
+def generate_music(midi_data, prompt_duration, musicgen_model, num_iterations, bpm):
+    audio_data, sample_rate = midi_data
 
-    song = song.to(device)
+    # Convert the audio data to a PyTorch tensor
+    song = torch.from_numpy(audio_data).to(device)
 
     # Use the user-provided BPM value for duration calculation
     duration = calculate_duration(bpm)
 
     # Create slices from the song using the user-provided BPM value
-    slices = create_slices(song, sr, 35, bpm, num_slices=5)
+    slices = create_slices(song, sample_rate, 35, bpm, num_slices=5)
 
     # Load the model
     model_name = musicgen_model.split(" ")[0]
@@ -160,10 +160,10 @@ def generate_music(midi_audio, prompt_duration, musicgen_model, num_iterations,
         
         print(f"Running iteration {i + 1} using slice {slice_idx}...")
         
-        prompt_waveform = slices[slice_idx][..., :int(prompt_duration * sr)]
+        prompt_waveform = slices[slice_idx][..., :int(prompt_duration * sample_rate)]
         prompt_waveform = preprocess_audio(prompt_waveform)
         
-        output = model_continue.generate_continuation(prompt_waveform, prompt_sample_rate=sr, progress=True)
+        output = model_continue.generate_continuation(prompt_waveform, prompt_sample_rate=sample_rate, progress=True)
         output = output.cpu()  # Move the output tensor back to CPU
         
         # Make sure the output tensor has at most 2 dimensions
@@ -184,10 +184,7 @@ def generate_music(midi_audio, prompt_duration, musicgen_model, num_iterations,
     combined_audio_filename = f"combined_audio_{random.randint(1, 10000)}.mp3"
     combined_audio.export(combined_audio_filename, format="mp3")
 
-    
     # Clean up temporary files
-    if not isinstance(midi_audio, tuple):
-        os.remove(temp_filename)
     for filename in all_audio_files:
         os.remove(filename)
 
@@ -253,6 +250,6 @@ with gr.Blocks() as iface:
             output_audio = gr.Audio(label="Generated Music")
 
     generate_midi_button.click(generate_midi, inputs=[seed, use_chords, chord_progression, bpm], outputs=midi_audio)
-    generate_music_button.click(generate_music, inputs=[midi_audio, prompt_duration, musicgen_model, num_iterations, bpm], outputs=output_audio)
+    generate_music_button.click(generate_music, inputs=[midi_audio[0], prompt_duration, musicgen_model, num_iterations, bpm], outputs=output_audio)
 
 iface.launch()