tango / audioldm /utils.py
deepanway's picture
add required files
6b448ad
raw
history blame
9.8 kB
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"]),
)
)