Audio postprocessing

app.py

@@ -2,6 +2,8 @@ import torch
 import gradio as gr
 from transformers import ViTImageProcessor, ViTModel
 from audiodiffusion import AudioDiffusionPipeline, ImageEncoder
+from pedalboard.io import AudioFile
+from pedalboard import Pedalboard, NoiseGate, Compressor, LowShelfFilter, Gain, HighShelfFilter, Reverb
 
 device = "cuda" if torch.cuda.is_available() else "cpu"
 generator1 = torch.Generator(device)
@@ -13,6 +15,16 @@ processor = ViTImageProcessor.from_pretrained('google/vit-base-patch16-224-in21k
 extractor = ViTModel.from_pretrained('google/vit-base-patch16-224-in21k')
 image_encoder = ImageEncoder(processor, extractor)
 
+board = Pedalboard([
+    NoiseGate(threshold_db=-60, ratio=10.0),
+    Compressor(threshold_db=60, ratio=1.0),
+    LowShelfFilter(cutoff_frequency_hz=220, gain_db=-10),
+    HighShelfFilter(cutoff_frequency_hz=1200, gain_db=-10),
+    Gain(gain_db=40),
+    Reverb(room_size=0.5),
+    
+])
+
 def _encode_image(image):
     return torch.unsqueeze(image_encoder.encode(image), axis=1).to(device)
 
@@ -28,9 +40,13 @@ def _generate_spectrogram(condition, steps, eta):
     )
     return images[0], (sample_rate, audios[0])
 
+def _denoise_audio(audio, sr):
+    return board(audio, sr)
+
 def run_generation(image, steps, eta):
     condition = _encode_image(image)
     spectrogram, (sr, audio) = _generate_spectrogram(condition, steps, eta)
+    audio = _denoise_audio(audio, sr)
     return spectrogram, (sr, audio)
 
 with gr.Blocks(title="Image-based soundtrack generation") as demo:
@@ -44,21 +60,21 @@ with gr.Blocks(title="Image-based soundtrack generation") as demo:
                 label="Conditioning image"
             )
             steps = gr.Slider(
-                minimum=1, 
+                minimum=10, 
                 maximum=1000, 
-                step=1, 
+                step=10, 
                 value=50,
                 label="Denoising steps"
             )
             eta = gr.Slider(
-                minimum=0.1, 
+                minimum=0.0, 
                 maximum=1.0, 
                 step=0.1, 
-                value=0.9,
+                value=0.6,
                 label="η"
             )
             gr.Markdown('''
-                Eta (η) is a variable that controls the level of interpolation between a deterministic DDIM (η=0.0) and a stochastic DDPM (η=1.0).
+                Eta (η) is a variable that controls the level of interpolation between deterministic (η=0.0) and stochastic (η=1.0) denoising schedule.
             ''')
             btn = gr.Button("Generate")
             clear = gr.ClearButton(image)

requirements.txt

@@ -4,4 +4,5 @@ transformers==4.35.2
 numpy==1.23.5
 Pillow==9.3.0
 diffusers==0.23.1
-librosa==0.10.1
+librosa==0.10.1
+pedalboard==0.8.6