Update app.py

app.py

@@ -30,8 +30,13 @@ guidance_scale = st.sidebar.slider("Guidance Scale", min_value=1.0, max_value=10
 random_seed = st.sidebar.number_input("Random Seed", min_value=0, value=42, step=1)
 latent_t_per_second = 12.8
 
+# Helper function: Plot linear STFT spectrogram
 # Helper function: Plot linear STFT spectrogram
 def plot_spectrogram(waveform, sample_rate, title):
+    # Ensure waveform is a 1D tensor
+    if len(waveform.shape) > 1:
+        waveform = waveform.squeeze()  # Remove extra dimensions
+    
     plt.figure(figsize=(10, 4))
     spectrogram = torch.stft(
         torch.tensor(waveform),
@@ -44,7 +49,7 @@ def plot_spectrogram(waveform, sample_rate, title):
         np.log1p(spectrogram),
         aspect="auto",
         origin="lower",
-        extent=[0, waveform.shape[-1] / sample_rate, 0, sample_rate / 2],
+        extent=[0, len(waveform) / sample_rate, 0, sample_rate / 2],
         cmap="viridis",
     )
     plt.colorbar(format="%+2.0f dB")
@@ -54,6 +59,7 @@ def plot_spectrogram(waveform, sample_rate, title):
     plt.tight_layout()
     st.pyplot(plt)
 
+
 # Process Button
 if uploaded_file and st.button("Enhance Audio"):
     st.write("Processing audio...")