include audio gen

gradio_components/prediction.py

@@ -8,13 +8,16 @@ import gradio as gr
 import torch
 from audiocraft.data.audio import audio_write
 from audiocraft.data.audio_utils import convert_audio
-from audiocraft.models import MusicGen
+from audiocraft.models import MusicGen, AudioGen
 from basic_pitch import ICASSP_2022_MODEL_PATH
 from transformers import AutoModelForSeq2SeqLM
 
 
 def load_model(version="facebook/musicgen-melody"):
-    return MusicGen.get_pretrained(version)
+    if version in ["facebook/audiogen-medium"]:
+        return AudioGen.get_pretrained(version)
+    else:
+        return MusicGen.get_pretrained(version)
 
 
 def _do_predictions(