Update app.py

app.py

@@ -145,19 +145,25 @@ def load_gan_model(generator, model_path, device):
     generator.eval()  # Set the model to evaluation mode
     return generator
 
-# Convert magnitude-only spectrogram to complex format by assuming zero phase
 def magnitude_to_complex_spectrogram(magnitude_spectrogram):
+    # Create a zero-phase tensor with the same shape as the magnitude spectrogram
     zero_phase = torch.zeros_like(magnitude_spectrogram)
-    complex_spectrogram = torch.stack([magnitude_spectrogram, zero_phase], dim=-1)
+    
+    # Create a complex-valued spectrogram using the magnitude and zero phase
+    complex_spectrogram = torch.complex(magnitude_spectrogram, zero_phase)
+    
     return complex_spectrogram
 
-# Convert spectrogram back to audio using inverse STFT
 def spectrogram_to_audio(magnitude_spectrogram):
-    magnitude_spectrogram = torch.expm1(magnitude_spectrogram)
+    # Convert magnitude spectrogram to complex spectrogram
     complex_spectrogram = magnitude_to_complex_spectrogram(magnitude_spectrogram)
+
+    # Inverse STFT to convert the spectrogram back to audio
     audio = torch.istft(complex_spectrogram, n_fft=n_fft, hop_length=hop_length)
+    
     return audio
 
+
 # Function to generate audio from an uploaded image
 def generate_audio_from_image(image):
     test_img = image_transform(image).unsqueeze(0).to(device)  # Preprocess image
@@ -179,12 +185,16 @@ def main():
     generator = Generator(output_time_frames).to(device)
 
     # Load the pre-trained model
-    model_path = './gan_model.pth'  # Change this path
+    model_path = './gan_model.pth'  # Ensure the model is in the correct relative path
     generator = load_gan_model(generator, model_path, device)
 
     # Gradio interface: allow users to upload an image and generate audio
-    iface = gr.Interface(fn=generate_audio_from_image, inputs=gr.Image(type="pil"), outputs=gr.Audio(type="numpy", label="Generated Audio"))
+    iface = gr.Interface(fn=generate_audio_from_image, 
+                         inputs=gr.Image(type="pil"), 
+                         outputs=gr.Audio(type="numpy", label="Generated Audio"))
+    
     iface.launch()
 
+
 if __name__ == "__main__":
     main()