refactor and improve apps

app.py

@@ -4,7 +4,16 @@ import gradio as gr
 
 from audiodiffusion import AudioDiffusion
 
-audio_diffusion = AudioDiffusion()
+
+def generate_spectrogram_audio_and_loop(model_id):
+    audio_diffusion = AudioDiffusion(model_id=model_id)
+    image, (sample_rate,
+            audio) = audio_diffusion.generate_spectrogram_and_audio()
+    loop = AudioDiffusion.loop_it(audio, sample_rate)
+    if loop is None:
+        loop = audio
+    return image, (sample_rate, audio), (sample_rate, loop)
+
 
 if __name__ == "__main__":
     parser = argparse.ArgumentParser()
@@ -13,14 +22,23 @@ if __name__ == "__main__":
     args = parser.parse_args()
 
     demo = gr.Interface(
-        fn=audio_diffusion.generate_spectrogram_and_audio,
+        fn=generate_spectrogram_audio_and_loop,
         title="Audio Diffusion",
         description="Generate audio using Huggingface diffusers.\
             This takes about 20 minutes without a GPU, so why not make yourself a cup of tea in the meantime?",
-        inputs=[],
+        inputs=[
+            gr.Dropdown(label="Model",
+                        choices=[
+                            "teticio/audio-diffusion-256",
+                            "teticio/audio-diffusion-breaks-256"
+                        ],
+                        value="teticio/audio-diffusion-256")
+        ],
         outputs=[
             gr.Image(label="Mel spectrogram", image_mode="L"),
             gr.Audio(label="Audio"),
+            gr.Audio(label="Loop"),
         ],
+        allow_flagging="never"
     )
     demo.launch(server_name=args.server or "0.0.0.0", server_port=args.port)

audiodiffusion/__init__.py

@@ -1,6 +1,8 @@
+import numpy as np
 from PIL import Image
 from torch import cuda
 from diffusers import DDPMPipeline
+from librosa.beat import beat_track
 
 from .mel import Mel
 
@@ -38,3 +40,12 @@ class AudioDiffusion:
         image = Image.fromarray(images[0][0])
         audio = self.mel.image_to_audio(image)
         return image, (self.mel.get_sample_rate(), audio)
+
+    @staticmethod
+    def loop_it(audio, sample_rate, loops=12):
+        tempo, beats = beat_track(y=audio, sr=sample_rate, units='samples')
+        if len(beats) > 8:
+            return np.tile(audio[beats[0]:beats[8]], loops)
+        if len(beats) > 4:
+            return np.tile(audio[beats[0]:beats[4]], loops)
+        return None

streamlit_app.py

@@ -2,16 +2,21 @@ from io import BytesIO
 import streamlit as st
 import soundfile as sf
 from librosa.util import normalize
+from librosa.beat import beat_track
 
 from audiodiffusion import AudioDiffusion
 
-audio_diffusion = AudioDiffusion()
-
 if __name__ == "__main__":
     st.header("Audio Diffusion")
     st.markdown("Generate audio using Huggingface diffusers.\
         This takes about 20 minutes without a GPU, so why not make yourself a cup of tea in the meantime?"
                 )
+
+    model_id = st.selectbox(
+        "Model",
+        ["teticio/audio-diffusion-256", "teticio/audio-diffusion-breaks-256"])
+    audio_diffusion = AudioDiffusion(model_id=model_id)
+
     if st.button("Generate"):
         st.markdown("Generating...")
         image, (sample_rate,
@@ -20,3 +25,10 @@ if __name__ == "__main__":
         buffer = BytesIO()
         sf.write(buffer, normalize(audio), sample_rate, format="WAV")
         st.audio(buffer, format="audio/wav")
+
+        audio = AudioDiffusion.loop_it(audio, sample_rate)
+        if audio is not None:
+            st.markdown("Loop")
+            buffer = BytesIO()
+            sf.write(buffer, normalize(audio), sample_rate, format="WAV")
+            st.audio(buffer, format="audio/wav")