File size: 3,473 Bytes
d8595ba |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 |
import streamlit as st
import torch
import os
import librosa
import librosa.display
import matplotlib.pyplot as plt
from audiosr import build_model, super_resolution, save_wave
import tempfile
import numpy as np
# Set MPS device if available (for Mac M-Series GPUs)
device = "cuda" if torch.cuda.is_available() else "cpu"
# Title and Description
st.title("AudioSR: Versatile Audio Super-Resolution")
st.write("""
Upload your low-resolution audio files, and AudioSR will enhance them to high fidelity!
Supports all types of audio (music, speech, sound effects, etc.) with arbitrary sampling rates.
""")
# Upload audio file
uploaded_file = st.file_uploader("Upload an audio file (WAV format)", type=["wav"])
# Model Parameters
st.sidebar.title("Model Parameters")
model_name = st.sidebar.selectbox("Select Model", ["basic", "speech"], index=0)
ddim_steps = st.sidebar.slider("DDIM Steps", min_value=10, max_value=100, value=50)
guidance_scale = st.sidebar.slider("Guidance Scale", min_value=1.0, max_value=10.0, value=3.5)
random_seed = st.sidebar.number_input("Random Seed", min_value=0, value=42, step=1)
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)
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')
plt.title(title)
plt.tight_layout()
return 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")
output_path = os.path.join(tmp_dir, "output.wav")
# Save uploaded file locally
with open(input_path, "wb") as f:
f.write(uploaded_file.read())
# Plot input spectrogram
st.write("Input Audio Spectrogram:")
input_spectrogram = plot_spectrogram(input_path, title="Input Audio Spectrogram")
st.pyplot(input_spectrogram)
# Build and load the model
audiosr = build_model(model_name=model_name, device=device)
# Perform super-resolution
waveform = super_resolution(
audiosr,
input_path,
seed=random_seed,
guidance_scale=guidance_scale,
ddim_steps=ddim_steps,
latent_t_per_second=latent_t_per_second,
)
# Save enhanced audio
save_wave(waveform, inputpath=input_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)
# Display audio players and download link
st.audio(input_path, format="audio/wav")
st.write("Original Audio:")
st.audio(output_path, format="audio/wav")
st.write("Enhanced Audio:")
st.download_button("Download Enhanced Audio", data=open(output_path, "rb").read(), file_name="enhanced_audio.wav")
# Footer
st.write("Built with [Streamlit](https://streamlit.io) and [AudioSR](https://audioldm.github.io/audiosr)")
|