move to gradio

app.py

@@ -1,15 +1,13 @@
-import streamlit as st
+import gradio as gr
 import os
 from transformer_wrapper import TransformerWrapper
 from omegaconf import OmegaConf
 
 
-@st.cache(show_spinner=False)
 def get_file_content_as_string(path):
     return open(path, "r", encoding="utf-8").read()
 
 
-@st.cache(show_spinner=True)
 def model_load():
     config = OmegaConf.load("config.yaml")
     wrapper = TransformerWrapper(config)
@@ -23,43 +21,68 @@ def model_load():
     return wrapper, model_id, config
 
 
-def main():
-
-    wrapper, model_id, config = model_load()
-    composers = list(config.composer_to_feature_token.keys())
-    dest_dir = "ytsamples"
-    os.makedirs(dest_dir, exist_ok=True)
-    composer = st.selectbox(label="Arranger", options=composers)
-    file_up = st.file_uploader("Upload an audio", type=["mp3", "wav"])
-
-    if st.button("convert"):
-
-        if file_up is not None:
-            bytes_data = file_up.getvalue()
-            target_file = f"{dest_dir}/{file_up.name}"
-            with open(target_file, "wb") as f:
-                f.write(bytes_data)
-
-            with st.spinner("Wait for it..."):
-                midi, arranger, mix_path, midi_path = wrapper.generate(
-                    audio_path=target_file,
-                    composer=composer,
-                    model=model_id,
-                    ignore_duplicate=True,
-                    show_plot=False,
-                    save_midi=True,
-                    save_mix=True,
-                )
-
-            with open(midi_path, "rb") as midi_f:
-                file_down = st.download_button(
-                    "Download midi",
-                    data=midi_f,
-                    file_name=os.path.basename(midi_path),
-                )
-            with open(mix_path, "rb") as audio_f:
-                st.audio(audio_f.read(), format="audio/wav")
-
-
-if __name__ == "__main__":
-    main()
+wrapper, model_id, config = model_load()
+composers = list(config.composer_to_feature_token.keys())
+dest_dir = "ytsamples"
+os.makedirs(dest_dir, exist_ok=True)
+
+
+def inference(file_up, composer):
+
+    midi, arranger, mix_path, midi_path = wrapper.generate(
+        audio_path=file_up,
+        composer=composer,
+        model=model_id,
+        ignore_duplicate=True,
+        show_plot=False,
+        save_midi=True,
+        save_mix=True,
+    )
+
+    return mix_path
+
+
+block = gr.Blocks()
+
+
+with block:
+    gr.HTML(
+        """
+            <div style="text-align: center; max-width: 700px; margin: 0 auto;">
+              <div
+                style="
+                  display: inline-flex;
+                  align-items: center;
+                  gap: 0.8rem;
+                  font-size: 1.75rem;
+                "
+              >
+                <h1 style="font-weight: 900; margin-bottom: 7px;">
+                  Pop2piano
+                </h1>
+              </div>
+              <p style="margin-bottom: 10px; font-size: 94%">
+                A demo for Pop2Piano:Pop Audio-based Piano Cover Generation. Please select the composer and upload the pop audio to submit.
+              </p>
+            </div>
+        """
+    )
+    with gr.Group():
+        with gr.Box():
+            with gr.Row().style(mobile_collapse=False, equal_height=True):
+                file_up = gr.Audio(label="Upload an audio", type="filepath")
+                composer = gr.Dropdown(label="Arranger", choices=composers, value="composer1")
+                btn = gr.Button("Convert")
+        out = gr.Audio(label="Output")
+
+        btn.click(inference, inputs=[file_up, composer], outputs=out)
+        gr.HTML(
+            """
+        <div class="footer">
+                    <p><a href="http://sweetcocoa.github.io/pop2piano_samples" style="text-decoration: underline;" target="_blank">Project Page</a>
+                    </p>
+        </div>
+        """
+        )
+
+block.launch(debug=True)

transformer_wrapper.py

@@ -155,9 +155,7 @@ class TransformerWrapper(pl.LightningModule):
 
         return relative_tokens, notes, pm
 
-    def prepare_inference_mel(
-        self, audio, beatstep, n_bars, padding_value, composer_value=None
-    ):
+    def prepare_inference_mel(self, audio, beatstep, n_bars, padding_value, composer_value=None):
         n_steps = n_bars * 4
         n_target_step = len(beatstep)
         sample_rate = self.config.dataset.sample_rate
@@ -240,9 +238,7 @@ class TransformerWrapper(pl.LightningModule):
             composer = random.sample(list(composer_to_feature_token.keys()), 1)[0]
 
         composer_value = composer_to_feature_token[composer]
-        mix_sample_rate = (
-            config.dataset.sample_rate if mix_sample_rate is None else mix_sample_rate
-        )
+        mix_sample_rate = config.dataset.sample_rate if mix_sample_rate is None else mix_sample_rate
 
         if not ignore_duplicate:
             if os.path.exists(midi_path):
@@ -295,8 +291,7 @@ class TransformerWrapper(pl.LightningModule):
             feature_tokens=fzs,
             audio=_audio,
             beatstep=beatsteps - beatsteps[0],
-            max_length=config.dataset.target_length
-            * max(1, (n_bars // config.dataset.n_bars)),
+            max_length=config.dataset.target_length * max(1, (n_bars // config.dataset.n_bars)),
             max_batch_size=max_batch_size,
             n_bars=n_bars,
             composer_value=composer_value,
@@ -311,22 +306,15 @@ class TransformerWrapper(pl.LightningModule):
                 y = librosa.core.resample(y, orig_sr=sr, target_sr=mix_sample_rate)
                 sr = mix_sample_rate
             if add_click:
-                clicks = (
-                    librosa.clicks(times=beatsteps, sr=sr, length=len(y)) * click_amp
-                )
+                clicks = librosa.clicks(times=beatsteps, sr=sr, length=len(y)) * click_amp
                 y = y + clicks
             pm_y = pm.fluidsynth(sr)
             stereo = get_stereo(y, pm_y, pop_scale=stereo_amp)
 
         if show_plot:
-            import IPython.display as ipd
-            from IPython.display import display
             import note_seq
 
-            display("Stereo MIX", ipd.Audio(stereo, rate=sr))
-            display("Rendered MIDI", ipd.Audio(pm_y, rate=sr))
-            display("Original Song", ipd.Audio(y, rate=sr))
-            display(note_seq.plot_sequence(note_seq.midi_to_note_sequence(pm)))
+            note_seq.plot_sequence(note_seq.midi_to_note_sequence(pm))
 
         if save_mix:
             sf.write(