app.py

import os
import glob
import yaml
import torch
import __main__
import numpy as np
import soundfile as sf
import librosa
import librosa.display
import matplotlib.pyplot as plt

import spaces
import gradio as gr

from src.model.nn.synthesizer import Synthesizer
from src.utils.misc import triangular, downsample
from src.utils.plot import state_video as plot_state_video
from src.utils.audio import mel_basis, state_to_wav
from src.utils.control import vibrato as control_vibrato

class ConfigArgument:
    def __getitem__(self,key):
        return getattr(self, key)
    def __setitem__(self,key,value):
        return setattr(self, key, value)
setattr(__main__, "ConfigArgument", ConfigArgument)

def filter_state_dict(ckpt):
    out_dict = {}
    for key in ckpt.keys():
        new_key = key[6:] if str(key)[:6] == 'model.' else key
        out_dict[new_key] = ckpt[key]
    return out_dict

def flush(directory):
    os.makedirs(directory, exist_ok=True)
    files = glob.glob(f'{directory}/*')
    for f in files:
        os.remove(f)

def add_glissando(f_0, Nt, sr, glissando, max_t):
    front = int(0.2 * np.random.rand() * sr * max_t)
    rear =  int((0.2 * np.random.rand() + 0.3) * sr * max_t)
    middle = max(0, len(f_0) - front - rear)
    ramp = glissando * torch.cat((torch.zeros(front), torch.linspace(0,1,middle), torch.ones(rear)), dim=-1)
    return f_0 * (1 + ramp)

def plot_spectrogram(path, x, n_fft=2048, hop_length=512, n_mel=256, samplerate=48000, max_duration=1):
    x_wave = np.zeros(int(max_duration * samplerate))
    x_wave[:len(x)] += x
    x_spec = librosa.stft(
        x_wave, n_fft=n_fft, hop_length=hop_length, win_length=n_fft, pad_mode='reflect')
    mag = np.abs(x_spec) # (n_frames, n_freq)
    mel_fbank = mel_basis(samplerate, n_fft, n_mel) # (n_mel, n_freq)
    mel = np.einsum('ij,jk->ik', mel_fbank, mag) # (n_frames, n_mel)
    
    plt.figure(figsize=(7,7))
    librosa.display.specshow(mel)
    plt.xticks([])
    plt.yticks([])
    plt.clim([0, 30])
    plt.tight_layout()
    plt.savefig(path, transparent=True)
    plt.close('all')
    plt.clf()

with open("ckpt/config.yaml") as stream:
    configs = yaml.safe_load(stream)

with open("ckpt/pitch.yaml") as stream:
    pitch_dict = yaml.safe_load(stream)

def get_data(duration, resolution, note, glissando, vibrato, stiffness, tension, pluck, amplitude):
    sr = configs['sr']
    Nt = int(duration * sr)
    Nx = int(resolution)

    xgrid = torch.linspace(0,1,Nx)
    tgrid = torch.arange(Nt) / sr
    pitch = pitch_dict[note]
  
    t60_min_1=20.; t60_max_1=30.; t60_min_2=30.; t60_max_2=30.
    t60_diff_max=5.
    T60 = torch.Tensor([[[1000., 25.],[100., 30.]]])

    Nw = int(Nt / configs['block_size']) + 1
    
    xg, tg = torch.meshgrid(xgrid, tgrid, indexing='ij')
    ka = torch.Tensor([stiffness]).view(-1,1) # (1,1)
    al = torch.Tensor([tension]).view(-1,1) # (1,1)
    f_0 = torch.ones(Nt) * pitch # (Nt)
    nx  = torch.Tensor([[[Nx]]]).float()
    p_x = torch.ones_like(nx) * pluck
    p_a = torch.ones_like(nx) * amplitude
    u_0 = triangular(Nx, nx, p_x, p_a) # (1, 1, Nx)
   
    f_0 = add_glissando(f_0, Nt, sr, glissando, Nt / sr)
    f_0 = f_0 + control_vibrato(f_0.view(1,-1), 1/sr, mf=[3.,5.], ma=vibrato)
    f_0 = downsample(f_0, factor=configs['block_size'])

    xg  = xg[:,0].view(-1,1) # (Nx, 1)
    tg  = tg                 # (Nx, Nt)
    ka  = ka.repeat(Nx,1)    # (Nx, 1)
    al  = al.repeat(Nx,1)    # (Nx, 1)
    T60 = T60                # (Nx, 1, 1)
    f_0 = f_0.repeat(Nx,1)   # (Nx, Nw)
    u_0 = u_0.repeat(Nx,1,1) # (Nx, 1, Nx)
    
    params = [xg, tg, ka, al, T60, None, None]
    return params, f_0, u_0
 
@spaces.GPU
def run(duration, resolution, pitch, glissando, vibrato, stiffness, tension, pluck, amplitude):
    checkpoint = torch.load('ckpt/dmsp.ckpt', map_location='cpu')
    checkpoint = filter_state_dict(checkpoint['state_dict'])
    model = Synthesizer(**configs)
    model.load_state_dict(checkpoint)
    if torch.cuda.is_available():
        model = model.cuda()
    
    params, f_0, u_0 = get_data( \
        duration, resolution, pitch, glissando, vibrato, stiffness, tension, pluck, amplitude)

    if torch.cuda.is_available():
        params = [p.cuda() if p is not None else p for p in params]
        f_0 = f_0.cuda()
        u_0 = u_0.cuda()

    with torch.no_grad():
        ut, mode_input, mode_output = model(params, f_0, u_0)
    ut = ut.detach().cpu() # (Nx, Nt)
    ut_wave = configs['gain'] * ut.mean(0)
    
    save_dir = 'results'
    prefix = 'dmsp'
    fname = 'output'
    flush(save_dir)
    audio_name = f'{save_dir}/{fname}.wav'
    video_name = f'{save_dir}/{prefix}-{fname}.mp4'
    spec_name  = f'{save_dir}/spec.png'

    ut = ut.numpy().T
    ut_wave = ut_wave.numpy()
    maxy = 0.022
    sf.write(audio_name, ut_wave, samplerate=configs['sr'])
    plot_spectrogram(spec_name, ut_wave, samplerate=configs['sr'])
    plot_state_video(save_dir, ut, configs['sr'], prefix=prefix, fname=fname, maxy=maxy)
    return spec_name, video_name

pitch_list = ["G2", "Ab2", "A2", "Bb2", "B2", "C3", "Db3", "D3", "Eb3", "E3", "F3", "Gb3", "G3", "Ab3", "A3", "Bb3", "B3", "C4", "Db4", "D4", "Eb4", "E4", "F4", "Gb4", "G4",]

duration   = gr.Slider(0.1, 1.0, value=1.0, label="Temporal Duration")
resolution = gr.Slider(128, 256, value=256, label="Spatial Resolution", info='Reduce to simulate faster. Recommended to leave it as 256.')
pitch      = gr.Dropdown(pitch_list, value="C3", label="Pitch", info="Specify the fundamental frequency as a musical note.")
glissando  = gr.Slider(-0.4, 0.4, value=0, label="Glissando", info='Set +/- to ascend (+) or descend (-) the pitch')
vibrato    = gr.Slider(0, 0.25, value=0, label="Vibrato", info='Set larger value to add more vibrato')
stiffness  = gr.Slider(0.011, 0.029, value=0.02, label="Stiffness", info='Stiffness can change the resulting pitch. Specify low values when tension is high')
tension    = gr.Slider(1.0,  25, value=4, label="Stiffness-Tension Ratio", info='Tension can introduce non-linear effects such as pitch glide. Specify low values when stiffness is high')
pluck      = gr.Slider(0.12, 0.5, value=0.2, label="Plucking Position", info='Peak position of an initial condition')
amplitude  = gr.Slider(0.001, 0.02, value=0.015, label="Plucking Amplitude", info='Peak amplitude of an initial condition')

demo = gr.Interface(
    fn=run,
    inputs=[
        duration, resolution, pitch, glissando, vibrato,
        stiffness, tension, pluck, amplitude,
    ],
    outputs=[
        gr.Image(),
        gr.Video(format='mp4', include_audio=True),
    ],
)
demo.launch()