Update app.py

app.py

@@ -14,6 +14,16 @@ from diffusers import StableDiffusionImg2ImgPipeline
 from converter import load_wav, mel_spectrogram, normalize_spectrogram, denormalize_spectrogram, Generator, get_mel_spectrogram_from_audio
 from utils import pad_spec, image_add_color, torch_to_pil, normalize, denormalize
 
+# ——
+
+def save_spectrogram_image(spectrogram, filename):
+    """Save a spectrogram as an image."""
+    plt.figure(figsize=(10, 4))
+    plt.imshow(spectrogram.squeeze(), aspect='auto', origin='lower', cmap='magma')
+    plt.axis('off')  # Hide axes for a cleaner image
+    plt.savefig(filename, bbox_inches='tight', pad_inches=0)
+    plt.close()
+
 def infer(prompt, progress=gr.Progress(track_tqdm=True)):
     
     prompt = prompt
@@ -106,8 +116,16 @@ def infer_img2img(prompt, audio_path, desired_strength, progress=gr.Progress(tra
 
     # Save as WAV
     sf.write("output.wav", concat_audio, 16000)
+    
+    # Save input spectrogram image
+    input_spec_image_path = "input_spectrogram.png"
+    raw_image.save(input_spec_image_path)
+    
+    # Save concatenated spectrogram image
+    output_spec_image_path = "output_spectrogram.png"
+    concat_image.save(output_spec_image_path)
 
-    return "output.wav"
+    return "output.wav", input_spec_image_path, output_spec_image_path
 
 css="""
 div#col-container{
@@ -161,6 +179,10 @@ with gr.Blocks(css=css) as demo:
             prompt_strength = gr.Slider(label="Prompt Strength", minimum=0.0, maximum=1.0, step=0.1, value=0.7)
             submit_btn_img2img = gr.Button("Submit")
             audio_out_img2img = gr.Audio(label="Audio Ressult")
+
+            with gr.Row():
+                input_spectrogram = gr.Image(label="Input Spectrogram")
+                output_spectrogram = gr.Image(label="Output Spectrogram")
     
             gr.Examples(
                 examples = [
@@ -174,7 +196,7 @@ with gr.Blocks(css=css) as demo:
             submit_btn_img2img.click(
                 fn = infer_img2img,
                 inputs = [prompt_img2img, audio_in_img2img, prompt_strength],
-                outputs = [audio_out_img2img]
+                outputs = [audio_out_img2img, input_spectrogram, output_spectrogram]
             )
 
 demo.queue().launch(show_api=False, show_error=True)