parson/audioEditing
1
1import contextlib2import importlib3 4from inspect import isfunction5import os6import soundfile as sf7import time8import wave9 10import urllib.request11import progressbar12 13CACHE_DIR = os.getenv(14 "AUDIOLDM_CACHE_DIR",15 os.path.join(os.path.expanduser("~"), ".cache/audioldm"))16 17def get_duration(fname):18 with contextlib.closing(wave.open(fname, 'r')) as f:19 frames = f.getnframes()20 rate = f.getframerate()21 return frames / float(rate)22 23def get_bit_depth(fname):24 with contextlib.closing(wave.open(fname, 'r')) as f:25 bit_depth = f.getsampwidth() * 826 return bit_depth27 28def get_time():29 t = time.localtime()30 return time.strftime("%d_%m_%Y_%H_%M_%S", t)31 32def seed_everything(seed):33 import random, os34 import numpy as np35 import torch36 37 random.seed(seed)38 os.environ["PYTHONHASHSEED"] = str(seed)39 np.random.seed(seed)40 torch.manual_seed(seed)41 torch.cuda.manual_seed(seed)42 torch.backends.cudnn.deterministic = True43 torch.backends.cudnn.benchmark = True44 45 46def save_wave(waveform, savepath, name="outwav"):47 if type(name) is not list:48 name = [name] * waveform.shape[0]49 50 for i in range(waveform.shape[0]):51 path = os.path.join(52 savepath,53 "%s_%s.wav"54 % (55 os.path.basename(name[i])56 if (not ".wav" in name[i])57 else os.path.basename(name[i]).split(".")[0],58 i,59 ),60 )61 print("Save audio to %s" % path)62 sf.write(path, waveform[i, 0], samplerate=16000)63 64 65def exists(x):66 return x is not None67 68 69def default(val, d):70 if exists(val):71 return val72 return d() if isfunction(d) else d73 74 75def count_params(model, verbose=False):76 total_params = sum(p.numel() for p in model.parameters())77 if verbose:78 print(f"{model.__class__.__name__} has {total_params * 1.e-6:.2f} M params.")79 return total_params80 81 82def get_obj_from_str(string, reload=False):83 module, cls = string.rsplit(".", 1)84 if reload:85 module_imp = importlib.import_module(module)86 importlib.reload(module_imp)87 return getattr(importlib.import_module(module, package=None), cls)88 89 90def instantiate_from_config(config):91 if not "target" in config:92 if config == "__is_first_stage__":93 return None94 elif config == "__is_unconditional__":95 return None96 raise KeyError("Expected key `target` to instantiate.")97 return get_obj_from_str(config["target"])(**config.get("params", dict()))98 99 100def default_audioldm_config(model_name="audioldm-s-full"): 101 basic_config = {102 "wave_file_save_path": "./output",103 "id": {104 "version": "v1",105 "name": "default",106 "root": "/mnt/fast/nobackup/users/hl01486/projects/general_audio_generation/AudioLDM-python/config/default/latent_diffusion.yaml",107 },108 "preprocessing": {109 "audio": {"sampling_rate": 16000, "max_wav_value": 32768},110 "stft": {"filter_length": 1024, "hop_length": 160, "win_length": 1024},111 "mel": {112 "n_mel_channels": 64,113 "mel_fmin": 0,114 "mel_fmax": 8000,115 "freqm": 0,116 "timem": 0,117 "blur": False,118 "mean": -4.63,119 "std": 2.74,120 "target_length": 1024,121 },122 },123 "model": {124 "device": "cuda",125 "target": "audioldm.pipline.LatentDiffusion",126 "params": {127 "base_learning_rate": 5e-06,128 "linear_start": 0.0015,129 "linear_end": 0.0195,130 "num_timesteps_cond": 1,131 "log_every_t": 200,132 "timesteps": 1000,133 "first_stage_key": "fbank",134 "cond_stage_key": "waveform",135 "latent_t_size": 256,136 "latent_f_size": 16,137 "channels": 8,138 "cond_stage_trainable": True,139 "conditioning_key": "film",140 "monitor": "val/loss_simple_ema",141 "scale_by_std": True,142 "unet_config": {143 "target": "audioldm.latent_diffusion.openaimodel.UNetModel",144 "params": {145 "image_size": 64,146 "extra_film_condition_dim": 512,147 "extra_film_use_concat": True,148 "in_channels": 8,149 "out_channels": 8,150 "model_channels": 128,151 "attention_resolutions": [8, 4, 2],152 "num_res_blocks": 2,153 "channel_mult": [1, 2, 3, 5],154 "num_head_channels": 32,155 "use_spatial_transformer": True,156 },157 },158 "first_stage_config": {159 "base_learning_rate": 4.5e-05,160 "target": "audioldm.variational_autoencoder.autoencoder.AutoencoderKL",161 "params": {162 "monitor": "val/rec_loss",163 "image_key": "fbank",164 "subband": 1,165 "embed_dim": 8,166 "time_shuffle": 1,167 "ddconfig": {168 "double_z": True,169 "z_channels": 8,170 "resolution": 256,171 "downsample_time": False,172 "in_channels": 1,173 "out_ch": 1,174 "ch": 128,175 "ch_mult": [1, 2, 4],176 "num_res_blocks": 2,177 "attn_resolutions": [],178 "dropout": 0.0,179 },180 },181 },182 "cond_stage_config": {183 "target": "audioldm.clap.encoders.CLAPAudioEmbeddingClassifierFreev2",184 "params": {185 "key": "waveform",186 "sampling_rate": 16000,187 "embed_mode": "audio",188 "unconditional_prob": 0.1,189 },190 },191 },192 },193 }194 195 if("-l-" in model_name):196 basic_config["model"]["params"]["unet_config"]["params"]["model_channels"] = 256197 basic_config["model"]["params"]["unet_config"]["params"]["num_head_channels"] = 64198 elif("-m-" in model_name):199 basic_config["model"]["params"]["unet_config"]["params"]["model_channels"] = 192200 basic_config["model"]["params"]["cond_stage_config"]["params"]["amodel"] = "HTSAT-base" # This model use a larger HTAST201 202 return basic_config203 204def get_metadata():205 return {206 "audioldm-s-full": {207 "path": os.path.join(208 CACHE_DIR,209 "audioldm-s-full.ckpt",210 ),211 "url": "https://zenodo.org/record/7600541/files/audioldm-s-full?download=1",212 },213 "audioldm-l-full": {214 "path": os.path.join(215 CACHE_DIR,216 "audioldm-l-full.ckpt",217 ),218 "url": "https://zenodo.org/record/7698295/files/audioldm-full-l.ckpt?download=1",219 },220 "audioldm-s-full-v2": {221 "path": os.path.join(222 CACHE_DIR,223 "audioldm-s-full-v2.ckpt",224 ),225 "url": "https://zenodo.org/record/7698295/files/audioldm-full-s-v2.ckpt?download=1",226 },227 "audioldm-m-text-ft": {228 "path": os.path.join(229 CACHE_DIR,230 "audioldm-m-text-ft.ckpt",231 ),232 "url": "https://zenodo.org/record/7813012/files/audioldm-m-text-ft.ckpt?download=1",233 },234 "audioldm-s-text-ft": {235 "path": os.path.join(236 CACHE_DIR,237 "audioldm-s-text-ft.ckpt",238 ),239 "url": "https://zenodo.org/record/7813012/files/audioldm-s-text-ft.ckpt?download=1",240 },241 "audioldm-m-full": {242 "path": os.path.join(243 CACHE_DIR,244 "audioldm-m-full.ckpt",245 ),246 "url": "https://zenodo.org/record/7813012/files/audioldm-m-full.ckpt?download=1",247 },248 }249 250class MyProgressBar():251 def __init__(self):252 self.pbar = None253 254 def __call__(self, block_num, block_size, total_size):255 if not self.pbar:256 self.pbar=progressbar.ProgressBar(maxval=total_size)257 self.pbar.start()258 259 downloaded = block_num * block_size260 if downloaded < total_size:261 self.pbar.update(downloaded)262 else:263 self.pbar.finish()264 265def download_checkpoint(checkpoint_name="audioldm-s-full"):266 meta = get_metadata()267 if(checkpoint_name not in meta.keys()):268 print("The model name you provided is not supported. Please use one of the following: ", meta.keys())269 270 if not os.path.exists(meta[checkpoint_name]["path"]) or os.path.getsize(meta[checkpoint_name]["path"]) < 2*10**9:271 os.makedirs(os.path.dirname(meta[checkpoint_name]["path"]), exist_ok=True)272 print(f"Downloading the main structure of {checkpoint_name} into {os.path.dirname(meta[checkpoint_name]['path'])}")273 274 urllib.request.urlretrieve(meta[checkpoint_name]["url"], meta[checkpoint_name]["path"], MyProgressBar())275 print(276 "Weights downloaded in: {} Size: {}".format(277 meta[checkpoint_name]["path"],278 os.path.getsize(meta[checkpoint_name]["path"]),279 )280 )281