Update app.py

app.py

@@ -1,9 +1,8 @@
 import streamlit as st
+import torchaudio
 import torch
-import os
-import librosa
-import librosa.display
 import matplotlib.pyplot as plt
+import soundfile as sf
 from audiosr import build_model, super_resolution, save_wave
 import tempfile
 import numpy as np
@@ -31,48 +30,51 @@ random_seed = st.sidebar.number_input("Random Seed", min_value=0, value=42, step
 latent_t_per_second = 12.8
 
 # Helper function to plot spectrogram
-def plot_spectrogram(audio_path, title):
-    y, sr = librosa.load(audio_path, sr=None)
-    S = librosa.feature.melspectrogram(y=y, sr=sr, n_mels=128, fmax=sr // 2)
-    S_dB = librosa.power_to_db(S, ref=np.max)
-    
+def plot_spectrogram(waveform, sample_rate, title):
     plt.figure(figsize=(10, 4))
-    librosa.display.specshow(S_dB, sr=sr, x_axis='time', y_axis='mel', fmax=sr // 2, cmap='viridis')
-    plt.colorbar(format='%+2.0f dB')
+    spectrogram = torchaudio.transforms.MelSpectrogram(
+        sample_rate=sample_rate, n_fft=2048, hop_length=512, n_mels=128
+    )(torch.tensor(waveform))
+    log_spectrogram = torchaudio.transforms.AmplitudeToDB()(spectrogram)
+    plt.imshow(log_spectrogram.numpy(), aspect="auto", origin="lower", cmap="viridis")
+    plt.colorbar(format="%+2.0f dB")
     plt.title(title)
+    plt.xlabel("Time")
+    plt.ylabel("Frequency")
     plt.tight_layout()
-    return plt
+    st.pyplot(plt)
 
 # Process Button
 if uploaded_file and st.button("Enhance Audio"):
     st.write("Processing audio...")
 
     # Create temp directory for saving files
-    with tempfile.TemporaryDirectory() as tmp_dir:
-        input_path = os.path.join(tmp_dir, "input.wav")
-        truncated_path = os.path.join(tmp_dir, "truncated.wav")
-        output_path = os.path.join(tmp_dir, "output.wav")
+    with tempfile.TemporaryDirectory() as temp_dir:
+        input_path = os.path.join(temp_dir, "input.wav")
+        truncated_path = os.path.join(temp_dir, "truncated.wav")
+        output_path = os.path.join(temp_dir, "output.wav")
 
         # Save uploaded file locally
         with open(input_path, "wb") as f:
             f.write(uploaded_file.read())
 
-        # Load and truncate the first 10 seconds
-        y, sr = librosa.load(input_path, sr=None)
-        max_samples = sr * 10  # First 10 seconds
-        y_truncated = y[:max_samples]
-        librosa.output.write_wav(truncated_path, y_truncated, sr)
+        # Load audio and truncate the first 10 seconds
+        waveform, sample_rate = torchaudio.load(input_path)
+        max_samples = sample_rate * 10  # First 10 seconds
+        if waveform.size(1) > max_samples:
+            waveform = waveform[:, :max_samples]
+            st.write("Truncated audio to the first 10 seconds.")
+        sf.write(truncated_path, waveform[0].numpy(), sample_rate)
 
         # Plot truncated spectrogram
         st.write("Truncated Input Audio Spectrogram (First 10 seconds):")
-        truncated_spectrogram = plot_spectrogram(truncated_path, title="Truncated Input Audio Spectrogram")
-        st.pyplot(truncated_spectrogram)
+        plot_spectrogram(waveform[0].numpy(), sample_rate, title="Truncated Input Audio Spectrogram")
 
         # Build and load the model
         audiosr = build_model(model_name=model_name, device=device)
 
         # Perform super-resolution
-        waveform = super_resolution(
+        waveform_sr = super_resolution(
             audiosr,
             truncated_path,
             seed=random_seed,
@@ -82,12 +84,11 @@ if uploaded_file and st.button("Enhance Audio"):
         )
 
         # Save enhanced audio
-        save_wave(waveform, inputpath=truncated_path, savepath=tmp_dir, name="output", samplerate=48000)
+        save_wave(waveform_sr, inputpath=truncated_path, savepath=tmp_dir, name="output", samplerate=48000)
 
         # Plot output spectrogram
         st.write("Enhanced Audio Spectrogram:")
-        output_spectrogram = plot_spectrogram(output_path, title="Enhanced Audio Spectrogram")
-        st.pyplot(output_spectrogram)
+        plot_spectrogram(waveform_sr.numpy(), 48000, title="Enhanced Audio Spectrogram")
 
         # Display audio players and download link
         st.audio(truncated_path, format="audio/wav")