add duration param

app.py

@@ -5,15 +5,11 @@ from matplotlib.animation import FuncAnimation
 import io
 import librosa
 import tempfile
-import os
 
-def extract_waveform_animation(audio_file, duration):
+def extract_waveform_animation(audio_file):
     # Load audio file
     y, sr = librosa.load(audio_file, sr=None)
-    
-    # Determine the number of frames needed for the specified duration
-    num_frames = int(duration * sr)
-    y = y[:num_frames]
+    duration = librosa.get_duration(y=y, sr=sr)
     
     # Create a figure and axis for the animation
     fig, ax = plt.subplots()
@@ -28,15 +24,17 @@ def extract_waveform_animation(audio_file, duration):
     
     # Function to update the animation frame
     def update(frame):
-        line.set_data(np.linspace(0, duration, num=len(y[:frame])), y[:frame])
+        start = frame * sr
+        end = min(start + sr, len(y))
+        line.set_data(np.linspace(0, duration, num=len(y[:end])), y[:end])
         return line,
     
     # Create the animation
-    ani = FuncAnimation(fig, update, frames=np.arange(0, num_frames, sr // 10), init_func=init, blit=True)
+    ani = FuncAnimation(fig, update, frames=np.arange(0, int(duration)), init_func=init, blit=True)
     
     # Save the animation to a temporary file
     with tempfile.NamedTemporaryFile(delete=False, suffix='.mp4') as tmpfile:
-        ani.save(tmpfile.name, writer='ffmpeg', fps=30)
+        ani.save(tmpfile.name, writer='ffmpeg', fps=1)
         video_path = tmpfile.name
     
     return video_path
@@ -44,12 +42,9 @@ def extract_waveform_animation(audio_file, duration):
 # Define the Gradio interface
 iface = gr.Interface(
     fn=extract_waveform_animation,
-    inputs=[
-        gr.Audio(type="filepath"),
-        gr.Number(label="Duration (seconds)", value=5)
-    ],
-    outputs=gr.Video(),
-    description="Upload an audio file to extract a video animation from its waveform. Specify the duration of the animation."
+    inputs=gr.Audio(source="upload", type="file"),
+    outputs=gr.Video(type="file"),
+    description="Upload an audio file to extract a video animation from its waveform."
 )
 
 # Launch the app