mustango / audioldm /utils.py
deepanway's picture
Uplaod files
f1069cc
import contextlib
import importlib
from inspect import isfunction
import os
import soundfile as sf
import time
import wave
import urllib.request
import progressbar
CACHE_DIR = os.getenv(
"AUDIOLDM_CACHE_DIR",
os.path.join(os.path.expanduser("~"), ".cache/audioldm"))
def get_duration(fname):
with contextlib.closing(wave.open(fname, 'r')) as f:
frames = f.getnframes()
rate = f.getframerate()
return frames / float(rate)
def get_bit_depth(fname):
with contextlib.closing(wave.open(fname, 'r')) as f:
bit_depth = f.getsampwidth() * 8
return bit_depth
def get_time():
t = time.localtime()
return time.strftime("%d_%m_%Y_%H_%M_%S", t)
def seed_everything(seed):
import random, os
import numpy as np
import torch
random.seed(seed)
os.environ["PYTHONHASHSEED"] = str(seed)
np.random.seed(seed)
torch.manual_seed(seed)
torch.cuda.manual_seed(seed)
torch.backends.cudnn.deterministic = True
torch.backends.cudnn.benchmark = True
def save_wave(waveform, savepath, name="outwav"):
if type(name) is not list:
name = [name] * waveform.shape[0]
for i in range(waveform.shape[0]):
path = os.path.join(
savepath,
"%s_%s.wav"
% (
os.path.basename(name[i])
if (not ".wav" in name[i])
else os.path.basename(name[i]).split(".")[0],
i,
),
)
print("Save audio to %s" % path)
sf.write(path, waveform[i, 0], samplerate=16000)
def exists(x):
return x is not None
def default(val, d):
if exists(val):
return val
return d() if isfunction(d) else d
def count_params(model, verbose=False):
total_params = sum(p.numel() for p in model.parameters())
if verbose:
print(f"{model.__class__.__name__} has {total_params * 1.e-6:.2f} M params.")
return total_params
def get_obj_from_str(string, reload=False):
module, cls = string.rsplit(".", 1)
if reload:
module_imp = importlib.import_module(module)
importlib.reload(module_imp)
return getattr(importlib.import_module(module, package=None), cls)
def instantiate_from_config(config):
if not "target" in config:
if config == "__is_first_stage__":
return None
elif config == "__is_unconditional__":
return None
raise KeyError("Expected key `target` to instantiate.")
return get_obj_from_str(config["target"])(**config.get("params", dict()))
def default_audioldm_config(model_name="audioldm-s-full"):
basic_config = {
"wave_file_save_path": "./output",
"id": {
"version": "v1",
"name": "default",
"root": "/mnt/fast/nobackup/users/hl01486/projects/general_audio_generation/AudioLDM-python/config/default/latent_diffusion.yaml",
},
"preprocessing": {
"audio": {"sampling_rate": 16000, "max_wav_value": 32768},
"stft": {"filter_length": 1024, "hop_length": 160, "win_length": 1024},
"mel": {
"n_mel_channels": 64,
"mel_fmin": 0,
"mel_fmax": 8000,
"freqm": 0,
"timem": 0,
"blur": False,
"mean": -4.63,
"std": 2.74,
"target_length": 1024,
},
},
"model": {
"device": "cuda",
"target": "audioldm.pipline.LatentDiffusion",
"params": {
"base_learning_rate": 5e-06,
"linear_start": 0.0015,
"linear_end": 0.0195,
"num_timesteps_cond": 1,
"log_every_t": 200,
"timesteps": 1000,
"first_stage_key": "fbank",
"cond_stage_key": "waveform",
"latent_t_size": 256,
"latent_f_size": 16,
"channels": 8,
"cond_stage_trainable": True,
"conditioning_key": "film",
"monitor": "val/loss_simple_ema",
"scale_by_std": True,
"unet_config": {
"target": "audioldm.latent_diffusion.openaimodel.UNetModel",
"params": {
"image_size": 64,
"extra_film_condition_dim": 512,
"extra_film_use_concat": True,
"in_channels": 8,
"out_channels": 8,
"model_channels": 128,
"attention_resolutions": [8, 4, 2],
"num_res_blocks": 2,
"channel_mult": [1, 2, 3, 5],
"num_head_channels": 32,
"use_spatial_transformer": True,
},
},
"first_stage_config": {
"base_learning_rate": 4.5e-05,
"target": "audioldm.variational_autoencoder.autoencoder.AutoencoderKL",
"params": {
"monitor": "val/rec_loss",
"image_key": "fbank",
"subband": 1,
"embed_dim": 8,
"time_shuffle": 1,
"ddconfig": {
"double_z": True,
"z_channels": 8,
"resolution": 256,
"downsample_time": False,
"in_channels": 1,
"out_ch": 1,
"ch": 128,
"ch_mult": [1, 2, 4],
"num_res_blocks": 2,
"attn_resolutions": [],
"dropout": 0.0,
},
},
},
"cond_stage_config": {
"target": "audioldm.clap.encoders.CLAPAudioEmbeddingClassifierFreev2",
"params": {
"key": "waveform",
"sampling_rate": 16000,
"embed_mode": "audio",
"unconditional_prob": 0.1,
},
},
},
},
}
if("-l-" in model_name):
basic_config["model"]["params"]["unet_config"]["params"]["model_channels"] = 256
basic_config["model"]["params"]["unet_config"]["params"]["num_head_channels"] = 64
elif("-m-" in model_name):
basic_config["model"]["params"]["unet_config"]["params"]["model_channels"] = 192
basic_config["model"]["params"]["cond_stage_config"]["params"]["amodel"] = "HTSAT-base" # This model use a larger HTAST
return basic_config
def get_metadata():
return {
"audioldm-s-full": {
"path": os.path.join(
CACHE_DIR,
"audioldm-s-full.ckpt",
),
"url": "https://zenodo.org/record/7600541/files/audioldm-s-full?download=1",
},
"audioldm-l-full": {
"path": os.path.join(
CACHE_DIR,
"audioldm-l-full.ckpt",
),
"url": "https://zenodo.org/record/7698295/files/audioldm-full-l.ckpt?download=1",
},
"audioldm-s-full-v2": {
"path": os.path.join(
CACHE_DIR,
"audioldm-s-full-v2.ckpt",
),
"url": "https://zenodo.org/record/7698295/files/audioldm-full-s-v2.ckpt?download=1",
},
"audioldm-m-text-ft": {
"path": os.path.join(
CACHE_DIR,
"audioldm-m-text-ft.ckpt",
),
"url": "https://zenodo.org/record/7813012/files/audioldm-m-text-ft.ckpt?download=1",
},
"audioldm-s-text-ft": {
"path": os.path.join(
CACHE_DIR,
"audioldm-s-text-ft.ckpt",
),
"url": "https://zenodo.org/record/7813012/files/audioldm-s-text-ft.ckpt?download=1",
},
"audioldm-m-full": {
"path": os.path.join(
CACHE_DIR,
"audioldm-m-full.ckpt",
),
"url": "https://zenodo.org/record/7813012/files/audioldm-m-full.ckpt?download=1",
},
}
class MyProgressBar():
def __init__(self):
self.pbar = None
def __call__(self, block_num, block_size, total_size):
if not self.pbar:
self.pbar=progressbar.ProgressBar(maxval=total_size)
self.pbar.start()
downloaded = block_num * block_size
if downloaded < total_size:
self.pbar.update(downloaded)
else:
self.pbar.finish()
def download_checkpoint(checkpoint_name="audioldm-s-full"):
meta = get_metadata()
if(checkpoint_name not in meta.keys()):
print("The model name you provided is not supported. Please use one of the following: ", meta.keys())
if not os.path.exists(meta[checkpoint_name]["path"]) or os.path.getsize(meta[checkpoint_name]["path"]) < 2*10**9:
os.makedirs(os.path.dirname(meta[checkpoint_name]["path"]), exist_ok=True)
print(f"Downloading the main structure of {checkpoint_name} into {os.path.dirname(meta[checkpoint_name]['path'])}")
urllib.request.urlretrieve(meta[checkpoint_name]["url"], meta[checkpoint_name]["path"], MyProgressBar())
print(
"Weights downloaded in: {} Size: {}".format(
meta[checkpoint_name]["path"],
os.path.getsize(meta[checkpoint_name]["path"]),
)
)