rearranged model files and added config

app.py

@@ -23,10 +23,8 @@ from hmmlearn import hmm
 import time
 import soundfile as sf
 
-pitch_path = '/network/scratch/n/nithya.shikarpur/checkpoints/pitch-diffusion/corrected-attention-v3/4833583'
-audio_path = '/network/scratch/n/nithya.shikarpur/checkpoints/pitch-diffusion/corrected-attention-v3/4835364'
-pitch_primes = '/network/scratch/n/nithya.shikarpur/pitch-diffusion/data/merged_data-final/listening_study_primes.npz'
-output_folder = '/network/scratch/n/nithya.shikarpur/pitch-diffusion/user-studies/listening-study-2/task-3'
+pitch_path = 'models/diffusion_pitch/'
+audio_path = 'models/pitch_to_audio/'
 device = 'cpu'
 
 global_ind = -1
@@ -232,11 +230,12 @@ def set_prime_and_generate(audio, full_pitch, full_audio, full_user):
     return audio, pitch, full_pitch, full_audio, full_user, fig
 
 def save_session(full_pitch, full_audio, full_user):
-    os.makedirs(output_folder, exist_ok=True)
-    filename = f'session-{time.time()}'
-    logging.log(logging.INFO, f"Saving session to {filename}")
-    pd.DataFrame({'pitch': full_pitch, 'time': np.arange(0, len(full_pitch)/100, 0.01), 'user': full_user}).to_csv(os.path.join(output_folder, filename + '.csv'), index=False)
-    sf.write(os.path.join(output_folder, filename + '.wav'), full_audio[1], 16000)
+    pass
+    # os.makedirs(output_folder, exist_ok=True)
+    # filename = f'session-{time.time()}'
+    # logging.log(logging.INFO, f"Saving session to {filename}")
+    # pd.DataFrame({'pitch': full_pitch, 'time': np.arange(0, len(full_pitch)/100, 0.01), 'user': full_user}).to_csv(os.path.join(output_folder, filename + '.csv'), index=False)
+    # sf.write(os.path.join(output_folder, filename + '.wav'), full_audio[1], 16000)
 
 with gr.Blocks() as demo:
     full_audio = gr.State((16000, np.array([])))

models/diffusion_pitch/config.gin

@@ -0,0 +1,69 @@
+from __gin__ import dynamic_registration
+from src import dataset
+from src import model
+from src import utils
+import torch
+
+# Macros:
+# ==============================================================================
+LR = 0.0001
+SEQ_LEN = 1200
+TRANSPOSE_VALUE = 400
+
+# Parameters for torch.optim.AdamW:
+# ==============================================================================
+torch.optim.AdamW.betas = (0.9, 0.99)
+torch.optim.AdamW.lr = %LR
+
+# Parameters for utils.build_warmed_exponential_lr_scheduler:
+# ==============================================================================
+utils.build_warmed_exponential_lr_scheduler.cycle_length = 200000
+utils.build_warmed_exponential_lr_scheduler.eta_max = %LR
+utils.build_warmed_exponential_lr_scheduler.eta_min = 0.1
+utils.build_warmed_exponential_lr_scheduler.peak_iteration = 10000
+utils.build_warmed_exponential_lr_scheduler.start_factor = 0.01
+
+# Parameters for model.UNetBase.configure_optimizers:
+# ==============================================================================
+model.UNetBase.configure_optimizers.optimizer_cls = @torch.optim.AdamW
+model.UNetBase.configure_optimizers.scheduler_cls = \
+    @utils.build_warmed_exponential_lr_scheduler
+
+# Parameters for dataset.pitch_read_w_downsample:
+# ==============================================================================
+dataset.pitch_read_w_downsample.add_noise_to_silence = True
+dataset.pitch_read_w_downsample.decoder_key = 'pitch'
+dataset.pitch_read_w_downsample.max_clip = 600
+dataset.pitch_read_w_downsample.min_clip = 200
+dataset.pitch_read_w_downsample.min_norm_pitch = -4915
+dataset.pitch_read_w_downsample.pitch_downsample = 10
+dataset.pitch_read_w_downsample.seq_len = %SEQ_LEN
+dataset.pitch_read_w_downsample.time_downsample = 2
+
+# Parameters for train/dataset.pitch_read_w_downsample:
+# ==============================================================================
+train/dataset.pitch_read_w_downsample.transpose_pitch = %TRANSPOSE_VALUE
+
+# Parameters for train/dataset.SequenceDataset:
+# ==============================================================================
+train/dataset.SequenceDataset.task_fn = @train/dataset.pitch_read_w_downsample
+
+# Parameters for val/dataset.SequenceDataset:
+# ==============================================================================
+val/dataset.SequenceDataset.task_fn = @dataset.pitch_read_w_downsample
+
+# Parameters for model.UNet:
+# ==============================================================================
+model.UNet.dropout = 0.3
+model.UNet.features = [512, 640, 1024]
+model.UNet.inp_dim = 1
+model.UNet.kernel_size = 5
+model.UNet.nonlinearity = 'mish'
+model.UNet.norm = True
+model.UNet.num_attns = 4
+model.UNet.num_convs = 4
+model.UNet.num_heads = 8
+model.UNet.project_dim = 256
+model.UNet.seq_len = %SEQ_LEN
+model.UNet.strides = [4, 2, 2]
+model.UNet.time_dim = 128

models/pitch_to_audio/config.gin

@@ -0,0 +1,91 @@
+from __gin__ import dynamic_registration
+from src import dataset
+from src import model
+from src import pitch_to_audio_utils
+from src import utils
+import torch
+
+# Macros:
+# ==============================================================================
+AUDIO_SEQ_LEN = 750
+LR = 0.0001
+NFFT = 1024
+NUM_MELS = 192
+SINGER_CONDITIONING = True
+SR = 16000
+
+# Parameters for torch.optim.AdamW:
+# ==============================================================================
+torch.optim.AdamW.betas = (0.9, 0.99)
+torch.optim.AdamW.lr = 0.0001
+
+# Parameters for utils.build_warmed_exponential_lr_scheduler:
+# ==============================================================================
+utils.build_warmed_exponential_lr_scheduler.cycle_length = 480000
+utils.build_warmed_exponential_lr_scheduler.eta_max = %LR
+utils.build_warmed_exponential_lr_scheduler.eta_min = 0.1
+utils.build_warmed_exponential_lr_scheduler.peak_iteration = 10000
+utils.build_warmed_exponential_lr_scheduler.start_factor = 0.01
+
+# Parameters for model.UNetBase.configure_optimizers:
+# ==============================================================================
+model.UNetBase.configure_optimizers.optimizer_cls = @torch.optim.AdamW
+model.UNetBase.configure_optimizers.scheduler_cls = \
+    @utils.build_warmed_exponential_lr_scheduler
+
+# Parameters for pitch_to_audio_utils.from_mels:
+# ==============================================================================
+pitch_to_audio_utils.from_mels.nfft = %NFFT
+pitch_to_audio_utils.from_mels.num_mels = %NUM_MELS
+pitch_to_audio_utils.from_mels.sr = %SR
+
+# Parameters for dataset.load_cached_dataset:
+# ==============================================================================
+dataset.load_cached_dataset.audio_len = %AUDIO_SEQ_LEN
+dataset.load_cached_dataset.return_singer = %SINGER_CONDITIONING
+
+# Parameters for pitch_to_audio_utils.normalized_mels_to_audio:
+# ==============================================================================
+pitch_to_audio_utils.normalized_mels_to_audio.n_iter = 100
+pitch_to_audio_utils.normalized_mels_to_audio.nfft = %NFFT
+pitch_to_audio_utils.normalized_mels_to_audio.num_mels = %NUM_MELS
+pitch_to_audio_utils.normalized_mels_to_audio.sr = %SR
+
+# Parameters for dataset.SequenceDataset:
+# ==============================================================================
+dataset.SequenceDataset.task_fn = @dataset.load_cached_dataset
+
+# Parameters for pitch_to_audio_utils.torch_gl:
+# ==============================================================================
+pitch_to_audio_utils.torch_gl.n_iter = 200
+pitch_to_audio_utils.torch_gl.nfft = %NFFT
+pitch_to_audio_utils.torch_gl.sr = %SR
+
+# Parameters for pitch_to_audio_utils.torch_istft:
+# ==============================================================================
+pitch_to_audio_utils.torch_istft.nfft = %NFFT
+
+# Parameters for model.UNetPitchConditioned:
+# ==============================================================================
+model.UNetPitchConditioned.audio_seq_len = %AUDIO_SEQ_LEN
+model.UNetPitchConditioned.cfg = True
+model.UNetPitchConditioned.cond_drop_prob = 0.2
+model.UNetPitchConditioned.dropout = 0.3
+model.UNetPitchConditioned.f0_dim = 128
+model.UNetPitchConditioned.features = [512, 640, 1024]
+model.UNetPitchConditioned.inp_dim = %NUM_MELS
+model.UNetPitchConditioned.kernel_size = 5
+model.UNetPitchConditioned.log_samples_every = 10
+model.UNetPitchConditioned.log_wandb_samples_every = 50
+model.UNetPitchConditioned.nonlinearity = 'mish'
+model.UNetPitchConditioned.norm = False
+model.UNetPitchConditioned.num_attns = 4
+model.UNetPitchConditioned.num_convs = 4
+model.UNetPitchConditioned.num_heads = 8
+model.UNetPitchConditioned.project_dim = 256
+model.UNetPitchConditioned.singer_conditioning = %SINGER_CONDITIONING
+model.UNetPitchConditioned.singer_dim = 128
+model.UNetPitchConditioned.singer_vocab = 55
+model.UNetPitchConditioned.sr = %SR
+model.UNetPitchConditioned.strides = [4, 2, 2]
+model.UNetPitchConditioned.time_dim = 128

requirements.txt

@@ -16,4 +16,6 @@ torchaudio==2.4.0
 tqdm==4.65.0
 wandb==0.15.4
 x_transformers==1.32.15
+crepe==0.0.15
+hmmlearn==0.3.2