| import torch | |
| from audiocraft.models import AudioGen | |
| from audiocraft.data.audio import audio_write | |
| class AudioCraftGenerator(torch.nn.Module): | |
| def __init__(self, model_name='facebook/audiogen-medium', duration=5): | |
| super(AudioCraftGenerator, self).__init__() | |
| self.model = AudioGen.get_pretrained(model_name) | |
| self.model.set_generation_params(duration=duration) | |
| def forward(self, descriptions): | |
| wav = self.model.generate(descriptions) | |
| return wav | |
| def save_wav(self, wav, idx): | |
| audio_write(f'{idx}', wav.cpu(), self.model.sample_rate, strategy="loudness", loudness_compressor=True) |