added other functionalities to the interface

app.py

@@ -35,9 +35,11 @@ from gamadhani.utils.generate_utils import load_pitch_fns, load_audio_fns
 import gamadhani.utils.pitch_to_audio_utils as p2a
 from gamadhani.utils.utils import get_device
 
+import copy
 
 logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s', force=True)
-pitch_path = 'models/diffusion_pitch/'
+pitch_paths = {'Diffusion Pitch Generator': 'models/diffusion_pitch/'}
+model_loaded = None
 audio_path = 'models/pitch_to_audio/'
 device = get_device()
 
@@ -96,7 +98,15 @@ def generate_pitch_reinterp(pitch, pitch_model, invert_pitch_fn, num_samples, nu
     noisy_pitch = torch.Tensor(pitch[:, :, -1200:]).to(pitch_model.device) + (torch.normal(mean=0.0, std=noise_std*torch.ones((1200)))).to(pitch_model.device)
     noisy_pitch = torch.clamp(noisy_pitch, -5.19, 5.19)     # clipping the pitch values to be within the range of the model
     samples = pitch_model.sample_sdedit(noisy_pitch, num_samples, num_steps)
-    inverted_pitches = [invert_pitch_fn(f0=samples.detach().cpu().numpy()[0])[0]]   # pitch values in Hz
+    inverted_pitches = invert_pitch_fn(f0=samples.detach().cpu().numpy()[0]).flatten()   # pitch values in Hz
+
+    return samples, inverted_pitches
+
+def generate_pitch_response(pitch, pitch_model, invert_pitch_fn, num_samples, num_steps):
+    '''Generate pitch values for the call and response task'''
+    pitch = torch.Tensor(pitch[:, :, -400:]).to(pitch_model.device)
+    samples = pitch_model.sample_fn(num_samples, num_steps, prime=pitch)
+    inverted_pitches = invert_pitch_fn(f0=samples.detach().cpu().numpy()[0]).flatten()   # pitch values in Hz
 
     return samples, inverted_pitches
 
@@ -109,11 +119,16 @@ def generate_audio(audio_model, f0s, invert_audio_fn, singers=[3], num_steps=100
     return audio
     
 @spaces.GPU(duration=150)
-def generate(pitch, num_samples=1, num_steps=100, singers=[3], outfolder='temp', audio_seq_len=750, pitch_qt=None ):
+def generate(pitch, num_samples=1, num_steps=100, singers=[3], outfolder='temp', audio_seq_len=750, pitch_qt=None, type='response', invert_pitch_fn=None):
     
     logging.log(logging.INFO, 'Generate function')
     logging.log(logging.INFO, 'Generating pitch')
-    pitch, inverted_pitch = generate_pitch_reinterp(pitch, pitch_model, invert_pitch_fn, num_samples=num_samples, num_steps=100)
+    if type == 'response':
+        pitch, inverted_pitch = generate_pitch_response(pitch, pitch_model, invert_pitch_fn, num_samples=num_samples, num_steps=100)
+    elif type == 'reinterp':
+        pitch, inverted_pitch = generate_pitch_reinterp(pitch, pitch_model, invert_pitch_fn, num_samples=num_samples, num_steps=100)
+    else:
+        raise ValueError(f'Invalid type: {type}')
     if pitch_qt is not None:
         # if there is not pitch quantile transformer, undo the default quantile transformation that occurs
         def undo_qt(x, min_clip=200):
@@ -129,40 +144,54 @@ def generate(pitch, num_samples=1, num_steps=100, singers=[3], outfolder='temp',
     audio = generate_audio(audio_model, interpolated_pitch, invert_audio_fn, singers=singers, num_steps=100)
     audio = audio.detach().cpu().numpy()
     pitch = pitch.detach().cpu().numpy()
-    pitch_vals = np.where(pitch[0][:, 0] == 0, np.nan, pitch[0].flatten())
-
     # generate plot of model output to display on interface
+    pdb.set_trace()
     model_output_plot = plt.figure()
-    plt.plot(pitch_vals, figure=model_output_plot, label='Model Output')
+    inverted_pitch = np.where(inverted_pitch == 0, np.nan, inverted_pitch)
+    plt.plot(inverted_pitch, figure=model_output_plot, label='Model Output')
     plt.close(model_output_plot)
-    return (16000, audio[0]), model_output_plot, pitch_vals
+    return (16000, audio[0]), model_output_plot # return audio and plot
 
-# pdb.set_trace()
-pitch_model, pitch_qt, pitch_task_fn, invert_pitch_fn, _ = load_pitch_fns(
-    os.path.join(pitch_path, 'last.ckpt'), \
-    model_type = 'diffusion', \
-    config_path = os.path.join(pitch_path, 'config.gin'), \
-    qt_path = os.path.join(pitch_path, 'qt.joblib'), \
-    device = device
-)
+pitch_model, pitch_qt, pitch_task_fn, invert_pitch_fn = None, None, None, None # initialize pitch model based on user preference
 audio_model, audio_qt, audio_seq_len, invert_audio_fn = load_audio_fns(
     os.path.join(audio_path, 'last.ckpt'),
     qt_path = os.path.join(audio_path, 'qt.joblib'),
     config_path = os.path.join(audio_path, 'config.gin'),
     device = device
 )
-partial_generate = partial(generate, num_samples=1, num_steps=100, singers=[3], outfolder=None, pitch_qt=pitch_qt)  # generate function with default arguments
 
-@spaces.GPU(duration=150)
-def set_guide_and_generate(audio):
-    global selected_prime, pitch_task_fn
-    
+
+def load_pitch_model(model_selection):
+    global device
+    pitch_path = pitch_paths[model_selection]
+    pitch_model, pitch_qt, pitch_task_fn, invert_pitch_fn, _ = load_pitch_fns(
+        os.path.join(pitch_path, 'last.ckpt'), \
+        model_type = 'diffusion', \
+        config_path = os.path.join(pitch_path, 'config.gin'), \
+        qt_path = os.path.join(pitch_path, 'qt.joblib'), \
+        device = device
+    )
+    return pitch_model, pitch_qt, pitch_task_fn, invert_pitch_fn
+
+def container_generate(model_selection, task_selection, audio):
+    global pitch_model, pitch_qt, pitch_task_fn, invert_pitch_fn, model_loaded
+    # load pitch model
+    if model_loaded is None or model_loaded != model_selection:
+        pitch_model, pitch_qt, pitch_task_fn, invert_pitch_fn = load_pitch_model(model_selection)
+        model_loaded = model_selection
+    else:
+        logging.log(logging.INFO, f'using existing model: {model_selection}')
+
+    # extract pitch from input
     if audio is None:
         return None, None
     sr, audio = audio
-    if len(audio) < 12*sr:
+    if len(audio) < 12*sr and task_selection == 'Melodic Reinterpretation':    
+        # make sure the audio is at least 12 s long
         audio = np.pad(audio, (0, 12*sr - len(audio)), mode='constant')
-    
+    if len(audio) < 4*sr and task_selection == 'Call and Response':     
+        # make sure the audio is at least 4 s long
+        audio = np.pad(audio, (4*sr - len(audio), 0), mode='constant')
     audio = audio.astype(np.float32)
     audio /= np.max(np.abs(audio))
     audio = librosa.resample(audio, orig_sr=sr, target_sr=16000) # convert only last 4 s
@@ -186,55 +215,73 @@ def set_guide_and_generate(audio):
     f0 = f0.reshape(1, 1, -1)
     f0 = torch.tensor(f0).to(pitch_model.device).float()
     logging.log(logging.INFO, 'Calling generate function')
-    audio, pitch, _ = partial_generate(f0)
     mic_f0 = np.where(mic_f0 == 0, np.nan, mic_f0)
     # plot user input
     user_input_plot = plt.figure()
     plt.plot(np.arange(0, len(mic_f0)), mic_f0, label='User Input', figure=user_input_plot)
     plt.close(user_input_plot)
-    return audio, user_input_plot, pitch
-
+    
+    if task_selection == 'Call and Response':
+        partial_generate = partial(generate, num_samples=1, num_steps=100, singers=[3], outfolder=None, pitch_qt=pitch_qt, type='response', invert_pitch_fn=invert_pitch_fn) 
+    else:
+        partial_generate = partial(generate, num_samples=1, num_steps=100, singers=[3], outfolder=None, pitch_qt=pitch_qt, type='reinterp', invert_pitch_fn=invert_pitch_fn) 
+    audio, output_plot = partial_generate(f0)
+    return audio, user_input_plot, output_plot
 
-gr.HTML("""
-<style>
+css = """
 .center-text {
     text-align: center;
 }
-</style>
-""")
-with gr.Blocks(theme=gr.themes.Glass()) as demo:
-    demo.title("GaMaDHaNi: Hierarchical Generative Modeling of Melodic Vocal Contours in Hindustani Classical Music")
-    demo.description("""
-                    :book: Read more about the project [here](https://arxiv.org/pdf/2408.12658) <br>
-                    :samples: Listen to the samples [here](https://snnithya.github.io/gamadhani-samples) <br>
-    """)
+.justify-text {
+    text-align: justify;
+}   
+"""
+
+with gr.Blocks(css=css) as demo:
+    gr.Markdown("# GaMaDHaNi: Hierarchical Generative Modeling of Melodic Vocal Contours in Hindustani Classical Music", elem_classes="center-text")
+    gr.Markdown("### Abstract", elem_classes="center-text")
+    gr.Markdown("""
+        Hindustani music is a performance-driven oral tradition that exhibits the rendition of rich melodic patterns. In this paper, we focus on generative modeling of singers' vocal melodies extracted from audio recordings, as the voice is musically prominent within the tradition. Prior generative work in Hindustani music models melodies as coarse discrete symbols which fails to capture the rich expressive melodic intricacies of singing. Thus, we propose to use a finely quantized pitch contour, as an intermediate representation for hierarchical audio modeling. We propose GaMaDHaNi, a modular two-level hierarchy, consisting of a generative model on pitch contours, and a pitch contour to audio synthesis model. We compare our approach to non-hierarchical audio models and hierarchical models that use a self-supervised intermediate representation, through a listening test and qualitative analysis. We also evaluate audio model's ability to faithfully represent the pitch contour input using Pearson correlation coefficient. By using pitch contours as an intermediate representation, we show that our model may be better equipped to listen and respond to musicians in a human-AI collaborative setting by highlighting two potential interaction use cases (1) primed generation, and (2) coarse pitch conditioning.
+    """, elem_classes="justify-text")
+    gr.Markdown("""
+                    📖 Read more about the project [here](https://arxiv.org/pdf/2408.12658) <br>
+                    🎧 Listen to the samples [here](https://snnithya.github.io/gamadhani-samples) <br>
+    """, elem_classes="center-text")
     with gr.Column():
         gr.Markdown("""
                     ## Instructions
                     In this demo you can interact with the model in two ways: 
-                    1. **Call and response**: The model will try to continue the idea that you input. This is similar to `primed generation' discussed in the paper.
-                    2. **Melodic reinterpretation**: Akin to the idea of `coarse pitch conditioning' presented in the paper, you can input a pitch contour and the model will generate audio that is similar to but not exactly the same. <br><br>
+                    1. **Call and response**: The model will try to continue the idea that you input. This is similar to 'primed generation' discussed in the paper. The last 4 s of the audio will be considered as a 'prime' for the model to continue. <br><br>
+                    2. **Melodic reinterpretation**: Akin to the idea of 'coarse pitch conditioning' presented in the paper, you can input a pitch contour and the model will generate audio that is similar to but not exactly the same. <br><br>
                     ### Upload an audio file or record your voice to get started!
-                    """, elem_classes="center-text")
+                    """)
         gr.Markdown("""
                     This is still a work in progress, so please feel free to share any weird or interesting examples, we would love to hear them! Contact us at [snnithya.mit.edu](mailto:snnithya.mit.edu).
                     """)
         gr.Markdown("""
-                *Note: If you see an error message on the screen after clicking 'Submit', please wait for five seconds and click 'Submit' again.*
+            *Note: If you see an error message on the screen after clicking 'Submit', please wait for five seconds and click 'Submit' again.*
         """)
-
-    with gr.Row(equal_heights=True):
-        
+        gr.Markdown("""
+            *Another note: The model may take around 40-60s to generate an output. Hang tight! But if you're left hanging for too long, let me know!*
+            """)
+        gr.Markdown("""
+            *Last note, I promise: There are some example audio samples at the bottom of the page. You can start with those if you'd like!*
+            """)
+    model_dropdown = gr.Dropdown(["Diffusion Pitch Generator"], label="Select a model type")
+    task_dropdown = gr.Dropdown(label="Select a task", choices=["Call and Response", "Melodic Reinterpretation"])
+    sbmt = gr.Button()
+    with gr.Row(equal_height=True):
         with gr.Column():    
             audio = gr.Audio(label="Input")
-            sbmt = gr.Button()
-            with gr.Accordion("View Pitch Plot"):
-                user_input = gr.Plot(label="User Input")
         with gr.Column():
             generated_audio = gr.Audio(label="Generated Audio")
+    with gr.Row():
+        with gr.Column():
+            with gr.Accordion("View Pitch Plot"):
+                user_input = gr.Plot(label="User Input")  
+        with gr.Column():
             with gr.Accordion("View Pitch Plot"):
                 generated_pitch = gr.Plot(label="Generated Pitch")
-    example_description = gr.Textbox(label="Example Description", interactive=False)
     examples = gr.Examples(
         examples=[
             ["examples/ex1.wav"],
@@ -245,8 +292,7 @@ with gr.Blocks(theme=gr.themes.Glass()) as demo:
         ],
         inputs=audio
     )
-
-    sbmt.click(set_guide_and_generate, inputs=[audio], outputs=[generated_audio, user_input, generated_pitch])
+    sbmt.click(container_generate, inputs=[model_dropdown, task_dropdown, audio], outputs=[generated_audio, user_input, generated_pitch])
 
 def main(argv):