Update app.py

app.py

@@ -127,6 +127,92 @@ 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)):
+
+    pretrained_model_name_or_path = "auffusion/auffusion-full-no-adapter"
+    dtype = torch.float16
+    device = "cuda"
+    
+    if not os.path.isdir(pretrained_model_name_or_path):
+        pretrained_model_name_or_path = snapshot_download(pretrained_model_name_or_path)     
+    
+    vocoder = Generator.from_pretrained(pretrained_model_name_or_path, subfolder="vocoder")
+    vocoder = vocoder.to(device=device, dtype=dtype)
+
+    pipe = StableDiffusionImg2ImgPipeline.from_pretrained(pretrained_model_name_or_path, torch_dtype=dtype)
+    pipe = pipe.to(device)
+
+    width_start, width = 256, 512
+    prompt = "A siren ringing with a vehicle speeding closer"
+    seed = 42
+
+    # 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))
+    
+    # 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)
+
+    # Generation
+    generator = torch.Generator(device=device).manual_seed(seed)
+    
+    with torch.autocast("cuda"):
+        output_spec = pipe(
+            prompt=prompt, image=norm_spec, mask_image=mask, num_inference_steps=100, generator=generator, height=256, width=1024, output_type="pt"
+        ).images[0]
+    
+    output_spec_image = torch_to_pil(output_spec)
+    color_output_spec_image = image_add_color(output_spec_image)
+
+    # Display audio result: raw audio, masked raw audio, generated audio
+    post_norm_spec = denormalize(norm_spec).to(device, dtype)
+    raw_chunk_spec = denormalize_spectrogram(post_norm_spec)
+    raw_chunk_audio = vocoder.inference(raw_chunk_spec)
+    
+    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)
+    
+    denorm_spec = denormalize_spectrogram(output_spec)
+    denorm_spec_audio = vocoder.inference(denorm_spec)
+
+    #———
+
+    # Ensure correct shape
+    denorm_spec_audio = denorm_spec_audio.flatten()  # Converts (1, N) → (N,)
+
+    # Normalize the audio to prevent clipping or excessive loudness
+    denorm_spec_audio = denorm_spec_audio / np.max(np.abs(denorm_spec_audio))  # Scale between -1 and 1
+
+    # Save as WAV
+    sf.write("output.wav", denorm_spec_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"
+
 css="""
 div#col-container{
     margin: 0 auto;
@@ -185,14 +271,7 @@ with gr.Blocks(css=css) as demo:
                     input_spectrogram = gr.Image(label="Input Spectrogram")
                     output_spectrogram = gr.Image(label="Output Spectrogram")
     
-            gr.Examples(
-                examples = [
-                    "Rolling thunder with lightning strikes",
-                    "Two gunshots followed by birds chirping",
-                    "A train whistle blowing in the distance"
-                ],
-                inputs = [prompt_img2img]
-            )
+            
 
             submit_btn_img2img.click(
                 fn = infer_img2img,
@@ -200,4 +279,23 @@ with gr.Blocks(css=css) as demo:
                 outputs = [audio_out_img2img, input_spectrogram, output_spectrogram]
             )
 
+        with gr.Tab("Audio InPainting"):
+            prompt_inp = gr.Textbox(label="Prompt")
+            audio_in_inp = gr.Audio(label="Audio Reference", type="filepath")
+            submit_btn_inp = gr.Button("Submit")
+            audio_out_inp = gr.Audio(label="Audio Ressult")
+
+            with gr.Accordion("Compare Spectrograms", open=False):
+                with gr.Column():
+                    input_spectrogram_inp = gr.Image(label="Input Spectrogram")
+                    output_spectrogram_inp = gr.Image(label="Output Spectrogram")
+    
+            
+
+            submit_btn_inp.click(
+                fn = infer_inp,
+                inputs = [prompt_inp, audio_in_inp],
+                outputs = [audio_out_inp]
+            )
+
 demo.queue().launch(show_api=False, show_error=True)