added conditional diffusion, descriptions, and examples

app.py

@@ -1,5 +1,6 @@
 # Imports
 import gradio as gr
+import os
 import matplotlib.pyplot as plt
 import torch
 import torchaudio
@@ -109,19 +110,40 @@ def load_checkpoint(model, ckpt_path) -> None:
 
 
 # Generate Samples
-def generate_samples(model_name, num_samples, num_steps, duration=32768):
+def generate_samples(model_name, num_samples, num_steps, init_audio=None, noise_level=0.7, duration=32768):
     # load_checkpoint
     ckpt_path = models[model_name]
     load_checkpoint(model, ckpt_path)
-
+    
     if num_samples > 1:
-        duration = duration / 2
+        duration = int(duration / 2)
 
+    # Generate samples
     with torch.no_grad():
-        all_samples = torch.zeros(2, 0) # initialize all samples
+        if init_audio:
+            # load audio sample
+            audio_sample = torch.tensor(init_audio[1].T, dtype=torch.float32).unsqueeze(0).to(model.device)
+            audio_sample = audio_sample / torch.max(torch.abs(audio_sample)) # normalize init_audio
+
+            # Trim audio
+            og_shape = audio_sample.shape
+            if duration < og_shape[2]:
+                audio_sample = audio_sample[:,:,:duration]
+            elif duration > og_shape[2]:
+                # Pad tensor with zeros to match sample length
+                audio_sample = torch.concat((audio_sample, torch.zeros(og_shape[0], og_shape[1], duration - og_shape[2]).to(model.device)), dim=2)
+
+        else:
+            audio_sample = torch.zeros((1, 2, int(duration)), device=model.device)
+            noise_level = 1.0
+
+        all_samples = torch.zeros(2, 0)
         for i in range(num_samples):
-            noise = torch.randn((1, 2, int(duration)), device=model.device) # [batch_size, in_channels, length]
-            generated_sample = model.model_ema.ema_model.sample(noise, num_steps=num_steps).squeeze(0).cpu() # Suggested num_steps 10-100
+            noise = torch.randn_like(audio_sample, device=model.device) * noise_level # [batch_size, in_channels, length]
+            audio = (audio_sample * abs(1-noise_level)) + noise # add noise
+
+            # generate samples
+            generated_sample = model.model_ema.ema_model.sample(audio, num_steps=num_steps).squeeze(0).cpu() # Suggested num_steps 10-100
             
             # concatenate all samples:
             all_samples = torch.concat((all_samples, generated_sample), dim=1)
@@ -133,6 +155,8 @@ def generate_samples(model_name, num_samples, num_steps, duration=32768):
 
     return (sr, all_samples.cpu().detach().numpy().T), fig # (sample rate, audio), plot
 
+
+# Define Constants & initialize model
 # load model & configs
 sr = 44100 # sampling rate
 config_path = "saved_models/config.yaml" # config path
@@ -147,19 +171,70 @@ models = {
     "Percussion": "saved_models/percussion/percussion_v0.ckpt"
 }
 
-demo = gr.Interface(
-    generate_samples,
-    inputs=[
-        gr.Dropdown(choices=list(models.keys()), value=list(models.keys())[3], label="Model"),
-        gr.Slider(1, 25, step=1, label="Number of Samples to Generate", value=3),
-        gr.Slider(1, 100, step=1, label="Number of Diffusion Steps", value=15)
-    ],
-    outputs=[
-        gr.Audio(label="Generated Audio Sample"),
-        gr.Plot(label="Generated Audio Spectrogram")
-    ]
-)
-
-if __name__ == "__main__":
-    demo.launch()
-    
+intro = """
+<h1 style="font-weight: 1400; text-align: center; margin-bottom: 6px;">
+    Tiny Audio Diffusion
+</h1>
+<h3 style="font-weight: 600; text-align: center;">
+    Christopher Landschoot - Audio waveform diffusion built to run on consumer-grade hardware (<2GB VRAM)
+</h3>
+<h4 style="text-align: center; margin-bottom: 6px;">
+    <a href="https://github.com/crlandsc/tiny-audio-diffusion" style="text-decoration: underline;" target="_blank">GitHub Repo</a> 
+    | <a href="https://www.youtube.com/watch?v=m6Eh2srtTro&t=3s" style="text-decoration: underline;" target="_blank">Repo Tutorial Video</a> 
+    | <a href="https://medium.com/towards-data-science/tiny-audio-diffusion-ddc19e90af9b" style="text-decoration: underline;" target="_blank">Towards Data Science Article</a>
+</h4>
+"""
+
+
+with gr.Blocks() as demo:
+    # Layout
+    gr.HTML(intro)
+
+    with gr.Row(equal_height=False):
+        with gr.Column():
+            # Inputs
+            model_name = gr.Dropdown(choices=list(models.keys()), value=list(models.keys())[3], label="Model")
+            num_samples = gr.Slider(1, 25, step=1, label="Number of Samples to Generate", value=3)
+            num_steps = gr.Slider(1, 100, step=1, label="Number of Diffusion Steps", value=15)
+            
+            # Conditioning Audio Input
+            with gr.Accordion("Input Audio (optional)", open=False):
+                init_audio_description = gr.HTML('Upload an audio file to perform conditional "style transfer" diffusion.<br>Leaving input audio blank results in unconditional generation.')
+                init_audio = gr.Audio(label="Input Audio Sample")
+                init_audio_noise = gr.Slider(0, 1, step=0.01, label="Noise to add to input audio", value=0.70)#, visible=True)
+
+                # Examples
+                gr.Examples(
+                    examples=[
+                        os.path.join(os.path.dirname(__file__), "samples", "guitar.wav"),
+                        os.path.join(os.path.dirname(__file__), "samples", "snare.wav"),
+                        os.path.join(os.path.dirname(__file__), "samples", "kick.wav"),
+                        os.path.join(os.path.dirname(__file__), "samples", "hihat.wav")
+                    ],
+                    inputs=init_audio,
+                    label="Example Audio Inputs"
+                )
+
+            # Buttons
+            with gr.Row():
+                with gr.Column():
+                    clear_button = gr.Button(value="Reset All")
+                with gr.Column():
+                    generate_btn = gr.Button("Generate Samples!")
+            
+        with gr.Column():
+            # Outputs
+            output_audio = gr.Audio(label="Generated Audio Sample")
+            output_plot = gr.Plot(label="Generated Audio Spectrogram")
+        
+    # Functionality
+    # Generate samples
+    generate_btn.click(fn=generate_samples, inputs=[model_name, num_samples, num_steps, init_audio, init_audio_noise], outputs=[output_audio, output_plot])
+
+    # clear_button button to reset everything
+    clear_button.click(fn=lambda: [3, 15, None, 0.70, None, None], outputs=[num_samples, num_steps, init_audio, init_audio_noise, output_audio, output_plot])
+
+
+
+if __name__ == "__main__":    
+    demo.launch()