updated the models

.gitattributes

@@ -5,3 +5,5 @@ models/pitch_to_audio/last.ckpt filter=lfs diff=lfs merge=lfs -text
 models/diffusion_pitch/qt.joblib filter=lfs diff=lfs merge=lfs -text
 models/pitch_to_audio/qt.joblib filter=lfs diff=lfs merge=lfs -text
 examples/** filter=lfs diff=lfs merge=lfs -text
+models/diffusion_pitch/model.ckpt filter=lfs diff=lfs merge=lfs -text
+models/transformer_pitch/model.ckpt filter=lfs diff=lfs merge=lfs -text

app.py

@@ -41,7 +41,10 @@ from gamadhani.utils.utils import get_device
 import copy
 
 logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s', force=True)
-pitch_paths = {'Diffusion Pitch Generator': 'models/diffusion_pitch/'}
+pitch_paths = {
+    'Diffusion Pitch Generator': ('diffusion', 'models/diffusion_pitch/'),
+    'Autoregressive Pitch Generator': ('transformer', 'models/transformer_pitch/')
+    }
 model_loaded = None
 audio_path = 'models/pitch_to_audio/'
 device = get_device()
@@ -110,10 +113,13 @@ def generate_pitch_reinterp(pitch, pitch_model, invert_pitch_fn, num_samples, nu
 
     return samples, inverted_pitches
 
-def generate_pitch_response(pitch, pitch_model, invert_pitch_fn, num_samples, num_steps):
+def generate_pitch_response(pitch, pitch_model, invert_pitch_fn, num_samples, num_steps, model_type='diffusion'):
     '''Generate pitch values for the call and response task'''
     pitch = pitch[:, :, -400:]   # consider only the last 4 s of the pitch contour
-    samples = pitch_model.sample_fn(num_samples, num_steps, prime=pitch)
+    if model_type == 'diffusion':
+        samples = pitch_model.sample_fn(num_samples, num_steps, prime=pitch)
+    else:
+        samples = pitch_model.sample_fn(batch_size=num_samples, seq_len=800, prime=pitch)
     inverted_pitches = invert_pitch_fn(f0=samples.clone().detach().cpu().numpy()[0]).flatten()   # pitch values in Hz
 
     return samples, inverted_pitches
@@ -127,7 +133,7 @@ def generate_audio(audio_model, f0s, invert_audio_fn, singers=[3], num_steps=100
     return audio
     
 @spaces.GPU(duration=30)
-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, t0=0.5):
+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, t0=0.5, model_type='diffusion'):
     global pitch_model, audio_model
     # move the models to device
     pitch_model = pitch_model.to(device)
@@ -135,10 +141,11 @@ def generate(pitch, num_samples=1, num_steps=100, singers=[3], outfolder='temp',
     logging.log(logging.INFO, 'Generate function')
     # load pitch values onto GPU
     pitch = torch.tensor(pitch).float().unsqueeze(0).unsqueeze(0).to(device)
-    pitch_qt = p2a.GPUQuantileTransformer(pitch_qt, device=device)
+    if pitch_qt is not None:
+        pitch_qt = p2a.GPUQuantileTransformer(pitch_qt, device=device)
     logging.log(logging.INFO, 'Generating pitch')
     if type == 'response':
-        pitch, inverted_pitch = generate_pitch_response(pitch, pitch_model, invert_pitch_fn, num_samples=num_samples, num_steps=100)
+        pitch, inverted_pitch = generate_pitch_response(pitch, pitch_model, invert_pitch_fn, num_samples=num_samples, num_steps=100, model_type=model_type)
     elif type == 'reinterp':
         pitch, inverted_pitch = generate_pitch_reinterp(pitch, pitch_model, invert_pitch_fn, num_samples=num_samples, num_steps=100, t0=t0)
     
@@ -151,9 +158,10 @@ def generate(pitch, num_samples=1, num_steps=100, singers=[3], outfolder='temp',
             pitch= pitch_qt.inverse_transform(x).squeeze(0) # qt transform expects shape (bs, seq_len, 1)
             pitch = torch.round(pitch) # round to nearest integer, done in preprocessing of pitch contour fed into model
             pitch[pitch < 200] = np.nan
+            pitch = pitch.unsqueeze(0)
             return pitch
         pitch = undo_qt(pitch)
-    interpolated_pitch = p2a.interpolate_pitch(pitch=pitch.unsqueeze(0), audio_seq_len=audio_seq_len).squeeze(0)    # interpolate pitch values to match the audio model's input size
+    interpolated_pitch = p2a.interpolate_pitch(pitch=pitch, audio_seq_len=audio_seq_len).squeeze(0)    # interpolate pitch values to match the audio model's input size
     interpolated_pitch = torch.nan_to_num(interpolated_pitch, nan=196)  # replace nan values with silent token
     interpolated_pitch = interpolated_pitch.squeeze(1) # to match input size by removing the extra dimension
     logging.log(logging.INFO, 'Generating audio')
@@ -178,12 +186,12 @@ audio_model, audio_qt, audio_seq_len, invert_audio_fn = load_audio_fns(
 
 def load_pitch_model(model_selection):
     global device
-    pitch_path = pitch_paths[model_selection]
+    model_type, 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', \
+        os.path.join(pitch_path, 'model.ckpt'), \
+        model_type = model_type, \
         config_path = os.path.join(pitch_path, 'config.gin'), \
-        qt_path = os.path.join(pitch_path, 'qt.joblib'), \
+        qt_path = os.path.join(pitch_path, 'qt.joblib') if model_type == 'diffusion' else None, \
         device = 'cpu'
     )
     return pitch_model, pitch_qt, pitch_task_fn, invert_pitch_fn
@@ -239,9 +247,9 @@ def container_generate(model_selection, task_selection, audio, singer_id, t0):
     elif singer_id == 'Singer 2':
         singer = [27]
     if task_selection == 'Call and Response':
-        partial_generate = partial(generate, num_samples=1, num_steps=100, singers=singer, outfolder=None, pitch_qt=pitch_qt, type='response', invert_pitch_fn=invert_pitch_fn) 
+        partial_generate = partial(generate, num_samples=1, num_steps=100, singers=singer, outfolder=None, pitch_qt=pitch_qt, type='response', invert_pitch_fn=invert_pitch_fn, model_type=model_selection) 
     else:
-        partial_generate = partial(generate, num_samples=1, num_steps=100, singers=singer, outfolder=None, pitch_qt=pitch_qt, type='reinterp', invert_pitch_fn=invert_pitch_fn, t0=t0) 
+        partial_generate = partial(generate, num_samples=1, num_steps=100, singers=singer, outfolder=None, pitch_qt=pitch_qt, type='reinterp', invert_pitch_fn=invert_pitch_fn, t0=t0, model_type=model_selection) 
     audio, output_plot = partial_generate(f0)
     return audio, user_input_plot, output_plot
 
@@ -260,6 +268,13 @@ def toggle_visibility(selection):
         return gr.update(visible=True)
     else:
         return gr.update(visible=False)
+    
+def toggle_options(selection, options = ['Call and Response', 'Melodic Reinterpretation']):
+    # Show element if selection is "Show", otherwise hide it
+    if selection == "Melodic Reinterpretation":
+        return gr.update(choices=options)
+    else:
+        return gr.update(choices=options[:-1])
 
 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")
@@ -291,8 +306,9 @@ with gr.Blocks(css=css) as demo:
         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")
+    model_dropdown = gr.Dropdown(["Diffusion Pitch Generator", "Autoregressive Pitch Generator"], label="Select a model type")
     task_dropdown = gr.Dropdown(label="Select a task", choices=["Call and Response", "Melodic Reinterpretation"])
+    model_dropdown.change(toggle_options, outputs=task_dropdown)
     t0 = gr.Slider(label="Faithfulness to the input (For melodic reinterpretation task only)", minimum=0.0, maximum=1.0, step=0.01, value=0.3, visible=False)
     task_dropdown.change(toggle_visibility, inputs=task_dropdown, outputs=t0)
     singer_dropdown = gr.Dropdown(label="Select a singer", choices=["Singer 1", "Singer 2"])

models/transformer_pitch/config.gin

@@ -0,0 +1,61 @@
+from __gin__ import dynamic_registration
+from gamadhani import src
+from gamadhani.src import dataset
+from gamadhani.src import model_transformer
+from gamadhani.src import task_functions
+from gamadhani.utils import utils
+import torch.optim
+
+MODEL_DIM = 512
+EMB_DIM = 512
+NUM_TOKENS = 7928
+NUM_QUANTIZERS = 1
+DROPOUT_RATE = 0.3
+NUM_HEADS = 8
+SEQ_LEN = 1200
+HEAD_DIM = 32
+NUM_LAYERS = 8
+LR = 1e-3
+
+model_transformer.XTransformerPrior:
+    num_tokens = %NUM_TOKENS
+    seq_len = %SEQ_LEN
+    model_dim = %MODEL_DIM
+    emb_dim = %EMB_DIM
+    head_dim = %HEAD_DIM
+    num_layers = %NUM_LAYERS
+    num_heads = %NUM_HEADS
+    dropout_rate = %DROPOUT_RATE
+    
+
+src.dataset.Task:
+    read_fn = @src.task_functions.pitch_read_downsample
+    invert_fn = @src.task_functions.invert_pitch_read_downsample
+    kwargs = {"seq_len": %SEQ_LEN, 
+                    "decoder_key": "pitch",
+                    "min_norm_pitch": -4915,
+                    "time_downsample": 2,
+                    "pitch_downsample": 10,
+                    "base_tonic": 440.}
+
+src.dataset.SequenceDataset:
+    task = @dataset.Task()
+    apply_transform = False
+
+model_transformer.XTransformerPrior.configure_optimizers:
+    optimizer_cls = @torch.optim.AdamW
+    scheduler_cls = @utils.build_warmed_exponential_lr_scheduler
+
+utils.build_warmed_exponential_lr_scheduler:
+    start_factor = .01
+    peak_iteration = 10000
+    cycle_length = 394600
+    eta_min = 0.1
+    eta_max = %LR
+
+utils.set_seed:
+    seed = 2023
+
+torch.optim.AdamW:
+    lr = %LR
+    betas = (.9, .98)

models/transformer_pitch/model.ckpt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:d849eaca79a0bc390d0550b8187d47a843bdb3a6c81b9401e5e925ae1220acc4
+size 356915980

requirements.txt

@@ -1,4 +1,4 @@
 crepe==0.0.15
 hmmlearn==0.3.2
 tensorflow==2.17.0
-GaMaDHaNi @ git+https://github.com/snnithya/GaMaDHaNi.git@c59e24e28ef7b80d43f56b39d3a9dd0563e01df6
+GaMaDHaNi @ git+https://github.com/snnithya/GaMaDHaNi.git@055df71380e0feced7e409470ffc8603f1cfa926