testing feasibility

app.py

@@ -0,0 +1,268 @@
+from gradio import Interface, Audio
+import gradio as gr
+import numpy as np
+import torch
+import subprocess
+import librosa
+import matplotlib.pyplot as plt
+import pandas as pd
+import os
+from functools import partial
+import gin
+import sys
+sys.path.append('./')
+from src.generate_utils import invert_pitch_read, load_pitch_model, load_audio_model
+import src.pitch_to_audio_utils as p2a
+import torchaudio
+from absl import app
+from torch.nn.functional import interpolate
+import pdb
+import logging
+import crepe
+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'
+device = 'cpu'
+
+global_ind = -1
+global_audios = np.array([0.0])
+global_pitches = np.array([0])
+singer = 3
+audio_components = []
+preprocessed_primes = []
+selected_prime = None
+
+
+
+def make_prime_npz(prime):
+    np.savez('./temp/prime.npz', concatenated_array=[[prime]])
+
+def load_pitch_fns():
+    pitch_model, pitch_qt, _, pitch_task_fn = load_pitch_model(
+        os.path.join(pitch_path, 'config.gin'), 
+        os.path.join(pitch_path, 'models', 'last.ckpt'), 
+        os.path.join(pitch_path, 'qt.joblib'),
+        device=device
+        )
+    invert_pitch_fn = partial(
+        invert_pitch_read,
+        min_norm_pitch=gin.query_parameter('dataset.pitch_read_w_downsample.min_norm_pitch'),
+        time_downsample=gin.query_parameter('dataset.pitch_read_w_downsample.time_downsample'),
+        pitch_downsample=gin.query_parameter('dataset.pitch_read_w_downsample.pitch_downsample'),
+        qt_transform=pitch_qt,
+        min_clip=gin.query_parameter('dataset.pitch_read_w_downsample.min_clip'),
+        max_clip=gin.query_parameter('dataset.pitch_read_w_downsample.max_clip')
+    )
+    return pitch_model, pitch_qt, pitch_task_fn, invert_pitch_fn
+
+def interpolate_pitch(pitch, audio_seq_len):
+    pitch = interpolate(pitch, size=audio_seq_len, mode='linear')
+    plt.plot(pitch[0].squeeze(0).detach().cpu().numpy())
+    plt.savefig(f"./temp/interpolated_pitch.png")
+    plt.close()
+    return pitch
+
+def load_audio_fns():
+    ckpt = os.path.join(audio_path, 'models', 'checkpoint-epoch=3279-val_cross_entropy=0.00-cross_entropy=0.00.ckpt')
+    config = os.path.join(audio_path, 'config.gin')
+    qt = os.path.join(db_path_audio, 'qt.joblib')
+
+    audio_model, audio_qt = load_audio_model(config, ckpt, qt, device=device)
+    audio_seq_len = gin.query_parameter('%AUDIO_SEQ_LEN')
+
+    invert_audio_fn = partial(
+        p2a.normalized_mels_to_audio,
+        qt=audio_qt,
+        n_iter=200
+    )
+
+    return audio_model, audio_qt, audio_seq_len, invert_audio_fn
+
+def predict_voicing(confidence):
+    # https://github.com/marl/crepe/pull/26
+    """
+    Find the Viterbi path for voiced versus unvoiced frames.
+    Parameters
+    ----------
+    confidence : np.ndarray [shape=(N,)]
+        voicing confidence array, i.e. the confidence in the presence of
+        a pitch
+    Returns
+    -------
+    voicing_states : np.ndarray [shape=(N,)]
+        HMM predictions for each frames state, 0 if unvoiced, 1 if
+        voiced
+    """
+    # uniform prior on the voicing confidence
+    starting = np.array([0.5, 0.5])
+
+    # transition probabilities inducing continuous voicing state
+    transition = np.array([[0.99, 0.01], [0.01, 0.99]])
+
+    # mean and variance for unvoiced and voiced states
+    means = np.array([[0.0], [1.0]])
+    variances = np.array([[0.25], [0.25]])
+
+    # fix the model parameters because we are not optimizing the model
+    model = hmm.GaussianHMM(n_components=2)
+    model.startprob_, model.covars_, model.transmat_, model.means_, \
+    model.n_features = starting, variances, transition, means, 1
+
+    # find the Viterbi path
+    voicing_states = model.predict(confidence.reshape(-1, 1), [len(confidence)])
+
+    return np.array(voicing_states)
+
+def extract_pitch(audio, unvoice=True, sr=16000, frame_shift_ms=10, log=True):
+    time, frequency, confidence, _ = crepe.predict(
+      audio, sr=sr,
+      viterbi=True,
+      step_size=frame_shift_ms,
+      verbose=0 if not log else 1)
+    f0 = frequency
+    if unvoice:
+      is_voiced = predict_voicing(confidence)
+      frequency_unvoiced = frequency * is_voiced
+      f0 = frequency_unvoiced
+
+    return time, f0, confidence
+
+def generate_pitch(pitch, pitch_model, invert_pitch_fn, num_samples, num_steps, outfolder=None, processed_primes=None):
+    noisy_pitch = torch.Tensor(pitch[:, :, :1200]).to(pitch_model.device) + (torch.normal(mean=0.0, std=0.4*torch.ones(( 1200)))).to(pitch_model.device)
+    noisy_pitch = torch.clamp(noisy_pitch, -5.19, 5.19)
+    samples = pitch_model.sample_sdedit(noisy_pitch, num_samples, num_steps)
+    inverted_pitches = [invert_pitch_fn(samples.detach().cpu().numpy()[0])[0]]
+
+    if outfolder is not None:
+        os.makedirs(outfolder, exist_ok=True)
+        # pdb.set_trace()
+        for i, pitch in enumerate(inverted_pitches):
+            flattened_pitch = pitch.flatten()
+            pd.DataFrame({'f0': flattened_pitch}).to_csv(f"{outfolder}/{i}.csv", index=False)
+            plt.plot(np.where(flattened_pitch == 0, np.nan, flattened_pitch))
+            plt.savefig(f"{outfolder}/{i}.png")
+            plt.close()
+    return samples, inverted_pitches
+
+def generate_audio(audio_model, f0s, invert_audio_fn, outfolder, singers=[3], num_steps=100):
+    singer_tensor = torch.tensor(np.repeat(singers, repeats=f0s.shape[0])).to(audio_model.device)
+    samples, _, singers = audio_model.sample_cfg(f0s.shape[0], f0=f0s, num_steps=num_steps, singer=singer_tensor, strength=3)
+    audio = invert_audio_fn(samples)
+    
+    if outfolder is not None:
+        os.makedirs(outfolder, exist_ok=True)
+        for i, a in enumerate(audio):
+            logging.log(logging.INFO, f"Saving audio {i}")
+            torchaudio.save(f"{outfolder}/{i}.wav", torch.tensor(a).detach().unsqueeze(0).cpu(), 16000)
+    return audio
+     
+def generate(pitch, num_samples=2, num_steps=100, singers=[3], outfolder='temp', audio_seq_len=750, pitch_qt=None ):
+    global global_ind, audio_components
+    global preprocessed_primes
+    # pdb.set_trace()
+    logging.log(logging.INFO, 'Generate function')
+    pitch, inverted_pitch = generate_pitch(pitch, pitch_model, invert_pitch_fn, 1, 100, outfolder=outfolder, processed_primes=selected_prime if global_ind != 0 else None)
+    if pitch_qt is not None:
+        def undo_qt(x, min_clip=200):
+            pitch= pitch_qt.inverse_transform(x.reshape(-1, 1)).reshape(1, -1)
+            pitch = np.around(pitch) # round to nearest integer, done in preprocessing of pitch contour fed into model
+            pitch[pitch < 200] = np.nan
+            return pitch
+        pitch = torch.tensor(np.array([undo_qt(x) for x in pitch.detach().cpu().numpy()])).to(pitch_model.device)
+    interpolated_pitch = interpolate_pitch(pitch=pitch, audio_seq_len=audio_seq_len)
+    interpolated_pitch = torch.nan_to_num(interpolated_pitch, nan=196)
+    interpolated_pitch = interpolated_pitch.squeeze(1) # to match input size by removing the extra dimension
+    audio = generate_audio(audio_model, interpolated_pitch, invert_audio_fn, singers=singers, num_steps=100, outfolder=outfolder)
+    # pdb.set_trace()
+    audio = audio.detach().cpu().numpy()[:, :]
+    pitch = pitch.detach().cpu().numpy()
+    # state = [(16000, audio[0]), (16000, audio[1])]
+    # pdb.set_trace()
+    pitch_vals = np.where(pitch[0][:, 0] == 0, np.nan, pitch[0].flatten())
+    fig1 = plt.figure()
+    # plt.plot(np.arange(0, 400), pitch_vals[:400], figure=fig1, label='User Input')
+    plt.plot(pitch_vals, figure=fig1, label='Pitch')
+    # plt.legend(fig1)
+    # state.append(fig1)
+    plt.close(fig1)
+    return (16000, audio[0]), fig1, pitch_vals
+
+pitch_model, pitch_qt, pitch_task_fn, invert_pitch_fn = load_pitch_fns()
+audio_model, audio_qt, audio_seq_len, invert_audio_fn = load_audio_fns()
+partial_generate = partial(generate, num_samples=1, num_steps=100, singers=[3], outfolder='temp', pitch_qt=pitch_qt)
+
+def set_prime_and_generate(audio, full_pitch, full_audio, full_user):
+    global selected_prime, pitch_task_fn
+    
+    if audio is None:
+        return None, None
+    sr, audio = audio
+    if len(audio) < 12*sr:
+        audio = np.pad(audio, (0, 12*sr - len(audio)), 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
+    mic_audio = audio.copy()
+    audio = audio[-12*16000:]
+    _, f0, _ = extract_pitch(audio)
+    mic_f0 = f0.copy()
+    f0 = pitch_task_fn({
+        'pitch': {
+            'data': f0,
+            'sampling_rate': 100
+        }
+    }, qt_transform=pitch_qt)
+    f0 = f0.reshape(1, 1, -1)
+    f0 = torch.tensor(f0).to(pitch_model.device).float()
+    audio, pitch, pitch_vals = partial_generate(f0)
+    # pdb.set_trace()
+    full_pitch = np.concatenate((full_pitch, mic_f0, pitch_vals))
+    full_user = np.concatenate((full_user, ['User'] * len(mic_f0), ['Model'] * len(pitch_vals)))
+    full_audio[1] = np.concatenate((full_audio[1], mic_audio, audio[1]))
+    # pdb.set_trace()
+    fig = plt.figure()
+    plt.plot(np.arange(0, len(mic_f0)), mic_f0, label='User Input', figure=fig)
+    plt.close(fig)
+    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)
+
+with gr.Blocks() as demo:
+    full_audio = gr.State((16000, np.array([])))
+    full_pitch = gr.State(np.array([]))
+    full_user = gr.State(np.array([]))
+    with gr.Row():
+        with gr.Column():    
+            audio = gr.Audio(label="Input")
+            sbmt = gr.Button()
+            user_input = gr.Plot(label="User Input")
+        with gr.Column():
+            generated_audio = gr.Audio(label="Generated Audio")
+            generated_pitch = gr.Plot(label="Generated Pitch")
+    sbmt.click(set_prime_and_generate, inputs=[audio, full_pitch, full_audio, full_user], outputs=[generated_audio, generated_pitch, full_pitch, full_audio, full_user, user_input])
+    save = gr.Button("Save Session")
+    save.click(save_session, inputs=[full_pitch, full_audio, full_user])
+
+
+
+def main(argv):
+    # audio = np.random.randint(0, high=128, size=(44100*5), dtype=np.int16)
+    # sr = 44100
+    # pdb.set_trace()
+    # p, a = set_prime_and_generate((sr, audio))
+    
+    demo.launch(share=True)
+
+if __name__ == '__main__':
+    app.run(main)

src/generate_utils.py

@@ -0,0 +1,88 @@
+import numpy as np
+from typing import Optional
+from sklearn.preprocessing import QuantileTransformer
+import sys
+import pdb
+sys.path.append('../pitch-diffusion')
+import torch
+import gin
+from src.model import UNet, UNetPitchConditioned
+from functools import partial
+import joblib
+from src.dataset import hz_to_cents, pitch_read_w_downsample
+
+def invert_pitch_read(pitch,
+                  min_norm_pitch: int,
+                  time_downsample: int,
+                  pitch_downsample: int,
+                  qt_transform: Optional[QuantileTransformer],
+                  min_clip: int,
+                  max_clip: int):
+    try:
+        pitch = pitch.detach().cpu().numpy()
+    except:
+        pass
+    if qt_transform is not None:
+        pitch = qt_transform.inverse_transform(pitch.reshape(-1, 1))
+        pitch.reshape(1, -1)
+    pitch[pitch < min_clip] = np.nan
+    pitch[~np.isnan(pitch)] = (pitch[~np.isnan(pitch)] - 1) * pitch_downsample
+    pitch[~np.isnan(pitch)] = pitch[~np.isnan(pitch)] + min_norm_pitch
+    pitch[~np.isnan(pitch)] = 440 * 2**(pitch[~np.isnan(pitch)] / 1200)
+    pitch[np.isnan(pitch)] = 0
+
+    return pitch, 200//time_downsample
+
+def invert_tonic(tonic: Optional[int] = None, 
+                      min_norm_pitch: int = 0,
+                      min_clip: int = 200,
+                      pitch_downsample: int = 1,
+                      ):
+    tonic += min_clip
+    tonic = pitch_downsample * (tonic - 1)
+    tonic += min_norm_pitch
+    tonic = 440 * 2**(tonic / 1200)
+   
+    return tonic
+
+def load_processed_pitch(pitch,
+                         audio_seq_len: int,
+                         min_norm_pitch: int,
+                         pitch_downsample: int,
+                         min_clip: int,
+                         max_clip: int,
+                        ):
+    # pdb.set_trace()
+    pitch = hz_to_cents(pitch, min_norm_pitch=min_norm_pitch, min_clip=min_clip, max_clip=max_clip, pitch_downsample=pitch_downsample, silence_token=min_clip-4)
+    pitch_inds = np.linspace(0, pitch.shape[0], num=audio_seq_len, endpoint=False)
+    pitch = np.interp(pitch_inds, np.arange(0, pitch.shape[0]), pitch)
+    return pitch
+
+def load_pitch_model(config, ckpt, qt = None, prime_file=None, device='cuda'):
+    gin.parse_config_file(config)
+    model = UNet()
+    model.load_state_dict(torch.load(ckpt)['state_dict'])  
+    model.to(device)
+    if qt is not None:
+        qt = joblib.load(qt)
+    if prime_file is not None:
+        with gin.config_scope('val'): # probably have to change this
+            with gin.unlock_config():
+                gin.bind_parameter('dataset.pitch_read_w_downsample.qt_transform', qt)
+        primes = np.load(prime_file, allow_pickle=True)['concatenated_array'][:, 0]
+    else:
+        primes = None
+        task_fn = None
+    task_fn = partial(pitch_read_w_downsample,
+    seq_len=None)
+    return model, qt, primes, task_fn
+
+def load_audio_model(config, ckpt, qt = None, device='cuda'):
+    gin.parse_config_file(config)
+    model = UNetPitchConditioned() # there are no gin parameters for some reason
+    model.load_state_dict(torch.load(ckpt)['state_dict'])  
+    model.to(device)
+    if qt is not None:
+        qt = joblib.load(qt)
+
+    return model, qt

src/pitch_to_audio_utils.py

@@ -0,0 +1,121 @@
+import math
+import librosa as li
+import torch
+from tqdm import tqdm
+import numpy as np
+import gin
+import logging
+
+import pdb
+
+@gin.configurable
+def torch_stft(x, nfft):
+    window = torch.hann_window(nfft).to(x)
+    x = torch.stft(
+        x,
+        n_fft=nfft,
+        hop_length=nfft // 4,
+        win_length=nfft,
+        window=window,
+        center=True,
+        return_complex=True,
+    )
+    x = 2 * x / torch.mean(window)
+    return x
+
+@gin.configurable
+def torch_istft(x, nfft):
+    # pdb.set_trace()
+    window = torch.hann_window(nfft).to(x.device)
+    x = x / 2 * torch.mean(window)
+    return torch.istft(
+        x,
+        n_fft=nfft,
+        hop_length=nfft // 4,
+        win_length=nfft,
+        window=window,
+        center=True,
+    )
+
+@gin.configurable
+def to_mels(stft, nfft, num_mels, sr, eps=1e-2):
+    mels = li.filters.mel(
+        sr=sr,
+        n_fft=nfft,
+        n_mels=num_mels,
+        fmin=40,
+    )
+    # pdb.set_trace()
+    mels = torch.from_numpy(mels).to(stft)
+    mel_stft = torch.einsum("mf,bft->bmt", mels, stft)
+    mel_stft = torch.log(mel_stft + eps)
+    return mel_stft
+
+@gin.configurable
+def from_mels(mel_stft, nfft, num_mels, sr, eps=1e-2):
+    mels = li.filters.mel(
+        sr=sr,
+        n_fft=nfft,
+        n_mels=num_mels,
+        fmin=40,
+    )
+    mels = torch.from_numpy(mels).to(mel_stft)
+    mels = torch.pinverse(mels)
+    mel_stft = torch.exp(mel_stft) - eps
+    stft = torch.einsum("fm,bmt->bft", mels, mel_stft)
+    return stft
+
+@gin.configurable
+def torch_gl(stft, nfft, sr, n_iter):
+
+    def _gl_iter(phase, xs, stft):
+        del xs
+        # pdb.set_trace()
+        c_stft = stft * torch.exp(1j * phase)
+        rec = torch_istft(c_stft, nfft)
+        r_stft = torch_stft(rec, nfft)
+        phase = torch.angle(r_stft)
+        return phase, None
+
+    phase = torch.rand_like(stft) * 2 * torch.pi
+
+    for _ in tqdm(range(n_iter)):
+        phase, _ = _gl_iter(phase, None, stft)
+
+    c_stft = stft * torch.exp(1j * phase)
+    audio = torch_istft(c_stft, nfft)
+
+    return audio
+
+@gin.configurable
+def normalize(x, qt=None):
+    x_flat = x.reshape(-1, 1)
+    if qt is None:
+        logging.warning('No quantile transformer found, returning input')
+        return x
+    return torch.Tensor(qt.transform(x_flat).reshape(x.shape))
+
+@gin.configurable
+def unnormalize(x, qt=None):
+    x_flat = x.reshape(-1, 1)
+    if qt is None:
+        logging.warning('No quantile transformer found, returning input')
+        return x
+    if isinstance(x_flat, torch.Tensor):
+        x_flat = x_flat.detach().cpu().numpy()
+    return torch.Tensor(qt.inverse_transform(x_flat).reshape(x.shape))
+
+@gin.configurable
+def audio_to_normalized_mels(x, nfft, num_mels, sr, qt):
+    # pdb.set_trace()
+    stfts = torch_stft(x, nfft=nfft).abs()[..., :-1]
+    mel_stfts = to_mels(stfts, nfft, num_mels, sr)
+    return normalize(mel_stfts, qt).to(x)
+
+@gin.configurable
+def normalized_mels_to_audio(x, nfft, num_mels, sr, qt, n_iter=20):
+    x = unnormalize(x, qt).to(x)
+    x = from_mels(x, nfft, num_mels, sr)
+    x = torch.clamp(x, 0, nfft)
+    x = torch_gl(x, nfft, sr, n_iter=n_iter)
+    return x