Update app.py

app.py

@@ -127,7 +127,10 @@ def infer_img2img(prompt, audio_path, desired_strength, progress=gr.Progress(tra
 
     return "output.wav", input_spec_image_path, output_spec_image_path
 
-def infer_inp(prompt, audio_path, progress=gr.Progress(track_tqdm=True)):
+def infer_inp(prompt, audio_path, spec_with_mask, progress=gr.Progress(track_tqdm=True)):
+
+    if spec_with_mask:
+        print(spec_with_mask)
 
     pretrained_model_name_or_path = "auffusion/auffusion-full-no-adapter"
     dtype = torch.float16
@@ -213,6 +216,18 @@ def infer_inp(prompt, audio_path, progress=gr.Progress(track_tqdm=True)):
 
     return "output.wav"
 
+def load_spec_for_manual_masking(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))
+    return raw_image
+    
+
 css="""
 div#col-container{
     margin: 0 auto;
@@ -282,6 +297,7 @@ 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")
+            spec_for_mask = gr.ImageMask(label="Draw Mask", type="pil", interactive=False)
             submit_btn_inp = gr.Button("Submit")
             audio_out_inp = gr.Audio(label="Audio Ressult")
 
@@ -290,11 +306,15 @@ 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_spec_for_manual_masking,
+                inputs = [audio_in_inp],
+                outputs = [spec_for_mask]
+            )
 
             submit_btn_inp.click(
                 fn = infer_inp,
-                inputs = [prompt_inp, audio_in_inp],
+                inputs = [prompt_inp, audio_in_inp, spec_for_mask],
                 outputs = [audio_out_inp]
             )