Update app.py

app.py

@@ -129,41 +129,58 @@ def load_gan_model(generator, model_path, device):
     return generator
 
 
-# Generator model class definitions remain the same as in your original code.
-
-# Convert magnitude-only spectrogram to complex format by assuming zero phase
 def magnitude_to_complex_spectrogram(magnitude_spectrogram):
+    # Clip values to avoid extreme values or potential invalid inputs
+    magnitude_spectrogram = torch.clamp(magnitude_spectrogram, min=1e-10, max=1e5)
+
     zero_phase = torch.zeros_like(magnitude_spectrogram)
     complex_spectrogram = torch.stack([magnitude_spectrogram, zero_phase], dim=-1)
+
+    # Check for NaNs in the complex spectrogram
+    if torch.isnan(complex_spectrogram).any():
+        raise ValueError("Complex spectrogram contains NaN values.")
+
     return complex_spectrogram
 
+
 def spectrogram_to_audio(magnitude_spectrogram):
     # Perform inverse log scaling
-    magnitude_spectrogram = torch.expm1(magnitude_spectrogram)
-    
+    try:
+        magnitude_spectrogram = torch.expm1(magnitude_spectrogram)
+        if torch.isnan(magnitude_spectrogram).any():
+            raise ValueError("NaN values found in magnitude_spectrogram after expm1.")
+    except Exception as e:
+        raise ValueError(f"Error in expm1 operation: {e}")
+
     # Convert magnitude-only spectrogram to complex format
-    complex_spectrogram = magnitude_to_complex_spectrogram(magnitude_spectrogram)
+    try:
+        complex_spectrogram = magnitude_to_complex_spectrogram(magnitude_spectrogram)
+        if torch.isnan(complex_spectrogram).any():
+            raise ValueError("Complex spectrogram contains NaN values after conversion.")
+    except Exception as e:
+        raise ValueError(f"Error in complex spectrogram creation: {e}")
 
     # Inverse STFT to convert the spectrogram back to audio
-    audio = torch.istft(complex_spectrogram, n_fft=n_fft, hop_length=hop_length)
+    try:
+        audio = torch.istft(complex_spectrogram, n_fft=n_fft, hop_length=hop_length)
+        if torch.isnan(audio).any():
+            raise ValueError("Generated audio contains NaN values after istft.")
+    except Exception as e:
+        raise ValueError(f"Error in istft operation: {e}")
 
     # Normalize audio to the range [-1, 1] (standard audio range)
-    if torch.max(torch.abs(audio)) != 0:
-        audio = audio / torch.max(torch.abs(audio))
+    try:
+        if torch.max(torch.abs(audio)) != 0:
+            audio = audio / torch.max(torch.abs(audio))
+    except Exception as e:
+        raise ValueError(f"Error in audio normalization: {e}")
 
     # 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()
 
-    # Convert audio tensor to numpy array for Gradio
-    audio = audio.cpu().numpy().astype(np.int16)
-
     return audio