Update app.py

app.py

@@ -206,6 +206,10 @@ def infer_inp(prompt, audio_path, mask_start_point, mask_end_point, progress=gr.
     post_masked_spec = denormalize(masked_spec).to(device, dtype)
     denorm_masked_spec = denormalize_spectrogram(post_masked_spec)
     denorm_masked_spec_audio = vocoder.inference(denorm_masked_spec)
+
+    # Rescale generated spectrogram to match original range
+    output_spec = (output_spec - output_spec.min()) / (output_spec.max() - output_spec.min())  # Normalize to [0,1]
+    output_spec = output_spec * (norm_spec.max() - norm_spec.min()) + norm_spec.min()  # Rescale to match input range
     
     denorm_spec = denormalize_spectrogram(output_spec)
     denorm_spec_audio = vocoder.inference(denorm_spec)
@@ -225,7 +229,9 @@ def infer_inp(prompt, audio_path, mask_start_point, mask_end_point, progress=gr.
     print("Output spectrogram min/max:", output_spec.min().item(), output_spec.max().item())
 
     # Save as WAV
-    sf.write("output.wav", denorm_spec_audio, 16000)
+    sf.write("raw_output.wav", raw_chunk_audio, 16000)
+    sf.write("masked_raw_output.wav", denorm_masked_spec_audio, 16000)
+    sf.write("generated_output.wav", denorm_spec_audio, 16000)
     
     # Save input spectrogram image
     input_spec_image_path = "input_spectrogram.png"
@@ -235,7 +241,7 @@ def infer_inp(prompt, audio_path, mask_start_point, mask_end_point, progress=gr.
     output_spec_image_path = "output_spectrogram.png"
     color_output_spec_image.save(output_spec_image_path)
 
-    return "output.wav", input_spec_image_path, color_output_spec_image
+    return "raw_output.wav", input_spec_image_path, color_output_spec_image, "raw_output.wav", "masked_raw_output.wav"
 
 def load_input_spectrogram(audio_path):
     # Loading
@@ -380,6 +386,11 @@ with gr.Blocks(css=css) as demo:
                 with gr.Column():
                     input_spectrogram_inp = gr.Image(label="Input Spectrogram")
                     output_spectrogram_inp = gr.Image(label="Output Spectrogram")
+            
+            with gr.Accordion("Raw Processed audio", open=False):
+                with gr.Column():
+                    raw_out_audio = gr.Audio(label="RAW Audio")
+                    raw_masked_out_audio = gr.Audio(label="RAW Masked Audio")
 
             gr.Examples(
                 examples = [
@@ -415,7 +426,7 @@ with gr.Blocks(css=css) as demo:
             submit_btn_inp.click(
                 fn = infer_inp,
                 inputs = [prompt_inp, audio_in_inp, mask_start_point, mask_end_point],
-                outputs = [audio_out_inp, input_spectrogram_inp, output_spectrogram_inp]
+                outputs = [audio_out_inp, input_spectrogram_inp, output_spectrogram_inp, raw_out_audio, raw_masked_out_audio]
             )
 
 demo.queue().launch(show_api=False, show_error=True)