Update app.py

app.py

@@ -213,7 +213,48 @@ def infer_inp(prompt, audio_path, mask_start_point, mask_end_point, progress=gr.
 
     return "output.wav", input_spec_image_path, color_output_spec_image
 
+def load_input_spectrogram(audio_path):
+    # Loading
+    audio, sampling_rate = load_wav(audio_path)
+    audio, spec = get_mel_spectrogram_from_audio(audio)
+    norm_spec = normalize_spectrogram(spec)
+    norm_spec = pad_spec(norm_spec, 1024)
+    norm_spec = normalize(norm_spec) # normalize to [-1, 1], because pipeline do not normalize for torch.Tensor input
+    
+    raw_image = image_add_color(torch_to_pil(norm_spec))
+    
+    # Save input spectrogram image
+    input_spec_image_path = "input_spectrogram.png"
+    raw_image.save(input_spec_image_path)
+
+def preview_masked_area(audio_path, mask_start_point, mask_end_point):
+    # Loading
+    audio, sampling_rate = load_wav(audio_path)
+    audio, spec = get_mel_spectrogram_from_audio(audio)
+    norm_spec = normalize_spectrogram(spec)
+    norm_spec = pad_spec(norm_spec, 1024)
+    norm_spec = normalize(norm_spec) # normalize to [-1, 1], because pipeline do not normalize for torch.Tensor input
+
+    # Add Mask
+    mask = torch.zeros_like(norm_spec)[:1,...]
+    mask[:, :, width_start:width_start+width] = 1
+    mask_image = torch_to_pil(mask)
+    
+    mask, masked_spec = prepare_mask_and_masked_image(norm_spec, mask)
+    masked_spec_image = torch_to_pil(masked_spec)
 
+    # color masked spec and paint masked area to black
+    color_masked_spec_image = image_add_color(masked_spec_image)
+    color_masked_spec_image = np.array(color_masked_spec_image)
+    color_masked_spec_image[:, width_start:width_start+width, :] = 0
+    color_masked_spec_image = Image.fromarray(color_masked_spec_image)
+
+    # Save the masked spectrogram image
+    masked_spec_image_path = "masked_spectrogram.png"
+    color_masked_spec_image.save(masked_spec_image_path)
+
+    return masked_spec_image_path
+    
 css="""
 div#col-container{
     margin: 0 auto;
@@ -283,9 +324,15 @@ with gr.Blocks(css=css) as demo:
         with gr.Tab("Audio InPainting"):
             prompt_inp = gr.Textbox(label="Prompt")
             audio_in_inp = gr.Audio(label="Audio Reference", type="filepath")
+            
+            audio_in_spec = gr.Image(label="Audio IN spectrogram")
             mask_start_point = gr.Slider(label="Mask Start point", minimum=0, maximum=1024, step=1, value=256)
             mask_end_point = gr.Slider(label="Mask End point", minimum=0, maximum=1024, step=1, value=768)
+            preview_mask_btn = gr.Button("Preview Mask")
+            
+            masked_spec_preview = gr.Image(label="Spectrogram Mask Preview")
             submit_btn_inp = gr.Button("Submit")
+            
             audio_out_inp = gr.Audio(label="Audio Ressult")
 
             with gr.Accordion("Compare Spectrograms", open=False):
@@ -293,10 +340,22 @@ with gr.Blocks(css=css) as demo:
                     input_spectrogram_inp = gr.Image(label="Input Spectrogram")
                     output_spectrogram_inp = gr.Image(label="Output Spectrogram")
 
+            audio_in_inp.upload(
+                fn = load_input_spectrogram,
+                inputs = [audio_in_inp],
+                outputs = [audio_in_spec]
+            )
+
+            preview_mask_btn.click(
+                fn = preview_masked_area,
+                inputs = [audio_in_inp, mask_start_point, mask_end_point],
+                outputs = [masked_spec_preview]
+            )
+            
             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, output_spectrogram]
+                outputs = [audio_out_inp, input_spectrogram_inp, output_spectrogram_inp]
             )
 
 demo.queue().launch(show_api=False, show_error=True)