from transformers import PretrainedConfig
from huggingface_integration import AudioCraftForHuggingFace

# Example usage:
model_name = 'facebook/audiogen-medium'
duration = 5

config = PretrainedConfig.from_pretrained(model_name)
audio_craft_model = AudioCraftForHuggingFace(config, model_name, duration)

descriptions = ['creepy cartoon dungeon music']
wav = audio_craft_model(descriptions)

for idx, one_wav in enumerate(wav):
    audio_craft_model.save_wav(one_wav, idx)