lambda/generative-music-visualizer

1import librosa2import numpy as np3import moviepy.editor as mpy4import random5import torch6from tqdm import tqdm7import dnnlib8import legacy9 10target_sr = 2205011 12def visualize(audio_file,13              network,14              truncation,15              tempo_sensitivity,16              jitter,17              frame_length,18              duration,19              ):20    print(audio_file)21 22    if audio_file:23        print('\nReading audio \n')24        audio, sr = librosa.load(audio_file, duration=duration)25    else:26        raise ValueError("you must enter an audio file name in the --song argument")27 28    # print(sr)29    # print(audio.dtype)30    # print(audio.shape)31    # if audio.shape[0] < duration * sr:32    #     duration = None33    # else:34    #     frames = duration * sr35    #     audio = audio[:frames]36    #37    # print(audio.dtype)38    # print(audio.shape)39    # if audio.dtype == np.int16:40    #     print(f'min: {np.min(audio)}, max: {np.max(audio)}')41    #     audio = audio.astype(np.float32, order='C') / 2**1542    # elif audio.dtype == np.int32:43    #     print(f'min: {np.min(audio)}, max: {np.max(audio)}')44    #     audio = audio.astype(np.float32, order='C') / 2**3145    # audio = audio.T46    # audio = librosa.to_mono(audio)47    # audio = librosa.resample(audio, orig_sr=sr, target_sr=target_sr, res_type="kaiser_best")48    # print(audio.dtype)49    # print(audio.shape)50 51 52 53    # TODO:54    batch_size = 155    resolution = 51256    outfile="output.mp4"57 58    tempo_sensitivity = tempo_sensitivity * frame_length / 51259 60    # Load pre-trained model61    device = torch.device('cuda')62    with dnnlib.util.open_url(network) as f:63        G = legacy.load_network_pkl(f)['G_ema'].to(device) # type: ignore64        G.eval()65 66    with torch.no_grad():67        z = torch.randn([1, G.z_dim]).cuda()    # latent codes68        c = None                                # class labels (not used in this example)69        img = G(z, c)                           # NCHW, float32, dynamic range [-1, +1], no truncation70 71    #set device72    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')73 74    #create spectrogram75    spec = librosa.feature.melspectrogram(y=audio, sr=target_sr, n_mels=512,fmax=8000, hop_length=frame_length)76 77    #get mean power at each time point78    specm=np.mean(spec,axis=0)79 80    #compute power gradient across time points81    gradm=np.gradient(specm)82 83    #set max to 184    gradm=gradm/np.max(gradm)85 86    #set negative gradient time points to zero87    gradm = gradm.clip(min=0)88 89    #normalize mean power between 0-190    specm=(specm-np.min(specm))/np.ptp(specm)91 92    #initialize first noise vector93    nv1 = torch.randn([G.z_dim]).cuda()94 95    #initialize list of class and noise vectors96    noise_vectors=[nv1]97 98    #initialize previous vectors (will be used to track the previous frame)99    nvlast=nv1100 101    #initialize the direction of noise vector unit updates102    update_dir=np.zeros(512)103    print(len(nv1))104    for ni,n in enumerate(nv1):105        if n<0:106            update_dir[ni] = 1107        else:108            update_dir[ni] = -1109 110    #initialize noise unit update111    update_last=np.zeros(512)112 113    #get new jitters114    def new_jitters(jitter):115        jitters=np.zeros(512)116        for j in range(512):117            if random.uniform(0,1)<0.5:118                jitters[j]=1119            else:120                jitters[j]=1-jitter121        return jitters122 123 124    #get new update directions125    def new_update_dir(nv2,update_dir):126        for ni,n in enumerate(nv2):127            if n >= 2*truncation - tempo_sensitivity:128                update_dir[ni] = -1129 130            elif n < -2*truncation + tempo_sensitivity:131                update_dir[ni] = 1132        return update_dir133 134    print('\nGenerating input vectors \n')135    for i in tqdm(range(len(gradm))):136 137        #update jitter vector every 100 frames by setting ~half of noise vector units to lower sensitivity138        if i%200==0:139            jitters=new_jitters(jitter)140 141        #get last noise vector142        nv1=nvlast143 144        #set noise vector update based on direction, sensitivity, jitter, and combination of overall power and gradient of power145        update = np.array([tempo_sensitivity for k in range(512)]) * (gradm[i]+specm[i]) * update_dir * jitters146 147        #smooth the update with the previous update (to avoid overly sharp frame transitions)148        update=(update+update_last*3)/4149 150        #set last update151        update_last=update152 153        #update noise vector154        nv2=nv1.cpu()+update155 156        #append to noise vectors157        noise_vectors.append(nv2)158 159        #set last noise vector160        nvlast=nv2161 162        #update the direction of noise units163        update_dir=new_update_dir(nv2,update_dir)164 165    noise_vectors = torch.stack([nv.cuda() for nv in noise_vectors])166 167 168    print('\n\nGenerating frames \n')169    frames = []170    for i in tqdm(range(noise_vectors.shape[0] // batch_size)):171 172        noise_vector=noise_vectors[i*batch_size:(i+1)*batch_size]173 174        c = None  # class labels (not used in this example)175        with torch.no_grad():176            img = np.array(G(noise_vector, c, truncation_psi=truncation, noise_mode='const').cpu())            # NCHW, float32, dynamic range [-1, +1], no truncation177            img = np.transpose(img, (0,2,3,1)) #CHW -> HWC178            img = np.clip((img * 127.5 + 128), 0, 255).astype(np.uint8)179 180        # add to frames181        for im in img:182            frames.append(im)183 184 185    #Save video186    aud = mpy.AudioFileClip(audio_file)187 188    if duration < aud.duration:189        aud.duration = duration190 191    fps = target_sr / frame_length192    clip = mpy.ImageSequenceClip(frames, fps=fps)193    clip = clip.set_audio(aud)194    clip.write_videofile(outfile, audio_codec='aac', ffmpeg_params=[195        # "-vf", "scale=-1:2160:flags=lanczos",196        "-bf", "2",197        "-g", f"{fps/2}",198        "-crf", "18",199        "-movflags", "faststart"200    ])201 202    return outfile