Update app.py

app.py

@@ -1,6 +1,5 @@
 import gradio as gr
 import torch, os
-from torchvision import transforms
 import numpy as np
 from PIL import Image
 import matplotlib.pyplot as plt
@@ -128,10 +127,7 @@ 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, spec_with_mask, progress=gr.Progress(track_tqdm=True)):
-
-    if spec_with_mask:
-        print(spec_with_mask)
+def infer_inp(prompt, audio_path, mask_start_point, mask_end_point, progress=gr.Progress(track_tqdm=True)):
 
     pretrained_model_name_or_path = "auffusion/auffusion-full-no-adapter"
     dtype = torch.float16
@@ -146,7 +142,7 @@ def infer_inp(prompt, audio_path, spec_with_mask, progress=gr.Progress(track_tqd
     pipe = StableDiffusionInpaintPipeline.from_pretrained(pretrained_model_name_or_path, torch_dtype=dtype)
     pipe = pipe.to(device)
 
-    width_start, width = 256, 512
+    width_start, width = mask_start_point, mask_end_point-mask_start_point
     prompt = "A siren ringing with a vehicle speeding closer"
     seed = 42
 
@@ -160,24 +156,11 @@ def infer_inp(prompt, audio_path, spec_with_mask, progress=gr.Progress(track_tqd
     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)
-
-    # Load the mask image (input from user)
-    mask_pil = spec_with_mask['layers'][0]
-    
-    # Convert to tensor and normalize
-    mask_tensor = transforms.ToTensor()(mask_pil)  # Shape: (1, H, W), values in [0, 1]
-    
-    # Ensure the shape matches expected input (add batch dimension if needed)
-    mask_tensor = mask_tensor[:1, :, :]  # Keep only one channel (grayscale)
-    mask_tensor = mask_tensor.to(device, dtype)  # Send to correct device and dtype
+    mask = torch.zeros_like(norm_spec)[:1,...]
+    mask[:, :, width_start:width_start+width] = 1
+    mask_image = torch_to_pil(mask)
     
-    # Convert to PIL image if needed for visualization
-    mask_image = torch_to_pil(mask_tensor)
-    
-    mask, masked_spec = prepare_mask_and_masked_image(norm_spec, mask_tensor)
+    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
@@ -221,47 +204,15 @@ def infer_inp(prompt, audio_path, spec_with_mask, progress=gr.Progress(track_tqd
     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"
-    #denorm_spec_audio.save(output_spec_image_path)
-
-    return "output.wav"
-
-def create_transparent_layer(image_path):
-    """Creates a transparent PNG with the same size as the background image."""
-    background = Image.open(image_path)
-    transparent_layer = Image.new("RGBA", background.size, (0, 0, 0, 0))
-    
-    layer_path = "layer_one.png"
-    transparent_layer.save(layer_path)
-    return layer_path
-
-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))
-
     input_spec_image_path = "input_spectrogram.png"
     raw_image.save(input_spec_image_path)
     
-    # Create transparent layer
-    layer_one_path = create_transparent_layer(input_spec_image_path)
-    
-    # Return as EditorValue
-    return {
-        "background": input_spec_image_path,
-        "layers": [layer_one_path],
-        "composite": None
-    }
-    
+    # Save output spectrogram image
+    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
+
 
 css="""
 div#col-container{
@@ -332,8 +283,8 @@ 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")
-            brush = gr.Brush(colors=["#000"], color_mode="fixed")
-            spec_for_mask = gr.ImageEditor(label="Draw Mask", type="pil", brush=brush, interactive=True, layers=False)
+            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)
             submit_btn_inp = gr.Button("Submit")
             audio_out_inp = gr.Audio(label="Audio Ressult")
 
@@ -341,23 +292,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")
-    
-            audio_in_inp.upload(
-                fn = load_spec_for_manual_masking,
-                inputs = [audio_in_inp],
-                outputs = [spec_for_mask]
-            )
-
-            spec_for_mask.clear(
-                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, spec_for_mask],
-                outputs = [audio_out_inp]
+                inputs = [prompt_inp, audio_in_inp, mask_start_point, mask_end_point],
+                outputs = [audio_out_inp, input_spectrogram, output_spectrogram]
             )
 
 demo.queue().launch(show_api=False, show_error=True)