mrtroydev's picture
Upload folder using huggingface_hub
3883c60 verified
import gc
import os.path
import diffusers
import torch.cuda
import transformers
import librosa
model: diffusers.AudioLDM2Pipeline = None
loaded = False
device: str = None
models = ['cvssp/audioldm2', 'cvssp/audioldm2-large', 'cvssp/audioldm2-music']
def create_model(pretrained='cvssp/audioldm2', map_device='cuda' if torch.cuda.is_available() else 'cpu'):
if is_loaded():
delete_model()
global model, loaded, device
try:
cache_dir = os.path.join('data', 'models', 'audioldm')
model = diffusers.AudioLDM2Pipeline.from_pretrained(pretrained, cache_dir=cache_dir).to(map_device)
device = map_device
loaded = True
except:
pass
def is_loaded():
return loaded
def delete_model():
global model, loaded, clap_model, processor, device
try:
del model
gc.collect()
torch.cuda.empty_cache()
loaded = False
device = None
except:
pass
def generate(prompt='', negative_prompt='', steps=10, duration=5.0, cfg=2.5, seed=-1, callback=None):
if is_loaded():
try:
sample_rate = 16000
seed = seed if seed >= 0 else torch.seed()
torch.manual_seed(seed)
output = model(prompt=prompt, negative_prompt=negative_prompt, audio_length_in_s=duration, num_inference_steps=steps, guidance_scale=cfg, callback=callback)
waveforms = output.audios
waveform = waveforms[0]
return seed, (sample_rate, waveform)
except Exception as e:
return f'An exception occurred: {str(e)}'
return 'No model loaded! Please load a model first.'