Update app.py

app.py

@@ -137,7 +137,6 @@ def magnitude_to_complex_spectrogram(magnitude_spectrogram):
     complex_spectrogram = torch.stack([magnitude_spectrogram, zero_phase], dim=-1)
     return complex_spectrogram
 
-    
 def spectrogram_to_audio(magnitude_spectrogram):
     # Perform inverse log scaling
     magnitude_spectrogram = torch.expm1(magnitude_spectrogram)
@@ -155,12 +154,16 @@ def spectrogram_to_audio(magnitude_spectrogram):
     # Clip the audio to ensure it fits in the range [-1, 1]
     audio = torch.clamp(audio, min=-1, max=1)
 
+    # Check for NaNs in the audio tensor
+    if torch.isnan(audio).any():
+        raise ValueError("Generated audio contains NaN values.")
+
     # Convert to 16-bit PCM format by scaling and casting to int16
     audio = (audio * 32767).short()
 
-    # Ensure the audio is in the correct format (int16)
+    # Convert audio tensor to numpy array for Gradio
     audio = audio.cpu().numpy().astype(np.int16)
-    
+
     return audio
 
 
@@ -173,9 +176,6 @@ def generate_audio_from_image(image):
 
     # Convert the generated spectrogram to audio
     generated_audio = spectrogram_to_audio(generated_spectrogram.squeeze(0).cpu())
-    if torch.isnan(generated_audio).any():
-        raise ValueError("Generated audio contains NaN values.")
-
 
     # Return the audio as (sample_rate, audio) tuple for Gradio to process
     return (sample_rate, generated_audio)